"""
Adaptive MLP training loop with convergence detection.

Trains seed batches per timeframe round-by-round, measures whether the winner
Calmar improved, and retires timeframes that have stopped benefiting.  When a
TF plateaus on all-regime seeds, it switches to --regime bull as an
intervention before retiring permanently.

State machine per TF:
    ACTIVE (phase=all)  →  no improvement × convergence_rounds  →  ACTIVE (phase=bull)
    ACTIVE (phase=bull) →  no improvement × convergence_rounds  →  RETIRED
    RETIRED             →  skipped for all future rounds

All console output is tee'd to results/adaptive_logs/run_<timestamp>.log so the
full run history is available for later review.

Usage:
    # Run with defaults (BTC all TFs, 3 seeds/round, up to 10 rounds)
    python3 tools/run_mlp_adaptive.py --notify

    # Specific TFs, smaller batches
    python3 tools/run_mlp_adaptive.py --tfs 6H 12H 1D --seeds-per-round 2 --notify

    # Preview plan without executing
    python3 tools/run_mlp_adaptive.py --dry-run
"""

from __future__ import annotations

import argparse
import concurrent.futures
import csv
import os
import subprocess
import sys
import time
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import IO

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"
WINNER_DIR  = REPO / "results" / "winners"
LOG_DIR     = REPO / "results" / "adaptive_logs"
VENV_PYTHON = REPO / ".venv" / "bin" / "python3"
PYTHON      = str(VENV_PYTHON) if VENV_PYTHON.exists() else sys.executable

_log_file: IO[str] | None = None


# ── Output ────────────────────────────────────────────────────────────────────

def tee(line: str = "") -> None:
    """Write a line to stdout and the current run log file."""
    print(line, flush=True)
    if _log_file is not None:
        try:
            _log_file.write(line + "\n")
            _log_file.flush()
        except (OSError, ValueError):
            pass


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


# ── File helpers ──────────────────────────────────────────────────────────────

def winner_path(asset: str, tf: str) -> Path:
    return WINNER_DIR / f"optimization_winner_strategy_mlp_scores_{asset}_{tf}.csv"


def _winner_artifact_compatible(winner_row: dict) -> bool:
    """Return False if the winner artifact's feature_cols don't match current FEATURE_COLS."""
    import json as _json
    from strategies.strategy_mlp_scores import FEATURE_COLS
    artifact_path = winner_row.get("mlp_weights_file", "")
    if not artifact_path:
        return True  # no artifact recorded — treat as compatible (missing)
    try:
        art = _json.loads(Path(artifact_path).read_text())
        art_cols = art.get("feature_cols", [])
        return art_cols == FEATURE_COLS
    except Exception:
        return True  # can't read artifact — don't block on uncertainty


def read_winner_calmar(asset: str, tf: str) -> float:
    p = winner_path(asset, tf)
    if not p.exists():
        return 0.0
    try:
        rows = list(csv.DictReader(p.open()))
        if not rows:
            return 0.0
        row = rows[0]
        if not _winner_artifact_compatible(row):
            print(f"  ⚠️  {tf}: winner artifact has incompatible feature_cols — resetting baseline to 0.0 (force-promote enabled)")
            return 0.0
        return float(row.get("Calmar Ratio", 0.0) or 0.0)
    except (ValueError, TypeError, OSError):
        return 0.0


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


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


def next_seed_start(asset: str, tf: str, tag: str) -> int:
    """First seed not yet trained for this asset/TF/tag combination."""
    seeds = existing_seeds(asset, tf, tag)
    return (max(seeds) + 101) if seeds else 101


def phase_tag(phase: str, base_tag: str = "bb50") -> str:
    return base_tag if phase == "all" else f"{base_tag}_{phase}"


def phase_regime(phase: str) -> str:
    return "all" if phase == "all" else phase


# ── TF state ──────────────────────────────────────────────────────────────────

@dataclass
class TFState:
    tf: str
    phase: str = "all"           # "all" | "bull" | "retired"
    no_improve_rounds: int = 0
    last_calmar: float = 0.0
    initial_calmar: float = 0.0  # Calmar at loop start — for final comparison
    phase_won: str = "—"         # phase that produced the best promotion
    rounds_used: int = 0         # total rounds this TF participated in
    next_seed: int = 101         # incremented after each training round


# ── Training phase ────────────────────────────────────────────────────────────

