"""
Trains one multi-label classifier per attribute type on the labeled dataset.

Approach
--------
- Features: character n-gram TF-IDF (char_wb, 2-5) over the raw description.
  Character n-grams generalize better than word n-grams on a 61-row dataset
  because they share sub-word signal across morphological variants
  ("embroidered"/"embroidery", "beaded"/"beading", "off-shoulder"/
  "off shoulder") without needing a stemmer.
- Label space: for each attribute type, the label space is exactly the
  canonical vocabulary in lexicon.py (colors use COLOR_VOCAB). This keeps the
  model closed-vocabulary and directly comparable to the rules baseline.
- Classifier: OneVsRestClassifier(LogisticRegression) per attribute type,
  i.e. independent binary classifiers per label, combined into a multi-label
  prediction per attribute (a description can have 0, 1, or many labels for
  a given attribute, e.g. two colors).
- Given the small dataset, model artifacts are ensembled at inference time
  with the deterministic rules extractor (see predict.py) rather than relied
  on in isolation.
"""
import json
import joblib
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.multiclass import OneVsRestClassifier
from sklearn.preprocessing import MultiLabelBinarizer

from lexicon import ATTRIBUTE_LEXICON, COLOR_VOCAB

ATTR_TYPES = ["silhouette", "fabric", "neckline", "sleeve", "length",
              "embellishment", "category", "color"]


def label_space(attr):
    if attr == "color":
        return COLOR_VOCAB
    return list(ATTRIBUTE_LEXICON[attr].keys())


def load_dataset(path="../data/dataset.json"):
    with open(path) as f:
        return json.load(f)


def train_all(dataset_path="../data/dataset.json", out_path="model.joblib"):
    data = load_dataset(dataset_path)
    texts = [d["text"] for d in data]

    vectorizer = TfidfVectorizer(analyzer="char_wb", ngram_range=(2, 5), min_df=1)
    X = vectorizer.fit_transform(texts)

    models = {}
    binarizers = {}

    for attr in ATTR_TYPES:
        classes = label_space(attr)
        mlb = MultiLabelBinarizer(classes=classes)
        y_labels = [d["labels"][attr] for d in data]
        Y = mlb.fit_transform(y_labels)

        # Skip labels with zero positive examples (OneVsRest can't fit them);
        # in practice every label in our lexicon appears >=1 time in the 61
        # rows, but guard anyway for robustness on future dataset edits.
        if Y.sum() == 0:
            models[attr] = None
            binarizers[attr] = mlb
            continue

        clf = OneVsRestClassifier(
            LogisticRegression(max_iter=1000, class_weight="balanced")
        )
        clf.fit(X, Y)
        models[attr] = clf
        binarizers[attr] = mlb

    artifact = {
        "vectorizer": vectorizer,
        "models": models,
        "binarizers": binarizers,
        "attr_types": ATTR_TYPES,
    }
    joblib.dump(artifact, out_path)
    print(f"Trained models for attrs: {ATTR_TYPES}")
    print(f"Saved artifact to {out_path}")
    return artifact


if __name__ == "__main__":
    train_all()