"""
MLP training orchestrator — trains N seeds per asset/TF combo, sweeps thresholds,
optionally regenerates Pine presets and sends ntfy notification.

Replaces the ad-hoc shell scripts written to /tmp/ each session.

Usage:
    # Train BTC all TFs, 3 seeds, sweep + promote + notify
    python3 tools/run_mlp_train.py --assets COINBASE_BTCUSD --seeds 3 --sweep --promote --notify

    # Train all 4 assets, 4H only, 6 seeds, then generate presets
    python3 tools/run_mlp_train.py --tfs 4H --seeds 6 --sweep --promote --presets --notify

    # Extra seeds 404-606 for BTC 6H and 8H (specific seeds)
    python3 tools/run_mlp_train.py --assets COINBASE_BTCUSD --tfs 6H 8H --seed-start 404 --seeds 3

    # Dry run — show what would be launched without running
    python3 tools/run_mlp_train.py --assets COINBASE_BTCUSD --seeds 3 --dry-run

    # All 4 assets, all TFs, 3 seeds, full pipeline
    python3 tools/run_mlp_train.py --seeds 3 --sweep --promote --presets --table --notify
"""

from __future__ import annotations

import argparse
import concurrent.futures
import json
import os
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO))

ALL_ASSETS = ["COINBASE_BTCUSD", "COINBASE_ETHUSD", "BINANCE_SOLUSD", "BINANCE_LINKUSD"]
ALL_TFS    = ["4H", "6H", "8H", "12H", "1D"]
TF_PERIOD  = {"4H": "240", "6H": "360", "8H": "480", "12H": "720", "1D": "1D"}

WEIGHTS_DIR = REPO / "strategies" / "params" / "mlp"
VENV_PYTHON = REPO / ".venv" / "bin" / "python3"
PYTHON      = str(VENV_PYTHON) if VENV_PYTHON.exists() else sys.executable


def log(msg: str) -> None:
    print(f"[{datetime.now(timezone.utc).strftime('%H:%M:%S')}] {msg}", flush=True)


def data_path(asset: str, tf: str) -> Path:
    period = TF_PERIOD[tf]
    return REPO / "data" / "mlp" / f"{asset}, {period}.csv"


def artifact_path(asset: str, tf: str, seed: int, tag: str = "bb50") -> Path:
    return WEIGHTS_DIR / f"mlp_weights_{asset}_{tf}_{tag}_seed{seed}.json"


def seed_range(seed_start: int, n_seeds: int) -> list[int]:
    """Generate seeds as seed_start, seed_start+101, seed_start+202, ..."""
    return [seed_start + i * 101 for i in range(n_seeds)]


def existing_seeds(asset: str, tf: str, tag: str = "bb50") -> list[int]:
    """Return sorted list of seeds already trained for this asset/TF/tag."""
    seeds = []
    for p in WEIGHTS_DIR.glob(f"mlp_weights_{asset}_{tf}_{tag}_seed*.json"):
        try:
            s = int(p.stem.split("_seed")[-1])
            seeds.append(s)
        except ValueError:
            pass
    return sorted(seeds)


def effective_artifact_tag(tag: str, regime: str, input_structure: str) -> str:
    """Keep non-dense experiments out of the production ``bb50`` namespace."""
    if input_structure == "hybrid_grouped" and tag.startswith((
        "grouped", "random_sparse", "bb50_grouped", "bb50_random_sparse",
    )):
        raise ValueError(
            f"Artifact tag '{tag}' is reserved by an existing input-structure experiment; "
            "use an isolated hybrid tag such as 'hybrid_grouped_v1'."
        )
    if input_structure != "dense" and tag == "bb50":
        tag = f"{tag}_{input_structure}"
    return tag if regime == "all" else f"{tag}_{regime}"