def train_round(
    asset: str,
    active_states: list[TFState],
    seeds_per_round: int,
    base_tag: str,
    hidden: list[int],
    fold_objective: str,
    l2: float,
    es_workers: int,
    es_generations: int,
    dry_run: bool,
) -> None:
    """Train seeds_per_round new seeds for each active TF (all in parallel).

    Prints one summary line per TF when all its seeds finish, not per seed.
    """
    jobs: list[tuple[str, int, str, str]] = []
    for state in active_states:
        tag    = phase_tag(state.phase, base_tag)
        regime = phase_regime(state.phase)
        for i in range(seeds_per_round):
            jobs.append((state.tf, state.next_seed + i * 101, tag, regime))

    if not jobs:
        return

    phase_summary = ", ".join(f"{s.tf}={s.phase}" for s in active_states)
    log(f"Training: {len(active_states)} TFs × {seeds_per_round} seeds = {len(jobs)} jobs "
        f"[{phase_summary}]")

    seeds_per_tf = defaultdict(int)
    for tf, *_ in jobs:
        seeds_per_tf[tf] += 1

    completed: dict[str, list[tuple[bool, int, int]]] = defaultdict(list)

    def _train_one(tf: str, seed: int, tag: str, regime: str) -> tuple[str, int, bool, int]:
        out  = WEIGHTS_DIR / f"mlp_weights_{asset}_{tf}_{tag}_seed{seed}.json"
        data = data_path(asset, tf)
        if not data.exists():
            return tf, seed, False, 0

        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 dry_run:
            tee(f"  [dry-run] {' '.join(cmd)}")
            return tf, seed, True, 0

        log_p = Path(f"/tmp/mlp_train_{asset}_{tf}_{tag}_seed{seed}.log")
        t0 = time.time()
        try:
            with open(log_p, "w") as fh:
                rc = subprocess.run(cmd, stdout=fh, stderr=fh, cwd=str(REPO)).returncode
            return tf, seed, rc == 0, int(time.time() - t0)
        except Exception:
            return tf, seed, False, int(time.time() - t0)

    with concurrent.futures.ThreadPoolExecutor(max_workers=max(1, len(jobs))) as pool:
        futures = {
            pool.submit(_train_one, tf, seed, tag, regime): (tf, seed)
            for tf, seed, tag, regime in jobs
        }
        for fut in concurrent.futures.as_completed(futures):
            tf, seed, ok, elapsed = fut.result()
            completed[tf].append((ok, seed, elapsed))
            if len(completed[tf]) == seeds_per_tf[tf]:
                results = sorted(completed[tf], key=lambda x: x[1])
                n_ok    = sum(1 for ok, _, _ in results if ok)
                ids     = ", ".join(f"seed{s}" for _, s, _ in results)
                times   = "+".join(f"{e}s" for _, _, e in results)
                icon    = "✅" if n_ok == seeds_per_tf[tf] else ("⚠️" if n_ok > 0 else "❌")
                log(f"  {icon} {tf}: {n_ok}/{seeds_per_tf[tf]} trained  ({ids}) [{times}]")

    for state in active_states:
        state.next_seed += seeds_per_round * 101


# ── Sweep phase ───────────────────────────────────────────────────────────────

def sweep_round(
    asset: str,
    active_states: list[TFState],
    base_tag: str,
    samples: int,
    dry_run: bool,
) -> None:
    """Sweep all active TFs in parallel, each with an equal share of CPU cores.

    Always passes all seed artifacts for this TF (all tags) so the best
    historical weight can still win with updated thresholds.
    """
    n_tfs = len(active_states)
    cpu_count = os.cpu_count() or 4
    workers_per_sweep = max(1, cpu_count // n_tfs)

    log(f"Sweeping: {n_tfs} TFs in parallel ({workers_per_sweep} workers each) ...")

    def _sweep_one(state: TFState) -> tuple[str, bool, str]:
        extras = sorted(WEIGHTS_DIR.glob(
            f"mlp_weights_{asset}_{state.tf}_{base_tag}*_seed*.json"
        ))
        if not extras:
            return state.tf, False, "no artifacts found"

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

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

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

    with concurrent.futures.ThreadPoolExecutor(max_workers=n_tfs) as pool:
        futures = {pool.submit(_sweep_one, state): state.tf for state in active_states}
        for fut in concurrent.futures.as_completed(futures):
            tf, ok, detail = fut.result()
            log(f"  {'✅' if ok else '❌'} sweep {tf}: {detail}")


# ── Summary helpers ───────────────────────────────────────────────────────────

def print_comparison_table(asset: str, states: dict[str, TFState]) -> None:
    """Print before/after Calmar comparison for each TF."""
    tee()
    tee(f"{'─'*65}")
    tee(f"{'TF':<6} {'Calmar Before':>14} {'Calmar After':>13} {'Delta':>8}  {'Phase Won':<12} Rounds")
    tee(f"{'─'*65}")
    any_improved = False
    for tf, state in states.items():
        after = state.last_calmar
        delta = after - state.initial_calmar
        delta_str = f"+{delta:.4f}" if delta > 0.0001 else f"{delta:+.4f}"
        improved  = delta > 0.0001
        icon      = "✅" if improved else " "
        if improved:
            any_improved = True
        tee(f"{icon} {tf:<5} {state.initial_calmar:>14.4f} {after:>13.4f} {delta_str:>8}  "
            f"{state.phase_won:<12} {state.rounds_used}")
    tee(f"{'─'*65}")
    if not any_improved:
        tee("  No Calmar improvements — all TFs at architecture/feature ceiling.")
    tee()


def run_presets(dry_run: bool) -> bool:
    if dry_run:
        tee("  [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,
    )
    if result.stdout.strip():
        tee(result.stdout.strip())
    if result.returncode != 0:
        print(result.stderr.strip(), file=sys.stderr)
    return result.returncode == 0


def run_results_table(asset: str) -> str:
    result = subprocess.run(
        [PYTHON, "tools/mlp_results_table.py", "--asset", asset],
        cwd=str(REPO), capture_output=True, text=True,
    )
    return result.stdout.strip()


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,
    )


