"""Shared safe training utilities. Real training is gated by dataset metadata."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import json, random
import numpy as np
import torch
from torch import Tensor, nn
from sklearn.model_selection import GroupShuffleSplit

class DatasetContractError(ValueError): pass

def seed_everything(seed: int = 17) -> None:
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)

def masked_mse(pred: Tensor, target: Tensor, valid: Tensor) -> Tensor:
    valid = valid.bool()
    if valid.numel() == 0 or not valid.any(): raise DatasetContractError("Refusing training: target validity mask has no validated labels.")
    if pred.shape != target.shape or valid.shape != target.shape:
        raise DatasetContractError(f"prediction, target, and valid mask must have identical shapes; got {tuple(pred.shape)}, {tuple(target.shape)}, {tuple(valid.shape)}")
    if not torch.isfinite(pred).all(): raise DatasetContractError("model prediction is non-finite")
    # Index first: invalid targets may safely be NaN, without contaminating autograd.
    p, y = pred[valid], target[valid]
    if not torch.isfinite(y).all(): raise DatasetContractError("a target marked valid is non-finite")
    return (p - y).square().mean()

@dataclass
class TrainOnlyScaler:
    mean: np.ndarray | None = None; scale: np.ndarray | None = None
    def fit(self, x: np.ndarray, feature_valid: np.ndarray) -> "TrainOnlyScaler":
        if x.shape != feature_valid.shape: raise DatasetContractError("feature_valid must exactly match x")
        if np.any(~np.isfinite(x[feature_valid])): raise DatasetContractError("observed input feature is non-finite")
        self.mean=np.zeros(x.shape[1],dtype=np.float64); self.scale=np.ones(x.shape[1],dtype=np.float64)
        for j in range(x.shape[1]):
            observed=x[feature_valid[:,j],j]
            if len(observed): self.mean[j]=observed.mean(); self.scale[j]=max(float(observed.std()),1.) if observed.std()<1e-8 else float(observed.std())
        return self
    def transform(self, x: np.ndarray, feature_valid: np.ndarray) -> np.ndarray:
        if self.mean is None or self.scale is None: raise RuntimeError("fit scaler on training split first")
        if x.shape != feature_valid.shape: raise DatasetContractError("feature_valid must exactly match x")
        return np.where(feature_valid, (x - self.mean) / self.scale, 0.).astype(np.float32)

def group_split(groups: np.ndarray, seed: int = 17) -> tuple[np.ndarray, np.ndarray]:
    if len(np.unique(groups)) < 2: raise DatasetContractError("Need at least two independent group IDs; random point/patch splits are prohibited.")
    return next(GroupShuffleSplit(n_splits=1, test_size=.25, random_state=seed).split(np.zeros(len(groups)), groups=groups))

def _require_contract(d: Any) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    keys = set(d.files)
    required = {"x", "feature_valid", "y", "training_valid", "group_id", "label_status"}
    if not required <= keys: raise DatasetContractError(f"Training NPZ requires {sorted(required)}; missing {sorted(required - keys)}")
    x, feature_valid, y, valid, groups = d["x"], d["feature_valid"].astype(bool), d["y"], d["training_valid"].astype(bool), d["group_id"]
    if x.ndim != 2 or feature_valid.shape != x.shape or y.ndim != 2 or y.shape[0] != x.shape[0] or valid.shape != y.shape or groups.shape != (x.shape[0],):
        raise DatasetContractError("contract requires x/feature_valid [samples,features], y/training_valid [samples,targets], and group_id [samples]")
    if not valid.any(): raise DatasetContractError("Refusing training: no externally validated real labels (training_valid is empty).")
    if not np.all(d["label_status"].astype(str) == "externally_validated"):
        raise DatasetContractError("Refusing training: label_status must be externally_validated for every sample.")
    return x, y, valid, groups, feature_valid

def fit_validated_npz(model: nn.Module, dataset: str | Path, checkpoint: str | Path, epochs: int = 5, seed: int = 17) -> dict[str, Any]:
    """Reference entry point for tabular/vector models with strict data contract."""
    seed_everything(seed); d = np.load(dataset, allow_pickle=False); x, y, valid, groups, feature_valid = _require_contract(d)
    train, test = group_split(groups, seed); scaler = TrainOnlyScaler().fit(x[train], feature_valid[train]); x = np.concatenate([scaler.transform(x, feature_valid), feature_valid.astype(np.float32)], axis=1)
    xt, yt, mt = (torch.as_tensor(x[train], dtype=torch.float32), torch.as_tensor(y[train], dtype=torch.float32), torch.as_tensor(valid[train]))
    opt = torch.optim.Adam(model.parameters(), 1e-3); model.train()
    for _ in range(epochs):
        opt.zero_grad(); loss = masked_mse(model(xt), yt, mt); loss.backward(); opt.step()
    path = Path(checkpoint); path.parent.mkdir(parents=True, exist_ok=True)
    torch.save({"state_dict": model.state_dict(), "scaler_mean": scaler.mean, "scaler_scale": scaler.scale, "train_groups": np.unique(groups[train]).tolist()}, path)
    # A held-out number is only lawful for labels marked valid in the held-out groups.
    model.eval()
    with torch.no_grad():
        pred = model(torch.as_tensor(x[test], dtype=torch.float32)); heldout = masked_mse(pred, torch.as_tensor(y[test], dtype=torch.float32), torch.as_tensor(valid[test]))
    return {"status": "trained", "train_samples": int(len(train)), "test_samples": int(len(test)), "train_loss": float(loss.detach()), "heldout_masked_mse": float(heldout)}

def validated_contract_dimensions(dataset: str | Path) -> tuple[int, int]:
    """Validate before exposing dimensions, so malformed NPZs yield a contract error."""
    d = np.load(dataset, allow_pickle=False)
    x, y, _, _, _ = _require_contract(d)
    return int(x.shape[1]), int(y.shape[1])