import os
from typing import Dict, Optional, Tuple
from pathlib import Path

import librosa
import numpy as np
import torch
import torch.nn as nn

# ============================================================
# Config inferred from the training pipeline
# ============================================================
SR = 16000
TARGET_SECONDS = 1.0
TARGET_LEN = int(SR * TARGET_SECONDS)
N_MELS = 40
HOP_LENGTH = int(SR * 0.01)   # 10 ms -> 160
WIN_LENGTH = int(SR * 0.025)  # 25 ms -> 400
N_FFT = 400
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

ORDERED_LABELS = [
    "five", "zero", "yes", "seven", "no", "nine",
    "down", "one", "go", "two", "stop", "six",
    "unknown", "silence",
]
IDX2LABEL = {i: label for i, label in enumerate(ORDERED_LABELS)}

BASE_DIR = Path(__file__).resolve().parent

BASELINE_CKPT = BASE_DIR / "baseline_cnn14_best.pth"
CRNN_CKPT = BASE_DIR / "crnn_specaugment14_best.pth"


# ============================================================
# Models copied from the PDF pipeline
# ============================================================
class CNN2DBaseline(nn.Module):
    def __init__(self, num_classes: int, dropout: float = 0.3):
        super().__init__()

        def block(in_ch: int, out_ch: int) -> nn.Sequential:
            return nn.Sequential(
                nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1, bias=False),
                nn.BatchNorm2d(out_ch),
                nn.ReLU(inplace=True),
                nn.MaxPool2d(2),
            )

        self.backbone = nn.Sequential(
            block(1, 16),
            block(16, 32),
            block(32, 64),
            block(64, 128),
        )
        self.dropout = nn.Dropout(dropout)
        self.fc = nn.Linear(128, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.backbone(x)
        x = x.mean(dim=[2, 3])
        x = self.dropout(x)
        return self.fc(x)


class CRNN(nn.Module):
    def __init__(self, num_classes: int, rnn_hidden: int = 128, dropout: float = 0.2):
        super().__init__()
        self.cnn = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1, bias=False),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d((2, 2)),
            nn.Conv2d(32, 64, 3, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d((2, 2)),
            nn.Conv2d(64, 128, 3, padding=1, bias=False),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d((2, 1)),
        )
        self.dropout = nn.Dropout(dropout)
        self.bigru = nn.GRU(
            input_size=128 * 5,
            hidden_size=rnn_hidden,
            num_layers=1,
            batch_first=True,
            bidirectional=True,
        )
        self.fc = nn.Linear(rnn_hidden * 2, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.cnn(x)
        b, c, freq, time = x.shape
        x = x.permute(0, 3, 1, 2).contiguous().view(b, time, c * freq)
        x, _ = self.bigru(x)
        x = x.mean(dim=1)
        x = self.dropout(x)
        return self.fc(x)


@torch.no_grad()
def load_model(model_name: str) -> nn.Module:
    if model_name == "Baseline CNN2D":
        model = CNN2DBaseline(num_classes=len(ORDERED_LABELS), dropout=0.2)
        ckpt_path = BASELINE_CKPT
    else:
        model = CRNN(num_classes=len(ORDERED_LABELS), rnn_hidden=128, dropout=0.3)
        ckpt_path = CRNN_CKPT

    if not os.path.exists(ckpt_path):
        raise FileNotFoundError(
            f"Khong tim thay checkpoint: {ckpt_path}. Hay dat file checkpoint cung thu muc voi app.py"
        )

    state_dict = torch.load(ckpt_path, map_location="cpu")
    model.load_state_dict(state_dict)
    model.eval().to(DEVICE)
    return model


MODELS: Dict[str, nn.Module] = {
    "Baseline CNN2D": load_model("Baseline CNN2D"),
    "CRNN + SpecAugment": load_model("CRNN + SpecAugment"),
}


def _to_mono_resampled(audio: np.ndarray, sample_rate: int) -> np.ndarray:
    audio = np.asarray(audio)
    if audio.ndim == 2:
        audio = audio.mean(axis=1)
    audio = audio.astype(np.float32)
    max_abs = float(np.max(np.abs(audio))) if audio.size else 1.0
    if max_abs > 1.5:
        audio = audio / 32768.0
    if sample_rate != SR:
        audio = librosa.resample(audio, orig_sr=sample_rate, target_sr=SR)
    return audio.astype(np.float32)


def _fix_length(wav: np.ndarray) -> np.ndarray:
    if wav.shape[0] < TARGET_LEN:
        wav = np.pad(wav, (0, TARGET_LEN - wav.shape[0]))
    elif wav.shape[0] > TARGET_LEN:
        wav = wav[:TARGET_LEN]
    return wav.astype(np.float32)


def _logmel(wav: np.ndarray) -> np.ndarray:
    mel = librosa.feature.melspectrogram(
        y=wav,
        sr=SR,
        n_fft=N_FFT,
        hop_length=HOP_LENGTH,
        win_length=WIN_LENGTH,
        n_mels=N_MELS,
        center=False,
        power=2.0,
    )
    logmel = librosa.power_to_db(mel, ref=1.0)
    return logmel.astype(np.float32)


def preprocess_numpy_audio(audio: np.ndarray, sample_rate: int) -> torch.Tensor:
    wav = _to_mono_resampled(audio, sample_rate)
    wav = _fix_length(wav)
    logmel = _logmel(wav)
    x = torch.from_numpy(logmel).unsqueeze(0).unsqueeze(0)  # [1, 1, 40, T]
    return x.to(DEVICE)


@torch.no_grad()
def run_prediction(audio_input: Optional[Tuple[int, np.ndarray]], model_name: str, top_k: int = 5) -> Tuple[str, float, Dict[str, float]]:
    if audio_input is None:
        return "", 0.0, {}

    sample_rate, audio = audio_input
    if audio is None or len(audio) == 0:
        return "", 0.0, {}

    model = MODELS[model_name]
    x = preprocess_numpy_audio(audio, sample_rate)
    logits = model(x)
    probs = torch.softmax(logits, dim=1).squeeze(0).cpu().numpy()

    order = np.argsort(probs)[::-1]
    top_k_results = {IDX2LABEL[int(i)]: float(probs[int(i)]) for i in order[:top_k]}
    pred_idx = int(order[0])
    pred_label = IDX2LABEL[pred_idx]
    pred_conf = float(probs[pred_idx])

    return pred_label, pred_conf, top_k_results