"""Smoke-validate data.py, graph.py, distributions.py against the real data."""
import sys
sys.path.insert(0, "/work")

import numpy as np
import torch

from heatingforecast import data, graph
from heatingforecast.distributions import JohnsonSU, QUANTILE_LEVELS

print("=== data ===")
wide, n_q3 = data.load_wide("/work/epronet_full_20260427.csv")
print(f"wide: {wide.shape}, dropped {n_q3} quality!=1 rows")
print(f"range: {wide.index.min()} .. {wide.index.max()}")
print(f"NaNs before bridging: {int(wide.isna().sum().sum())}")
wide, obs_mask = data.bridge_gaps(wide)
print(f"NaNs after bridging: {int(wide.isna().sum().sum())}, "
      f"obs_mask mean {obs_mask.mean().mean():.4f}")
tr, va, te = data.splits(wide)
for name, m in [("train", tr), ("val", va), ("test", te)]:
    share = wide[m].isna().mean()
    print(f"{name}: NaN share per consumer: max {share.max():.2f} "
          f"({share.idxmax()}), median {share.median():.4f}")
feats = data.exog_features("/work/weather/weather_memmingen_hourly.csv",
                           "/work/weather/holidays_bavaria.csv", wide.index)
print(f"exog: {feats.shape}, NaN {int(feats.isna().sum().sum())}")
tr, va, te = data.splits(wide)
print(f"split slots: train {tr.sum()}, val {va.sum()}, test {te.sum()}")
print(f"per-consumer kWh median: min {wide.median().min():.3f}, "
      f"max {wide.median().max():.2f}")

print("\n=== graph ===")
g = graph.load_graph("/work/Memmingen_pressure.yaml",
                     "/work/graph/consumer_mapping.csv")
print(f"nodes {g['n_nodes']}, edges {g['edge_index'].shape[1]}, "
      f"edge_attr {g['edge_attr'].shape}")
print("consumers ordered:", g["consumers"][:6], "...")
assert g["n_nodes"] == 42 and g["edge_attr"].shape[0] == g["edge_index"].shape[1]
assert g["consumers"] == list(wide.columns), "graph consumers != wide columns"

print("\n=== JohnsonSU ===")
torch.manual_seed(0)
loc = torch.randn(5)
d = JohnsonSU(loc=loc, scale=torch.rand(5) + 0.5, skew=torch.randn(5) * 0.5,
              tail=torch.rand(5) + 0.3)
q = d.quantile(torch.tensor(QUANTILE_LEVELS, dtype=torch.float32))
assert torch.all(torch.diff(q, dim=0) > 0), "quantiles not monotone"
med = d.quantile(torch.tensor([0.5]))
assert torch.allclose(med[0], loc - d.scale * torch.sinh(d.skew / d.tail),
                      atol=1e-5), "median formula mismatch"
s = d.rsample(torch.Size([200_000]))
assert torch.isfinite(d.log_prob(s)).all()
emp = s.mean(dim=0)
print("mean: closed form vs empirical:",
      [f"{a:.3f}/{b:.3f}" for a, b in zip(d.mean.tolist(), emp.tolist())])
print("OK: monotone quantiles, median formula, finite log-prob, mean matches")