"""Pre-flight training time / cost estimate from a *measured* throughput.

Runs the real training loop (same data pipeline, model, micro-batch, AMP,
compile) for `bench_seconds` on the current machine, takes the steady-state
samples/s (first logging window dropped: compile / warm-up), then extrapolates.
Run it on the rented instance itself for GPU numbers - measuring locally and
dividing by a guessed speed-up is exactly the kind of invented number to avoid.
"""
from __future__ import annotations

import csv
import copy
import tempfile
import time
from pathlib import Path

from minijev.model import human_params


def measure_throughput(cfg: dict, seconds: float = 60.0) -> dict:
    from minijev.training.supervised import train

    cfg = copy.deepcopy(cfg)
    data_dir = Path(cfg["data"]["dir"])
    tmp = Path(tempfile.mkdtemp(prefix="minijev-est-"))
    if not (data_dir / "train.npz").exists():
        # same generator settings, small size: per-row encoding cost is identical
        gen = dict(cfg["data"].get("generate", {}))
        gen["samples"] = 30_000
        cfg["data"] = {"dir": str(tmp / "data"), "generate": gen}
    tcfg = cfg["training"]
    tcfg["log_every"] = 10
    tcfg["eval_every"] = 10**9
    tcfg["checkpoint_every"] = 10**9
    s = train(cfg, run_dir=tmp / "run", max_minutes=seconds / 60, quiet=True, final_eval=False)
    rows = list(csv.DictReader(open(tmp / "run" / "metrics.csv")))
    rates = [float(r["samples_per_sec"]) for r in rows if r.get("samples_per_sec")]
    steady = rates[1:] or rates
    sps = sorted(steady)[len(steady) // 2] if steady else s["throughput_samples_per_sec"]
    return {"samples_per_sec": sps, "runtime": s["runtime"], "params": s["params"], "train_rows": s["train_rows"],
            "windows": len(rates)}


def measure_generation(cfg: dict, n: int = 5000) -> float | None:
    """Rows/s of oracle data generation on this machine (all cores)."""
    gen = cfg["data"].get("generate")
    if not gen:
        return None
    import importlib
    mod = importlib.import_module(f"environments.{cfg['task']['name']}.dataset")
    g = {k: v for k, v in gen.items() if k not in ("samples", "seed")}
    t = time.time()
    mod.generate(n, **g)
    return n / (time.time() - t)


def estimate(cfg: dict, hourly_cost: float, throughput: float | None = None, bench_seconds: float = 60,
             gpu_name: str | None = None, budget: float | None = None, overhead_minutes: float = 3.0) -> dict:
    from minijev.device import resolve_runtime
    from minijev.tasks import get_task
    from minijev.training.supervised import make_model

    task = get_task(cfg["task"]["name"], cfg["task"])
    params = make_model(cfg, task).num_parameters()
    data_dir = Path(cfg["data"]["dir"])
    if (data_dir / "train.npz").exists():
        import numpy as np
        rows = int(len(np.load(data_dir / "train.npz")["target"]))
        gen_rows = 0
    else:
        total = int(cfg["data"].get("generate", {}).get("samples", 0))
        rows, gen_rows = int(total * 0.9), total
    if cfg["data"].get("train_limit"):
        rows = min(rows, int(cfg["data"]["train_limit"]))
    epochs = float(cfg["training"].get("epochs", 1))
    prof = resolve_runtime(cfg["training"], params)
    measured = throughput is None
    if measured:
        print(f"[estimate] measuring real throughput for {bench_seconds:.0f}s on {prof.device} ({prof.device_name}) ...")
        m = measure_throughput(cfg, bench_seconds)
        throughput = m["samples_per_sec"]
    gen_rate = measure_generation(cfg) if gen_rows else None
    train_s = rows * epochs / throughput
    evals = max(1, int(rows * epochs / cfg["training"].get("batch_size", 256) / cfg["training"].get("eval_every", 1000)))
    eval_rows = min(cfg["training"].get("eval_max_rows", 20000), int(rows / 18))
    eval_s = evals * eval_rows / (3 * throughput)  # inference ~3x cheaper than a training step
    gen_s = gen_rows / gen_rate if gen_rate else 0.0
    total_s = train_s + eval_s + gen_s + overhead_minutes * 60
    hours = total_s / 3600
    cost = hours * hourly_cost
    out = {"model": cfg.get("name"), "params": params, "train_rows": rows, "epochs": epochs,
           "device": gpu_name or prof.device_name, "throughput_samples_per_sec": round(throughput, 1),
           "throughput_measured": measured, "train_minutes": round(train_s / 60, 1),
           "data_generation_minutes": round(gen_s / 60, 1), "eval_minutes": round(eval_s / 60, 1),
           "overhead_minutes": overhead_minutes, "total_minutes": round(total_s / 60, 1),
           "gpu_hours": round(hours, 3), "hourly_cost": hourly_cost, "estimated_cost": round(cost, 3)}
    print()
    print(f"Model: MiniJEV {cfg.get('name')}")
    print(f"Parameters: {human_params(params)}")
    print(f"Dataset: {rows:,} training examples" + (f" (+ generating {gen_rows:,})" if gen_rows else ""))
    print(f"Epochs: {epochs:g}")
    print()
    print(f"Device: {out['device']}  (micro-batch {prof.micro_batch} x accum {prof.grad_accum}, amp={prof.amp_dtype})")
    print(f"{'Measured' if measured else 'Given'} throughput: {throughput:,.0f} samples/sec")
    print()
    print(f"Estimated duration: {_fmt_min(total_s / 60)}  (train {_fmt_min(train_s / 60)}, "
          f"data gen {_fmt_min(gen_s / 60)}, eval {_fmt_min(eval_s / 60)}, overhead {overhead_minutes:g} min)")
    print(f"Estimated GPU hours: {hours:.3f}")
    print(f"Estimated cost: €{cost:.2f}  (at €{hourly_cost:.3f}/h)")
    if budget is not None and cost > budget:
        scale = budget / cost
        print(f"\n!! Over budget (€{budget:.2f}). Reduce e.g. epochs to {epochs * scale:.2f} "
              f"or train rows to {int(rows * scale):,}.")
    return out


def _fmt_min(m: float) -> str:
    return f"{m:.1f} min" if m < 120 else f"{m / 60:.1f} h"