import os
os.environ.setdefault("USE_TF", "0")
os.environ.setdefault("USE_TORCH", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import json
import time

import numpy as np
import pandas as pd
import trackio

def _trackio_safe(fn):
    def wrapper(*a, **k):
        try:
            return fn(*a, **k)
        except Exception as e:
            print(f"trackio.{fn.__name__} failed (non-fatal): {e}")
    return wrapper
trackio.log = _trackio_safe(trackio.log)
trackio.init = _trackio_safe(trackio.init)
trackio.finish = _trackio_safe(trackio.finish)

from datasets import Dataset
from huggingface_hub import HfApi
from sklearn.metrics import f1_score
from sklearn.model_selection import train_test_split
from setfit import SetFitModel, TrainingArguments, Trainer

SMOKE = os.environ.get("SMOKE", "0") == "1"
PARQUET_URL = "https://huggingface.co/datasets/GIZ/vulnerability_training_data_full/resolve/refs%2Fconvert%2Fparquet/default/train/0000.parquet"
LABELS = [
    "Agricultural communities", "Coastal communities", "Ethnic, racial or other minorities",
    "Fishery communities", "Informal sector workers", "Members of indigenous and local communities",
    "Migrants and displaced persons", "Older persons", "Other", "Persons living in poverty",
    "Persons with disabilities", "Persons with pre-existing health conditions",
    "Residents of drought-prone regions", "Rural populations", "Sexual minorities (LGBTQI+)",
    "Urban populations", "Women and other genders",
]
REPO_ID = "peter2000/setfit-vulnerability-groups"

def load_data():
    df = pd.read_parquet(PARQUET_URL)
    assert len(df) == 475, f"expected 475 rows, got {len(df)}"
    Y = df[LABELS].values.astype(np.int64)
    nlab = Y.sum(1)
    idx_tr, idx_te = train_test_split(
        np.arange(len(df)), test_size=0.2, random_state=42, stratify=np.minimum(nlab, 3)
    )
    return df, Y, np.asarray(idx_tr), np.asarray(idx_te)

def ece(conf, correct, n_bins=15):
    conf = np.asarray(conf, dtype=np.float64)
    corr = np.asarray(correct, dtype=np.float64)
    bins = np.linspace(0.0, 1.0, n_bins + 1)
    e = 0.0
    for lo, hi in zip(bins[:-1], bins[1:]):
        m = (conf > lo) & (conf <= hi)
        if m.sum() > 0:
            e += m.mean() * abs(corr[m].mean() - conf[m].mean())
    return float(e)

def evaluate(Y_true, P_pred, threshold=0.5):
    pred = (P_pred >= threshold).astype(int)
    per_label = f1_score(Y_true, pred, average=None, zero_division=0)
    macro = float(f1_score(Y_true, pred, average="macro", zero_division=0))
    micro = float(f1_score(Y_true, pred, average="micro", zero_division=0))
    conf = np.where(pred == 1, P_pred, 1.0 - P_pred)
    corr = (pred == Y_true).astype(np.float64)
    return {
        "macro_f1": macro,
        "micro_f1": micro,
        "ece": ece(conf, corr),
        "subset_accuracy": float(((pred == Y_true).all(axis=1)).mean()),
        "per_label_f1": {LABELS[i]: round(float(per_label[i]), 4) for i in range(len(LABELS))},
    }

def main():
    df, Y, idx_tr, idx_te = load_data()
    texts = df["text"].tolist()
    X_tr = [texts[i] for i in idx_tr]
    X_te = [texts[i] for i in idx_te]
    Y_tr, Y_te = Y[idx_tr], Y[idx_te]
    print(f"train={len(X_tr)} test={len(X_te)} labels={len(LABELS)}")

    if SMOKE:
        X_tr, Y_tr = X_tr[:40], Y_tr[:40]
        X_te, Y_te = X_te[:10], Y_te[:10]

    model = SetFitModel.from_pretrained(
        "sentence-transformers/paraphrase-mpnet-base-v2",
        multi_target_strategy="one-vs-rest",
        labels=LABELS,
    )
    model.model_body.max_seq_length = 256

    train_ds = Dataset.from_dict({
        "text": X_tr,
        "label": [list(map(int, y)) for y in Y_tr],
    })

    args = TrainingArguments(
        batch_size=16,
        num_epochs=1,
        num_iterations=20 if not SMOKE else 2,
        seed=42,
    )
    trainer = Trainer(model=model, args=args, train_dataset=train_ds)
    trainer.train()

    P_te = np.asarray(model.predict_proba(X_te))
    metrics = evaluate(Y_te, P_te)
    print(json.dumps(metrics, indent=2))

    trackio.init(project="vulnerability-multilabel-classifier", space_id="peter2000/vulnerability-multilabel-classifier-trackio")
    trackio.log({"setfit_macro_f1": metrics["macro_f1"], "setfit_micro_f1": metrics["micro_f1"], "setfit_ece": metrics["ece"]}, step=1)
    trackio.log({f"setfit_f1/{k}": v for k, v in metrics["per_label_f1"].items()}, step=1)

    trainer.push_to_hub(repo_id=REPO_ID)
    api = HfApi(token=os.environ.get("HF_TOKEN"))
    api.upload_file(
        path_or_fileobj=json.dumps(metrics, indent=2).encode(),
        path_in_repo="metrics.json",
        repo_id=REPO_ID,
        repo_type="model",
        commit_message="Add evaluation metrics (macro-F1, micro-F1, per-label F1, ECE)",
    )
    trackio.finish()
    print("DONE")

if __name__ == "__main__":
    t0 = time.time()
    main()
    print(f"elapsed {time.time()-t0:.0f}s")