"""
EigenWave-ASR Model — Internal Architecture (HIDDEN from users)
================================================================

This file contains the full model architecture.
Users never import or see this directly — they use EigenWaveASR class only.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio
import math
import os
import time
import json
import warnings
warnings.filterwarnings('ignore')


# ============================================================================
# INTERNAL MODEL ARCHITECTURE (users never see this)
# ============================================================================

class _MultiScaleRobinFeatures(nn.Module):
    """Multi-scale Robin differential operator feature extractor."""
    def __init__(self, n_mels=80, d_model=384):
        super().__init__()
        self.mel = torchaudio.transforms.MelSpectrogram(
            sample_rate=16000, n_fft=400, hop_length=160,
            n_mels=n_mels, normalized=True
        )
        self.scales = [1, 3, 5]
        self.alphas = nn.ParameterList([nn.Parameter(torch.ones(n_mels) * 0.6) for _ in self.scales])
        self.betas = nn.ParameterList([nn.Parameter(torch.ones(n_mels) * 0.25) for _ in self.scales])
        self.gammas = nn.ParameterList([nn.Parameter(torch.ones(n_mels) * 0.15) for _ in self.scales])
        self.scale_conv = nn.Conv1d(n_mels * len(self.scales), d_model, 1)
        self.conv = nn.Sequential(
            nn.Conv1d(d_model, d_model, 3, stride=2, padding=1),
            nn.BatchNorm1d(d_model),
            nn.GELU(),
            nn.Conv1d(d_model, d_model, 3, padding=1),
            nn.BatchNorm1d(d_model),
            nn.GELU(),
        )
        self.norm = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(0.1)

    def _robin(self, mel, scale, alpha, beta, gamma):
        pad = F.pad(mel, (scale, scale), mode='replicate')
        d1 = (pad[..., 2*scale:] - pad[..., :-2*scale]) / (2 * scale)
        d2 = pad[..., 2*scale:] - 2*pad[..., scale:-scale] + pad[..., :-2*scale]
        d2 = d2 / (scale * scale)
        d1, d2 = d1[..., :mel.size(-1)], d2[..., :mel.size(-1)]
        a, b, g = alpha.view(1,-1,1), beta.view(1,-1,1), gamma.view(1,-1,1)
        return a * mel + b * d1 + g * d2

    def forward(self, wav):
        if wav.dim() == 2: wav = wav.unsqueeze(1)
        mel = torch.log(self.mel(wav).squeeze(1).clamp(min=1e-5))
        feats = [self._robin(mel, s, self.alphas[i], self.betas[i], self.gammas[i])
                 for i, s in enumerate(self.scales)]
        out = self.conv(self.scale_conv(torch.cat(feats, dim=1)))
        return self.dropout(self.norm(out.transpose(1, 2)))


class _RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_len=80000):
        super().__init__()
        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)
        self._build_cache(max_len)

    def _build_cache(self, max_len):
        t = torch.arange(max_len, device=self.inv_freq.device).float()
        freqs = torch.einsum('i,j->ij', t, self.inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        self.register_buffer('cos_cached', emb.cos().unsqueeze(0), persistent=False)
        self.register_buffer('sin_cached', emb.sin().unsqueeze(0), persistent=False)

    def forward(self, x, seq_len):
        if seq_len > self.cos_cached.size(1):
            self._build_cache(seq_len + 1000) # Rebuild with some buffer
        return self.cos_cached[:, :seq_len], self.sin_cached[:, :seq_len]


def _rotate_half(x):
    x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
    return torch.cat([-x2, x1], dim=-1)


class _ConvModule(nn.Module):
    def __init__(self, d_model, kernel_size=31, dropout=0.1):
        super().__init__()
        self.ln = nn.LayerNorm(d_model)
        self.pw1 = nn.Linear(d_model, d_model * 2)
        self.dw = nn.Conv1d(d_model, d_model, kernel_size, padding=kernel_size//2, groups=d_model)
        self.bn = nn.BatchNorm1d(d_model)
        self.pw2 = nn.Linear(d_model, d_model)
        self.drop = nn.Dropout(dropout)

    def forward(self, x):
        x = self.ln(x)
        x = self.pw1(x)
        x, g = x.chunk(2, dim=-1)
        x = x * torch.sigmoid(g)
        x = F.silu(self.bn(self.dw(x.transpose(1,2)))).transpose(1,2)
        return self.drop(self.pw2(x))


class _TransformerBlock(nn.Module):
    def __init__(self, d_model, n_heads, dropout=0.1):
        super().__init__()
        self.d_model, self.n_heads = d_model, n_heads
        self.d_head = d_model // n_heads
        self.norm1 = nn.LayerNorm(d_model)
        self.qkv = nn.Linear(d_model, d_model * 3)
        self.out_proj = nn.Linear(d_model, d_model)
        self.attn_drop = nn.Dropout(dropout)
        self.rope = _RotaryEmbedding(self.d_head)
        self.conv = _ConvModule(d_model, dropout=dropout)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_model * 4), nn.GELU(),
            nn.Dropout(dropout), nn.Linear(d_model * 4, d_model), nn.Dropout(dropout)
        )
        self.scale = 1.0 / math.sqrt(self.d_head)

    def forward(self, x, mask=None):
        B, T, D = x.shape
        h = self.norm1(x)
        qkv = self.qkv(h).reshape(B, T, 3, self.n_heads, self.d_head)
        q, k, v = [t.transpose(1,2) for t in qkv.unbind(2)]
        cos, sin = self.rope(x, T)
        cos, sin = cos.unsqueeze(1), sin.unsqueeze(1)
        q = q * cos + _rotate_half(q) * sin
        k = k * cos + _rotate_half(k) * sin
        attn = F.softmax(torch.matmul(q, k.transpose(-2,-1)) * self.scale, dim=-1)
        attn = self.attn_drop(attn)
        x = x + self.out_proj(torch.matmul(attn, v).transpose(1,2).reshape(B, T, D))
        x = x + self.conv(x)
        return x + self.ffn(self.norm2(x))


class _EigenWaveCore(nn.Module):
    """Core model — never exposed to users."""
    def __init__(self, d_model=384, n_layers=12, n_heads=8, vocab_size=29, dropout=0.0):
        super().__init__()
        self.frontend = _MultiScaleRobinFeatures(80, d_model)
        self.encoder = nn.ModuleList([_TransformerBlock(d_model, n_heads, dropout) for _ in range(n_layers)])
        self.norm = nn.LayerNorm(d_model)
        self.out = nn.Linear(d_model, vocab_size)

    def forward(self, wav):
        x = self.frontend(wav)
        for layer in self.encoder:
            x = layer(x)
        return self.out(self.norm(x))


# ============================================================================
# VOCAB (internal)
# ============================================================================

class _Vocab:
    def __init__(self):
        self.chars = " abcdefghijklmnopqrstuvwxyz'"
        self.c2i = {c: i+1 for i, c in enumerate(self.chars)}
        self.c2i["<b>"] = 0
        self.i2c = {v: k for k, v in self.c2i.items()}
        self.size = len(self.c2i)

    def decode_greedy(self, ids):
        out, prev = [], -1
        for i in ids:
            if i == 0 or i == prev:
                prev = i
                continue
            if i in self.i2c and self.i2c[i] != "<b>":
                out.append(self.i2c[i])
            prev = i
        return ''.join(out)

    def get_labels(self):
        labels = []
        for idx in range(self.size):
            c = self.i2c.get(idx, "")
            labels.append("" if c == "<b>" else c)
        return labels


# ============================================================================
# PUBLIC API — This is what users see and use
# ============================================================================

class EigenWaveASR:
    """
    EigenWave-ASR: High-Performance Speech Recognition Model.
    
    Uses Multi-Scale Robin Features with mathematical differential operators
    for state-of-the-art speech recognition on English audio.
    
    Example:
        >>> model = EigenWaveASR.from_pretrained("sakibhasan/eigenwave-asr")
        >>> text = model.transcribe("path/to/audio.wav")
        >>> print(text)
        "hello world this is a test"
        
        >>> # Batch transcription
        >>> texts = model.transcribe(["audio1.wav", "audio2.wav"])
        
        >>> # From numpy array / tensor (16kHz mono)
        >>> import torch
        >>> audio = torch.randn(16000 * 5)  # 5 seconds
        >>> text = model.transcribe(audio)
    """

    # Best hyperparameters (Optuna-optimized on LibriSpeech dev-clean)
    DEFAULT_CONFIG = {
        "alpha": 0.9268,        # LM weight
        "beta": 0.061,          # Word insertion bonus  
        "temperature": 0.536,   # Softmax temperature
        "beam_width": 50,       # Beam search width
        "sample_rate": 16000,   # Expected sample rate
        "d_model": 384,
        "n_layers": 12,
        "n_heads": 8,
        "vocab_size": 29,
    }

    def __init__(self, model=None, decoder=None, vocab=None, config=None, device=None):
        self.config = config or self.DEFAULT_CONFIG.copy()
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self._model = model
        self._decoder = decoder
        self._vocab = vocab or _Vocab()

    @classmethod
    def from_pretrained(cls, path_or_repo, device=None, use_lm=True, lm_path=None):
        """
        Load a pre-trained EigenWave-ASR model.
        
        Args:
            path_or_repo: Local path to checkpoint (.pt file) OR directory containing it
            device: 'cpu', 'cuda', or None (auto-detect)
            use_lm: Whether to use KenLM language model for decoding (better accuracy)
            lm_path: Path to KenLM .arpa or .bin file (auto-downloads if None)
            
        Returns:
            EigenWaveASR instance ready for transcription
            
        Example:
            >>> model = EigenWaveASR.from_pretrained("./my_checkpoint.pt")
            >>> model = EigenWaveASR.from_pretrained("./model_dir/")
            >>> model = EigenWaveASR.from_pretrained("sakibhasan/eigenwave-asr")
        """
        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        config = cls.DEFAULT_CONFIG.copy()
        vocab = _Vocab()

        # ---- Find checkpoint ----
        checkpoint_path = None
        if os.path.isfile(path_or_repo) and path_or_repo.endswith('.pt'):
            checkpoint_path = path_or_repo
        elif os.path.isdir(path_or_repo):
            # Look for .pt files in directory
            for f in sorted(os.listdir(path_or_repo)):
                if f.endswith('.pt'):
                    checkpoint_path = os.path.join(path_or_repo, f)
                    break
            # Also check for config.json
            config_path = os.path.join(path_or_repo, "config.json")
            if os.path.exists(config_path):
                with open(config_path) as f:
                    config.update(json.load(f))
        else:
            # Try HuggingFace Hub download
            try:
                from huggingface_hub import hf_hub_download
                checkpoint_path = hf_hub_download(
                    repo_id=path_or_repo,
                    filename="model.pt",
                )
                try:
                    config_file = hf_hub_download(repo_id=path_or_repo, filename="config.json")
                    with open(config_file) as f:
                        config.update(json.load(f))
                except Exception:
                    pass
            except ImportError:
                raise ValueError(
                    f"Cannot find model at '{path_or_repo}'. "
                    "For HuggingFace Hub models, install: pip install huggingface_hub"
                )

        if checkpoint_path is None:
            raise FileNotFoundError(f"No .pt checkpoint found at '{path_or_repo}'")

        # ---- Build model ----
        model = _EigenWaveCore(
            d_model=config["d_model"],
            n_layers=config["n_layers"],
            n_heads=config["n_heads"],
            vocab_size=config["vocab_size"],
            dropout=0.0,
        ).to(device)

        # Load weights
        ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
        state_dict = ckpt.get('model', ckpt.get('state_dict', ckpt))
        if isinstance(state_dict, dict) and any(k.startswith('module.') for k in state_dict):
            state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
        model.load_state_dict(state_dict, strict=False)
        model.eval()

        n_params = sum(p.numel() for p in model.parameters())
        print(f"[OK] EigenWave-ASR loaded ({n_params/1e6:.1f}M parameters)")

        # ---- Build decoder ----
        decoder = None
        if use_lm:
            try:
                from pyctcdecode import build_ctcdecoder
                
                # Find or download KenLM
                if lm_path is None:
                    lm_path = cls._find_or_download_lm()
                
                if lm_path and os.path.exists(lm_path):
                    labels = vocab.get_labels()
                    decoder = build_ctcdecoder(
                        labels=labels,
                        kenlm_model_path=lm_path,
                        alpha=config["alpha"],
                        beta=config["beta"],
                    )
                    print(f"[OK] Language model loaded (KenLM beam search)")
                else:
                    print(f"[WARN] KenLM not found — using greedy decoding")
            except ImportError:
                print("[WARN] pyctcdecode not installed — using greedy decoding")
                print("   Install for better accuracy: pip install pyctcdecode kenlm")

        return cls(model=model, decoder=decoder, vocab=vocab, config=config, device=device)

    @staticmethod
    def _find_or_download_lm():
        """Find or download KenLM language model."""
        search_paths = [
            "kenlm_models/4-gram_lower.arpa",
            "kenlm_models/4-gram.arpa",
            "kenlm_models/4-gram.bin",
            "/kaggle/input/librispeech-lm/4-gram.arpa",
        ]
        for p in search_paths:
            if os.path.exists(p):
                return p
        return None

    def _forward_chunked(self, wav, chunk_sec=10):
        """Process long audio in chunks to prevent Out-Of-Memory (OOM) crashes."""
        chunk_len = int(chunk_sec * self.config["sample_rate"])
        
        # If audio is shorter than the chunk size, process normally
        if wav.size(-1) <= chunk_len:
            try:
                return self._model(wav.unsqueeze(0))
            except RuntimeError as e:
                if 'out of memory' in str(e).lower():
                    # If it somehow OOMs on short audio, try CPU
                    self.to("cpu")
                    return self._model(wav.to("cpu").unsqueeze(0))
                raise e
            
        # Otherwise, process in chunks and concatenate logits
        all_logits = []
        for i in range(0, wav.size(-1), chunk_len):
            chunk = wav[..., i:i+chunk_len]
            # Pad last chunk if it's too short (less than 1 sec)
            if chunk.size(-1) < 16000:
                chunk = F.pad(chunk, (0, 16000 - chunk.size(-1)))
            
            try:
                chunk_logits = self._model(chunk.unsqueeze(0))
            except RuntimeError as e:
                if 'out of memory' in str(e).lower():
                    # Rescue strategy for OOMs mid-way
                    import gc; gc.collect()
                    if torch.cuda.is_available(): torch.cuda.empty_cache()
                    self.to("cpu") # Force fallback to CPU
                    chunk = chunk.to("cpu")
                    chunk_logits = self._model(chunk.unsqueeze(0))
                else:
                    raise e
                    
            # Move chunk logits to CPU instantly to save VRAM
            all_logits.append(chunk_logits.cpu())
            
            # Aggressive cleanup
            import gc; gc.collect()
            if torch.cuda.is_available(): torch.cuda.empty_cache()
            
        return torch.cat(all_logits, dim=1).to(self.device)

    def transcribe(self, audio, beam_width=None, temperature=None):
        """
        Transcribe audio to text.
        
        Args:
            audio: One of:
                - str: Path to audio file (.wav, .flac, .mp3, etc.)
                - list[str]: List of audio file paths
                - torch.Tensor: Raw waveform (16kHz, mono)
                - numpy.ndarray: Raw waveform (16kHz, mono)
            beam_width: Override beam search width (default: 50)
            temperature: Override softmax temperature (default: 0.536)
            
        Returns:
            str or list[str]: Transcribed text(s)
            
        Example:
            >>> text = model.transcribe("speech.wav")
            >>> texts = model.transcribe(["a.wav", "b.wav"])
            >>> text = model.transcribe(torch.randn(48000))  # 3 sec
        """
        bw = beam_width or self.config["beam_width"]
        temp = temperature or self.config["temperature"]

        # Handle batch
        if isinstance(audio, list):
            return [self.transcribe(a, bw, temp) for a in audio]

        # Load audio
        wav = self._load_audio(audio)
        wav = wav.to(self.device)

        # Inference with automatic chunking for long audio
        with torch.no_grad():
            logits = self._forward_chunked(wav)

            if self._decoder is not None:
                # Beam search + KenLM (best quality)
                log_probs = F.log_softmax(logits[0] / temp, dim=-1).cpu().numpy()
                text = self._decoder.decode(log_probs, beam_width=bw)
            else:
                # Greedy decode (fallback)
                ids = logits[0].argmax(dim=-1).cpu().tolist()
                text = self._vocab.decode_greedy(ids)

        return text.strip()

    def transcribe_with_details(self, audio, beam_width=None, temperature=None):
        """
        Transcribe with additional details (timing, confidence, etc.)
        
        Returns:
            dict with 'text', 'duration', 'processing_time', 'rtf'
        """
        bw = beam_width or self.config["beam_width"]
        temp = temperature or self.config["temperature"]

        wav = self._load_audio(audio)
        duration = len(wav) / self.config["sample_rate"]
        wav = wav.to(self.device)

        t0 = time.time()
        with torch.no_grad():
            logits = self._forward_chunked(wav)
            if self._decoder is not None:
                log_probs = F.log_softmax(logits[0] / temp, dim=-1).cpu().numpy()
                text = self._decoder.decode(log_probs, beam_width=bw)
            else:
                ids = logits[0].argmax(dim=-1).cpu().tolist()
                text = self._vocab.decode_greedy(ids)
        proc_time = time.time() - t0

        return {
            "text": text.strip(),
            "duration": round(duration, 2),
            "processing_time": round(proc_time, 3),
            "rtf": round(proc_time / duration, 3),
            "real_time": proc_time < duration,
        }

    def _load_audio(self, audio):
        """Load and preprocess audio from various formats."""
        if isinstance(audio, str):
            wav, sr = torchaudio.load(audio)
            wav = wav.mean(0) if wav.shape[0] > 1 else wav.squeeze(0)
            if sr != 16000:
                wav = torchaudio.transforms.Resample(sr, 16000)(wav)
            return wav
        elif isinstance(audio, torch.Tensor):
            if audio.dim() > 1:
                audio = audio.mean(0) if audio.shape[0] > 1 else audio.squeeze(0)
            return audio.float()
        else:
            # numpy array
            return torch.from_numpy(audio).float()

    def to(self, device):
        """Move model to specified device."""
        self.device = str(device)
        self._model = self._model.to(device)
        return self

    @property
    def parameters_count(self):
        """Number of model parameters."""
        return sum(p.numel() for p in self._model.parameters())

    def __repr__(self):
        return (
            f"EigenWaveASR(\n"
            f"  parameters={self.parameters_count/1e6:.1f}M,\n"
            f"  device='{self.device}',\n"
            f"  beam_width={self.config['beam_width']},\n"
            f"  has_lm={self._decoder is not None}\n"
            f")"
        )