def train_one(asset: str, tf: str, seed: int, hidden: list[int],
              fold_objective: str, l2: float, es_workers: int,
              es_generations: int, tag: str, dry_run: bool,
              regime: str = "all", input_structure: str = "dense",
              mask_seed: int | None = None, temporal_features: bool = False) -> tuple[str, bool, str]:
    out = artifact_path(asset, tf, seed, tag)
    data = data_path(asset, tf)
    label = f"{asset} {tf} seed{seed}"

    if not data.exists():
        return label, False, f"data file missing: {data}"

    cmd = [
        PYTHON, "tools/train_mlp.py",
        "--data", str(data),
        "--asset", asset, "--tf", tf,
        "--fold-objective", fold_objective,
        "--l2", str(l2),
        "--hidden", *[str(h) for h in hidden],
        "--seed", str(seed),
        "--es-workers", str(es_workers),
        "--es-generations", str(es_generations),
        "--out", str(out),
    ]
    if regime != "all":
        cmd += ["--regime", regime]
    if input_structure != "dense":
        cmd += ["--input-structure", input_structure]
    if input_structure == "random_sparse":
        if mask_seed is None:
            return label, False, "random_sparse requires a mask seed"
        cmd += ["--mask-seed", str(mask_seed)]
    if temporal_features:
        cmd += ["--temporal-features"]
    if dry_run:
        print(f"  [dry-run] {' '.join(cmd)}")
        return label, True, "dry-run"

    log_path = Path(f"/tmp/mlp_train_{asset}_{tf}_seed{seed}.log")
    t0 = time.time()
    try:
        with open(log_path, "w") as f:
            result = subprocess.run(cmd, stdout=f, stderr=f, cwd=str(REPO))
        elapsed = int(time.time() - t0)
        if result.returncode != 0:
            return label, False, f"exit {result.returncode} after {elapsed}s — see {log_path}"
        return label, True, f"{elapsed}s → {out.name}"
    except Exception as e:
        return label, False, str(e)


def sweep_one(asset: str, tf: str, tag: str, samples: int, promote: bool,
              dry_run: bool) -> tuple[str, bool, str]:
    label = f"{asset} {tf}"
    extras = list(WEIGHTS_DIR.glob(f"mlp_weights_{asset}_{tf}_{tag}_seed*.json"))
    if not extras:
        return label, False, "no artifacts to sweep"

    cmd = [
        PYTHON, "tools/run_mlp_deep_sweep.py",
        "--asset", asset,
        "--timeframes", tf,
        "--samples", str(samples),
        "--extra-weights", *[str(e) for e in extras],
    ]
    if promote:
        cmd.append("--promote")

    if dry_run:
        print(f"  [dry-run] {' '.join(cmd)}")
        return label, True, "dry-run"

    log_path = Path(f"/tmp/mlp_sweep_{asset}_{tf}.log")
    t0 = time.time()
    try:
        with open(log_path, "w") as f:
            result = subprocess.run(cmd, stdout=f, stderr=f, cwd=str(REPO))
        elapsed = int(time.time() - t0)
        if result.returncode != 0:
            return label, False, f"exit {result.returncode} after {elapsed}s — see {log_path}"
        return label, True, f"{elapsed}s"
    except Exception as e:
        return label, False, str(e)


def run_presets(dry_run: bool) -> bool:
    if dry_run:
        print("  [dry-run] python3 tools/generate_pine_mlp_presets.py")
        return True
    result = subprocess.run(
        [PYTHON, "tools/generate_pine_mlp_presets.py"],
        cwd=str(REPO), capture_output=True, text=True
    )
    print(result.stdout.strip())
    if result.returncode != 0:
        print(result.stderr.strip(), file=sys.stderr)
    return result.returncode == 0


def run_table(assets: list[str], dry_run: bool) -> str:
    if dry_run:
        print("  [dry-run] python3 tools/mlp_results_table.py --save")
        return ""
    lines = []
    for asset in assets:
        result = subprocess.run(
            [PYTHON, "tools/mlp_results_table.py", "--asset", asset],
            cwd=str(REPO), capture_output=True, text=True
        )
        lines.append(result.stdout.strip())
    return "\n".join(lines)


def send_notify(title: str, msg: str) -> None:
    subprocess.run(
        [PYTHON, "tools/ntfy.py", msg, "--title", title, "--priority", "high"],
        cwd=str(REPO), capture_output=True
    )


