"""
Run MLP and MLP+attention on test_data.csv; report metrics for both + ensemble.
Requires: manifests/test_data.csv with nativity_status, both models in models/
"""
import os
import sys
import pandas as pd
import numpy as np
import torch
import librosa
import requests
import io
from transformers import Wav2Vec2Model, AutoFeatureExtractor
from tqdm import tqdm
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, confusion_matrix

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from model import NativityClassifier, NativityClassifierWithAttention
from extract_embeddings import get_mfcc_from_audio, get_pitch_stats, PROSODY_EMBED_DIM, EMBED_DIM_WITH_MFCC

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
input_csv = "manifests/test_data.csv"
HIDDEN_DIMS = (384, 256, 128)
DROPOUT = 0.35
EVAL_THRESHOLD_DEFAULT = 0.35  # fallback if best_threshold.txt missing
PROSODY_LAYERS = (4, 6, 8, 10)
BEST_THRESHOLD_PATH = "models/best_threshold.txt"
CHUNK_SIZE = 160000
EMBED_DIM_FILE = "manifests/embedding_dim.txt"
THRESHOLD_SWEEP_TEST = [0.25, 0.28, 0.30, 0.32]  # report test F1 at these thresholds

# Load wav2vec once
model_name = "facebook/wav2vec2-large-xlsr-53"
feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
wav2vec = Wav2Vec2Model.from_pretrained(model_name).to(device)
wav2vec.eval()


def get_embedding(audio_np, embed_dim=PROSODY_EMBED_DIM):
    """Return 2051-d (mean+std+pitch) or 2090-d (+ MFCC) to match training."""
    with torch.no_grad():
        input_values = feature_extractor(
            audio_np, return_tensors="pt", sampling_rate=16000
        ).input_values.to(device)
        outputs = wav2vec(input_values, output_hidden_states=True)
        h = outputs.hidden_states
        num_layers = len(h) - 1
        layer_idx = [min(max(l, 1), num_layers) for l in PROSODY_LAYERS]
        embs = [h[i] for i in layer_idx]
        time_seq = torch.stack(embs, dim=0).mean(dim=0).squeeze(0)  # (T, 1024)
        emb_mean = time_seq.mean(dim=0).cpu().numpy().astype(np.float32)
        emb_std = time_seq.std(dim=0).cpu().numpy().astype(np.float32)
        np.nan_to_num(emb_std, copy=False, nan=0.0, posinf=0.0, neginf=0.0)
    pitch = get_pitch_stats(audio_np, sr=16000)
    out = np.concatenate([emb_mean, emb_std, pitch]).astype(np.float32)
    if embed_dim == EMBED_DIM_WITH_MFCC:
        mfcc_39 = get_mfcc_from_audio(audio_np, sr=16000)
        out = np.concatenate([out, mfcc_39]).astype(np.float32)
    n = np.linalg.norm(out)
    if n > 1e-8:
        out = out / n
    return out


def load_audio(url):
    """Load audio from URL; return None if URL returns HTML or unreadable format."""
    try:
        r = requests.get(url, timeout=15)
        r.raise_for_status()
        raw = r.content
        if not raw or len(raw) < 100:
            return None
        ct = (r.headers.get("Content-Type") or "").lower()
        if "text/html" in ct or raw.lstrip()[:50].startswith(b"<!") or raw.lstrip()[:20].startswith(b"<?"):
            return None
        buf = io.BytesIO(raw)
        y, _ = librosa.load(buf, sr=16000, mono=True)
        return y
    except Exception as e:
        return None


def predict_file(embs, model_mlp, model_att, lgb_model, threshold):
    """embs: (N, D) with D=2051 or 2090. Returns (pred_mlp, pred_att, pred_ens, pred_lgb). Ensemble = avg of MLP + attention + LGB (if present)."""
    with torch.no_grad():
        embs_t = torch.tensor(embs, dtype=torch.float32).to(device)
        logits_mlp = model_mlp(embs_t)
        logits_att = model_att(embs_t)
        prob_mlp = torch.softmax(logits_mlp, dim=1).cpu().numpy()
        prob_att = torch.softmax(logits_att, dim=1).cpu().numpy()
    p1_mlp = prob_mlp[:, 1].mean()
    p1_att = prob_att[:, 1].mean()
    pred_mlp = 1 if p1_mlp >= threshold else 0
    pred_att = 1 if p1_att >= threshold else 0

    mean_emb = embs.mean(axis=0).astype(np.float32).reshape(1, -1)
    p1_lgb = None
    pred_lgb = None
    if lgb_model is not None:
        prob_lgb = lgb_model.predict(mean_emb)  # (1, 2)
        p1_lgb = float(prob_lgb[0, 1])
        pred_lgb = 1 if p1_lgb >= threshold else 0

    probs_list = [p1_mlp, p1_att]
    if p1_lgb is not None:
        probs_list.append(p1_lgb)
    p1_ens = np.mean(probs_list)
    pred_ens = 1 if p1_ens >= threshold else 0
    return pred_mlp, pred_att, pred_ens, pred_lgb, p1_ens


