#!/usr/bin/env python3
"""Fit the WDL-sigmoid scale K on the TRAIN split only, then label
every non-exact position in both train and holdout with:
  - p_search: sigmoid(score / K) -- the search-score-derived win
    probability estimate.
  - r_outcome: that position's own game's actual final result,
    converted to side-to-move-relative (matching the score's own
    convention: 1.0 = the side to move at this FEN eventually won
    their game, 0.0 = eventually lost, 0.5 = drawn).

Does NOT blend p_search/r_outcome into a single training target --
that mixing weight (lambda) is left as a free parameter to sweep
later; storing both components separately means that sweep never
needs to re-derive anything from here.

K is fit by ternary search minimizing MSE(sigmoid(score/K), r_outcome)
over every non-exact TRAIN position (not holdout -- fitting K on
holdout data would leak information into what's supposed to be an
untouched generalization check). The same fitted K is then applied
(never refit) to label holdout positions too.

Usage:
    python3 fit_k_and_label.py --train runs/sd10_twic_2M/games.train.jsonl \
        --holdout runs/sd10_twic_2M/games.holdout.jsonl \
        --train-out runs/sd10_twic_2M/games.train.labeled.jsonl \
        --holdout-out runs/sd10_twic_2M/games.holdout.labeled.jsonl
"""
import argparse
import json
import math


def result_to_white_score(result):
    """1.0/0.5/0.0 from White's perspective, or None if unknown ('*')."""
    return {"1-0": 1.0, "0-1": 0.0, "1/2-1/2": 0.5}.get(result)


def stm_relative_outcome(white_score, fen):
    """Flip a White-perspective game result to the side-to-move's own
    perspective at this FEN, matching the score field's own
    side-to-move-relative convention."""
    stm_is_white = fen.split()[1] == "w"
    return white_score if stm_is_white else (1.0 - white_score)


def collect_calibration_pairs(games_path):
    """(score, r_outcome) for every non-exact position across all
    games in this file."""
    pairs = []
    n_unknown_result = 0
    with open(games_path) as f:
        for line in f:
            g = json.loads(line)
            white_score = result_to_white_score(g["result"])
            if white_score is None:
                n_unknown_result += 1
                continue
            for pos in g["positions"]:
                if pos.get("exact"):
                    continue
                r = stm_relative_outcome(white_score, pos["fen"])
                pairs.append((pos["score"], r))
    return pairs, n_unknown_result


def mse(pairs, k):
    err = 0.0
    for score, r in pairs:
        p = 1.0 / (1.0 + math.exp(-score / k))
        err += (r - p) ** 2
    return err / max(len(pairs), 1)


def fit_k(pairs, lo=0.05, hi=20.0, iterations=40):
    for _ in range(iterations):
        m1 = lo + (hi - lo) / 3
        m2 = hi - (hi - lo) / 3
        if mse(pairs, m1) < mse(pairs, m2):
            hi = m2
        else:
            lo = m1
    return (lo + hi) / 2


def label_file(in_path, out_path, k):
    n_games = 0
    n_positions = 0
    n_exact_skipped = 0
    n_unknown_result = 0
    with open(in_path) as fin, open(out_path, "w") as fout:
        for line in fin:
            g = json.loads(line)
            white_score = result_to_white_score(g["result"])
            if white_score is None:
                n_unknown_result += 1
                fout.write(json.dumps(g) + "\n")
                continue
            for pos in g["positions"]:
                n_positions += 1
                if pos.get("exact"):
                    n_exact_skipped += 1
                    continue
                r = stm_relative_outcome(white_score, pos["fen"])
                p_search = 1.0 / (1.0 + math.exp(-pos["score"] / k))
                pos["p_search"] = p_search
                pos["r_outcome"] = r
            fout.write(json.dumps(g) + "\n")
            n_games += 1
    return n_games, n_positions, n_exact_skipped, n_unknown_result


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--train", required=True)
    ap.add_argument("--holdout", required=True)
    ap.add_argument("--train-out", required=True)
    ap.add_argument("--holdout-out", required=True)
    args = ap.parse_args()

    print(f"Collecting calibration pairs from {args.train} (TRAIN only) ...")
    pairs, n_unknown = collect_calibration_pairs(args.train)
    print(f"  {len(pairs)} (score, outcome) pairs "
          f"({n_unknown} games skipped: unknown result)")

    print("Fitting K (ternary search, minimizing sigmoid-MSE against "
          "actual game outcomes) ...")
    k = fit_k(pairs)
    final_mse = mse(pairs, k)
    print(f"  K = {k:.2f}  (MSE at fit = {final_mse:.5f})")

    print(f"\nLabeling {args.train} -> {args.train_out} ...")
    n_g, n_p, n_exact, n_unk = label_file(args.train, args.train_out, k)
    print(f"  {n_g} games, {n_p} positions, {n_exact} exact (unlabeled), "
          f"{n_unk} unknown-result games (unlabeled)")

    print(f"Labeling {args.holdout} -> {args.holdout_out} (same K, not refit) ...")
    n_g, n_p, n_exact, n_unk = label_file(args.holdout, args.holdout_out, k)
    print(f"  {n_g} games, {n_p} positions, {n_exact} exact (unlabeled), "
          f"{n_unk} unknown-result games (unlabeled)")

    print(f"\nFitted K = {k:.4f} -- record this; applying a different K "
          f"later (e.g. after labeling more data) is not the same model.")


if __name__ == "__main__":
    main()