# -*- coding: utf-8 -*-
"""P3-S.18 — LOGISTIC BASELINE.

Simple linear probabilistic baseline (LogisticRegression, L2, C=1.0, no
search). Standard scaler fitted on TRAIN ONLY. Binary target: WIN vs LOSS
(leads only; UNRESOLVED/AMBIGUOUS excluded from fit, retained in evaluation
as separate columns). Writes predictions CSV columns for evaluate_baselines.
"""
import csv
import json
import os
import sys

import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler

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 main():
    os.makedirs(OUT, exist_ok=True)
    ctx = PD.load_verified_population()
    rows = PD.build_rows(ctx)
    # binary subset = leads with WIN/LOSS
    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_ids = set()
    tr, va, te, _ = PD.temporal_split(rows)
    tr_ids = {id(r) for r in tr}
    # rebuild split membership by setup_id (stable)
    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)
    clf = LogisticRegression(C=1.0, max_iter=5000, random_state=SEED)
    clf.fit(Xs[splits == "train"], y[splits == "train"])
    p = clf.predict_proba(Xs)[:, 1]

    rows_out = []
    for i, r in enumerate(bin_rows):
        rows_out.append({
            "setup_id": r["setup_id"], "split": splits[i],
            "outcome": r["outcome"], "y": int(y[i]),
            "pred_logistic": float(p[i]),
        })
    with open(os.path.join(OUT, "p3_s18_predictions_logistic.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)

    coef = {PD.FEATURE_COLS[i]: float(clf.coef_[0][i])
            for i in range(len(PD.FEATURE_COLS))}
    with open(os.path.join(OUT, "p3_s18_logistic_model.json"), "w",
              encoding="utf-8") as f:
        json.dump({
            "model": "LogisticRegression(C=1.0, L2, standardized-train-only)",
            "feature_sha16": PD.make_manifest(
                rows, tr, va, te, True)["feature_sha16"],
            "seed": SEED, "n_train": int((splits == "train").sum()),
            "n_val": int((splits == "val").sum()),
            "n_test": int((splits == "test").sum()),
            "class_counts": {
                "train_win": int(((splits == "train") & (y == 1)).sum()),
                "train_loss": int(((splits == "train") & (y == 0)).sum()),
                "val_win": int(((splits == "val") & (y == 1)).sum()),
                "val_loss": int(((splits == "val") & (y == 0)).sum()),
                "test_win": int(((splits == "test") & (y == 1)).sum()),
                "test_loss": int(((splits == "test") & (y == 0)).sum()),
            },
            "coefficients": coef,
        }, f, indent=2)
    print("[saved] p3_s18_predictions_logistic.csv / logistic_model.json")
    return 0


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