"""
Run inference on test CSV: download audio, chunk, embed with Wav2Vec (middle layers), classify.
Arabic Native vs Non-Native. Uses language-agnostic middle-layer embeddings.
"""
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,
)

# Import shared classifier (same architecture as training)
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from model import NativityClassifier
from extract_embeddings import get_mfcc_from_audio, get_pitch_stats, PROSODY_EMBED_DIM, EMBED_DIM_WITH_MFCC

# -----------------------------------------------------------------------------
# Config
# -----------------------------------------------------------------------------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_path = "models/nativity_model.pth"
input_csv = "manifests/test_data.csv"
output_csv = "final_predictions.csv"
EMBED_DIM_FILE = "manifests/embedding_dim.txt"

# Model arch must match train_classifier_compare / model.py
HIDDEN_DIMS = (256, 128)  # Must match train_classifier_compare.py
DROPOUT = 0.4
PCA_COMPONENTS = 256
PCA_PATH = "models/pca_transform.pkl"
EVAL_THRESHOLD_DEFAULT = 0.35
BEST_THRESHOLD_PATH = "models/best_threshold.txt"
PROSODY_LAYERS = (4, 6, 8, 10)  # must match extract_embeddings.py (prosody-focused)

# Embedding dim: 2051 (prosody) or 2090 (prosody+MFCC)
embedding_dim = PROSODY_EMBED_DIM
if os.path.isfile(EMBED_DIM_FILE):
    with open(EMBED_DIM_FILE) as f:
        embedding_dim = int(f.read().strip())

# Use validation-picked threshold if available (from train_classifier_compare.py)
if os.path.isfile(BEST_THRESHOLD_PATH):
    with open(BEST_THRESHOLD_PATH) as f:
        EVAL_THRESHOLD = float(f.read().strip())
else:
    EVAL_THRESHOLD = EVAL_THRESHOLD_DEFAULT

# -----------------------------------------------------------------------------
# Load models
# -----------------------------------------------------------------------------
import joblib

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

pca = joblib.load(PCA_PATH)

classifier = NativityClassifier(
    input_dim=PCA_COMPONENTS,
    hidden_dims=HIDDEN_DIMS,
    output_dim=2,
    dropout=DROPOUT,
).to(device)
classifier.load_state_dict(torch.load(model_path, map_location=device))
classifier.eval()


def get_embedding(audio_np, embed_dim=None):
    """Prosody embedding: mean+std+pitch (2051-d) or +MFCC (2090-d); L2 norm to match training."""
    if embed_dim is None:
        embed_dim = embedding_dim
    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
    out = out.reshape(1, -1)
    # Apply PCA to match training pipeline (2051-d -> 256-d)
    out = pca.transform(out).astype(np.float32)
    return torch.tensor(out, dtype=torch.float32, device=device)


def process_file_and_predict(url):
    """
    Download, chunk (10s), embed each chunk, then aggregate with confidence-weighted voting.
    Returns (label_str, confidence).
    """
    try:
        response = requests.get(url, timeout=15)
        response.raise_for_status()
        raw = response.content
        if not raw or len(raw) < 100:
            return "Error", 0.0
        ct = (response.headers.get("Content-Type") or "").lower()
        if "text/html" in ct or raw.lstrip()[:50].startswith(b"<!") or raw.lstrip()[:20].startswith(b"<?"):
            return "Error", 0.0
        audio_bytes = io.BytesIO(raw)
        y, sr = librosa.load(audio_bytes, sr=16000, mono=True)
    except Exception:
        return "Error", 0.0

    chunk_size = 160000  # 10 sec at 16 kHz
    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:
        return "Error", 0.0

    chunk_preds = []   # 0 = Native, 1 = Non-Native
    chunk_confs = []   # probability of predicted class

    with torch.no_grad():
        for chunk in chunks:
            emb = get_embedding(chunk, embed_dim=embedding_dim)
            logits = classifier(emb)
            probs = torch.softmax(logits, dim=1)
            p1 = probs[0, 1].item()
            pred = 1 if p1 >= EVAL_THRESHOLD else 0
            conf = p1 if pred == 1 else (1.0 - p1)
            chunk_preds.append(pred)
            chunk_confs.append(conf)

    # Confidence-weighted voting: weight each chunk by its confidence
    weights = np.array(chunk_confs, dtype=np.float64)
    weights = weights / (weights.sum() + 1e-10)
    weighted_score = np.sum(np.array(chunk_preds) * weights)
    final_label = "Non-Native" if weighted_score >= EVAL_THRESHOLD else "Native"
    # Final confidence = weighted average of confidences (or max)
    final_conf = float(np.average(chunk_confs, weights=weights))

    return final_label, final_conf


def main():
    df = pd.read_csv(input_csv)
    results = []
    print(f"Predicting for {len(df)} files (device={device})...")
    for _, row in tqdm(df.iterrows(), total=len(df)):
        label, conf = process_file_and_predict(row["audio_url"])
        results.append({
            "dp_id": row["dp_id"],
            "prediction": label,
            "confidence_score": round(conf, 4),
        })
    out_df = pd.DataFrame(results)
    out_df.to_csv(output_csv, index=False)
    n_ok = (out_df["prediction"] != "Error").sum()
    n_fail = len(out_df) - n_ok
    print(f"Done. Saved to {output_csv} ({n_ok} predicted, {n_fail} failed to load)")

    # If ground-truth labels exist (nativity_status), print test metrics like validation
    if "nativity_status" in df.columns:
        mask = out_df["prediction"] != "Error"
        y_true_str = df.loc[mask, "nativity_status"].values
        y_pred_str = out_df.loc[mask, "prediction"].values
        map_fn = lambda s: 1 if s == "Non-Native" else 0
        y_true = np.array([map_fn(s) for s in y_true_str])
        y_pred = np.array([map_fn(s) for s in y_pred_str])

        acc = accuracy_score(y_true, y_pred)
        f1 = f1_score(y_true, y_pred, average="macro", zero_division=0)
        prec = precision_score(y_true, y_pred, average="macro", zero_division=0)
        rec = recall_score(y_true, y_pred, average="macro", zero_division=0)
        cm = confusion_matrix(y_true, y_pred)

        print("\nTest metrics (from test_data.csv, using nativity_status):")
        print(f"  Accuracy:  {acc:.4f}")
        print(f"  F1 macro:  {f1:.4f}")
        print(f"  Precision: {prec:.4f}")
        print(f"  Recall:    {rec:.4f}")
        print("  Confusion matrix (rows=true, cols=pred):")
        print(f"    {cm}")
        print("  (Native=0, Non-Native=1)")


if __name__ == "__main__":
    main()