import os
import sys
import ast
import torch
import itertools
import collections

sys.path.append(os.getcwd())

from main.library.speaker_diarization.speechbrain import if_main_process, ddp_barrier
from main.library.speaker_diarization.features import register_checkpoint_hooks, mark_as_saver, mark_as_loader

DEFAULT_UNK = "<unk>"
DEFAULT_BOS = "<bos>"
DEFAULT_EOS = "<eos>"
DEFAULT_BLANK = "<blank>"

@register_checkpoint_hooks
class CategoricalEncoder:
    VALUE_SEPARATOR = " => "
    EXTRAS_SEPARATOR = "================\n"

    def __init__(self, starting_index=0, **special_labels):
        self.lab2ind = {}
        self.ind2lab = {}
        self.starting_index = starting_index
        self.handle_special_labels(special_labels)

    def handle_special_labels(self, special_labels):
        if "unk_label" in special_labels: self.add_unk(special_labels["unk_label"])

    def __len__(self):
        return len(self.lab2ind)

    @classmethod
    def from_saved(cls, path):
        obj = cls()
        obj.load(path)
        return obj

    def update_from_iterable(self, iterable, sequence_input=False):
        label_iterator = itertools.chain.from_iterable(iterable) if sequence_input else iter(iterable)
        for label in label_iterator:
            self.ensure_label(label)

    def update_from_didataset(self, didataset, output_key, sequence_input=False):
        with didataset.output_keys_as([output_key]):
            self.update_from_iterable((data_point[output_key] for data_point in didataset), sequence_input=sequence_input)

    def limited_labelset_from_iterable(self, iterable, sequence_input=False, n_most_common=None, min_count=1):
        label_iterator = itertools.chain.from_iterable(iterable) if sequence_input else iter(iterable)
        counts = collections.Counter(label_iterator)

        for label, count in counts.most_common(n_most_common):
            if count < min_count: break
            self.add_label(label)

        return counts

    def load_or_create(self, path, from_iterables=[], from_didatasets=[], sequence_input=False, output_key=None, special_labels={}):
        try:
            if if_main_process():
                if not self.load_if_possible(path):
                    for iterable in from_iterables:
                        self.update_from_iterable(iterable, sequence_input)

                    for didataset in from_didatasets:
                        if output_key is None: raise ValueError
                        self.update_from_didataset(didataset, output_key, sequence_input)

                    self.handle_special_labels(special_labels)
                    self.save(path)
        finally:
            ddp_barrier()
            self.load(path)

    def add_label(self, label):
        if label in self.lab2ind: raise KeyError
        index = self._next_index()

        self.lab2ind[label] = index
        self.ind2lab[index] = label

        return index

    def ensure_label(self, label):
        if label in self.lab2ind: return self.lab2ind[label]
        else: return self.add_label(label)

    def insert_label(self, label, index):
        if label in self.lab2ind: raise KeyError
        else: self.enforce_label(label, index)

    def enforce_label(self, label, index):
        index = int(index)

        if label in self.lab2ind:
            if index == self.lab2ind[label]: return
            else: del self.ind2lab[self.lab2ind[label]]

        if index in self.ind2lab:
            saved_label = self.ind2lab[index]
            moving_other = True
        else: moving_other = False

        self.lab2ind[label] = index
        self.ind2lab[index] = label

        if moving_other:
            new_index = self._next_index()
            self.lab2ind[saved_label] = new_index
            self.ind2lab[new_index] = saved_label

    def add_unk(self, unk_label=DEFAULT_UNK):
        self.unk_label = unk_label
        return self.add_label(unk_label)

    def _next_index(self):
        index = self.starting_index
        while index in self.ind2lab:
            index += 1

        return index

    def is_continuous(self):
        indices = sorted(self.ind2lab.keys())
        return self.starting_index in indices and all(j - i == 1 for i, j in zip(indices[:-1], indices[1:]))

    def encode_label(self, label, allow_unk=True):
        self._assert_len()

        try:
            return self.lab2ind[label]
        except KeyError:
            if hasattr(self, "unk_label") and allow_unk: return self.lab2ind[self.unk_label]
            elif hasattr(self, "unk_label") and not allow_unk: raise KeyError
            elif not hasattr(self, "unk_label") and allow_unk: raise KeyError
            else: raise KeyError

    def encode_label_torch(self, label, allow_unk=True):
        return torch.LongTensor([self.encode_label(label, allow_unk)])

    def encode_sequence(self, sequence, allow_unk=True):
        self._assert_len()
        return [self.encode_label(label, allow_unk) for label in sequence]

    def encode_sequence_torch(self, sequence, allow_unk=True):
        return torch.LongTensor([self.encode_label(label, allow_unk) for label in sequence])

    def decode_torch(self, x):
        self._assert_len()
        decoded = []

        if x.ndim == 1:  
            for element in x:
                decoded.append(self.ind2lab[int(element)])
        else:
            for subtensor in x:
                decoded.append(self.decode_torch(subtensor))

        return decoded

    def decode_ndim(self, x):
        self._assert_len()
        try:
            decoded = []
            for subtensor in x:
                decoded.append(self.decode_ndim(subtensor))

            return decoded
        except TypeError:  
            return self.ind2lab[int(x)]

    @mark_as_saver
    def save(self, path):
        self._save_literal(path, self.lab2ind, self._get_extras())

    def load(self, path):
        lab2ind, ind2lab, extras = self._load_literal(path)
        self.lab2ind = lab2ind
        self.ind2lab = ind2lab
        self._set_extras(extras)

    @mark_as_loader
    def load_if_possible(self, path, end_of_epoch=False):
        del end_of_epoch

        try:
            self.load(path)
        except FileNotFoundError:
            return False
        except (ValueError, SyntaxError):
            return False
        
        return True 

    def expect_len(self, expected_len):
        self.expected_len = expected_len

    def ignore_len(self):
        self.expected_len = None

    def _assert_len(self):
        if hasattr(self, "expected_len"):
            if self.expected_len is None: return
            if len(self) != self.expected_len: raise RuntimeError
        else:
            self.ignore_len()
            return

    def _get_extras(self):
        extras = {"starting_index": self.starting_index}
        if hasattr(self, "unk_label"): extras["unk_label"] = self.unk_label

        return extras

    def _set_extras(self, extras):
        if "unk_label" in extras: self.unk_label = extras["unk_label"]
        self.starting_index = extras["starting_index"]

    @staticmethod
    def _save_literal(path, lab2ind, extras):
        with open(path, "w", encoding="utf-8") as f:
            for label, ind in lab2ind.items():
                f.write(repr(label) + CategoricalEncoder.VALUE_SEPARATOR + str(ind) + "\n")

            f.write(CategoricalEncoder.EXTRAS_SEPARATOR)

            for key, value in extras.items():
                f.write(repr(key) + CategoricalEncoder.VALUE_SEPARATOR + repr(value) + "\n")

            f.flush()

    @staticmethod
    def _load_literal(path):
        lab2ind, ind2lab, extras = {}, {}, {}

        with open(path, encoding="utf-8") as f:
            for line in f:
                if line == CategoricalEncoder.EXTRAS_SEPARATOR: break
                literal, ind = line.strip().split(CategoricalEncoder.VALUE_SEPARATOR, maxsplit=1)
                label = ast.literal_eval(literal)
                lab2ind[label] = int(ind)
                ind2lab[ind] = label

            for line in f:
                literal_key, literal_value = line.strip().split(CategoricalEncoder.VALUE_SEPARATOR, maxsplit=1)
                extras[ast.literal_eval(literal_key)] = ast.literal_eval(literal_value)
                
        return lab2ind, ind2lab, extras