"""
Two-phase trainer for the MLP activation-score strategy (strategy_mlp_scores).

Phase 1 — supervised pretrain (PyTorch):
    Target: y_t = tanh( scale * (k-bar forward log return / k) / rolling_vol_t )
    where rolling_vol is past-only (shifted) so there is no lookahead leakage.
    Temporal split: train on bars up to the last WFO fold's IS end, early-stop
    on the last fold's OOS span. (Full per-epoch expanding-window CV across all
    folds was deliberately simplified away — Phase 2 does the real fold-averaged
    model selection on the actual backtest metric.)

Phase 2 — CMA-ES fine-tune (cma):
    Parameter vector = flattened MLP weights + [entry, exit, exit_conf, trail]
    thresholds. Fitness = -(mean per-fold OOS Calmar across config.WFO_FOLDS,
    invalid folds scoring 0, >= WFO_MIN_VALID_FOLDS valid folds required)
    + lambda * mean squared deviation from the pretrained weights (L2 anchor).
    `--fold-objective robust` keeps the same WFO-only training boundary but
    subtracts penalties for high fold drawdown, zero P&L/DD, low fold trade
    coverage, negative folds, and threshold fragility.
    Per-fold Calmar matches the existing WFO rescore convention in
    auto_optimize_loop.py (composite_score's 1000% P&L gate is meaningless on
    ~1-year folds, so Calmar is the per-fold metric).

Usage:
    python3 tools/train_mlp.py --data data/COINBASE_BTCUSD-6H.csv \
        [--hidden 16 8] [--target-k 10] [--phase1-epochs 300] \
        [--es-generations 300] [--es-popsize 32] [--es-sigma 0.05] [--l2 1e-3] \
        [--fold-objective mean|robust] [--es-workers 0] \
        [--seed 42] [--smoke] [--skip-phase2] [--out PATH]

Output: strategies/params/mlp/mlp_weights_{ASSET}_{TF}.json
"""
import argparse
import concurrent.futures
import hashlib
import json
import multiprocessing as mp
import os
import re
import sys
import time
from datetime import datetime, timezone

import numpy as np
import pandas as pd

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
import config  # noqa: E402
from strategies.strategy_mlp_scores import (  # noqa: E402
    ACTIVATION_NAME,
    FEATURE_COLS,
    _score_to_signals,
    calculate_metrics,
    mlp_forward,
    save_mlp_artifact,
)
from strategies.strategy_activation_scores import _prepare_features  # noqa: E402
from strategies.mlp_feature_groups import (  # noqa: E402
    grouped_first_layer_mask,
    grouped_structure_metadata,
    hybrid_grouped_first_layer_mask,
    hybrid_grouped_structure_metadata,
    random_sparse_first_layer_mask,
    random_sparse_structure_metadata,
)
from strategies.mlp_temporal_features import inject_temporal_features  # noqa: E402
from tools.worker_utils import DEFAULT_WORKER_FRACTION, resolve_worker_count  # noqa: E402

VOL_WINDOW = 60          # bars of past 1-bar log returns for the vol normaliser
TARGET_SCALE = 2.0       # spreads targets across tanh's dynamic range
THRESHOLD_SCALE = 100.0  # thresholds enter the ES vector divided by this
ROBUST_DRAWDOWN_TARGET = 50.0
ROBUST_FRAGILITY_TIGHT_MARGIN = 0.05
ROBUST_FRAGILITY_WIDE_MARGIN = 1.0
_FITNESS_CONTEXT = None


# ---------------------------------------------------------------------------
# Data / target construction
# ---------------------------------------------------------------------------

def derive_asset_tf(data_path):
    m = re.match(r"([A-Z0-9_!]+)-([0-9]+[HDW])\.csv", os.path.basename(data_path))
    if not m:
        raise ValueError(f"Cannot derive ASSET/TF from filename: {data_path}")
    return m.group(1), m.group(2)


def load_data(data_path):
    df = pd.read_csv(data_path)
    df.columns = df.columns.str.lower().str.strip()
    df["time"] = pd.to_datetime(df["time"], utc=True).dt.tz_localize(None)
    mask = (df["time"] >= config.TRAIN_START) & (df["time"] <= config.TRAIN_END)
    return df.loc[mask].reset_index(drop=True)