# ── Main loop ─────────────────────────────────────────────────────────────────

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Adaptive MLP training loop with convergence detection",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument("--assets", nargs="+", default=["COINBASE_BTCUSD"], choices=ALL_ASSETS,
                        metavar="ASSET",
                        help="Asset to train (default: COINBASE_BTCUSD; multi-asset not yet supported)")
    parser.add_argument("--tfs", nargs="+", default=ALL_TFS, choices=ALL_TFS, metavar="TF",
                        help="Timeframes to include (default: all 5)")
    parser.add_argument("--seeds-per-round", type=int, default=3,
                        help="New seeds to train per TF per round (default: 3)")
    parser.add_argument("--max-rounds", type=int, default=10,
                        help="Hard cap on total rounds (default: 10)")
    parser.add_argument("--improvement-threshold", type=float, default=0.02,
                        help="Min Calmar improvement to count as progress (default: 0.02)")
    parser.add_argument("--convergence-rounds", type=int, default=3,
                        help="Consecutive non-improving rounds before phase change (default: 2)")
    parser.add_argument("--tag", default="bb50",
                        help="Artifact tag prefix; bull suffix auto-appended (default: bb50)")
    parser.add_argument("--samples", type=int, default=80000,
                        help="Sweep candidate samples per TF (default: 80000)")
    # 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 workers per training job (default: 2)")
    parser.add_argument("--es-generations", type=int, default=300)
    # Output
    parser.add_argument("--presets", action="store_true", default=True,
                        help="Regenerate Pine presets after any promotion round (default: on)")
    parser.add_argument("--no-presets", dest="presets", action="store_false",
                        help="Skip Pine preset regeneration (use for intermediate training passes)")
    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 executing")
    args = parser.parse_args()

    if len(args.assets) > 1:
        print("ERROR: multi-asset not yet supported — pass a single --asset", file=sys.stderr)
        sys.exit(1)
    asset = args.assets[0]

    # ── Log file setup ────────────────────────────────────────────────────────
    global _log_file
    if not args.dry_run:
        LOG_DIR.mkdir(parents=True, exist_ok=True)
        run_ts   = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H%M%S")
        log_path = LOG_DIR / f"run_{run_ts}.log"
        _log_file = open(log_path, "w")
        tee(f"Log: {log_path.relative_to(REPO)}")

    # ── Initialise TF states ──────────────────────────────────────────────────
    states: dict[str, TFState] = {}
    for tf in args.tfs:
        calmar = read_winner_calmar(asset, tf)
        nxt    = next_seed_start(asset, tf, args.tag)
        states[tf] = TFState(
            tf=tf, phase="all",
            last_calmar=calmar, initial_calmar=calmar,
            next_seed=nxt,
        )

    log(f"Adaptive MLP — {asset} — {len(states)} TFs")
    log(f"Config: {args.seeds_per_round} seeds/round, max {args.max_rounds} rounds, "
        f"convergence={args.convergence_rounds} rounds, threshold={args.improvement_threshold}")
    log("Starting Calmar per TF:")
    for tf, s in states.items():
        bull_n = len(existing_seeds(asset, tf, f"{args.tag}_bull"))
        log(f"  {tf}: Calmar={s.last_calmar:.4f}  next_seed={s.next_seed} "
            f"(bull seeds already trained: {bull_n})")

    any_promoted = False
    round_num    = 0
    t_start      = time.time()

    while round_num < args.max_rounds:
        active = [s for s in states.values() if s.phase != "retired"]
        if not active:
            log("All TFs retired — stopping.")
            break

        round_num += 1
        log(f"\n{'═'*60}")
        log(f"Round {round_num}/{args.max_rounds} — {len(active)} active TFs: "
            f"{', '.join(s.tf for s in active)}")
        log(f"{'═'*60}")

        pre_calmars = {tf: read_winner_calmar(asset, tf) for tf in states}

        # ── Training ──────────────────────────────────────────────────────────
        train_round(
            asset, active, args.seeds_per_round, args.tag,
            args.hidden, args.fold_objective, args.l2,
            args.es_workers, args.es_generations, args.dry_run,
        )

        # ── Sweep ─────────────────────────────────────────────────────────────
        sweep_round(asset, active, args.tag, args.samples, args.dry_run)

        # ── Improvement check & state transitions ─────────────────────────────
        log("Improvement check:")
        promoted_this_round = False
        for state in active:
            state.rounds_used += 1
            post_calmar = read_winner_calmar(asset, state.tf) if not args.dry_run else state.last_calmar
            delta       = post_calmar - pre_calmars[state.tf]
            improved    = delta >= args.improvement_threshold

            if improved:
                state.no_improve_rounds = 0
                state.last_calmar       = post_calmar
                state.phase_won         = state.phase
                promoted_this_round     = True
                any_promoted            = True
                log(f"  ✅ {state.tf}: {pre_calmars[state.tf]:.4f} → {post_calmar:.4f} "
                    f"(+{delta:.4f})  [phase={state.phase}]")
            else:
                state.no_improve_rounds += 1
                icon = "⚠️" if state.no_improve_rounds >= args.convergence_rounds else "—"
                log(f"  {icon} {state.tf}: {pre_calmars[state.tf]:.4f} → {post_calmar:.4f} "
                    f"({delta:+.4f})  no_improve={state.no_improve_rounds}/{args.convergence_rounds} "
                    f"[phase={state.phase}]")

                if state.no_improve_rounds >= args.convergence_rounds:
                    if state.phase == "all":
                        state.phase           = "bull"
                        state.no_improve_rounds = 0
                        state.next_seed       = next_seed_start(asset, state.tf, f"{args.tag}_bull")
                        bull_n                = len(existing_seeds(asset, state.tf, f"{args.tag}_bull"))
                        log(f"  🔄 {state.tf}: plateau in all-regime → switching to bull "
                            f"(existing bull seeds: {bull_n}, next: seed{state.next_seed})")
                    else:
                        state.phase = "retired"
                        log(f"  🚫 {state.tf}: RETIRED (converged in all + bull phases)")

        if promoted_this_round and args.presets and not args.dry_run:
            log("Regenerating Pine presets...")
            ok = run_presets(args.dry_run)
            log(f"  {'✅' if ok else '❌'} Pine presets")

    # ── Final summary ─────────────────────────────────────────────────────────
    elapsed = int(time.time() - t_start)
    log(f"\n{'═'*60}")
    log(f"ADAPTIVE LOOP COMPLETE — {round_num} rounds in {elapsed//60}m{elapsed%60}s")
    log(f"{'═'*60}")

    tee("\nCalmar improvement summary:")
    print_comparison_table(asset, states)

    # mlp_results table — always shown at end (captures IS + OOS metrics)
    if not args.dry_run:
        log("Running mlp_results_table...")
        table_output = run_results_table(asset)
        tee()
        for line in table_output.splitlines():
            tee(line)
        tee()

    # ── Notify ────────────────────────────────────────────────────────────────
    if args.notify and not args.dry_run:
        improved_tfs = [
            f"{tf} {s.initial_calmar:.3f}→{s.last_calmar:.3f} ({s.phase_won})"
            for tf, s in states.items()
            if s.last_calmar - s.initial_calmar > 0.001
        ]
        msg_lines = [
            f"Adaptive loop done — {round_num} rounds, {elapsed//60}m",
        ]
        if improved_tfs:
            msg_lines.append("Promoted:")
            msg_lines.extend(f"  {line}" for line in improved_tfs)
        else:
            msg_lines.append("No promotions — all TFs at ceiling.")
        send_notify("TradingBot25 MLP Adaptive", "\n".join(msg_lines))
        log("ntfy sent")

    if _log_file is not None:
        log_path_display = Path(_log_file.name).relative_to(REPO)
        _log_file.close()
        _log_file = None
        tee(f"\nFull log: {log_path_display}")

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


if __name__ == "__main__":
    main()