import os
import sys
import pickle
import numpy as np
from torch.utils.data import DataLoader
from stl import mesh as mesh_stl
from pathlib import Path
from data.utils import Frame


class Accumulator():
    def __init__(self):
        self.dict = dict()

    def update(self, dict):
        for key in dict.keys():
            if key in self.dict.keys():
                self.dict[key].append(dict[key])
            else:
                self.dict[key] = [dict[key]]

    def get(self):
        result = dict()
        for key in self.dict.keys():
            result[key] = np.mean(self.dict[key])
        return result

    def reset(self):
        self.dict = dict()


class InfiniteDataLoader(DataLoader):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # Initialize an iterator over the dataset.
        self.dataset_iterator = super().__iter__()

    def __iter__(self):
        return self

    def __next__(self):
        try:
            batch = next(self.dataset_iterator)
        except StopIteration:
            # Dataset exhausted, use a new fresh iterator.
            self.dataset_iterator = super().__iter__()
            batch = next(self.dataset_iterator)
        return batch

def text_to_html(lines):  
    html = ''
    for line in lines:
        html += "<pre>" + line + "</pre>"
    return html



def save_stl(file, mesh, faces):
    obj = mesh_stl.Mesh(np.zeros(faces.shape[0], dtype=mesh_stl.Mesh.dtype))
    for i, f in enumerate(faces):
        for j in range(3):
            obj.vectors[i][j] = mesh[f[j],:]    

    obj.save(file)


def create_temp_frames(folder: str, frame: Frame):
    with open(Path(folder) / 'frames_temp.txt', 'w') as f:
        f.writelines([f'{frame}\n' for _ in range(4)])


def pickle_load(file):
    """
    Load a pickle file.
    """
    with open(file, 'rb') as f:
        loadout = pickle.load(f)

    return loadout


def pickle_dump(loadout, file):
    """
    Dump a pickle file. Create the directory if it does not exist.
    """
    os.makedirs(os.path.dirname(str(file)), exist_ok=True)

    with open(file, 'wb') as f:
        pickle.dump(loadout, f)