src/run_full.py

"""
End-to-end runner, rebuilt against explicit contracts (see provenance.py):

  1. Data:      get_dataset() [REAL, hard-fails] or get_synthetic_dataset()
                [explicit opt-in] — never a silent fallback between them.
  2. Validation: trajectory lengths checked BEFORE training starts.
  3. Dataset reuse: DatasetRegistry.claim() blocks retraining on a dataset
                already CONSUMED by a prior run — required because
                multiple contributors will supply datasets over time.
  4. Checkpointing: CheckpointStore — content-addressed
                (sha256 of config+code+dataset), atomic write, and a
                human-readable meta.json sidecar.

Run (real data required by default):
  python -m src.run_full --data-root ./data/real

Run against synthetic data (explicit opt-in, for smoke-testing only):
  python -m src.run_full --synthetic
"""
from __future__ import annotations
import os, sys, argparse
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import torch
from torch.utils.data import DataLoader
import numpy as np

from src.data_real import get_dataset, get_synthetic_dataset
from src.data_pbdb import get_pbdb_dataset, DEFAULT_TAXON_GROUPS
from src.normalization import FieldNormalizer
from src.model import MultiScaleEncoder, HierarchicalHyperbolicPredictor, HyperbolicCritic
from src.physics_losses import combined_physics_loss
from src.env import MultiStepPoincareEnv
from src.ppo import PoincareActor, PPOTrainer
from src.provenance import (
    DatasetRegistry,
    CheckpointStore,
    hash_dataset,
    hash_code,
    validate_trajectory_lengths,
    DatasetAlreadyUsedError,
    DatasetInProgressError,
)
from src.config import BEST_HPARAMS as BEST, WINDOW

SRC_DIR = os.path.dirname(os.path.abspath(__file__))


def collate(batch):
    return torch.stack([b["fields"] for b in batch])


