from dataclasses import dataclass
import inspect

import numpy as np
import sentencepiece
import torch
import torchaudio
from transformers import PreTrainedTokenizerFast
from tokenizers import AddedToken

from .modules.featurizer import Featurizer
from .modules.encoder import Encoder
from .modules.decoder import Decoder
from .modules.decoder_v2 import DecoderV2
from .modules.augmenter import SpectrogramAugmentation


class Tokenizer(object):
    def __init__(self, filepath: str, is_sentencepiece: bool = True):
        super(Tokenizer, self).__init__()
        self.is_sentencepiece = is_sentencepiece
        if filepath.endswith(".json"):
            # this is using the huggingface tokenizer, not sentencepiece
            self.is_sentencepiece = False
        if self.is_sentencepiece:
            self._tokenizer = sentencepiece.SentencePieceProcessor()
            self._tokenizer.Load(filepath)
            self._n_vocab = self._tokenizer.vocab_size()
        else:
            self._tokenizer = PreTrainedTokenizerFast(
                tokenizer_file=filepath,
                unk_token="<unk>",
                pad_token="<pad>",
            )
            self._tokenizer.add_special_tokens(
                {"additional_special_tokens": [AddedToken("\n")]}
            )
            self._n_vocab = self._tokenizer.vocab_size

    @property
    def n_vocab(self):
        return self._n_vocab

    @property
    def vocab(self):
        if self.is_sentencepiece:
            return {
                self._tokenizer.id_to_piece(_id): _id
                for _id in range(self._tokenizer.get_piece_size())
            }
        else:
            return self._tokenizer.get_vocab()

    @property
    def inv_vocab(self):
        return {v: k for k, v in self.vocab.items()}

    def encode(self, text: str) -> np.ndarray:
        if self.is_sentencepiece:
            return np.array(self._tokenizer.encode(text))
        else:
            # not sure i understand this tokenizer...
            tokens = self._tokenizer(text, add_special_tokens=False)["input_ids"]
            if any(t >= self._n_vocab for t in tokens):
                # prrint(tokens)
                # Ensure tokens are within vocabulary range
                tokens = [t for t in tokens if t < self._n_vocab]
            return np.array(tokens)

    def decode(self, encoded_text) -> str:
        if not len(encoded_text):
            return ""
        text_to_decode = None
        if type(encoded_text) is list:
            if type(encoded_text[0]) is int:
                text_to_decode = encoded_text
            else:
                text_to_decode = [int(i) for i in encoded_text.tolist()]
        elif isinstance(encoded_text, (torch.Tensor, np.ndarray)):
            text_to_decode = [int(i) for i in encoded_text.tolist()]
        else:
            raise ValueError(
                f"Unsupported array type for decoding: {type(encoded_text)}"
            )
        if text_to_decode is not None:
            if self.is_sentencepiece:
                return self._tokenizer.decode(text_to_decode)
            else:
                return self._tokenizer.decode(text_to_decode, skip_special_tokens=False)
        else:
            raise ValueError(
                f"Unsupported array type for decoding: {type(encoded_text)}"
            )

    def decode_pieces(self, pieces: list[str]) -> str:
        if self.is_sentencepiece:
            return self._tokenizer.decode_pieces(pieces)
        else:
            # TO be worked on: sth weird with asian languages
            return "".join(pieces)

    def decode_logits(self, logits, prior_text=None, return_timing_indices=False):
        """Decode the logits to text, optionally return the timing indices.

        Args:
            logits: The logits to decode.
            prior_text: The prior text to use for decoding.
            return_timing_indices: Whether to return the timing indices.

        Returns:
            The decoded text.
            If return_timing_indices is True, the timing indices are returned.
            The format is a list of lists,
            where each inner list contains the start and end indices of the token.
        """
        # from time import time
        # start_time = time()
        if prior_text is not None:
            # print(prior_text)
            # we want to find the tokens from the prior text and only look for those
            # note that if the true lyrics are not contained within the prior text
            # this method won't work!
            expected_tokens = self.encode(prior_text.lower())
            set_of_tokens = set(token for token in expected_tokens)
            set_of_tokens.add(self.n_vocab)
            unique_tokens = sorted(list(set_of_tokens))
            # print(f"finish enocding {time() - start_time:.3f} seconds.")
            # only need to work on this subset of axis -- reduces the search dim
            sublogits = logits.take(indices=unique_tokens, axis=1)
            # print(sublogits.shape)
            sub_logit_id_to_token_map = {}
            for i in range(sublogits.shape[1]):
                sub_logit_id_to_token_map[i] = unique_tokens[i]
            # print(sub_logit_id_to_token_map)
        else:
            sublogits = logits
        # print(f"finish enocding and masking {time() - start_time:.3f} seconds.")
        max_prob_idxs = np.argmax(sublogits, axis=1)
        if prior_text is not None:
            max_prob_idxs = [sub_logit_id_to_token_map[idx] for idx in max_prob_idxs]
        # print(f"finish argmax {time() - start_time:.3f} seconds.")
        tokens = []
        token_indices = []
        prev_token = None
        for idx, current_token in enumerate(max_prob_idxs):
            # print(idx, current_token)
            if current_token == prev_token:
                continue
            if current_token != self.n_vocab:
                tokens.append(int(current_token))
                token_indices.append([idx, idx])
            elif current_token == self.n_vocab:
                if token_indices and token_indices[-1][1] == token_indices[-1][0]:
                    token_indices[-1][1] = idx
            prev_token = current_token
        # print(f"finish everything {time() - start_time:.3f} seconds.")
        if return_timing_indices:
            return tokens, token_indices
        else:
            return self.decode(tokens)


