# =========================================================
# IMPORT
# =========================================================
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio

# =========================================================
# CONFIG (GIỐNG TRAIN + OFFLINE TEST)
# =========================================================
SR = 16000
N_MELS = 80
N_FFT = 1024
HOP = 256

MAX_LEN_CONF = 160
CROP_SEC = 2.5
NUM_CROPS = 5          # 🔥 GIỐNG OFFLINE
TH_CONF = 0.5          # 🔥 GIỐNG OFFLINE

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# =========================================================
# MEL
# =========================================================
mel_extractor = torchaudio.transforms.MelSpectrogram(
    sample_rate=SR,
    n_fft=N_FFT,
    hop_length=HOP,
    n_mels=N_MELS,
    power=2.0
).to(DEVICE)

db_transform = torchaudio.transforms.AmplitudeToDB(stype="power")

# =========================================================
# MODEL
# =========================================================
class ConvSubsampling(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 64, 3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 64, 3, stride=2, padding=1),
            nn.ReLU(),
        )

    def forward(self, x):
        x = self.conv(x.unsqueeze(1))
        b, c, f, t = x.shape
        return x.view(b, c * f, t).transpose(1, 2)


class ConformerBlock(nn.Module):
    def __init__(self, dim, heads=4, dropout=0.3):
        super().__init__()

        self.ff1 = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, dim * 4),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(dim * 4, dim),
        )

        self.attn = nn.MultiheadAttention(
            dim, heads, dropout=dropout, batch_first=True
        )
        self.norm_attn = nn.LayerNorm(dim)

        self.conv = nn.Sequential(
            nn.Conv1d(dim, dim, 3, padding=1, groups=dim),
            nn.BatchNorm1d(dim),
            nn.GELU(),
            nn.Conv1d(dim, dim, 1),
        )
        self.norm_conv = nn.LayerNorm(dim)

        self.ff2 = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, dim * 4),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(dim * 4, dim),
        )

        self.norm_out = nn.LayerNorm(dim)

    def forward(self, x):
        x = x + 0.5 * self.ff1(x)
        attn, _ = self.attn(x, x, x)
        x = self.norm_attn(x + attn)
        conv = self.conv(x.transpose(1, 2)).transpose(1, 2)
        x = self.norm_conv(x + conv)
        x = x + 0.5 * self.ff2(x)
        return self.norm_out(x)


class AttentivePooling(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.att = nn.Linear(dim, 1)

    def forward(self, x):
        w = torch.softmax(self.att(x), dim=1)
        return (x * w).sum(dim=1)


ENC_DIM = 64 * (N_MELS // 4)

class AntiDeepfakeConformer(nn.Module):
    def __init__(self):
        super().__init__()
        self.subsample = ConvSubsampling()
        self.encoder = nn.Sequential(
            ConformerBlock(ENC_DIM),
            ConformerBlock(ENC_DIM),
            ConformerBlock(ENC_DIM),
        )
        self.pool = AttentivePooling(ENC_DIM)
        self.fc = nn.Sequential(
            nn.Linear(ENC_DIM, 128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, 1)
        )

    def forward(self, x):
        x = self.subsample(x)
        x = self.encoder(x)
        x = self.pool(x)
        return self.fc(x).squeeze(1)

# =========================================================
# LOAD MODEL
# =========================================================
conformer_model = AntiDeepfakeConformer().to(DEVICE)
conformer_model.load_state_dict(
    torch.load("anti_ai_conformer_best.pt", map_location=DEVICE)
)
conformer_model.eval()

# =========================================================
# UTILS (GIỐNG OFFLINE)
# =========================================================
def normalize_soft(audio, target_rms=0.08):
    rms = np.sqrt(np.mean(audio ** 2) + 1e-9)
    scale = min(target_rms / rms, 3.0)
    return audio * scale


def split_crops(audio, sr=SR, crop_sec=CROP_SEC, num_crops=NUM_CROPS):
    crop_len = int(sr * crop_sec)

    if len(audio) <= crop_len:
        return [audio]

    step = max((len(audio) - crop_len) // (num_crops - 1), 1)

    crops = []
    for i in range(num_crops):
        start = i * step
        seg = audio[start:start + crop_len]
        if len(seg) == crop_len:
            crops.append(seg)

    return crops

# =========================================================
# PREDICT (CHUẨN – GIỐNG OFFLINE)
# =========================================================
def predict_conformer(audio: np.ndarray):
    audio = normalize_soft(audio)

    crops = split_crops(audio)
    if len(crops) == 0:
        return 0.0, "UNKNOWN"

    probs = []

    with torch.no_grad():
        for seg in crops:
            seg = torch.tensor(seg, dtype=torch.float32).to(DEVICE)

            mel = db_transform(mel_extractor(seg))

            if mel.shape[1] < MAX_LEN_CONF:
                mel = F.pad(mel, (0, MAX_LEN_CONF - mel.shape[1]))
            else:
                mel = mel[:, :MAX_LEN_CONF]

            mel = mel.unsqueeze(0)
            prob = torch.sigmoid(conformer_model(mel)).item()
            probs.append(prob)

    final_score = float(np.mean(probs))
    label = "AI" if final_score >= TH_CONF else "HUMAN"

    return final_score, label