#!/usr/bin/env python3
"""Quick classification test of ethicalabs/Echo-DSRN-v0.1.3-Embed-Intent on MASSIVE.

Loads the HF model, encodes MASSIVE test split, runs 1-NN + logistic regression.
"""

import argparse
import numpy as np
from datasets import load_dataset
from sentence_transformers import SentenceTransformer
from sklearn.linear_model import SGDClassifier
from sklearn.metrics import accuracy_score


MASSIVE_REVISION = "refs/convert/parquet"


def main():
    parser = argparse.ArgumentParser(description="Test Echo-DSRN intent classifier on MASSIVE")
    parser.add_argument(
        "--model_id", default="ethicalabs/Echo-DSRN-v0.1.3-Embed-Intent",
        help="HF model repo to load",
    )
    parser.add_argument("--max_samples", type=int, default=0, help="Cap test samples (0=all)")
    parser.add_argument("--device", default="cuda")
    parser.add_argument(
        "--task", choices=["intent", "scenario", "both"], default="intent",
        help="Which classification label to test",
    )
    args = parser.parse_args()

    print(f"Loading model: {args.model_id}")
    model = SentenceTransformer(args.model_id, trust_remote_code=True, device=args.device)
    print(f"  Loaded. Pooling: {model.get_sentence_embedding_dimension()}-dim")

    print("Loading MASSIVE test split...")
    ds = load_dataset("AmazonScience/massive", split="test", revision=MASSIVE_REVISION)
    if args.max_samples > 0:
        ds = ds.shuffle(seed=42).select(range(min(args.max_samples, len(ds))))

    utts = ds["utt"]
    intents = np.array(ds["intent"])
    scenarios = np.array(ds["scenario"])
    locales = ds["locale"]

    n_intents = len(set(intents))
    n_scenarios = len(set(scenarios))
    print(f"  Samples: {len(ds)} | Intents: {n_intents} | Scenarios: {n_scenarios} | Locales: {len(set(locales))}")

    print("Encoding...")
    embeddings = model.encode(utts, batch_size=64, show_progress_bar=True, convert_to_numpy=True)

    # Normalise for cosine distance
    embeddings = embeddings / (np.linalg.norm(embeddings, axis=1, keepdims=True) + 1e-9)

    # Split: first 80% per locale as train, last 20% as test (simulates MTEB protocol)
    rng = np.random.default_rng(42)
    train_mask = np.zeros(len(ds), dtype=bool)
    for loc in sorted(set(locales)):
        loc_idx = np.where(np.array(locales) == loc)[0]
        rng.shuffle(loc_idx)
        split = int(0.8 * len(loc_idx))
        train_mask[loc_idx[:split]] = True
    test_mask = ~train_mask

    if args.task in ("intent", "both"):
        _evaluate("Intent (60-class)", embeddings, intents, n_intents, train_mask, test_mask)

    if args.task in ("scenario", "both"):
        _evaluate("Scenario (17-class)", embeddings, scenarios, n_scenarios, train_mask, test_mask)


def _evaluate(name, embeddings, labels, n_classes, train_mask, test_mask):
    X_train, y_train = embeddings[train_mask], labels[train_mask]
    X_test, y_test = embeddings[test_mask], labels[test_mask]

    # 1-NN (cosine)
    sim = X_test @ X_train.T
    nn_preds = y_train[np.argmax(sim, axis=1)]
    nn_acc = accuracy_score(y_test, nn_preds)

    # Logistic regression (SGD for memory safety on large datasets)
    lr = SGDClassifier(loss="log_loss", max_iter=1000, tol=1e-3, random_state=42)
    lr.fit(X_train, y_train)
    lr_preds = lr.predict(X_test)
    lr_acc = accuracy_score(y_test, lr_preds)

    print(f"\n{'='*50}")
    print(f"  {name}")
    print(f"  Train: {len(y_train)}  |  Test: {len(y_test)}")
    print(f"  1-NN accuracy:  {nn_acc:.4f}")
    print(f"  LR  accuracy:   {lr_acc:.4f}")


if __name__ == "__main__":
    main()