"""
train_model.py
--------------
Train a LightGBM binary risk classifier on the training dataset.

Outputs:
  scoring/model.lgb         — trained LightGBM model
  scoring/holdout_split.json — held-out test rows for evaluation (Section 6)

Usage:
    python scoring/train_model.py
"""

import json
import os
import pickle
import sys

import lightgbm as lgb
import numpy as np
from sklearn.model_selection import train_test_split

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
TRAIN_DIR = os.path.join(BASE_DIR, "data", "synthetic_train")
SCORING_DIR = os.path.join(BASE_DIR, "scoring")

MODEL_PATH   = os.path.join(SCORING_DIR, "model.lgb")
HOLDOUT_PATH = os.path.join(SCORING_DIR, "holdout_split.json")

# ---------------------------------------------------------------------------
# Feature columns (must match scoring/features.py)
# ---------------------------------------------------------------------------
FEATURE_COLUMNS = [
    "size",
    "strong_link_count",
    "weak_link_count",
    "strong_link_density",
    "has_same_contact_edge",
    "order_count",
    "total_amount",
    "avg_order_amount",
    "return_rate",
    "chargeback_rate",
    "shared_coupon_count",
    "time_span_days",
    "order_velocity",
    "formation_type_strong",
]

RANDOM_STATE = 42


def load_data():
    """Load clusters, orders, ground truth, and graph from the training directory."""
    sys.path.insert(0, BASE_DIR)
    from scoring.features import extract_all, _load_json, _load_graph

    clusters_path = os.path.join(TRAIN_DIR, "candidate_clusters.json")
    orders_path   = os.path.join(TRAIN_DIR, "orders.json")
    gt_path       = os.path.join(TRAIN_DIR, "ground_truth.json")
    graph_path    = os.path.join(TRAIN_DIR, "graph.gpickle")

    clusters     = _load_json(clusters_path)
    orders       = _load_json(orders_path)
    ground_truth = _load_json(gt_path)
    graph        = _load_graph(graph_path) if os.path.exists(graph_path) else None

    rows = extract_all(clusters, orders, ground_truth=ground_truth, graph=graph)
    return rows


def train():
    print("Loading training data...")
    rows = load_data()
    print(f"  {len(rows)} candidate clusters loaded")

    X = np.array([[r[f] for f in FEATURE_COLUMNS] for r in rows])
    y = np.array([r["label"] for r in rows])
    cluster_ids = [r["cluster_id"] for r in rows]
    feature_rows = [{f: r[f] for f in FEATURE_COLUMNS} for r in rows]

    pos = y.sum()
    neg = len(y) - pos
    print(f"  Labels — positive: {int(pos)}, negative: {int(neg)}")

    # 5.2 — stratified 80/20 split
    (X_train, X_test, y_train, y_test,
     ids_train, ids_test,
     feat_train, feat_test) = train_test_split(
        X, y, cluster_ids, feature_rows,
        test_size=0.20,
        stratify=y,
        random_state=RANDOM_STATE,
    )
    print(f"  Train: {len(X_train)}  Test (holdout): {len(X_test)}")
    print(f"  Holdout positives: {y_test.sum():.0f}  negatives: {(len(y_test)-y_test.sum()):.0f}")

    # 5.3 — simple LightGBM model (small dataset, keep it simple)
    model = lgb.LGBMClassifier(
        num_leaves=7,
        max_depth=4,
        min_child_samples=5,
        n_estimators=100,
        learning_rate=0.05,
        objective="binary",
        random_state=RANDOM_STATE,
        verbose=-1,
    )
    print("Training LightGBM classifier...")
    model.fit(X_train, y_train)
    print("  Training complete.")

    # 5.4 — save model
    model.booster_.save_model(MODEL_PATH)
    print(f"  Model saved to {MODEL_PATH}")

    # Verify reload
    loaded = lgb.Booster(model_file=MODEL_PATH)
    test_pred = loaded.predict(X_test)
    print(f"  Model reload verified (test pred shape: {test_pred.shape})")

    # 5.5 — feature importances
    importances = model.feature_importances_
    importance_pairs = sorted(
        zip(FEATURE_COLUMNS, importances),
        key=lambda x: x[1], reverse=True
    )
    print("\nFeature importances (gain):")
    for feat, imp in importance_pairs:
        print(f"  {feat:30s}: {imp}")
    if importance_pairs[0][0] == "size" and importance_pairs[0][1] > sum(importances) * 0.5:
        print("  [NOTE] 'size' dominates feature importance — worth flagging in summary")

    # Save holdout split (cluster IDs + features + labels + probabilities)
    holdout = []
    probs = loaded.predict(X_test)
    for i in range(len(ids_test)):
        holdout.append({
            "cluster_id": ids_test[i],
            "features": feat_test[i],
            "label": int(y_test[i]),
            "risk_probability": float(probs[i]),
        })

    with open(HOLDOUT_PATH, "w", encoding="utf-8") as f:
        json.dump(holdout, f, indent=2, ensure_ascii=False)
    print(f"  Holdout split saved to {HOLDOUT_PATH} ({len(holdout)} rows)")

    return model, holdout, importance_pairs


if __name__ == "__main__":
    train()