"""Per-step marginal KL(real||gen) per target dim, 10 vs 20 ODE steps.

Streaming histogram over shared bin edges (computed from the true distribution),
so the two runs are directly comparable dim-by-dim.
"""

import numpy as np
import polars as pl

from giant.analysis import RAW_TARGET_NAMES, _kl_from_counts, _raw_dim_expr

FILES = {
    "10-step": "/home/lars/Programming/giant/9879e806-5e88-4b06-b1fa-0e61de9cda6f.parquet",
    "20-step": "/home/lars/Programming/giant/0e236919-65ae-4b9b-9957-d31a8211aca4.parquet",
}
BINS = 100

# Fixed shared edges from the true distribution (identical across both files), using
# robust quantiles to avoid a few outliers dominating the range.
base = FILES["10-step"]
edges = {}
for j, name in enumerate(RAW_TARGET_NAMES):
    lo, hi = (
        pl.scan_parquet(base)
        .select(
            _raw_dim_expr("true", j).quantile(0.001).alias("lo"),
            _raw_dim_expr("true", j).quantile(0.999).alias("hi"),
        )
        .collect(engine="streaming")
        .row(0)
    )
    if not (hi - lo > 1e-9):
        lo, hi = lo - 0.5, hi + 0.5
    edges[name] = np.linspace(lo, hi, BINS + 1)


def counts(file, prefix, j, e):
    vals = (
        pl.scan_parquet(file)
        .select(_raw_dim_expr(prefix, j).alias("v"))
        .collect(engine="streaming")["v"]
        .to_numpy()
    )
    c, _ = np.histogram(vals, bins=e)
    return c


kls = {name: {} for name in FILES}
for j, name in enumerate(RAW_TARGET_NAMES):
    e = edges[name]
    real_c = counts(base, "true", j, e)  # identical true dist across files
    for run, f in FILES.items():
        gen_c = counts(f, "pred", j, e)
        kls[run][name] = _kl_from_counts(real_c, gen_c)

print(f"{'dim':<14}{'KL 10-step':>14}{'KL 20-step':>14}{'ratio 20/10':>14}")
print("-" * 56)
tot = {"10-step": 0.0, "20-step": 0.0}
for name in RAW_TARGET_NAMES:
    a, b = kls["10-step"][name], kls["20-step"][name]
    tot["10-step"] += a
    tot["20-step"] += b
    print(f"{name:<14}{a:>14.5f}{b:>14.5f}{b / a if a else float('nan'):>14.2f}")
print("-" * 56)
print(
    f"{'SUM':<14}{tot['10-step']:>14.5f}{tot['20-step']:>14.5f}"
    f"{tot['20-step'] / tot['10-step']:>14.2f}"
)
print(f"{'MEAN':<14}{tot['10-step'] / 9:>14.5f}{tot['20-step'] / 9:>14.5f}")