"""
Train the alternatives & insurance transaction classifier.

Model: word (1-2 gram) + character (2-5 gram) TF-IDF features into a multinomial
logistic regression. Character n-grams handle abbreviations ("DIST ROC", "MGMT FEE",
"PREM PMT") and typos; word n-grams carry phrase meaning. Predicted probabilities
drive a confidence threshold: items above it are classified straight through, items
below it go to a person.

Two rounds are trained to show the review loop:
  v1  - training data only
  v2  - training data + 180 items labeled by reviewers from v1's exception queue
The published model.skops is v2.

Run: python data/generate.py && python train.py
"""

import json
import zipfile
from pathlib import Path

import numpy as np
import pandas as pd
import skops.io as sio
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, classification_report, f1_score
from sklearn.pipeline import FeatureUnion, Pipeline

ROOT = Path(__file__).parent
SEED = 42
TARGET_PRECISION = 0.99


def build_pipeline():
    features = FeatureUnion([
        ("word", TfidfVectorizer(analyzer="word", ngram_range=(1, 2), lowercase=True, sublinear_tf=True,
                                 token_pattern=r"(?u)\b[a-zA-Z][a-zA-Z&/]+\b", min_df=2)),
        ("char", TfidfVectorizer(analyzer="char_wb", ngram_range=(2, 5), lowercase=True, sublinear_tf=True,
                                 min_df=3, max_features=30000)),
    ])
    return Pipeline([("features", features),
                     ("clf", LogisticRegression(C=8.0, max_iter=3000, random_state=SEED))])


def straight_through(proba, y_true, classes, target=TARGET_PRECISION):
    """Lowest confidence threshold whose auto-classified items meet the target precision."""
    conf, pred = proba.max(axis=1), classes[proba.argmax(axis=1)]
    for t in np.round(np.arange(0.30, 1.0, 0.01), 2):
        m = conf >= t
        if m.sum() and (pred[m] == y_true[m]).mean() >= target:
            return {"threshold": float(t), "auto_rate": round(float(m.mean()), 4),
                    "auto_precision": round(float((pred[m] == y_true[m]).mean()), 4),
                    "routed_to_review": round(float(1 - m.mean()), 4)}
    return None


def evaluate(model, test):
    classes = model.classes_
    proba = model.predict_proba(test.text)
    pred = classes[proba.argmax(axis=1)]
    y = test.label.to_numpy()
    out = {"accuracy": round(accuracy_score(y, pred), 4),
           "macro_f1": round(f1_score(y, pred, average="macro"), 4),
           "straight_through": straight_through(proba, y, classes), "by_condition": {}}
    for c in ("unseen_entities", "unseen_phrasing"):
        m = (test.condition == c).to_numpy()
        out["by_condition"][c] = {"n": int(m.sum()), "accuracy": round(accuracy_score(y[m], pred[m]), 4),
                                  "macro_f1": round(f1_score(y[m], pred[m], average="macro"), 4),
                                  "straight_through": straight_through(proba[m], y[m], classes)}
    return out, pred, y


if __name__ == "__main__":
    train = pd.read_csv(ROOT / "data" / "train.csv")
    test = pd.read_csv(ROOT / "data" / "test.csv")
    feedback = pd.read_csv(ROOT / "data" / "review_feedback.csv")

    v1 = build_pipeline().fit(train.text, train.label)
    m1, _, _ = evaluate(v1, test)

    v2 = build_pipeline().fit(pd.concat([train.text, feedback.text]), pd.concat([train.label, feedback.label]))
    m2, pred, y = evaluate(v2, test)

    report = classification_report(y, pred, digits=4, output_dict=True)
    wrong = test[pred != y].assign(predicted=pred[pred != y])
    metrics = {
        "train_size": len(train), "review_feedback_size": len(feedback), "test_size": len(test),
        "num_labels": len(v2.classes_), "target_precision": TARGET_PRECISION,
        "v1_train_only": m1, "v2_with_review_feedback": m2,
        "v2_per_label_f1": {k: round(v["f1-score"], 4) for k, v in report.items() if k in set(v2.classes_)},
        "v2_sample_errors": wrong[["text", "label", "predicted"]].head(12).to_dict("records"),
    }

    # Drop vocabulary-pruning artifacts; they are only for introspection and bloat the file.
    for _, vec in v2.named_steps["features"].transformer_list:
        if hasattr(vec, "stop_words_"):
            del vec.stop_words_
    sio.dump(v2, ROOT / "model.skops", compression=zipfile.ZIP_DEFLATED, compresslevel=9)
    (ROOT / "metrics.json").write_text(json.dumps(metrics, indent=2))

    print(classification_report(y, pred, digits=4))
    for name in ("v1_train_only", "v2_with_review_feedback"):
        print(name, json.dumps(metrics[name], indent=1))