def build_target(df, k):
    """tanh-squashed vol-normalised k-bar forward log return. Returns (y, valid)."""
    log_close = np.log(df["close"].values.astype(np.float64))
    ret1 = pd.Series(np.diff(log_close, prepend=log_close[0]))
    # Past-only vol: window ends at t-1 (shifted) -> no lookahead.
    vol = ret1.rolling(VOL_WINDOW).std().shift(1).values
    fwd = np.full(len(df), np.nan)
    fwd[:-k] = (log_close[k:] - log_close[:-k]) / k
    y = np.tanh(TARGET_SCALE * fwd / (vol + 1e-6))
    valid = ~np.isnan(y)
    valid &= (df["time"] >= config.SCORE_START).values
    return y, valid


# ---------------------------------------------------------------------------
# Phase 1 — supervised pretrain
# ---------------------------------------------------------------------------

def pretrain(X, y, valid, times, hidden, epochs, seed, val_start, val_end, log,
             first_layer_mask=None):
    import torch

    device = ("cuda" if torch.cuda.is_available()
              else "mps" if torch.backends.mps.is_available()
              else "cpu")
    torch.manual_seed(seed)
    np.random.seed(seed)

    arch = [X.shape[1]] + list(hidden) + [1]
    layers_t = []
    modules = []
    for i in range(len(arch) - 1):
        lin = torch.nn.Linear(arch[i], arch[i + 1])
        modules += [lin, torch.nn.Tanh()]
        layers_t.append(lin)
    model = torch.nn.Sequential(*modules).to(device)
    mask_t = None
    if first_layer_mask is not None:
        mask_t = torch.tensor(first_layer_mask, dtype=torch.bool, device=device)
        with torch.no_grad():
            layers_t[0].weight.masked_fill_(~mask_t, 0.0)

    val_mask = valid & (times >= val_start).values & (times <= val_end).values
    train_mask = valid & (times < val_start).values
    log(f"phase1: device={device} arch={arch} train_bars={train_mask.sum()} "
        f"val_bars={val_mask.sum()}")

    def tens(mask_arr):
        return (torch.tensor(X[mask_arr], dtype=torch.float32, device=device),
                torch.tensor(y[mask_arr], dtype=torch.float32, device=device).unsqueeze(1))

    Xtr, ytr = tens(train_mask)
    Xva, yva = tens(val_mask)

    opt = torch.optim.Adam(model.parameters(), lr=1e-2, weight_decay=1e-4)
    loss_fn = torch.nn.MSELoss()
    best_val, best_state, patience, since_best = np.inf, None, 20, 0

    for epoch in range(epochs):
        model.train()
        opt.zero_grad()
        loss = loss_fn(model(Xtr), ytr)
        loss.backward()
        opt.step()
        if mask_t is not None:
            with torch.no_grad():
                layers_t[0].weight.masked_fill_(~mask_t, 0.0)
        train_loss = loss.item()
        model.eval()
        with torch.no_grad():
            val_loss = float(loss_fn(model(Xva), yva)) if len(Xva) else train_loss
        if val_loss < best_val - 1e-6:
            best_val, since_best = val_loss, 0
            best_state = [(lin.weight.detach().cpu().double().numpy().copy(),
                           lin.bias.detach().cpu().double().numpy().copy())
                          for lin in layers_t]
        else:
            since_best += 1
            if since_best >= patience:
                log(f"phase1: early stop at epoch {epoch} (best val MSE {best_val:.5f})")
                break
        if epoch % 50 == 0:
            log(f"phase1: epoch {epoch} train={train_loss:.5f} val={val_loss:.5f}")

    return best_state, {"device": device, "arch": arch, "epochs_run": epoch + 1,
                        "best_val_mse": best_val,
                        "train_bars": int(train_mask.sum()),
                        "val_bars": int(val_mask.sum())}


# ---------------------------------------------------------------------------
# Phase 2 — CMA-ES on fold-mean Calmar
# ---------------------------------------------------------------------------

