import json
import os

from transformers import PreTrainedTokenizer


_RNA_VOCAB = {
    "<cls>": 0, "<pad>": 1, "<eos>": 2, "<unk>": 3,
    "A": 4, "C": 5, "G": 6, "U": 7,
    "R": 8, "Y": 9, "K": 10, "M": 11,
    "S": 12, "W": 13, "B": 14, "D": 15,
    "H": 16, "V": 17, "N": 18, "-": 19,
    "<null_1>": 20, "<null_2>": 21, "<null_3>": 22, "<null_4>": 23,
    "<mask>": 24,
}

_MRNA_VOCAB = {
    "<cls>": 0, "<pad>": 1, "<eos>": 2, "<unk>": 3,
    "GAG": 4, "AAG": 5, "GAA": 6, "CUG": 7, "CAG": 8, "GAU": 9,
    "AAA": 10, "GUG": 11, "GAC": 12, "AUG": 13, "GCC": 14, "AAC": 15,
    "GCU": 16, "AAU": 17, "AUC": 18, "UUC": 19, "GGA": 20, "AUU": 21,
    "GGC": 22, "UUU": 23, "CCA": 24, "AGC": 25, "GCA": 26, "UCU": 27,
    "CUC": 28, "ACC": 29, "CAA": 30, "CCU": 31, "UCC": 32, "ACA": 33,
    "UUG": 34, "GUU": 35, "CUU": 36, "UAC": 37, "ACU": 38, "CCC": 39,
    "UCA": 40, "GUC": 41, "GGU": 42, "CAC": 43, "AGU": 44, "UAU": 45,
    "AGA": 46, "CAU": 47, "GGG": 48, "UGG": 49, "UGC": 50, "AGG": 51,
    "UGU": 52, "AUA": 53, "CGC": 54, "UUA": 55, "GCG": 56, "CGG": 57,
    "CCG": 58, "GUA": 59, "CUA": 60, "ACG": 61, "UCG": 62, "CGA": 63,
    "CGU": 64, "UGA": 65, "UAA": 66, "UAG": 67,
    "<null_1>": 68, "<null_2>": 69, "<null_3>": 70, "<null_4>": 71,
    "<mask>": 72,
}


class RnaFmTokenizer(PreTrainedTokenizer):
    vocab_files_names = {"vocab_file": "vocab.json"}
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(
        self,
        vocab_file=None,
        k_mer: int = 1,
        cls_token="<cls>",
        pad_token="<pad>",
        eos_token="<eos>",
        unk_token="<unk>",
        mask_token="<mask>",
        **kwargs,
    ):
        self.k_mer = k_mer
        if vocab_file and os.path.isfile(vocab_file):
            with open(vocab_file) as f:
                self._vocab = json.load(f)
        else:
            self._vocab = dict(_MRNA_VOCAB if k_mer == 3 else _RNA_VOCAB)
        self._ids_to_tokens = {v: k for k, v in self._vocab.items()}
        super().__init__(
            cls_token=cls_token,
            pad_token=pad_token,
            eos_token=eos_token,
            unk_token=unk_token,
            mask_token=mask_token,
            k_mer=k_mer,
            **kwargs,
        )

    @property
    def vocab_size(self):
        return len(self._vocab)

    def get_vocab(self):
        return dict(self._vocab)

    def _tokenize(self, text):
        if self.k_mer == 1:
            return list(text)
        return [text[i:i + self.k_mer] for i in range(0, len(text), self.k_mer)]

    def _convert_token_to_id(self, token):
        return self._vocab.get(token, self._vocab["<unk>"])

    def _convert_id_to_token(self, index):
        return self._ids_to_tokens.get(index, "<unk>")

    def save_vocabulary(self, save_directory, filename_prefix=None):
        os.makedirs(save_directory, exist_ok=True)
        fname = (filename_prefix + "-" if filename_prefix else "") + "vocab.json"
        path = os.path.join(save_directory, fname)
        with open(path, "w") as f:
            json.dump(self._vocab, f, indent=2)
        return (path,)

    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
        cls = [self.cls_token_id]
        eos = [self.eos_token_id]
        if token_ids_1 is None:
            return cls + token_ids_0 + eos
        return cls + token_ids_0 + eos + cls + token_ids_1 + eos

    def get_special_tokens_mask(self, token_ids_0, token_ids_1=None, already_has_special_tokens=False):
        if already_has_special_tokens:
            return super().get_special_tokens_mask(token_ids_0, token_ids_1, already_has_special_tokens=True)
        mask = [1] + [0] * len(token_ids_0) + [1]
        if token_ids_1 is not None:
            mask += [1] + [0] * len(token_ids_1) + [1]
        return mask

    def create_token_type_ids_from_sequences(self, token_ids_0, token_ids_1=None):
        if token_ids_1 is None:
            return [0] * (len(token_ids_0) + 2)
        return [0] * (len(token_ids_0) + 2) + [0] * (len(token_ids_1) + 2)

    @staticmethod
    def _extract_cds(sequence, cds):
        """Extract CDS region from a sequence, trimmed to a multiple of 3."""
        import numpy as np
        if sum(cds) == 0:
            return sequence[:len(sequence) - (len(sequence) % 3)]
        first = int(np.argmax(cds == 1))
        last = int(len(cds) - 1 - np.argmax(np.flip(cds) == 1)) + 2
        region = sequence[first:last + 1]
        if len(region) % 3 != 0:
            region = region[:-(len(region) % 3)]
        return region

    def batch_encode_with_cds(self, sequences, cds, max_length=None, **kwargs):
        """Encode sequences with CDS extraction (k_mer=3 / mRNA-FM only).

        Applies T->U, extracts the CDS region, chunks to max_length nucleotides
        (aligned to codon boundaries), and encodes each chunk.

        Args:
            sequences: List of raw nucleotide strings (T or U).
            cds: List of numpy arrays marking CDS codon start positions.
            max_length: Nucleotide budget per chunk (defaults to
                (model_max_length - 2) * k_mer).
            **kwargs: Forwarded to batch_encode_plus (e.g. return_tensors,
                padding, add_special_tokens).

        Returns:
            Tuple of (BatchEncoding, chunk_counts) where chunk_counts[i] is the
            number of chunks produced for sequences[i].
        """
        if self.k_mer != 3:
            raise ValueError("batch_encode_with_cds requires k_mer=3 (mRNA-FM tokenizer)")

        budget = max_length if max_length is not None else (self.model_max_length - 2) * self.k_mer
        budget = (budget // self.k_mer) * self.k_mer

        all_chunks = []
        chunk_counts = []

        for seq, c in zip(sequences, cds):
            seq = seq.replace("T", "U").replace("t", "u")
            seq = self._extract_cds(seq, c)
            raw_chunks = [seq[i:i + budget] for i in range(0, max(len(seq), 1), budget)]
            chunks = []
            for chunk in raw_chunks:
                if len(chunk) % self.k_mer != 0:
                    chunk = chunk[:-(len(chunk) % self.k_mer)]
                if chunk:
                    chunks.append(chunk)
            if not chunks:
                chunks = ["AUG"]
            all_chunks.extend(chunks)
            chunk_counts.append(len(chunks))

        enc = self.batch_encode_plus(all_chunks, **kwargs)
        return enc, chunk_counts