def main() -> None:
    parser = argparse.ArgumentParser(
        description="MLP training orchestrator",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument("--assets", nargs="+", default=ALL_ASSETS, choices=ALL_ASSETS,
                        metavar="ASSET", help="Assets to train (default: all 4)")
    parser.add_argument("--tfs", nargs="+", default=ALL_TFS, choices=ALL_TFS,
                        metavar="TF", help="Timeframes (default: all 5)")
    parser.add_argument("--seeds", type=int, default=3,
                        help="Number of seeds per combo (default: 3)")
    parser.add_argument("--seed-start", type=int, default=101,
                        help="First seed value; subsequent seeds += 101 (default: 101)")
    parser.add_argument("--skip-existing", action="store_true",
                        help="Skip training if artifact already exists for that seed")
    parser.add_argument("--tag", default="bb50",
                        help="Artifact name tag, e.g. 'bb50' (default: bb50)")
    # Architecture / training
    parser.add_argument("--hidden", type=int, nargs="+", default=[16, 8])
    parser.add_argument("--fold-objective", default="robust",
                        choices=["mean", "min", "mean_min", "robust"])
    parser.add_argument("--l2", type=float, default=0.01)
    parser.add_argument("--es-workers", type=int, default=2,
                        help="CMA-ES worker processes per training job (default: 2, "
                             "keeps headroom for parallel jobs)")
    parser.add_argument("--es-generations", type=int, default=300)
    parser.add_argument("--parallel-jobs", type=int, default=0,
                        help="Max parallel training jobs (default: 0 = all combos at once)")
    # Post-training
    parser.add_argument("--sweep", action="store_true",
                        help="Run threshold sweep after training")
    parser.add_argument("--samples", type=int, default=80000,
                        help="Sweep samples per TF (default: 80000)")
    parser.add_argument("--promote", action="store_true",
                        help="Promote sweep winners (requires --sweep)")
    parser.add_argument("--presets", action="store_true",
                        help="Regenerate Pine presets after sweep")
    parser.add_argument("--table", action="store_true",
                        help="Print regression table after sweep")
    parser.add_argument("--notify", action="store_true",
                        help="Send ntfy notification when done")
    parser.add_argument("--dry-run", action="store_true",
                        help="Print commands without running them")
    parser.add_argument("--regime", choices=["bull", "bear", "sideways", "all"], default="all",
                        help="Filter Phase-1 training to MVRV regime bars "
                             "(requires mvrv_regime column). Auto-appends regime suffix to tag.")
    parser.add_argument("--input-structure", choices=["dense", "grouped", "hybrid_grouped", "random_sparse"], default="dense",
                        help="Pass the first-layer connectivity experiment to each training run.")
    parser.add_argument("--random-mask-seed-start", type=int,
                        help="Required with random_sparse; one explicit mask seed per training seed, "
                             "incremented by 101 in the same order as training seeds.")
    parser.add_argument("--temporal-features", action="store_true",
                        help="Experimental causal EMA companions; cannot sweep, promote, or generate Pine presets.")
    args = parser.parse_args()

    if args.input_structure == "random_sparse" and args.random_mask_seed_start is None:
        parser.error("--random-mask-seed-start is required with --input-structure random_sparse")
    if args.input_structure == "hybrid_grouped" and (args.sweep or args.promote or args.presets):
        parser.error("hybrid_grouped does not support sweep, promote, or presets; use canonical replay only")
    if args.temporal_features and (args.sweep or args.promote or args.presets):
        parser.error("temporal-features does not support sweep, promote, or presets; use canonical replay only")
    if args.temporal_features and args.tag == "bb50":
        args.tag = "temporal_ema_v1"

    # Keep experimental structures out of the production namespace, then suffix
    # regime so bull-trained artifacts cannot overwrite all-regime candidates.
    effective_tag = effective_artifact_tag(args.tag, args.regime, args.input_structure)
    if effective_tag != args.tag:
        log(f"regime={args.regime}: tag auto-suffixed '{args.tag}' → '{effective_tag}'")

    combos = [(a, tf) for a in args.assets for tf in args.tfs]
    seeds  = seed_range(args.seed_start, args.seeds)
    mask_seeds = (seed_range(args.random_mask_seed_start, len(seeds))
                  if args.input_structure == "random_sparse" else [None] * len(seeds))
    jobs = [
        (a, tf, seed, mask_seed)
        for a, tf in combos for seed, mask_seed in zip(seeds, mask_seeds)
    ]

    log(f"Plan: {len(args.assets)} assets × {len(args.tfs)} TFs × {len(seeds)} seeds "
        f"= {len(jobs)} training jobs")
    if args.skip_existing:
        jobs = [(a, tf, s, mask_seed) for a, tf, s, mask_seed in jobs
                if not artifact_path(a, tf, s, effective_tag).exists()]
        log(f"  After skipping existing: {len(jobs)} jobs remaining")

    if args.dry_run:
        log("DRY RUN — no commands executed")

    # ── Phase 1: Training ─────────────────────────────────────────────────────
    t_train = time.time()
    max_workers = args.parallel_jobs if args.parallel_jobs > 0 else max(1, len(jobs))

    failures: list[str] = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(train_one, a, tf, s, args.hidden, args.fold_objective,
                        args.l2, args.es_workers, args.es_generations, effective_tag,
                        args.dry_run, args.regime, args.input_structure, mask_seed, args.temporal_features): (a, tf, s)
            for a, tf, s, mask_seed in jobs
        }
        for fut in concurrent.futures.as_completed(futures):
            label, ok, detail = fut.result()
            status = "✅" if ok else "❌"
            log(f"  {status} {label}: {detail}")
            if not ok:
                failures.append(f"{label}: {detail}")

    log(f"Training done in {int(time.time()-t_train)}s — "
        f"{len(jobs)-len(failures)}/{len(jobs)} succeeded")

    if failures:
        log(f"Failures:\n" + "\n".join(f"  {f}" for f in failures))

    # ── Phase 2: Sweep ────────────────────────────────────────────────────────
    if args.sweep:
        log("Starting threshold sweeps (sequential per combo)...")
        sweep_failures = []
        for asset in args.assets:
            for tf in args.tfs:
                label, ok, detail = sweep_one(asset, tf, effective_tag,
                                              args.samples, args.promote, args.dry_run)
                status = "✅" if ok else "❌"
                log(f"  {status} sweep {label}: {detail}")
                if not ok:
                    sweep_failures.append(f"{label}: {detail}")
        if sweep_failures:
            log("Sweep failures:\n" + "\n".join(f"  {f}" for f in sweep_failures))

    # ── Phase 3: Pine presets ─────────────────────────────────────────────────
    if args.presets:
        log("Regenerating Pine presets...")
        ok = run_presets(args.dry_run)
        log(f"  {'✅' if ok else '❌'} Pine presets")

    # ── Phase 4: Regression table + performance artifact ─────────────────────
    table_output = ""
    if args.table:
        log("Running regression table...")
        table_output = run_table(args.assets, args.dry_run)
        print(table_output)
    elif not args.dry_run and args.input_structure == "dense":
        # Always refresh the tracked artifact even without --table
        for asset in args.assets:
            subprocess.run(
                [PYTHON, "tools/mlp_results_table.py", "--asset", asset],
                cwd=str(REPO), capture_output=True,
            )

    # ── Notify ────────────────────────────────────────────────────────────────
    if args.notify and not args.dry_run:
        n_ok   = len(jobs) - len(failures)
        msg    = f"{n_ok}/{len(jobs)} trained"
        if args.sweep:
            msg += " + swept"
        if args.presets:
            msg += " + presets"
        if table_output:
            # Pull OOS line from table for quick summary
            oos_lines = [l for l in table_output.splitlines() if "OOS" in l or "%" in l]
            if oos_lines:
                msg += "\n" + "\n".join(oos_lines[:6])
        send_notify("TradingBot25 MLP", msg)
        log("ntfy sent")

    log("Done.")
    os._exit(0)


if __name__ == "__main__":
    main()