def collate_fn(arr_list, pad_id=0, fixed_len=None):
    arr_len_list = [
        torch.tensor(arr.shape[-1], dtype=torch.long, device=arr.device)
        for arr in arr_list
    ]
    padded_arr_len_list = []
    padded_arr_list = []
    max_arr_len = fixed_len if fixed_len is not None else max(arr_len_list).item()
    for arr, arr_len in zip(arr_list, arr_len_list):
        pad_len = max_arr_len - arr_len.item()
        if pad_len >= 0:
            padded_arr_list.append(
                torch.nn.functional.pad(
                    arr, (0, max_arr_len - arr_len.item()), value=pad_id
                )
            )
            padded_arr_len_list.append(arr_len)
        else:
            padded_arr_list.append(arr[..., :max_arr_len])
            padded_arr_len_list.append(
                torch.tensor(max_arr_len, dtype=torch.long, device=arr.device)
            )
    padded_arrays = torch.stack(padded_arr_list)
    padded_arr_len = torch.stack(padded_arr_len_list)
    return padded_arrays, padded_arr_len


@dataclass
class HootConfig:
    n_layers: int = 18
    n_embd: int = 512
    n_classes: int = 1024
    n_augment_freq_masks: int = 2
    decoder_type: str = "v1"
    max_text_len: int = 2000  # controls how long the text tokens


