import itertools
import fix_ulimit

fix_ulimit.fix_ulimit()

import torch
import logging
from data import Vocab, VocabError
from mpire.exception import StopWorker
from mpire import WorkerPool
from dataclasses import dataclass, field
from typing import List, Optional, Any
from collections.abc import Iterator
import copy
import traceback


class AugmentError(Exception):
    pass


@dataclass
class RawExample:
    # Persistent identifier for the example (e.g. database primary key)
    # Can be any integer, used only debugging/logging/QA
    id: int

    # List of descriptions of the example
    # A description may be any object that the data augmentation pipeline can
    # convert to a string
    descs: List[Any]

    # Example as JSON-serialized MIDI notes
    #
    # One of these notes may be marked descriptionAnchor: true, in which
    # case the system will assume descriptions apply starting at this note
    # (this defaults to the first note if not present)
    example: List[dict]

    # List of possible accompaniments as JSON-serialized MIDI notess
    accomps: List[List[dict]]

    # List of example tags (optional - used by data augmentation)
    example_tags: List[str] = field(default_factory=list)

    # List of lists of accompaniment tags (optional - used by data augmentation)
    accomp_tagss: List[List[str]] = field(default_factory=list)


@dataclass
class AugmentedExample:
    # Persistent identifier for the example, should match the corresponding
    # RawExample
    id: int

    # Description of the example, fed directly to the model
    desc: str

    # Example as JSON-serialized MIDI notes, as in RawExample
    example: List[dict]

    # Accompaniment for the example, as JSON-serialized MIDI notes
    accomp: List[dict]


class HeritableDynamicInstancing:
    @classmethod
    def _inheritors(cls):
        subclasses = set()
        work = [cls]
        while work:
            parent = work.pop()
            for child in parent.__subclasses__():
                if child not in subclasses:
                    subclasses.add(child)
                    work.append(child)
        return subclasses

    @classmethod
    def _class_from_config(cls, config, subkey):
        my_config = copy.deepcopy(config[subkey])
        class_name = my_config.pop("name", None)
        if class_name is None:
            raise ValueError(f"no 'name' key in {subkey} config")
        candidates = [
            subclass
            for subclass in cls._inheritors()
            if subclass.__name__ == class_name
        ]
        if len(candidates) == 0:
            raise ValueError(f"no subclass named {class_name} found")
        if len(candidates) > 1:
            raise ValueError(f"multiple subclasses named {class_name} found")
        return candidates[0], my_config


class Dataset(HeritableDynamicInstancing):
    # returns a Dataset subclass instance from a configuration dictionary subkey,
    # where config[subkey].name is the subclass name
    @classmethod
    def dynamic_from_config(cls, config, subkey):
        subcls, my_config = cls._class_from_config(config, subkey)
        return subcls.from_config(config, my_config)

    # returns a Dataset object from a configuration dictionary
    @classmethod
    def from_config(cls, common_config, my_config):
        raise NotImplementedError("from_config must be implemented by subclass")

    # returns a dictionary { split_number: number_of_examples }
    def num_examples(self) -> dict[int, int]:
        raise NotImplementedError("num_examples must be implemented by subclass")

    # returns an iterator over RawExample objects
    def stream_examples(
        self,
        split: Optional[int],  # split number (None means all splits)
        ranksize: tuple[int, int] = (
            0,
            1,
        ),  # (rank, size) tuple for distributed training
        order: str = "random",  # order of examples, either "random" or "id", if "id" then the order should be deterministic
    ) -> Iterator[RawExample]:
        return self.stream_examples_impl(split, ranksize, order)

    # performs any mutation needed to shuffle the examples in the dataset
    # such that stream_examples will return a different order when called with order="random"
    def shuffle(self):
        raise NotImplementedError("shuffle must be implemented by subclass")

    def stream_examples_impl(
        self, split: Optional[int], ranksize: tuple[int, int], order: str
    ) -> Iterator[RawExample]:
        raise NotImplementedError(
            "stream_examples_impl must be implemented by subclass"
        )