def main():
    df = pd.read_csv(input_csv)
    if "nativity_status" not in df.columns:
        print("test_data.csv must have 'nativity_status' column")
        return

    # Match training: 2051 (prosody) or 2090 (prosody+MFCC)
    embed_dim = PROSODY_EMBED_DIM
    if os.path.isfile(EMBED_DIM_FILE):
        with open(EMBED_DIM_FILE) as f:
            embed_dim = int(f.read().strip())
    print(f"Embedding dimension: {embed_dim}")

    model_mlp = NativityClassifier(
        input_dim=embed_dim, hidden_dims=HIDDEN_DIMS, output_dim=2, dropout=DROPOUT
    ).to(device)
    model_mlp.load_state_dict(torch.load("models/nativity_model.pth", map_location=device))
    model_mlp.eval()

    model_att = NativityClassifierWithAttention(
        input_dim=embed_dim, hidden_dims=HIDDEN_DIMS, output_dim=2, dropout=DROPOUT,
        final_dropout=0.65, reduction_ratio=8
    ).to(device)
    model_att.load_state_dict(torch.load("models/nativity_model_attention.pth", map_location=device))
    model_att.eval()

    # Best threshold from validation (saved by train_classifier_compare.py)
    threshold = EVAL_THRESHOLD_DEFAULT
    if os.path.isfile(BEST_THRESHOLD_PATH):
        with open(BEST_THRESHOLD_PATH) as f:
            threshold = float(f.read().strip())
        print(f"Using threshold from validation: {threshold}")
    else:
        print(f"Using default threshold: {threshold}")

    lgb_model = None
    if os.path.isfile("models/lgbm_nativity.txt"):
        try:
            import lightgbm as lgb
            lgb_model = lgb.Booster(model_file="models/lgbm_nativity.txt")
        except Exception:
            pass

    preds_mlp, preds_att, preds_ens, preds_lgb, preds_probs_ens = [], [], [], [], []
    valid_idx = []
    skipped = 0
    print("Running on test set...")
    for idx, row in tqdm(df.iterrows(), total=len(df)):
        y = load_audio(row["audio_url"])
        if y is None or len(y) == 0:
            skipped += 1
            continue
        chunks = [
            y[i : i + CHUNK_SIZE]
            for i in range(0, len(y), CHUNK_SIZE)
            if len(y[i : i + CHUNK_SIZE]) == CHUNK_SIZE
        ]
        if not chunks:
            continue
        embs = np.stack([get_embedding(c, embed_dim=embed_dim) for c in chunks])
        p_mlp, p_att, p_ens, p_lgb, p1_ens = predict_file(embs, model_mlp, model_att, lgb_model, threshold)
        preds_mlp.append(p_mlp)
        preds_att.append(p_att)
        preds_ens.append(p_ens)
        preds_probs_ens.append(p1_ens)
        if p_lgb is not None:
            preds_lgb.append(p_lgb)
        valid_idx.append(idx)

    if not valid_idx:
        print("No valid test samples.")
        return
    if skipped > 0:
        print(f"Skipped {skipped} URL(s) (non-audio or load failed).")

    y_true = np.array([1 if df.loc[i, "nativity_status"] == "Non-Native" else 0 for i in valid_idx])
    preds_mlp = np.array(preds_mlp)
    preds_att = np.array(preds_att)
    preds_ens = np.array(preds_ens)
    probs_ens = np.array(preds_probs_ens)
    preds_lgb = np.array(preds_lgb) if len(preds_lgb) == len(y_true) else None

    n_ens = 2 + (1 if preds_lgb is not None else 0)

    def row(name, pred):
        return {
            "name": name,
            "acc": accuracy_score(y_true, pred),
            "f1": f1_score(y_true, pred, average="macro", zero_division=0),
            "prec": precision_score(y_true, pred, average="macro", zero_division=0),
            "rec": recall_score(y_true, pred, average="macro", zero_division=0),
            "cm": confusion_matrix(y_true, pred),
        }

    results = [
        row("MLP (baseline)", preds_mlp),
        row("MLP + attention (prosody)", preds_att),
        row(f"Ensemble ({n_ens} models, t={threshold})", preds_ens),
    ]
    if preds_lgb is not None:
        results.append(row("LightGBM", preds_lgb))

    print("\n" + "=" * 60)
    print("COMPARISON ON TEST DATA")
    print("=" * 60)
    print(f"Test samples: {len(y_true)} (Native=0, Non-Native=1)")
    print(f"{'Model':<30} {'Accuracy':>10} {'F1 macro':>10} {'Precision':>10} {'Recall':>10}")
    print("-" * 72)
    for r in results:
        print(f"{r['name']:<30} {r['acc']:>10.4f} {r['f1']:>10.4f} {r['prec']:>10.4f} {r['rec']:>10.4f}")
    print("\nConfusion matrices (rows=true, cols=pred):")
    for r in results:
        print(f"  {r['name']}:")
        print(f"    {r['cm']}")
    best = max(results, key=lambda x: x["f1"])
    print(f"\nBest on test (F1): {best['name']} (F1={best['f1']:.4f})")

    # Threshold sweep on test (ensemble): see which threshold gives best test F1
    print("\nEnsemble threshold sweep on test:")
    print(f"  {'t':>6} {'Accuracy':>10} {'F1 macro':>10}")
    best_t, best_t_f1 = threshold, 0.0
    for t in THRESHOLD_SWEEP_TEST:
        pred_t = (probs_ens >= t).astype(int)
        acc_t = accuracy_score(y_true, pred_t)
        f1_t = f1_score(y_true, pred_t, average="macro", zero_division=0)
        print(f"  {t:>6.2f} {acc_t:>10.4f} {f1_t:>10.4f}")
        if f1_t > best_t_f1:
            best_t_f1 = f1_t
            best_t = t
    print(f"  Best on test: t={best_t} (F1={best_t_f1:.4f})")


if __name__ == "__main__":
    main()