def flatten_layers(layers, weight_masks=None):
    """Flatten trainable weights, omitting structurally blocked connections."""
    if weight_masks is None:
        weight_masks = [np.ones_like(W, dtype=bool) for W, _ in layers]
    return np.concatenate([
        np.concatenate([np.asarray(W)[mask], np.asarray(b).ravel()])
        for (W, b), mask in zip(layers, weight_masks)
    ])


def unflatten_layers(theta, arch, weight_masks=None):
    layers, i = [], 0
    for li in range(len(arch) - 1):
        n_out, n_in = arch[li + 1], arch[li]
        mask = (np.ones((n_out, n_in), dtype=bool) if weight_masks is None
                else np.asarray(weight_masks[li], dtype=bool))
        if mask.shape != (n_out, n_in):
            raise ValueError(f"Weight mask {li} has shape {mask.shape}; expected {(n_out, n_in)}")
        W = np.zeros((n_out, n_in), dtype=np.float64)
        n_active = int(mask.sum())
        W[mask] = theta[i:i + n_active]; i += n_active
        b = theta[i:i + n_out]; i += n_out
        layers.append((W, b))
    if i != len(theta):
        raise ValueError("Parameter vector length does not match architecture and masks")
    return layers


def aggregate_fold_score(calmars, objective):
    if objective == "mean":
        return float(np.mean(calmars))
    if objective == "min":
        return float(np.min(calmars))
    if objective == "mean_min":
        return 0.7 * float(np.mean(calmars)) + 0.3 * float(np.min(calmars))
    raise ValueError(f"Unknown fold objective: {objective}")


def threshold_fragility_count(scores, thresholds, margin):
    thresholds = np.asarray(thresholds[:3], dtype=np.float64)
    if len(thresholds) == 0:
        return 0
    scores = np.asarray(scores, dtype=np.float64)
    dist = np.min(np.abs(scores[:, None] - thresholds[None, :]), axis=1)
    return int(np.sum(dist < margin))


def robust_objective_score(calmars, fold_metrics, scores, thresholds, min_fold_trades):
    base = aggregate_fold_score(calmars, "mean_min")
    trade_target = max(min_fold_trades + 5, int(np.ceil(min_fold_trades * 1.5)))
    n_folds = max(1, len(fold_metrics))

    drawdown_penalty = 0.0
    pnl_dd_penalty = 0.0
    trade_penalty = 0.0
    neg_fold_penalty = 0.0
    for metrics in fold_metrics:
        trades = int(metrics.get("Total Trades", 0))
        drawdown = abs(float(metrics.get("Max Drawdown %", 0.0)))
        pnl_dd = float(metrics.get("P&L/DD Ratio", 0.0))
        calmar = float(metrics.get("Calmar Ratio", -10.0))

        drawdown_penalty += max(0.0, drawdown - ROBUST_DRAWDOWN_TARGET) / ROBUST_DRAWDOWN_TARGET
        if pnl_dd <= 0.0:
            pnl_dd_penalty += 1.0
        trade_penalty += max(0.0, trade_target - trades) / trade_target
        neg_fold_penalty += max(0.0, min(1.0, -calmar))

    drawdown_penalty /= n_folds
    pnl_dd_penalty /= n_folds
    trade_penalty /= n_folds
    neg_fold_penalty /= n_folds

    fragile_tight = threshold_fragility_count(scores, thresholds, ROBUST_FRAGILITY_TIGHT_MARGIN)
    fragile_wide = threshold_fragility_count(scores, thresholds, ROBUST_FRAGILITY_WIDE_MARGIN)
    fragility_penalty = (0.05 * fragile_tight) + min(1.0, fragile_wide / 500.0)

    penalties = {
        "drawdown_penalty": drawdown_penalty,
        "pnl_dd_penalty": pnl_dd_penalty,
        "trade_penalty": trade_penalty,
        "negative_fold_penalty": neg_fold_penalty,
        "fragility_penalty": fragility_penalty,
        "fragile_0_05": fragile_tight,
        "fragile_1_0": fragile_wide,
        "trade_target": trade_target,
    }
    penalty = (
        0.8 * drawdown_penalty
        + 0.8 * pnl_dd_penalty
        + 0.25 * trade_penalty
        + 0.30 * neg_fold_penalty
        + fragility_penalty
    )
    return base - penalty, penalties


