# -*- coding: utf-8 -*-
"""
Sanskrit Sandhi Splitter — Transformer Seq2Seq Model (Inference Module)
Version: 0.1.0
Author: Tanuj Saxena (2025)
------------------------------------------------------------
Usage:
    from vedika.sandhi_splitter import split_word, SanskritSplit
"""

import os
import math
import torch
import torch.nn as nn

# ============================================================
# CONFIGURATION
CHECKPOINT_PATH = os.path.join(
    os.path.dirname(__file__),
    "data",
    "tf_checkpoint_split.pth"
)
DEFAULT_CONFIG = {
    "d_model": 256,
    "nhead": 8,
    "num_layers": 4,
    "ffn_mult": 4,
    "dropout": 0.1,
    "max_len": 220,
    "sos_token": "<SOS>",
    "eos_token": "<EOS>",
    "pad_token": "<PAD>",
    "unk_token": "<UNK>",
}


# ============================================================
# POSITIONAL ENCODING
# ============================================================
class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_len: int = 2000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer("pe", pe, persistent=False)

    def forward(self, x):
        return x + self.pe[:, : x.size(1)]


# ============================================================
# TRANSFORMER MODEL
# ============================================================
class TFSeq2Seq(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        d_model: int,
        nhead: int,
        num_layers: int,
        ffn_mult: int,
        dropout: float,
        pad_idx: int,
    ):
        super().__init__()
        self.pad_idx = pad_idx
        self.src_emb = nn.Embedding(vocab_size, d_model, padding_idx=pad_idx)
        self.tgt_emb = nn.Embedding(vocab_size, d_model, padding_idx=pad_idx)
        self.pos = SinusoidalPositionalEncoding(d_model)
        self.tf = nn.Transformer(
            d_model=d_model,
            nhead=nhead,
            num_encoder_layers=num_layers,
            num_decoder_layers=num_layers,
            dim_feedforward=ffn_mult * d_model,
            dropout=dropout,
            batch_first=True,
        )
        self.out = nn.Linear(d_model, vocab_size)

    def forward(self, src, tgt):
        src_emb = self.pos(self.src_emb(src))
        tgt_emb = self.pos(self.tgt_emb(tgt))
        src_kpm = (src == self.pad_idx)
        tgt_kpm = (tgt == self.pad_idx)
        causal_mask = torch.triu(
            torch.ones(tgt.size(1), tgt.size(1), device=tgt.device), diagonal=1
        ).bool()
        out = self.tf(
            src=src_emb,
            tgt=tgt_emb,
            src_key_padding_mask=src_kpm,
            tgt_key_padding_mask=tgt_kpm,
            memory_key_padding_mask=src_kpm,
            tgt_mask=causal_mask,
        )
        return self.out(out)


# ============================================================
# SANSKRIT SPLITTER SINGLETON
# ============================================================
class SanskritSplit:
    """
    Sanskrit Sandhi Splitter — Transformer-based
    Loads a fixed checkpoint and performs inference.
    Singleton pattern for efficiency.
    """
    _instance = None

    def __init__(self):
        if not os.path.exists(CHECKPOINT_PATH):
            raise FileNotFoundError(f"Checkpoint not found at {CHECKPOINT_PATH}")

        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"[INFO] Loading Sanskrit Splitter from: {CHECKPOINT_PATH}")
        ckpt = torch.load(CHECKPOINT_PATH, map_location=self.device)

        self.config = ckpt.get("config", DEFAULT_CONFIG)

        # --- Handle vocab safely ---
        if "vocab" in ckpt:
            self.vocab = ckpt["vocab"]
        else:
            print("[WARN] No vocab found — using auto-generated Unicode vocab.")
            base_chars = [chr(c) for c in range(2304, 2432)]
            self.vocab = {tok: i for i, tok in enumerate(
                [self.config["pad_token"], self.config["sos_token"],
                 self.config["eos_token"], self.config["unk_token"]] + base_chars
            )}

        self.inv_vocab = {v: k for k, v in self.vocab.items()}
        self.pad_idx = self.vocab[self.config["pad_token"]]

        # --- Model ---
        self.model = TFSeq2Seq(
            vocab_size=len(self.vocab),
            d_model=self.config["d_model"],
            nhead=self.config["nhead"],
            num_layers=self.config["num_layers"],
            ffn_mult=self.config["ffn_mult"],
            dropout=self.config["dropout"],
            pad_idx=self.pad_idx,
        ).to(self.device)

        # --- Safe weight loading ---
        state_dict = ckpt["model_state_dict"]
        model_vocab_size = state_dict["src_emb.weight"].shape[0]
        current_vocab_size = len(self.vocab)

        if model_vocab_size != current_vocab_size:
            print(f"[WARN] Vocab size mismatch: checkpoint={model_vocab_size}, current={current_vocab_size}.")
            print("[INFO] Auto-adjusting embedding/output layers to match checkpoint.")
            
            # Rebuild model layers to checkpoint vocab size
            self.model = TFSeq2Seq(
                vocab_size=model_vocab_size,
                d_model=self.config["d_model"],
                nhead=self.config["nhead"],
                num_layers=self.config["num_layers"],
                ffn_mult=self.config["ffn_mult"],
                dropout=self.config["dropout"],
                pad_idx=self.pad_idx,
            ).to(self.device)

        # Finally load the checkpoint weights
        self.model.load_state_dict(state_dict, strict=False)
        self.model.eval()


    # Singleton access
    @classmethod
    def get_instance(cls):
        if cls._instance is None:
            cls._instance = SanskritSplit()
        return cls._instance

    # --------------------------------------------------------
    def _encode(self, text: str) -> torch.Tensor:
        tokens = [self.config["sos_token"]] + list(text) + [self.config["eos_token"]]
        token_ids = [self.vocab.get(t, self.vocab[self.config["unk_token"]]) for t in tokens]
        return torch.tensor(token_ids, dtype=torch.long, device=self.device).unsqueeze(0)

    def _decode(self, token_ids: torch.Tensor) -> str:
        tokens = [self.inv_vocab[int(t)] for t in token_ids]
        clean = [
            t for t in tokens if t not in {
                self.config["sos_token"], self.config["eos_token"], self.config["pad_token"]
            }
        ]
        return "".join(clean)

    # --------------------------------------------------------
    @torch.no_grad()
    def split_word(self, text: str, max_len: int = 220) -> str:
        src = self._encode(text)
        tgt = torch.tensor([[self.vocab[self.config["sos_token"]]]],
                           dtype=torch.long, device=self.device)

        for _ in range(max_len):
            out = self.model(src, tgt)
            next_tok = out[:, -1, :].argmax(-1, keepdim=True)
            tgt = torch.cat([tgt, next_tok], dim=1)
            if next_tok.item() == self.vocab[self.config["eos_token"]]:
                break

        return self._decode(tgt.squeeze(0))

    # --------------------------------------------------------
    def split_sentence(self, sentence: str) -> dict:
        words = sentence.strip().split()
        return {w: self.split_word(w) for w in words}


# ============================================================
# EASY IMPORT FUNCTION
# ============================================================
def split_word(text: str) -> str:
    splitter = SanskritSplit.get_instance()
    return splitter.split_word(text)


def split_sentence(sentence: str) -> dict:
    splitter = SanskritSplit.get_instance()
    return splitter.split_sentence(sentence)


# ============================================================
# SELF TEST
# ============================================================
if __name__ == "__main__":
    print("🧩 Sanskrit Sandhi Splitter Test:")
    print(split_word("रामःसीतायैवनम्"))