class Augmenter(HeritableDynamicInstancing):
    # returns an Augmenter subclass instance from a configuration dictionary subkey,
    # where config[subkey].name is the subclass name
    @classmethod
    def dynamic_from_config(cls, vocab, config, subkey):
        subcls, my_config = cls._class_from_config(config, subkey)
        return subcls.from_config(vocab, config, my_config)

    # returns an Augmenter object from a configuration dictionary
    @classmethod
    def from_config(cls, vocab, common_config, my_config):
        raise NotImplementedError("from_config must be implemented by subclass")

    # can raise AugmentError
    def augment(self, example: RawExample) -> AugmentedExample:
        raise NotImplementedError("augment must be implemented by subclass")


class TrivialAugmenter(Augmenter):
    @classmethod
    def from_config(cls, vocab, common_config, my_config):
        return cls()

    def augment(self, example: RawExample) -> AugmentedExample:
        return AugmentedExample(
            id=example.id,
            desc=" ".join(str(d) for d in example.descs),
            example=example.example,
            accomp=list(itertools.chain.from_iterable(example.accomps)),
        )


class Batch:
    def __init__(
        self,
        vocab,
        xs,
        to_device,
        texts=None,
        ids=None,
        hashes=None,
        encoder_input_ids=None,
        encoder_attention_mask=None,
    ):
        pad_mask = xs[:, :-1, 0] != vocab.pad.index
        self.ntokens = pad_mask.data.sum()
        self.fill = self.ntokens / (xs.shape[0] * xs.shape[1])
        self.length = xs.shape[1]

        # copy to device
        xs = xs.to(to_device, copy=True)
        self.encoder_input_ids = (
            encoder_input_ids.to(to_device, copy=True)
            if encoder_input_ids is not None
            else None
        )
        self.encoder_attention_mask = (
            encoder_attention_mask.to(to_device, copy=True)
            if encoder_attention_mask is not None
            else None
        )
        self.tgt = xs[:, :-1].contiguous()
        self.tgt_y = xs[:, 1:, 0].contiguous()

        self.texts = texts
        self.ids = ids
        self.hashes = hashes


def worker_prepare_example(worker_state, raw_example: RawExample):
    try:
        state = worker_state["state"]

        try:
            aug_example = state.augmenter.augment(raw_example)
            vec = state.vocab.midi_to_tensor(
                aug_example.example,
                accompany=aug_example.accomp,
                strict=True,
            )
        except (VocabError, AugmentError) as e:
            logging.debug(f"failed to convert example to tensor: {str(e)}")
            return id, False, None, None
        example_hash = state.vocab.hash_tensor(vec) if vec is not False else None
        return (
            aug_example.id,
            vec,
            example_hash,
            aug_example.desc,
        )
    except Exception as e:
        if not isinstance(e, StopWorker):
            traceback.print_exception(e)  # MPIRE mangles exceptions, so print them here
        raise e


class WorkerState:
    def __init__(self, vocab, augmenter):
        self.vocab = vocab
        self.augmenter = augmenter

    def __call__(self, worker_state):
        from util import configure_logging

        configure_logging()
        worker_state["state"] = self