def make_fitness_context(df, X, arch, theta0_w, folds, min_fold_trades, min_valid_folds, l2,
                         fold_objective="mean", weight_masks=None):
    return {
        "base_df": df[["time", "close"]].copy(),
        "times": df["time"].copy(),
        "close_arr": df["close"].values.astype(np.float64),
        "high_arr": df["high"].values.astype(np.float64),
        "stoch_peak_arr": X[:, 9],
        "X": X,
        "arch": list(arch),
        "theta0_w": np.asarray(theta0_w, dtype=np.float64),
        "folds": list(folds),
        "min_fold_trades": min_fold_trades,
        "min_valid_folds": min_valid_folds,
        "l2": l2,
        "fold_objective": fold_objective,
        "weight_masks": weight_masks,
    }


def evaluate_theta(theta, context):
    theta = np.asarray(theta, dtype=np.float64)
    theta0_w = context["theta0_w"]
    n_w = len(theta0_w)
    layers = unflatten_layers(theta[:n_w], context["arch"], context["weight_masks"])
    thresholds = np.clip(
        theta[n_w:] * THRESHOLD_SCALE,
        [-900.0, -900.0, -900.0, 0.0],
        [900.0, 900.0, 900.0, 30.0],
    )

    scores = mlp_forward(context["X"], layers)
    d = context["base_df"].copy()
    d["activation_score"] = scores
    d["high"] = context["high_arr"]
    params = {
        "i_long_entry_activation_threshold": thresholds[0],
        "i_long_exit_activation_threshold": thresholds[1],
        "i_long_exit_activation_confirmation_threshold": thresholds[2],
        "i_use_long_exit_confirmation": 1.0,
        "i_use_long_entry_confirmation": False,
        "i_trailing_stop_threshold": max(0.0, thresholds[3]),
    }
    d = _score_to_signals(
        d,
        params,
        context["stoch_peak_arr"],
        context["close_arr"],
        context["high_arr"],
    )

    calmars, fold_metrics, valid = [], [], 0
    times = context["times"]
    min_fold_trades = context["min_fold_trades"]
    for _, oos_start, oos_end in context["folds"]:
        sl = d.loc[(times >= oos_start) & (times <= oos_end)]
        if len(sl) < 10:
            fold_metrics.append({"Total Trades": 0, "Calmar Ratio": -10.0,
                                 "P&L/DD Ratio": 0.0, "Max Drawdown %": 0.0})
            calmars.append(0.0)
            continue
        metrics = calculate_metrics(sl, min_trades=min_fold_trades)
        fold_metrics.append(metrics)
        if metrics["Calmar Ratio"] <= -10.0 or metrics["Total Trades"] < min_fold_trades:
            calmars.append(0.0)
        else:
            calmars.append(metrics["Calmar Ratio"])
            valid += 1

    diagnostics = {"valid_folds": valid}
    fold_objective = context["fold_objective"]
    if fold_objective == "robust":
        score, penalties = robust_objective_score(
            calmars,
            fold_metrics,
            scores,
            thresholds,
            min_fold_trades,
        )
        diagnostics.update(penalties)
        diagnostics["base_fold_score"] = aggregate_fold_score(calmars, "mean_min")
    elif valid < context["min_valid_folds"]:
        score = 0.0
    else:
        score = aggregate_fold_score(calmars, fold_objective)

    anchor = context["l2"] * float(np.mean((theta[:n_w] - theta0_w) ** 2))
    return -(score) + anchor, calmars, thresholds, diagnostics


def init_fitness_worker(context):
    global _FITNESS_CONTEXT
    _FITNESS_CONTEXT = context


def evaluate_theta_worker(theta):
    return evaluate_theta(theta, _FITNESS_CONTEXT)


def make_fitness(df, X, arch, theta0_w, folds, min_fold_trades, min_valid_folds, l2,
                 fold_objective="mean", weight_masks=None):
    context = make_fitness_context(
        df, X, arch, theta0_w, folds, min_fold_trades, min_valid_folds, l2,
        fold_objective, weight_masks,
    )

    def fitness(theta):
        fit, calmars, thresholds, diagnostics = evaluate_theta(theta, context)
        fitness.last_diagnostics = diagnostics
        return fit, calmars, thresholds

    fitness.context = context
    fitness.last_diagnostics = {}
    return fitness


