# -*- coding: utf-8 -*-
"""P3-S.18 — TREE BASELINE (shallow decision tree + constrained boosting).

Intentionally small: DecisionTree(max_depth=3) and GradientBoosting
(n_estimators=40, max_depth=2, lr=0.03). No broad hyperparameter search.
Binary target: WIN vs LOSS (leads only). Same split/standardization policy;
trees are scale-invariant, scaler applied for consistency of the pipeline.
"""
import csv
import json
import os
import sys

import numpy as np
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.tree import DecisionTreeClassifier

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "output")
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.normpath(os.path.join(HERE, "..", "setup_dataset")))

import prepare_dataset as PD          # noqa: E402

SEED = 42


def _setup():
    ctx = PD.load_verified_population()
    rows = PD.build_rows(ctx)
    bin_rows = [r for r in rows if r["lead"] and r["outcome"] in PD.BINARY_CLASSES]
    bin_rows.sort(key=lambda r: (r["creation_bar"], r["setup_id"]))
    tr, va, te, _ = PD.temporal_split(rows)
    split_map = {}
    for part, tag in ((tr, "train"), (va, "val"), (te, "test")):
        for r in part:
            split_map[r["setup_id"]] = tag
    X = np.asarray([[r["feature_" + k] for k in PD.FEATURE_COLS]
                    for r in bin_rows], dtype=float)
    y = np.asarray([1.0 if r["outcome"] == "WIN" else 0.0
                    for r in bin_rows], dtype=float)
    splits = np.asarray([split_map[r["setup_id"]] for r in bin_rows])
    scaler = StandardScaler().fit(X[splits == "train"])
    Xs = scaler.transform(X)
    return Xs, y, splits, bin_rows, tr, va, te, rows


def main():
    os.makedirs(OUT, exist_ok=True)
    Xs, y, splits, bin_rows, tr, va, te, rows = _setup()
    models = {
        "tree_depth3": DecisionTreeClassifier(max_depth=3, random_state=SEED,
                                              min_samples_leaf=5),
        "boost_small": GradientBoostingClassifier(
            n_estimators=40, max_depth=2, learning_rate=0.03,
            random_state=SEED),
    }
    preds = {}
    for name, clf in models.items():
        clf.fit(Xs[splits == "train"], y[splits == "train"])
        preds[name] = clf.predict_proba(Xs)[:, 1]

    rows_out = []
    for i, r in enumerate(bin_rows):
        d = {"setup_id": r["setup_id"], "split": splits[i],
             "outcome": r["outcome"], "y": int(y[i])}
        for name in models:
            d["pred_" + name] = float(preds[name][i])
        rows_out.append(d)
    with open(os.path.join(OUT, "p3_s18_predictions_tree.csv"), "w",
              newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=list(rows_out[0].keys()))
        w.writeheader()
        for r in rows_out:
            w.writerow(r)

    with open(os.path.join(OUT, "p3_s18_tree_model.json"), "w",
              encoding="utf-8") as f:
        json.dump({
            "models": list(models.keys()),
            "config": {"tree": "max_depth=3,min_samples_leaf=5",
                       "boost": "n_estimators=40,max_depth=2,lr=0.03"},
            "feature_sha16": PD.make_manifest(rows, tr, va, te,
                                              True)["feature_sha16"],
            "seed": SEED,
            "feature_importance": {
                "tree_depth3": dict(zip(PD.FEATURE_COLS, models[
                    "tree_depth3"].feature_importances_.tolist()))},
        }, f, indent=2)
    print("[saved] p3_s18_predictions_tree.csv / tree_model.json")
    return 0


if __name__ == "__main__":
    sys.exit(main())