class DataGenerator:
    def __init__(
        self,
        dataset: Dataset,
        vocab: Vocab,
        split: Optional[int],
        ranksize: tuple[int, int],
        batch_size: int,
        seq_len: int,
        seq_len_min: int,
        seq_len_max: Optional[int],
        train_target: str,
        parallelism: int,
        to_device: Any,
        text_tokenize: Any,
        example_continuous: bool,
        pack_batch: bool,
        augmenter: Augmenter,
    ):
        self.dataset = dataset
        self.vocab = vocab
        self.split = split
        self.ranksize = ranksize
        self.batch_size = batch_size
        self.seq_len = seq_len
        self.seq_len_min = seq_len_min
        self.seq_len_max = seq_len_max
        self.train_target = train_target
        self.parallelism = parallelism
        self.to_device = to_device
        self.text_tokenize = text_tokenize
        self.example_continuous = example_continuous
        self.pack_batch = pack_batch
        self.augmenter = augmenter

        if self.text_tokenize is not None:
            assert (
                not self.pack_batch
            ), "full batch packing not supported with text prompts"

    def shuffle(self):
        self.dataset.shuffle()

    def num_examples(self):
        split_examples = self.dataset.num_examples()
        if self.split is None:
            return sum(split_examples.values())
        return split_examples.get(self.split, 0)

    def generate(
        self,
        force_batches=None,
        order="random",
        seq_len_max=None,
    ):
        seq_len_max = self.seq_len_max if self.seq_len_max is not None else self.seq_len

        logging.info(
            f"start DataGenerator.generate split={self.split} ranksize={self.ranksize}"
        )

        with WorkerPool(
            n_jobs=self.parallelism, start_method="spawn", use_worker_state=True
        ) as pool:
            batch = torch.zeros(self.batch_size, self.seq_len, 9, dtype=torch.long)
            batch[:, :, 0] = self.vocab.pad.index
            insert_pos = 0
            insert_b = 0
            maxlen = 0
            num_batches = 0
            batch_ids = []
            batch_hashes = []
            batch_texts = []

            def write_batch():
                nonlocal num_batches, batch, insert_pos, insert_b, maxlen, batch_ids, batch_hashes, batch_texts

                # encoder batch size must be the same as the decoder batch size
                # when there are too few examples to fill the batch, pad with tokenized empty strings
                encoder_input_ids, encoder_attention_mask = (
                    (None, None)
                    if self.text_tokenize is None
                    else self.text_tokenize(
                        batch_texts + ["" for _ in range(self.batch_size - insert_b)]
                    )
                )

                batch_out = Batch(
                    self.vocab,
                    batch,
                    self.to_device,
                    texts=batch_texts,
                    ids=batch_ids,
                    hashes=batch_hashes,
                    encoder_input_ids=encoder_input_ids,
                    encoder_attention_mask=encoder_attention_mask,
                )
                batch = torch.zeros(self.batch_size, self.seq_len, 9, dtype=torch.long)
                batch[:, :, 0] = self.vocab.pad.index
                insert_pos = 0
                insert_b = 0
                maxlen = 0
                batch_ids = []
                batch_hashes = []
                batch_texts = []
                num_batches += 1
                yield batch_out

            map_fn = pool.imap if order == "random" else pool.imap_unordered
            mapped_iterator = map_fn(
                worker_prepare_example,
                self.dataset.stream_examples(
                    split=self.split, ranksize=self.ranksize, order=order
                ),
                chunk_size=8,
                worker_init=WorkerState(self.vocab, self.augmenter),
            )
            continuing_examples = []
            while True:
                if force_batches is not None and num_batches >= force_batches:
                    return

                continued = False
                if len(continuing_examples) > 0:
                    id, vec, hash, text = continuing_examples.pop()
                    continued = True
                else:
                    try:
                        id, vec, hash, text = next(mapped_iterator)
                    except StopIteration:
                        break

                if vec is False:
                    continue
                try:
                    if not continued:
                        if (
                            vec.shape[0] > seq_len_max
                            or vec.shape[0] < self.seq_len_min
                        ):
                            continue

                    insert_len = min(vec.shape[0], self.seq_len - insert_pos)
                    if self.example_continuous and insert_len < vec.shape[0]:
                        continuing_examples.append(
                            (id, vec[insert_len - 1 :], hash, text)
                        )

                    batch[insert_b, insert_pos : insert_pos + insert_len] = vec[
                        :insert_len
                    ]
                    batch_ids.append(id)
                    batch_hashes.append(hash)
                    batch_texts.append(text)
                    insert_pos += insert_len
                    maxlen = max(maxlen, insert_pos)
                    if self.pack_batch:
                        if insert_pos >= self.seq_len:
                            insert_b += 1
                            insert_pos = 0
                    else:
                        insert_b += 1
                        insert_pos = 0

                finally:
                    del vec

                if insert_b == self.batch_size:
                    yield from write_batch()

            if force_batches is None:
                if insert_b > 0 or insert_pos > 0:
                    yield from write_batch()
            else:
                while num_batches < force_batches:
                    yield from write_batch()