def run_cma(fitness, theta0, sigma0, popsize, generations, seed, log, workers=1):
    import cma

    es = cma.CMAEvolutionStrategy(theta0, sigma0,
                                  {"popsize": popsize, "seed": seed, "verbose": -9})
    workers = max(1, int(workers))
    best = {"fit": np.inf, "theta": theta0, "calmars": None, "thr": None,
            "diagnostics": {}, "workers": workers}

    def observe_solution(solution, result, fits):
        f, calmars, thr, diagnostics = result
        fits.append(f)
        nonlocal best
        if f < best["fit"]:
            best = {
                "fit": f,
                "theta": np.asarray(solution).copy(),
                "calmars": calmars,
                "thr": thr,
                "diagnostics": dict(diagnostics),
                "workers": workers,
            }

    def run_generation(gen, evaluator):
        solutions = [np.asarray(s, dtype=np.float64) for s in es.ask()]
        results = evaluator(solutions)
        fits = []
        for solution, result in zip(solutions, results):
            observe_solution(solution, result, fits)
        es.tell(solutions, fits)
        if gen % 10 == 0 or gen == generations - 1:
            log(f"phase2: gen {gen}/{generations} best_fitness={best['fit']:.4f} "
                f"fold_calmars={[round(c, 2) for c in (best['calmars'] or [])]}")

    if workers == 1:
        log("phase2: evaluating CMA-ES population sequentially")

        def sequential_evaluator(solutions):
            out = []
            for solution in solutions:
                f, calmars, thr = fitness(solution)
                out.append((f, calmars, thr,
                            dict(getattr(fitness, "last_diagnostics", {}))))
            return out

        for gen in range(generations):
            run_generation(gen, sequential_evaluator)
            if es.stop():
                log(f"phase2: CMA-ES converged at gen {gen}")
                break
        return best

    context = getattr(fitness, "context", None)
    if context is None:
        raise ValueError("Parallel CMA-ES requires a fitness context")

    log(f"phase2: evaluating CMA-ES population with {workers} worker processes")
    mp_context = mp.get_context("spawn")
    try:
        with concurrent.futures.ProcessPoolExecutor(
            max_workers=workers,
            mp_context=mp_context,
            initializer=init_fitness_worker,
            initargs=(context,),
        ) as pool:
            for gen in range(generations):
                run_generation(gen, lambda solutions: list(pool.map(evaluate_theta_worker, solutions)))
                if es.stop():
                    log(f"phase2: CMA-ES converged at gen {gen}")
                    break
    except (OSError, PermissionError) as exc:
        log(f"phase2: worker pool unavailable ({exc}); falling back to sequential evaluation")
        return run_cma(fitness, theta0, sigma0, popsize, generations, seed, log, workers=1)
    return best


# ---------------------------------------------------------------------------
# Approach B — position feature injection
# ---------------------------------------------------------------------------

