"""Train a pathogenicity classifier on ClinVar labels (Pathogenic/Likely pathogenic vs Benign/Likely benign).

Label leakage warning: CLIN_SIG must never be a feature. This is a learning exercise, not a clinical model.
Usage: python -m rarelens_ml.train --tsv results/clinvar.vep.tsv
"""
import argparse

import lightgbm as lgb
import mlflow
import mlflow.lightgbm
import pandas as pd
from sklearn.metrics import average_precision_score, roc_auc_score
from sklearn.model_selection import train_test_split

from rarelens_ml.features import build

POS = {"Pathogenic", "Likely_pathogenic", "Pathogenic/Likely_pathogenic"}
NEG = {"Benign", "Likely_benign", "Benign/Likely_benign"}


def load(tsv: str) -> tuple[pd.DataFrame, pd.Series]:
    df = pd.read_csv(tsv, sep="\t", comment="#", header=None, dtype=str)
    with open(tsv) as fh:
        df.columns = next(l for l in fh if l.startswith("#Uploaded")).lstrip("#").rstrip().split("\t")
    df = df.rename(columns={"IMPACT": "impact", "Consequence": "consequence", "gnomADe_AF": "gnomad_af",
                            "CADD_PHRED": "cadd_phred"})
    y = df["CLIN_SIG"].map(lambda s: 1 if s in POS else 0 if s in NEG else None)
    keep = y.notna()
    return build(df[keep]), y[keep].astype(int)


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--tsv", required=True)
    p.add_argument("--register", action="store_true")
    a = p.parse_args()

    X, y = load(a.tsv)
    Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.2, stratify=y, random_state=42)

    mlflow.set_experiment("rarelens-pathogenicity")
    with mlflow.start_run():
        params = {"n_estimators": 400, "learning_rate": 0.05, "num_leaves": 31, "class_weight": "balanced"}
        mlflow.log_params(params)
        model = lgb.LGBMClassifier(**params).fit(Xtr, ytr)
        proba = model.predict_proba(Xte)[:, 1]
        mlflow.log_metrics({"auroc": roc_auc_score(yte, proba), "auprc": average_precision_score(yte, proba)})
        mlflow.lightgbm.log_model(
            model, "model",
            registered_model_name="rarelens-pathogenicity" if a.register else None,
        )


if __name__ == "__main__":
    main()