# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Chronological partitions, spatial-holdout flag and stratified LLM evaluation subsamples.

* ``assign_splits``: every sale is assigned to train / validation / test / holdout by its sale date using the
  closed windows of ``configs/data.yaml``; sales outside every window get ``split = "unused"`` (reported).
  ``spatial_holdout_region`` flags the regions withheld in the spatial-generalisation experiment.
* ``stratified_subsample``: deterministic stratified random subsample of each evaluation partition
  (strata = region_group × property_class × price_tercile, proportional allocation with largest-remainder
  rounding, seeded draws). Baselines are evaluated on the full partitions; the LLM arms on these subsamples.
"""
from __future__ import annotations

import hashlib
import json
import time
from datetime import date
from pathlib import Path
from typing import Any

import numpy as np
import polars as pl

from ..config import PROCESSED, RESULTS, SEED, SYNTHETIC, load_config

SPLIT_ORDER = ("train", "validation", "test", "holdout")


def split_windows(cfg: dict[str, Any]) -> dict[str, tuple[date, date]]:
    out: dict[str, tuple[date, date]] = {}
    for name in SPLIT_ORDER:
        w = cfg["splits"][name]
        out[name] = (date.fromisoformat(w["start"]), date.fromisoformat(w["end"]))
    return out


def check_windows(cfg: dict[str, Any]) -> None:
    """Windows must be ordered, non-overlapping and contiguous (no sale can fall in two partitions)."""
    w = split_windows(cfg)
    prev_end: date | None = None
    for name in SPLIT_ORDER:
        start, end = w[name]
        if start > end:
            raise ValueError(f"split {name}: start {start} after end {end}")
        if prev_end is not None and start <= prev_end:
            raise ValueError(f"split {name} overlaps the previous window (starts {start} <= {prev_end})")
        prev_end = end


def assign_splits(analytic: pl.DataFrame, cfg: dict[str, Any] | None = None) -> pl.DataFrame:
    """→ (tx_id, split, spatial_holdout_region)."""
    cfg = cfg or load_config("data")
    check_windows(cfg)
    w = split_windows(cfg)
    expr: pl.Expr = pl.lit("unused")
    for name in reversed(SPLIT_ORDER):
        start, end = w[name]
        expr = pl.when((pl.col("date") >= start) & (pl.col("date") <= end)).then(pl.lit(name)).otherwise(expr)
    regions = list(cfg.get("spatial_holdout", {}).get("regions", []))
    return analytic.select(
        "tx_id",
        expr.alias("split"),
        pl.col("region_name").is_in(regions).alias("spatial_holdout_region"),
    )


def _stratum_label(cols: list[str]) -> pl.Expr:
    return pl.concat_str([pl.col(c).cast(pl.String) for c in cols], separator="|")


def proportional_allocation(counts: dict[str, int], n: int) -> dict[str, int]:
    """Largest-remainder (Hamilton) allocation of ``n`` draws across strata proportional to their sizes, capped at
    the stratum size; the shortfall of capped strata is redistributed among the others."""
    keys = sorted(counts)
    total = sum(counts.values())
    if total == 0 or n <= 0:
        return {k: 0 for k in keys}
    n = min(n, total)
    alloc = {k: 0 for k in keys}
    remaining = n
    open_keys = list(keys)
    while remaining > 0 and open_keys:
        pool = sum(counts[k] - alloc[k] for k in open_keys)
        quotas = {k: remaining * (counts[k] - alloc[k]) / pool for k in open_keys}
        floors = {k: int(np.floor(quotas[k])) for k in open_keys}
        left = remaining - sum(floors.values())
        order = sorted(open_keys, key=lambda k: (-(quotas[k] - floors[k]), k))
        for k in order[:left]:
            floors[k] += 1
        for k in open_keys:
            alloc[k] += floors[k]
        # cap and reopen
        over = {k: alloc[k] - counts[k] for k in open_keys if alloc[k] > counts[k]}
        for k, o in over.items():
            alloc[k] = counts[k]
        remaining = n - sum(alloc.values())
        open_keys = [k for k in open_keys if alloc[k] < counts[k]]
    return alloc


def _derive_seed(seed: int, split: str, stratum: str) -> int:
    h = hashlib.sha256(f"{seed}|{split}|{stratum}".encode()).hexdigest()
    return int(h[:8], 16)


def stratified_subsample(
    analytic: pl.DataFrame,
    splits: pl.DataFrame,
    cfg: dict[str, Any] | None = None,
    n_by_split: dict[str, int] | None = None,
    seed: int | None = None,
) -> pl.DataFrame:
    """→ (tx_id, split, stratum, draw_seed). Deterministic given (table, config, seed)."""
    cfg = cfg or load_config("data")
    sub = cfg["llm_eval_subsample"]
    strata_cols = list(sub["strata"])
    n_by_split = n_by_split or {"validation": sub["validation_n"], "test": sub["test_n"], "holdout": sub["holdout_n"]}
    seed = cfg.get("seed", SEED) if seed is None else seed
    df = analytic.select(["tx_id", *strata_cols]).join(splits.select("tx_id", "split"), on="tx_id", how="inner")
    df = df.with_columns(_stratum_label(strata_cols).alias("stratum"))
    frames: list[pl.DataFrame] = []
    for split, n in n_by_split.items():
        part = df.filter(pl.col("split") == split)
        counts = {k: int(v) for k, v in part.group_by("stratum").len().iter_rows()}
        alloc = proportional_allocation(counts, int(n))
        for stratum in sorted(alloc):
            k = alloc[stratum]
            if k == 0:
                continue
            ids = np.array(sorted(part.filter(pl.col("stratum") == stratum)["tx_id"].to_list()))
            s = _derive_seed(seed, split, stratum)
            rng = np.random.default_rng(s)
            chosen = rng.choice(ids, size=k, replace=False)
            frames.append(pl.DataFrame({"tx_id": sorted(chosen.tolist()), "split": split, "stratum": stratum, "draw_seed": s}, schema={"tx_id": pl.String, "split": pl.String, "stratum": pl.String, "draw_seed": pl.Int64}))
    if not frames:
        return pl.DataFrame(schema={"tx_id": pl.String, "split": pl.String, "stratum": pl.String, "draw_seed": pl.Int64})
    out = pl.concat(frames)
    return out.sort(["split", "stratum", "tx_id"])


def run(synthetic: bool = False) -> dict[str, Any]:
    cfg = load_config("data")
    t0 = time.perf_counter()
    base = SYNTHETIC if synthetic else PROCESSED
    analytic = pl.read_parquet(base / "analytic.parquet")
    splits = assign_splits(analytic, cfg)
    n_by = {"validation": 40, "test": 60, "holdout": 30} if synthetic else None
    subs = stratified_subsample(analytic, splits, cfg, n_by_split=n_by)
    splits.write_parquet(base / "splits.parquet")
    subs.write_parquet(base / "llm_subsamples.parquet")
    counts = {k: int(v) for k, v in splits.group_by("split").len().iter_rows()}
    sub_counts = {k: int(v) for k, v in subs.group_by("split").len().iter_rows()}
    summary = {
        "synthetic": synthetic,
        "splits": {k: counts.get(k, 0) for k in (*SPLIT_ORDER, "unused")},
        "spatial_holdout_rows": int(splits["spatial_holdout_region"].sum()),
        "spatial_holdout_regions": cfg["spatial_holdout"]["regions"],
        "windows": {k: [a.isoformat(), b.isoformat()] for k, (a, b) in split_windows(cfg).items()},
        "llm_subsamples": sub_counts,
        "llm_subsample_strata": int(subs["stratum"].n_unique()),
        "seed": cfg.get("seed", SEED),
        "seconds": round(time.perf_counter() - t0, 2),
    }
    out = RESULTS / "data" / ("splits_summary_synthetic.json" if synthetic else "splits_summary.json")
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(summary, indent=1, ensure_ascii=False))
    return summary