def _inject_position_features(df, asset, tf, args, log):
    """Compute in_long_position + bars_held_norm from the current winner policy.

    One-cycle circularity (documented in roadmap): the position state used for
    training comes from the OLD model's decisions, not the new model.  This is
    acceptable for the first retrain; a second retrain with the new policy's
    state would converge.

    Mutates df in-place, adding columns 'in_long_position' and 'bars_held_norm'.
    Returns the extended feature_cols list.
    """
    from strategies.strategy_mlp_scores import generate_signals as _gen_signals

    params_path = (
        args.approach_b_params
        or f"results/winners/optimization_winner_strategy_mlp_scores_{asset}_{tf}.csv"
    )
    log(f"approach-b: loading winner params from {params_path}")
    params_df = pd.read_csv(params_path)
    params = params_df.iloc[0].to_dict()

    log("approach-b: running generate_signals with old policy to get position state ...")
    # Use the same IS-filtered df — position always starts flat at TRAIN_START.
    signals = _gen_signals(df.copy(), **params)

    position = signals["position"].values  # int array: 1 if in long, 0 otherwise
    in_long = position.astype(np.float64)

    bars_held = np.zeros(len(position), dtype=np.float64)
    bars_in = 0
    for i in range(len(position)):
        if position[i] > 0:
            bars_held[i] = min(bars_in / 100.0, 1.0)
            bars_in += 1
        else:
            bars_in = 0

    df["in_long_position"] = in_long
    df["bars_held_norm"] = bars_held

    n_in_pos = int(in_long.sum())
    max_held = float(bars_held.max())
    log(f"approach-b: position features injected — {n_in_pos}/{len(df)} bars in-position, "
        f"max bars_held_norm={max_held:.3f} (={max_held*100:.0f} bars)")

    return FEATURE_COLS + ["in_long_position", "bars_held_norm"]


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    p = argparse.ArgumentParser(description="Two-phase MLP trainer")
    p.add_argument("--data", required=True)
    p.add_argument("--hidden", type=int, nargs="+", default=[16, 8])
    p.add_argument("--target-k", type=int, default=10)
    p.add_argument("--phase1-epochs", type=int, default=300)
    p.add_argument("--es-generations", type=int, default=300)
    p.add_argument("--es-popsize", type=int, default=32)
    p.add_argument("--es-sigma", type=float, default=0.05)
    p.add_argument("--l2", type=float, default=1e-3)
    p.add_argument("--fold-objective", choices=["mean", "min", "mean_min", "robust"],
                   default="mean",
                   help="Phase-2 fold aggregation: mean matches legacy behavior; "
                        "min emphasizes worst fold; mean_min blends 70%% mean / 30%% min; "
                        "robust subtracts drawdown, P&L/DD, trade-count, negative-fold, "
                        "and threshold-fragility penalties.")
    p.add_argument("--es-workers", type=int, default=0,
                   help="Worker processes for parallel CMA-ES population evaluation. "
                        "Use 0 for auto, currently 50%% of logical CPUs capped by popsize; "
                        "use 1 for sequential evaluation.")
    p.add_argument("--worker-fraction", type=float, default=DEFAULT_WORKER_FRACTION,
                   help="Logical CPU fraction used when --es-workers=0.")
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--smoke", action="store_true",
                   help="Tiny run on recent bars for plumbing validation")
    p.add_argument("--skip-phase2", action="store_true")
    p.add_argument("--out", help="Artifact path override")
    p.add_argument("--asset", help="Asset override (e.g. COINBASE_BTCUSD) when filename format doesn't match")
    p.add_argument("--tf",    help="Timeframe override (e.g. 4H) when filename format doesn't match")
    p.add_argument("--approach-b", action="store_true",
                   help="Probe: inject in_long_position + bars_held_norm (cols 50-51) computed "
                        "from the current winner artifact's position state (one-cycle circularity "
                        "per roadmap). Output artifact will have 52 features.")
    p.add_argument("--approach-b-params",
                        help="Winner CSV to derive position state from (default: auto-detect from asset/tf)")
    p.add_argument("--temporal-features", action="store_true",
                   help="Experiment: append causal 3- and 12-bar EMA companions for ten fast signals.")
    p.add_argument("--input-structure", choices=["dense", "grouped", "hybrid_grouped", "random_sparse"], default="dense",
                   help="First-layer connectivity: dense is the production baseline; grouped "
                   "routes each of the 16 first-layer units to one logical signal family; "
                        "hybrid_grouped keeps eight family specialists plus eight dense rows; "
                        "random_sparse is the matched seeded control.")
    p.add_argument("--mask-seed", type=int,
                   help="Required with random_sparse; seed for the auditable first-layer mask.")
    p.add_argument("--regime", choices=["bull", "bear", "sideways", "all"], default="all",
                   help="Filter Phase-1 supervised training to MVRV regime bars only "
                        "(requires mvrv_regime column). Phase-2 CMA-ES always uses all IS bars.")
    args = p.parse_args()

    if args.input_structure == "random_sparse" and args.mask_seed is None:
        p.error("--mask-seed is required with --input-structure random_sparse")
    if args.input_structure == "hybrid_grouped" and args.out is None:
        p.error("--out is required with --input-structure hybrid_grouped to protect live artifacts")
    if args.temporal_features and args.out is None:
        p.error("--temporal-features requires --out to protect live artifacts")

    t0 = time.time()

    def log(msg):
        print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)

    if args.asset and args.tf:
        asset, tf = args.asset, args.tf
    else:
        asset, tf = derive_asset_tf(args.data)
    df = load_data(args.data)

    if args.smoke:
        df = df.tail(800).reset_index(drop=True)
        if args.input_structure == "dense":
            args.hidden = [8, 4]
        args.phase1_epochs = min(args.phase1_epochs, 20)
        args.es_generations = min(args.es_generations, 5)
        args.es_popsize = min(args.es_popsize, 8)
        # Single pseudo-fold over the last 25% of the smoke window.
        t75 = df["time"].iloc[int(len(df) * 0.75)]
        folds = [(None, str(t75.date()), str(df["time"].iloc[-1].date()))]
        min_fold_trades, min_valid_folds = 1, 1
        val_start = t75
    else:
        folds = config.WFO_FOLDS
        min_fold_trades = config.WFO_MIN_OOS_TRADES
        min_valid_folds = config.WFO_MIN_VALID_FOLDS
        val_start = pd.Timestamp(folds[-1][1])  # last fold's OOS span = val set
    args.es_workers = resolve_worker_count(args.es_workers, args.es_popsize, args.worker_fraction)

    val_end = df["time"].iloc[-1]
    log(f"asset={asset} tf={tf} bars={len(df)} window="
        f"{df['time'].iloc[0].date()}..{df['time'].iloc[-1].date()} hidden={args.hidden}")

    # --- Approach B: inject position-awareness features (probe branch) ---
    if args.approach_b:
        feature_cols = _inject_position_features(df, asset, tf, args, log)
    else:
        feature_cols = FEATURE_COLS
    if args.temporal_features:
        feature_cols = [*feature_cols, *inject_temporal_features(df)]
        log(f"temporal-features: appended {len(feature_cols) - len(FEATURE_COLS)} causal EMA companions")

    first_layer_mask = None
    if args.input_structure == "grouped":
        first_layer_mask = grouped_first_layer_mask(feature_cols, args.hidden[0])
        log("input-structure=grouped: first-layer units are restricted to logical signal families")
    elif args.input_structure == "hybrid_grouped":
        first_layer_mask = hybrid_grouped_first_layer_mask(feature_cols, args.hidden[0])
        log("input-structure=hybrid_grouped: eight family specialists plus eight dense rows "
            f"(active_edges={int(first_layer_mask.sum())})")
    elif args.input_structure == "random_sparse":
        first_layer_mask = random_sparse_first_layer_mask(
            feature_cols, args.hidden[0], args.mask_seed,
        )
        log(f"input-structure=random_sparse: seeded mask={args.mask_seed} "
            f"active_edges={int(first_layer_mask.sum())}")

    X = _prepare_features(df, feature_cols)
    y, valid = build_target(df, args.target_k)

    # --- Regime filter (Phase 1 only) ---
    _REGIME_INT = {"bull": 1, "bear": -1, "sideways": 0}
    if args.regime != "all":
        regime_int = _REGIME_INT[args.regime]
        if "mvrv_regime" in df.columns:
            regime_mask = df["mvrv_regime"].values == regime_int
            before = int(valid.sum())
            valid = valid & regime_mask
            pct = int(valid.sum()) / max(1, before) * 100
            log(f"regime={args.regime}: Phase-1 training bars {before} → {int(valid.sum())} ({pct:.0f}%)")
        else:
            log(f"WARNING: --regime={args.regime} requested but mvrv_regime not in data; ignoring")

    # --- Phase 1 ---
    layers, p1_meta = pretrain(X, y, valid, df["time"], args.hidden,
                               args.phase1_epochs, args.seed, val_start, val_end, log,
                               first_layer_mask)
    arch = p1_meta["arch"]

    training_meta = {
        "phase1": {**p1_meta, "target_k": args.target_k,
                   "vol_window": VOL_WINDOW, "target_scale": TARGET_SCALE,
                   "val_span": [str(val_start), str(val_end)]},
        "data_file": args.data,
        "data_sha256": hashlib.sha256(open(args.data, "rb").read()).hexdigest(),
        "train_window": [str(df["time"].iloc[0].date()), str(df["time"].iloc[-1].date())],
        "seed": args.seed,
        "regime_filter": args.regime,
        "input_structure": (
            grouped_structure_metadata(feature_cols, args.hidden[0])
            if args.input_structure == "grouped"
            else hybrid_grouped_structure_metadata(feature_cols, args.hidden[0])
            if args.input_structure == "hybrid_grouped"
            else random_sparse_structure_metadata(feature_cols, args.hidden[0], args.mask_seed)
            if args.input_structure == "random_sparse"
            else {"name": "dense"}
        ),
        "created": datetime.now(timezone.utc).isoformat(timespec="seconds"),
    }
    # Seed thresholds from the pretrained net's actual score distribution —
    # fixed defaults (100/140) can sit entirely outside the score range, giving
    # zero trades on every fold and a flat fitness landscape CMA-ES can't escape.
    scores0 = mlp_forward(X, layers)
    thresholds = np.array([
        np.percentile(scores0, 60),   # entry: crossed downward often enough to trade
        np.percentile(scores0, 75),   # exit
        np.percentile(scores0, 30),   # exit confirmation
        0.0,                          # trailing stop off
    ])
    log(f"phase2: score range [{scores0.min():.1f}, {scores0.max():.1f}], "
        f"seed thresholds entry={thresholds[0]:.1f} exit={thresholds[1]:.1f} "
        f"conf={thresholds[2]:.1f}")

    # --- Phase 2 ---
    if not args.skip_phase2:
        weight_masks = None
        if first_layer_mask is not None:
            weight_masks = [first_layer_mask] + [
                np.ones((arch[i + 1], arch[i]), dtype=bool) for i in range(1, len(arch) - 1)
            ]
        theta0_w = flatten_layers(layers, weight_masks)
        theta0 = np.concatenate([theta0_w, thresholds / THRESHOLD_SCALE])
        fitness = make_fitness(df, X, arch, theta0_w, folds,
                               min_fold_trades, min_valid_folds, args.l2,
                               args.fold_objective, weight_masks)
        f0, calmars0, _ = fitness(theta0)
        log(f"phase2: pretrain baseline fitness={f0:.4f} "
            f"fold_calmars={[round(c, 2) for c in calmars0]}")
        best = run_cma(fitness, theta0, args.es_sigma, args.es_popsize,
                       args.es_generations, args.seed, log, args.es_workers)
        n_w = len(theta0_w)
        layers = unflatten_layers(best["theta"][:n_w], arch, weight_masks)
        thresholds = best["thr"]
        training_meta["phase2"] = {
            "es_popsize": args.es_popsize, "es_sigma": args.es_sigma,
            "es_generations": args.es_generations, "l2_anchor": args.l2,
            "es_workers": best.get("workers", args.es_workers),
            "fold_objective": args.fold_objective,
            "best_fitness": best["fit"],
            "fold_calmars": best["calmars"],
            "fold_mean_calmar": float(np.mean(best["calmars"])),
            "fold_min_calmar": float(np.min(best["calmars"])),
        }
        if args.fold_objective == "robust":
            training_meta["phase2"]["robust_diagnostics"] = best.get("diagnostics", {})

    training_meta["recommended_thresholds"] = {
        "i_long_entry_activation_threshold": float(thresholds[0]),
        "i_long_exit_activation_threshold": float(thresholds[1]),
        "i_long_exit_activation_confirmation_threshold": float(thresholds[2]),
        "i_trailing_stop_threshold": float(thresholds[3]),
        "i_use_long_exit_confirmation": 1.0,
        "i_use_long_entry_confirmation": False,
    }

    out = args.out or os.path.join("strategies", "params", "mlp",
                                   f"mlp_weights_{asset}_{tf}.json")
    save_mlp_artifact(out, layers, asset=asset, timeframe=tf,
                      training=training_meta, feature_cols=feature_cols)
    log(f"artifact written: {out}  ({time.time() - t0:.0f}s total)")
    log(f"recommended thresholds: {training_meta['recommended_thresholds']}")


if __name__ == "__main__":
    main()