def supervised_pretrain(model, norm, ds, device, epochs=5):
    loader = DataLoader(ds, batch_size=BEST["batch_size"], shuffle=True, collate_fn=collate)
    opt = torch.optim.Adam(model.parameters(), lr=BEST["lr"])
    ps = BEST["pred_steps"]
    w = BEST["w_phys"]
    for ep in range(epochs):
        total, n = 0.0, 0
        for batch in loader:
            B, T, C, H, W = batch.shape
            batch = batch.to(device)
            flat = norm.transform(batch.view(B * T, C, H, W)).view(B, T, C, H, W)
            x = flat[:, :WINDOW]
            with torch.no_grad():
                tgt = torch.stack([model.encode(flat[:, WINDOW + s]) for s in range(ps)], 1)
            pred = model(x)
            loss = model.hyperbolic_loss(pred, tgt) + combined_physics_loss(
                flat[:, : WINDOW + ps], w_smooth=w, w_temp=w, w_cons=w * 0.5
            )
            if not torch.isfinite(loss):
                raise RuntimeError(
                    f"[NON_FINITE_LOSS] loss became {loss.item()} at epoch {ep+1}; "
                    f"stopping rather than silently continuing with a corrupted model."
                )
            opt.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            total += loss.item()
            n += 1
        print(f"  Pretrain epoch {ep+1}/{epochs} loss={total/max(n,1):.4f}")
    return model


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-root", action="append", default=None,
                         help="Directory containing real .hdf5/.h5 Well files. "
                              "May be repeated. Default: ./data/real, ./data/well")
    parser.add_argument("--synthetic", action="store_true",
                         help="Explicit opt-in to synthetic data (smoke test only).")
    parser.add_argument("--pbdb", action="store_true",
                         help="Explicit opt-in to real PBDB fossil-occurrence data "
                              "(spatiotemporal occurrence/diversity density fields). "
                              "Requires network access to paleobiodb.org.")
    parser.add_argument("--pbdb-taxa", nargs="+", default=None,
                         help=f"Taxon groups to fetch from PBDB. Default: {list(DEFAULT_TAXON_GROUPS)}")
    parser.add_argument("--experiment-id", default=None,
                         help="Human label for this run. Default: auto-generated.")
    parser.add_argument("--allow-dataset-reuse", action="store_true",
                         help="Explicit override to retrain on an already-CONSUMED "
                              "dataset. Off by default — reuse is blocked.")
    args = parser.parse_args()

    device = "cpu"
    print("=" * 64)
    print("Full pipeline: data -> hierarchical Poincare -> physics -> PPO")
    print("Optuna best HPs:", BEST)
    print("=" * 64)

    # ---- 1. Data: explicit, no silent fallback -----------------------
    if args.synthetic:
        ds, provenance = get_synthetic_dataset(max_samples=128, n_steps=14)
    elif args.pbdb:
        ds, provenance = get_pbdb_dataset(
            taxon_groups=args.pbdb_taxa or DEFAULT_TAXON_GROUPS,
        )
    else:
        ds, provenance = get_dataset(
            max_samples=128, n_steps=14, search_roots=args.data_root,
        )
    print(f"[data] provenance={provenance} size={len(ds)}")

    # ---- 2. Validate BEFORE training, not mid-loop --------------------
    required_length = WINDOW + BEST["pred_steps"]
    validate_trajectory_lengths(ds, required_length=required_length)
    print(f"[validate] all sampled trajectories >= {required_length} steps: OK")

    # ---- 3. Dataset-reuse registry -------------------------------------
    dataset_hash = hash_dataset(ds, sample_cap=64)
    code_hash = hash_code(SRC_DIR)
    experiment_id = args.experiment_id or f"run_full:{dataset_hash[:8]}:{code_hash[:8]}"
    registry = DatasetRegistry(registry_dir="registry/datasets")

    if args.allow_dataset_reuse:
        status = registry.status(dataset_hash)
        if status and status["status"] == "CONSUMED":
            print(f"[registry] WARNING: explicit override — retraining on "
                  f"already-CONSUMED dataset {dataset_hash[:12]}")
            registry.allow_retry(dataset_hash)

    try:
        registry.claim(dataset_hash, experiment_id)
    except (DatasetAlreadyUsedError, DatasetInProgressError) as e:
        print(f"[registry] BLOCKED: {e}")
        raise

    try:
        # ---- 4. Normalizer ---------------------------------------------
        samples = []
        for i in range(min(48, len(ds))):
            item = ds[i]
            samples.append(item["fields"] if isinstance(item, dict) else item)
        data = torch.stack(samples)
        norm = FieldNormalizer(mode="zscore").fit(data)
        print("[norm] fitted")

        # ---- 5. Hierarchical model --------------------------------------
        enc = MultiScaleEncoder(hidden=BEST["hidden"], out_dim=8)
        model = HierarchicalHyperbolicPredictor(
            enc, c=BEST["curvature"], pred_steps=BEST["pred_steps"], levels=BEST["levels"]
        ).to(device)

        print("\n--- Supervised pre-training with physics priors ---")
        model = supervised_pretrain(model, norm, ds, device, epochs=4)

        # ---- 6. PPO with hyperbolic critic -------------------------------
        print("\n--- PPO fine-tuning with hyperbolic critic ---")
        env = MultiStepPoincareEnv(
            dataset=ds,
            normalizer=norm,
            encoder=model.encoder,
            poincare_module=model.poincare,
            window=WINDOW,
            horizon=BEST["pred_steps"],
            device=device,
        )
        actor = PoincareActor(obs_dim=8, action_dim=8, hidden=64)
        critic = HyperbolicCritic(c=BEST["curvature"])
        ppo = PPOTrainer(actor, critic, model.poincare, lr=BEST["lr"], device=device)

        returns = []
        for update in range(12):
            rollout = ppo.collect_rollout(env, n_steps=48)
            loss = ppo.update(rollout, n_epochs=3, batch_size=16)
            ep_ret = float(np.sum(rollout["rewards"]))
            returns.append(ep_ret)
            if (update + 1) % 3 == 0:
                print(f"  PPO update {update+1}/12  loss={loss:.4f}  rollout_return={ep_ret:.3f}")
        print(f"  Mean return (last 4): {np.mean(returns[-4:]):.3f}")

        # ---- 7. Content-addressed, atomic checkpoint ---------------------
        store = CheckpointStore(checkpoints_dir="checkpoints")
        result = store.save(
            model_state={
                "model": model.state_dict(),
                "actor": actor.state_dict(),
                "critic": critic.state_dict(),
            },
            config=BEST,
            dataset_hash=dataset_hash,
            code_hash=code_hash,
            data_provenance=provenance,
            extra={
                "normalizer": norm.state_dict(),
                "ppo_returns": returns,
                "experiment_id": experiment_id,
            },
        )
        print(f"\n[checkpoint] {result['outcome_code']} -> {result['path']}")

        registry.mark_consumed(dataset_hash)
        print(f"[registry] dataset {dataset_hash[:12]} marked CONSUMED "
              f"(future runs on this exact data will be blocked by default)")

    except Exception as e:
        registry.mark_failed(dataset_hash, error_detail=str(e))
        print(f"[registry] dataset {dataset_hash[:12]} marked FAILED: {e}")
        raise

    print("\nDone.")


if __name__ == "__main__":
    main()