class Hoot(torch.nn.Module):
    def __init__(self, config=None):
        super(Hoot, self).__init__()
        self.config = config if config is not None else HootConfig()
        self.featurizer = Featurizer(preemph=None)
        self.encoder = Encoder(self.config.n_layers, self.config.n_embd)
        if self.config.decoder_type == "v1":
            self.decoder = Decoder(self.config.n_embd, self.config.n_classes)
        elif self.config.decoder_type == "v2":
            self.decoder = DecoderV2(self.config.n_embd, self.config.n_classes)
        else:
            raise ValueError(f"Invalid decoder type: {self.config.decoder_type}")
        self.loss = torch.nn.CTCLoss(
            blank=self.config.n_classes, reduction="none", zero_infinity=False
        )
        self.spec_augmentation = SpectrogramAugmentation(
            freq_masks=self.config.n_augment_freq_masks
        )
        # init all weights
        self.apply(self._init_weights)
        # for pn, p in self.named_parameters():
        #     if pn.endswith("c_proj.weight"):
        #         torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))

    def forward(
        self,
        signals,
        signals_len,
        targets=None,
        targets_len=None,
        return_decoded=False,
        return_encoded_embedding=False,
        text_input=None,
    ):
        assert (targets is None) == (targets_len is None)
        processed_signal, processed_signal_length = self.featurizer(
            input_signal=signals, length=signals_len
        )
        if self.spec_augmentation is not None and self.training:
            # print("augment")
            new_signal = self.spec_augmentation(
                input_spec=processed_signal, length=signals_len
            )
            # print("post-aug check", torch.equal(new_signal, processed_signal))
            processed_signal = new_signal
        encoded, encoded_len = self.encoder(
            audio_signal=processed_signal, length=processed_signal_length
        )
        if return_encoded_embedding:
            return encoded
        if self.config.decoder_type == "v1":
            decoded = self.decoder(encoder_output=encoded)
        elif self.config.decoder_type == "v2":
            if text_input is None:
                text_input = torch.zeros(
                    encoded.shape[0], 1, device=encoded.device, dtype=torch.long
                )
            else:
                # Ensure text_input values are less than vocab size (n_classes)
                text_input = (
                    text_input[:, : self.config.max_text_len].to(encoded.device).long()
                )
                # Clamp values to be within valid vocabulary range
                # TODO: some tokenizers are very weird...
                max_token_id = text_input.max().item()
                if max_token_id > self.config.n_classes:
                    print(
                        f"Clamping token IDs: max={max_token_id}, n_classes={self.config.n_classes}"
                    )
                    text_input = torch.clamp(text_input, max=self.config.n_classes)
            decoded = self.decoder(encoder_output=encoded, text_input=text_input)
        if targets is None:
            return decoded, encoded_len
        loss = (
            self.loss(decoded.swapaxes(0, 1), targets, encoded_len, targets_len).sum()
            / targets_len.sum()  # mean_volume, longer samples weigh more
        )
        if not return_decoded:
            return loss
        else:
            return loss, decoded

    def infer(self, filepath):
        reset_to_train = False
        if self.training:
            self.eval()
            reset_to_train = True
        arr, sr = torchaudio.load(filepath)
        assert arr.shape[0] in (1, 2)
        arr = torchaudio.functional.resample(arr.sum(axis=0), sr, 16_000)
        arr = arr.to(next(self.parameters()).device)
        arrays, arrays_len = collate_fn([arr])
        with torch.no_grad():
            decoded, _ = self.forward(arrays, arrays_len)
        # no need to unpad cause batch 1
        logits = decoded[0].detach().cpu().numpy()
        if reset_to_train:
            self.train()
        return logits

    @classmethod
    def from_nemo(cls, checkpoint_fp):
        model = cls(HootConfig())
        d = torch.load(checkpoint_fp, map_location="cpu", weights_only=False)
        d["featurizer._mel_spec_extractor.spectrogram.window"] = d[
            "preprocessor.featurizer.window"
        ]
        del d["preprocessor.featurizer.window"]
        d["featurizer._mel_spec_extractor.mel_scale.fb"] = torch.swapaxes(
            d["preprocessor.featurizer.fb"][0], 0, 1
        )
        del d["preprocessor.featurizer.fb"]
        model.load_state_dict(d)
        return model

    @classmethod
    def from_checkpoint(cls, checkpoint_fp):
        d = torch.load(checkpoint_fp, map_location="cpu", weights_only=False)
        model = cls(HootConfig(**d["model_args"]))
        model.load_state_dict(d["model"])
        return model

    def _init_weights(self, module):
        if isinstance(module, torch.nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, torch.nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def configure_optimizers(
        self, weight_decay, learning_rate, betas, device_type, use_fused=True
    ):
        # start with all of the candidate parameters
        param_dict = {pn: p for pn, p in self.named_parameters()}
        # filter out those that do not require grad
        param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
        # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
        # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
        optim_groups = [
            {"params": decay_params, "weight_decay": weight_decay},
            {"params": nodecay_params, "weight_decay": 0.0},
        ]
        num_decay_params = sum(p.numel() for p in decay_params)
        num_nodecay_params = sum(p.numel() for p in nodecay_params)
        print(
            f"num decayed parameter tensors: {len(decay_params)},"
            f" with {num_decay_params:,} parameters"
        )
        print(
            f"num non-decayed parameter tensors: {len(nodecay_params)},"
            f" with {num_nodecay_params:,} parameters"
        )
        # Create AdamW optimizer and use the fused version if it is available
        fused_available = "fused" in inspect.signature(torch.optim.AdamW).parameters
        extra_args = (
            dict(fused=True)
            if fused_available and device_type == "cuda" and use_fused
            else dict()
        )
        optimizer = torch.optim.AdamW(
            optim_groups, lr=learning_rate, betas=betas, **extra_args
        )
        print(f"using fused AdamW: {use_fused}")
        return optimizer