#!/usr/bin/env python3
"""
100optimization.py — safe Hugging Face CRT strategy optimizer.

This script is designed for a Hugging Face Space. It:

1. Calls load_dataset("Mikecode123/volatility_100_index") with remote dataset
   code disabled.
2. Locates volatility_100_index.zip in the dataset snapshot.
3. Inspects the archive and extracts CSV files only. Scripts, executables,
   symlinks, traversal paths, and oversized archives are rejected.
4. Loads M1.csv, M5.csv, M15.csv, M30.csv, H1.csv, and H4.csv case-insensitively.
5. Implements a clearly documented Candle Range Theory (CRT) rule:
      - use a completed H1 or H4 reference candle;
      - wait for a sweep of its high/low;
      - require a close back inside the reference range;
      - require the next-bar confirmation through the sweep candle extreme;
      - enter at confirmation close;
      - target the opposite edge of the reference range and evaluate staged
        ATR-based TP/SL outcomes.
6. Uses a bounded deterministic candidate search. The supervisor model can
   select only from supplied candidate IDs; it cannot generate code or change
   the evaluation rules.
7. Saves every new best result and milestone result (20%, 30%, 80%) to the
   Space checkpoint directory, then creates a downloadable ZIP.

There is no guarantee that 80% accuracy is achievable. The target is a stopping
criterion, not a promise. Validation is used for optimization; the final test
period remains untouched until the end.

Supervisor model (pinned and approved):
    google/flan-t5-small
    revision=0fc9ddf78a1e988dac52e2dac162b0ede4fd74ab
    trust_remote_code=False

Expected Space requirements:
    datasets
    huggingface_hub
    transformers
    torch
    pandas
    numpy

Research sources for the CRT specification:
    https://innercircletrader.net/tutorials/candle-range-theory-crt/
    https://tradingwyckoff.com/en/crt/

The public CRT descriptions are educational retail sources, not peer-reviewed
proof of profitability. The implementation therefore reports validation and
test results separately and records coverage beside accuracy.
"""

from __future__ import annotations

import gc
import itertools
import json
import math
import os
import random
import re
import shutil
import stat
import time
import zipfile
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

import numpy as np
import pandas as pd

try:
    from datasets import load_dataset
except ImportError as exc:  # pragma: no cover - dependency supplied by the Space
    raise RuntimeError(
        "Install the Space requirements first: datasets, huggingface_hub, "
        "transformers, torch, pandas, numpy"
    ) from exc


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

DATASET_ID = os.environ.get("HF_DATASET_ID", "Mikecode123/volatility_100_index")
SUPERVISOR_MODEL_ID = "google/flan-t5-small"
SUPERVISOR_REVISION = "0fc9ddf78a1e988dac52e2dac162b0ede4fd74ab"

TF_MINUTES = {
    "M1": 1,
    "M5": 5,
    "M15": 15,
    "M30": 30,
    "H1": 60,
    "H4": 240,
}

CLASS_NAMES = {
    0: "NO_TRADE",
    1: "SL",
    2: "BE",
    3: "TP1",
    4: "TP2",
    5: "TP3",
}
LOCAL_R_VALUE = np.array([0.0, -1.0, 0.0, 2.0, 4.0, 6.0], dtype=np.float32)
LOCAL_RANK = np.array([1, 0, 2, 3, 4, 5], dtype=np.int8)


@dataclass
class Config:
    dataset_id: str = DATASET_ID
    seed: int = int(os.environ.get("OPTIMIZER_SEED", "42"))
    target_accuracy: float = float(os.environ.get("TARGET_ACCURACY", "0.80"))
    max_trials: int = int(os.environ.get("MAX_TRIALS", "30"))
    min_trades: int = int(os.environ.get("MIN_TRADES", "100"))
    min_coverage: float = float(os.environ.get("MIN_COVERAGE", "0.001"))
    max_rows: int = int(os.environ.get("MAX_ROWS", "0"))
    reference_timeframe: str = os.environ.get("CRT_REFERENCE_TF", "H1").upper()
    confirm_bars: int = int(os.environ.get("CRT_CONFIRM_BARS", "1"))
    atr_period: int = int(os.environ.get("ATR_PERIOD", "14"))
    default_sl_atr_mult: float = float(os.environ.get("SL_ATR_MULT", "1.0"))
    default_horizon: int = int(os.environ.get("MAX_HORIZON", "60"))
    checkpoint_dir: Path = Path(
        os.environ.get("CHECKPOINT_DIR", "/data/100optimization_checkpoints")
    )
    data_dir: Path = Path(
        os.environ.get("DATA_EXTRACT_DIR", "/tmp/volatility_100_index_data")
    )


CFG = Config()


# ---------------------------------------------------------------------------
# Logging and filesystem helpers
# ---------------------------------------------------------------------------


def log(message: str, *args: Any, level: str = "INFO") -> None:
    if args:
        message = message.format(*args)
    stamp = datetime.now().strftime("%H:%M:%S")
    print(f"[{stamp}] [{level}] {message}", flush=True)


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


def ensure_checkpoint_dir() -> Path:
    try:
        CFG.checkpoint_dir.mkdir(parents=True, exist_ok=True)
        return CFG.checkpoint_dir
    except PermissionError:
        fallback = Path("./100optimization_checkpoints")
        fallback.mkdir(parents=True, exist_ok=True)
        log(
            "Cannot write to {}; using {} instead",
            CFG.checkpoint_dir,
            fallback,
            level="WARN",
        )
        CFG.checkpoint_dir = fallback
        return fallback


# ---------------------------------------------------------------------------
# Hugging Face dataset discovery and safe archive handling
# ---------------------------------------------------------------------------


def _snapshot_dataset_files() -> Path:
    """Download only dataset files, never executable model or code files."""
    from huggingface_hub import snapshot_download

    return Path(
        snapshot_download(
            repo_id=CFG.dataset_id,
            repo_type="dataset",
            allow_patterns=["*.zip", "*.ZIP", "*.csv", "*.CSV"],
        )
    )


def locate_zip(snapshot_dir: Path) -> Path:
    candidates = sorted(
        p for p in snapshot_dir.rglob("*")
        if p.is_file() and p.suffix.lower() == ".zip"
    )
    if not candidates:
        raise FileNotFoundError(
            f"No ZIP file was found in the dataset snapshot {snapshot_dir}."
        )

    exact = [p for p in candidates if p.name.lower() == "volatility_100_index.zip"]
    if exact:
        return exact[0]
    if len(candidates) == 1:
        log("Using the only ZIP in the dataset snapshot: {}", candidates[0], level="WARN")
        return candidates[0]
    raise FileNotFoundError(
        "Multiple ZIP files were found and none was named "
        f"volatility_100_index.zip: {candidates}"
    )


def inspect_zip(zip_path: Path) -> list[zipfile.ZipInfo]:
    """Inspect the archive before extraction; allow data CSVs only."""
    max_entries = 2_000
    max_file_size = 25 * 1024 * 1024
    max_total_size = 250 * 1024 * 1024

    with zipfile.ZipFile(zip_path, "r") as archive:
        infos = archive.infolist()
        if len(infos) > max_entries:
            raise ValueError(f"Archive contains too many entries: {len(infos)}")

        total_size = 0
        files: list[zipfile.ZipInfo] = []
        log("Inspecting archive {}", zip_path)
        for info in infos:
            name = Path(info.filename.replace("\\", "/"))
            if name.is_absolute() or ".." in name.parts:
                raise ValueError(f"Unsafe archive path rejected: {info.filename}")

            file_mode = (info.external_attr >> 16) & 0o170000
            if file_mode == stat.S_IFLNK:
                raise ValueError(f"Symlink rejected: {info.filename}")

            if info.is_dir():
                log("  [directory] {}", info.filename)
                continue

            if name.suffix.lower() != ".csv":
                raise ValueError(
                    f"Non-CSV/script/configuration member rejected: {info.filename}. "
                    "This optimizer accepts a data-only archive."
                )
            if info.file_size > max_file_size:
                raise ValueError(f"Archive member is too large: {info.filename}")

            total_size += info.file_size
            if total_size > max_total_size:
                raise ValueError("Archive expanded size exceeds safety limit")

            log("  [csv] {} ({:,} bytes)", info.filename, info.file_size)
            files.append(info)

        if not files:
            raise ValueError("Archive contains no CSV files")
        return files


def safe_extract_zip(zip_path: Path, destination: Path) -> Path:
    infos = inspect_zip(zip_path)
    if destination.exists():
        shutil.rmtree(destination)
    destination.mkdir(parents=True, exist_ok=True)

    root = destination.resolve()
    with zipfile.ZipFile(zip_path, "r") as archive:
        for info in infos:
            target = (destination / info.filename).resolve()
            if target != root and root not in target.parents:
                raise ValueError(f"Extraction escaped destination: {info.filename}")
            target.parent.mkdir(parents=True, exist_ok=True)
            with archive.open(info, "r") as source, target.open("wb") as sink:
                shutil.copyfileobj(source, sink, length=1024 * 1024)
    return destination


def find_timeframe_file(directory: Path, timeframe: str) -> Optional[Path]:
    wanted = f"{timeframe}.csv".lower()
    for candidate in directory.iterdir():
        if candidate.is_file() and candidate.name.lower() == wanted:
            return candidate
    return None


def find_data_directory(root: Path) -> Path:
    for m1 in sorted(root.rglob("*.csv")):
        if m1.is_file() and m1.stem.lower() == "m1":
            required = ("M1", "M5", "M15", "M30", "H1", "H4")
            missing = [tf for tf in required if find_timeframe_file(m1.parent, tf) is None]
            if not missing:
                return m1.parent
            log("Ignoring {} because it is missing {}", m1.parent, missing, level="WARN")
    raise FileNotFoundError(
        f"Could not find a directory containing M1/M5/M15/M30/H1/H4 CSVs below {root}"
    )


def load_dataset_archive() -> Path:
    """Call load_dataset as requested, then use the ZIP snapshot safely."""
    try:
        loaded = load_dataset(CFG.dataset_id, trust_remote_code=False)
        if hasattr(loaded, "keys"):
            log("load_dataset opened splits: {}", list(loaded.keys()))
        else:
            log("load_dataset opened {} rows", len(loaded))
    except Exception as exc:
        # A repository containing only a ZIP may not have a Datasets builder.
        # This fallback downloads data files only and still inspects the ZIP
        # before extraction.
        log(
            "load_dataset could not parse the ZIP-only repository: {}. "
            "Using the data-file snapshot fallback.",
            exc,
            level="WARN",
        )

    snapshot = _snapshot_dataset_files()
    archive = locate_zip(snapshot)
    log("Dataset archive: {}", archive)
    return archive


# ---------------------------------------------------------------------------
# OHLC loading and CRT feature construction
# ---------------------------------------------------------------------------

COLUMN_ALIASES = {
    "timestamp": ["timestamp", "time", "date", "datetime", "open_time"],
    "open": ["open", "o"],
    "high": ["high", "h"],
    "low": ["low", "l"],
    "close": ["close", "c", "adj_close"],
    "volume": ["volume", "vol", "v", "tick_volume"],
}


def resolve_column(columns: list[str], aliases: list[str]) -> Optional[str]:
    mapping = {column.lower(): column for column in columns}
    for alias in aliases:
        if alias in mapping:
            return mapping[alias]
    return None


def load_ohlcv(path: Path) -> pd.DataFrame:
    header = pd.read_csv(path, nrows=0)
    columns = list(header.columns)
    resolved = {
        key: resolve_column(columns, aliases)
        for key, aliases in COLUMN_ALIASES.items()
    }
    required = ("timestamp", "open", "high", "low", "close")
    missing = [key for key in required if resolved[key] is None]
    if missing:
        raise ValueError(f"{path} is missing required columns: {missing}")

    usecols = [resolved[key] for key in required]
    if resolved["volume"] is not None:
        usecols.append(resolved["volume"])
    dtype = {
        resolved[key]: np.float32
        for key in ("open", "high", "low", "close")
    }
    if resolved["volume"] is not None:
        dtype[resolved["volume"]] = np.float32

    raw = pd.read_csv(
        path,
        usecols=list(dict.fromkeys(usecols)),
        dtype=dtype,
        parse_dates=[resolved["timestamp"]],
    )
    out = pd.DataFrame({
        "timestamp": raw[resolved["timestamp"]],
        "open": raw[resolved["open"]],
        "high": raw[resolved["high"]],
        "low": raw[resolved["low"]],
        "close": raw[resolved["close"]],
    })
    out["volume"] = (
        raw[resolved["volume"]]
        if resolved["volume"] is not None
        else np.float32(0.0)
    )
    out = (
        out.dropna(subset=["timestamp", "open", "high", "low", "close"])
        .sort_values("timestamp")
        .drop_duplicates("timestamp", keep="last")
        .reset_index(drop=True)
    )
    invalid = (
        (out["high"] < out["low"])
        | (out[["open", "high", "low", "close"]] <= 0).any(axis=1)
    )
    if invalid.any():
        log("{}: dropping {} invalid rows", path.name, int(invalid.sum()), level="WARN")
        out = out.loc[~invalid].reset_index(drop=True)
    return out


def compute_atr(df: pd.DataFrame, period: int) -> pd.Series:
    previous_close = df["close"].shift(1)
    true_range = pd.concat(
        [
            df["high"] - df["low"],
            (df["high"] - previous_close).abs(),
            (df["low"] - previous_close).abs(),
        ],
        axis=1,
    ).max(axis=1)
    return true_range.rolling(period, min_periods=period).mean()


def closed_reference_features(
    base: pd.DataFrame,
    reference: pd.DataFrame,
    timeframe: str,
) -> pd.DataFrame:
    """Attach only fully closed reference-candle levels to M1 rows."""
    ref = reference[["timestamp", "open", "high", "low", "close"]].copy()
    ref["timestamp"] = ref["timestamp"] + pd.Timedelta(
        minutes=TF_MINUTES[timeframe]
    )
    ref = ref.rename(columns={
        "open": "crt_open",
        "high": "crt_high",
        "low": "crt_low",
        "close": "crt_close",
    })
    return pd.merge_asof(
        base.sort_values("timestamp"),
        ref.sort_values("timestamp"),
        on="timestamp",
        direction="backward",
    )


# ---------------------------------------------------------------------------
# Staged outcomes used for CRT backtesting
# ---------------------------------------------------------------------------


def simulate_direction(
    close: np.ndarray,
    high: np.ndarray,
    low: np.ndarray,
    risk: np.ndarray,
    direction: int,
    horizon: int,
) -> tuple[np.ndarray, np.ndarray]:
    """Simulate SL -> BE -> TP1-lock -> TP3 for one direction."""
    n = len(close)
    sign = 1.0 if direction > 0 else -1.0
    entry = close
    stop = entry - sign * risk
    tp1 = entry + sign * risk * 2.0
    tp2 = entry + sign * risk * 4.0
    tp3 = entry + sign * risk * 6.0
    lock = tp1

    stage = np.zeros(n, dtype=np.int8)
    resolved = np.zeros(n, dtype=bool)
    outcome = np.zeros(n, dtype=np.int8)
    valid = np.isfinite(risk)

    for step in range(1, horizon + 1):
        owners = np.arange(n)
        future = owners + step
        active = (~resolved) & valid & (future < n)
        if not active.any():
            continue
        owners = owners[active]
        future = future[active]
        snapshot = stage[owners].copy()

        for stage_id, favorable, adverse in (
            (0, tp1, stop),
            (1, tp2, entry),
            (2, tp3, lock),
        ):
            mask = snapshot == stage_id
            if not mask.any():
                continue
            rows = owners[mask]
            hi = high[future[mask]]
            lo = low[future[mask]]
            if direction > 0:
                touched_favorable = hi >= favorable[rows]
                touched_adverse = lo <= adverse[rows]
            else:
                touched_favorable = lo <= favorable[rows]
                touched_adverse = hi >= adverse[rows]

            adverse_rows = rows[touched_adverse]
            favorable_rows = rows[touched_favorable & ~touched_adverse]
            if stage_id == 0:
                outcome[adverse_rows] = 1
                resolved[adverse_rows] = True
                stage[favorable_rows] = 1
            elif stage_id == 1:
                outcome[adverse_rows] = 2
                resolved[adverse_rows] = True
                stage[favorable_rows] = 2
            else:
                outcome[adverse_rows] = 3
                resolved[adverse_rows] = True
                final_rows = rows[touched_favorable & ~touched_adverse]
                outcome[final_rows] = 5
                resolved[final_rows] = True

    open_rows = (~resolved) & valid
    outcome[open_rows & (stage == 0)] = 0
    outcome[open_rows & (stage == 1)] = 3
    outcome[open_rows & (stage == 2)] = 4
    return outcome, LOCAL_R_VALUE[outcome]


def build_outcomes(
    m1: pd.DataFrame,
    atr_period: int,
    sl_atr_mult: float,
    horizon: int,
) -> dict[str, np.ndarray]:
    atr = compute_atr(m1, atr_period).to_numpy(dtype=np.float32)
    risk = np.float32(sl_atr_mult) * atr
    close = m1["close"].to_numpy(dtype=np.float32)
    high = m1["high"].to_numpy(dtype=np.float32)
    low = m1["low"].to_numpy(dtype=np.float32)

    long_outcome, long_r = simulate_direction(close, high, low, risk, 1, horizon)
    short_outcome, short_r = simulate_direction(close, high, low, risk, -1, horizon)
    long_rank = LOCAL_RANK[long_outcome]
    short_rank = LOCAL_RANK[short_outcome]
    choose_long = (long_rank > short_rank) | (
        (long_rank == short_rank) & (long_r >= short_r)
    )
    chosen_outcome = np.where(choose_long, long_outcome, short_outcome).astype(np.int8)
    chosen_direction = np.where(
        chosen_outcome == 0,
        0,
        np.where(choose_long, 1, -1),
    ).astype(np.int8)

    return {
        "long_outcome": long_outcome,
        "short_outcome": short_outcome,
        "long_r": long_r,
        "short_r": short_r,
        "chosen_outcome": chosen_outcome,
        "chosen_direction": chosen_direction,
    }


# ---------------------------------------------------------------------------
# CRT signals and metrics
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class Candidate:
    candidate_id: int
    reference_timeframe: str
    min_sweep_atr: float
    use_reference_bias: bool
    confirmation_window: int
    cooldown_bars: int
    sl_atr_mult: float
    horizon: int


def make_candidates(seed: int) -> list[Candidate]:
    rng = random.Random(seed)
    candidates: list[Candidate] = []
    candidate_id = 0
    for values in itertools.product(
        ("H1", "H4"),
        (0.00, 0.05, 0.10, 0.25),
        (False, True),
        (1, 2),
        (0, 5, 15),
        (0.75, 1.00, 1.50),
        (30, 60),
    ):
        candidates.append(Candidate(candidate_id=candidate_id, **dict(zip(
            (
                "reference_timeframe",
                "min_sweep_atr",
                "use_reference_bias",
                "confirmation_window",
                "cooldown_bars",
                "sl_atr_mult",
                "horizon",
            ),
            values,
        ))))
        candidate_id += 1
    rng.shuffle(candidates)
    return candidates


def generate_crt_signals(
    m1: pd.DataFrame,
    reference: pd.DataFrame,
    atr: pd.Series,
    candidate: Candidate,
) -> np.ndarray:
    base = m1[["timestamp", "open", "high", "low", "close"]].copy()
    merged = closed_reference_features(base, reference, candidate.reference_timeframe)
    atr_values = atr.to_numpy(dtype=np.float32)

    bullish_sweep = (
        (merged["low"].to_numpy() < merged["crt_low"].to_numpy())
        & (merged["close"].to_numpy() > merged["crt_low"].to_numpy())
        & (
            (merged["crt_low"].to_numpy() - merged["low"].to_numpy())
            >= np.float32(candidate.min_sweep_atr) * atr_values
        )
    )
    bearish_sweep = (
        (merged["high"].to_numpy() > merged["crt_high"].to_numpy())
        & (merged["close"].to_numpy() < merged["crt_high"].to_numpy())
        & (
            (merged["high"].to_numpy() - merged["crt_high"].to_numpy())
            >= np.float32(candidate.min_sweep_atr) * atr_values
        )
    )

    ref_open = merged["crt_open"].to_numpy()
    ref_close = merged["crt_close"].to_numpy()
    if candidate.use_reference_bias:
        bullish_sweep &= ref_close >= ref_open
        bearish_sweep &= ref_close <= ref_open

    highs = merged["high"].to_numpy()
    lows = merged["low"].to_numpy()
    closes = merged["close"].to_numpy()
    signals = np.zeros(len(merged), dtype=np.int8)

    # Confirmation is a close through the sweep candle's opposite extreme.
    for delay in range(1, candidate.confirmation_window + 1):
        prior_bull = np.zeros(len(merged), dtype=bool)
        prior_bear = np.zeros(len(merged), dtype=bool)
        if delay < len(merged):
            prior_bull[delay:] = bullish_sweep[:-delay]
            prior_bear[delay:] = bearish_sweep[:-delay]
            signals[delay:][prior_bull[delay:] & (closes[delay:] > highs[:-delay])] = 1
            signals[delay:][prior_bear[delay:] & (closes[delay:] < lows[:-delay])] = -1

    if candidate.cooldown_bars > 0:
        last_signal = -candidate.cooldown_bars - 1
        for i in range(len(signals)):
            if signals[i] != 0:
                if i - last_signal <= candidate.cooldown_bars:
                    signals[i] = 0
                else:
                    last_signal = i
    return signals


def metrics_for_slice(
    signals: np.ndarray,
    outcomes: dict[str, np.ndarray],
    start: int,
    end: int,
    min_trades: int,
    min_coverage: float,
) -> dict[str, Any]:
    signal = signals[start:end]
    chosen_direction = outcomes["chosen_direction"][start:end]
    long_r = outcomes["long_r"][start:end]
    short_r = outcomes["short_r"][start:end]
    long_outcome = outcomes["long_outcome"][start:end]
    short_outcome = outcomes["short_outcome"][start:end]

    trades = signal != 0
    n_rows = len(signal)
    n_trades = int(trades.sum())
    coverage = float(n_trades / n_rows) if n_rows else 0.0
    if n_trades:
        realized_r = np.where(signal > 0, long_r, short_r)
        realized_outcome = np.where(signal > 0, long_outcome, short_outcome)
        trade_accuracy = float((realized_r[trades] > 0).mean())
        direction_accuracy = float(
            (signal[trades] == chosen_direction[trades]).mean()
        )
        tp1_rate = float(np.isin(realized_outcome[trades], [3, 4, 5]).mean())
        tp2_rate = float(np.isin(realized_outcome[trades], [4, 5]).mean())
        tp3_rate = float((realized_outcome[trades] == 5).mean())
        sl_rate = float((realized_outcome[trades] == 1).mean())
        expectancy = float(realized_r[trades].mean())
        total_r = float(realized_r[trades].sum())
        positive = float(realized_r[trades][realized_r[trades] > 0].sum())
        negative = float(-realized_r[trades][realized_r[trades] < 0].sum())
        profit_factor = positive / negative if negative > 0 else math.inf
    else:
        trade_accuracy = direction_accuracy = 0.0
        tp1_rate = tp2_rate = tp3_rate = sl_rate = 0.0
        expectancy = total_r = 0.0
        profit_factor = 0.0

    eligible = n_trades >= min_trades and coverage >= min_coverage
    return {
        "rows": n_rows,
        "trades": n_trades,
        "coverage": coverage,
        "eligible": eligible,
        "trade_accuracy": trade_accuracy,
        "direction_accuracy": direction_accuracy,
        "tp1_or_better_rate": tp1_rate,
        "tp2_or_better_rate": tp2_rate,
        "tp3_rate": tp3_rate,
        "sl_rate": sl_rate,
        "expectancy_R": expectancy,
        "profit_factor": profit_factor,
        "total_R": total_r,
    }


def objective(metrics: dict[str, Any]) -> tuple[float, float, float]:
    if not metrics["eligible"]:
        return (-1.0, metrics["coverage"], metrics["expectancy_R"])
    return (
        metrics["trade_accuracy"],
        metrics["coverage"],
        metrics["expectancy_R"],
    )


# ---------------------------------------------------------------------------
# Pinned FLAN-T5 supervisor
# ---------------------------------------------------------------------------


class Supervisor:
    def __init__(self) -> None:
        from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
        import torch

        self.torch = torch
        log(
            "Loading pinned supervisor {} at revision {}",
            SUPERVISOR_MODEL_ID,
            SUPERVISOR_REVISION,
        )
        self.tokenizer = AutoTokenizer.from_pretrained(
            SUPERVISOR_MODEL_ID,
            revision=SUPERVISOR_REVISION,
            trust_remote_code=False,
        )
        self.model = AutoModelForSeq2SeqLM.from_pretrained(
            SUPERVISOR_MODEL_ID,
            revision=SUPERVISOR_REVISION,
            trust_remote_code=False,
        )
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)
        self.model.eval()

    def choose(
        self,
        candidates: list[Candidate],
        history: list[dict[str, Any]],
    ) -> tuple[Optional[int], str]:
        if not candidates:
            return None, "no candidates"

        candidate_text = "\n".join(
            f"ID {c.candidate_id}: ref={c.reference_timeframe}, "
            f"sweep_atr={c.min_sweep_atr}, bias={c.use_reference_bias}, "
            f"confirm={c.confirmation_window}, cooldown={c.cooldown_bars}, "
            f"sl={c.sl_atr_mult}, horizon={c.horizon}"
            for c in candidates[:24]
        )
        history_text = "\n".join(
            f"trial {row['trial']}: candidate={row['candidate_id']}, "
            f"accuracy={row.get('validation', {}).get('trade_accuracy', 0):.4f}, "
            f"coverage={row.get('validation', {}).get('coverage', 0):.4f}"
            for row in history[-8:]
        ) or "No previous trials."
        prompt = (
            "You are a constrained trading-strategy supervisor. Choose one "
            "candidate ID from the list. Do not invent an ID. Prefer enough "
            "coverage and realistic validation accuracy. Reply exactly as "
            "CANDIDATE_ID=<integer> followed by one short reason.\n\n"
            f"Candidates:\n{candidate_text}\n\nHistory:\n{history_text}"
        )
        inputs = self.tokenizer(
            prompt,
            return_tensors="pt",
            truncation=True,
            max_length=768,
        ).to(self.device)
        with self.torch.no_grad():
            output = self.model.generate(**inputs, max_new_tokens=48)
        reply = self.tokenizer.decode(output[0], skip_special_tokens=True)
        match = re.search(r"CANDIDATE_ID\s*=\s*(\d+)", reply)
        selected = int(match.group(1)) if match else None
        valid_ids = {candidate.candidate_id for candidate in candidates}
        if selected not in valid_ids:
            selected = None
        return selected, reply


# ---------------------------------------------------------------------------
# Checkpointing and optimization loop
# ---------------------------------------------------------------------------


def save_json(path: Path, value: Any) -> None:
    path.write_text(json.dumps(value, indent=2, default=str), encoding="utf-8")


def save_checkpoint(
    checkpoint_dir: Path,
    candidate: Candidate,
    result: dict[str, Any],
    milestone: Optional[int] = None,
) -> None:
    prefix = "best" if milestone is None else f"milestone_{milestone:02d}pct"
    payload = {
        "saved_at": utc_now(),
        "candidate": asdict(candidate),
        "result": result,
        "supervisor_model": {
            "id": SUPERVISOR_MODEL_ID,
            "revision": SUPERVISOR_REVISION,
            "trust_remote_code": False,
        },
        "crt_sources": [
            "https://innercircletrader.net/tutorials/candle-range-theory-crt/",
            "https://tradingwyckoff.com/en/crt/",
        ],
    }
    save_json(checkpoint_dir / f"{prefix}_checkpoint.json", payload)
    log("Saved {} checkpoint at validation accuracy {:.2%}", prefix, result["validation"]["trade_accuracy"])


def package_checkpoints(checkpoint_dir: Path) -> Path:
    archive_path = checkpoint_dir.parent / "100optimization_checkpoints.zip"
    if archive_path.exists():
        archive_path.unlink()
    with zipfile.ZipFile(archive_path, "w", zipfile.ZIP_DEFLATED) as archive:
        for path in sorted(checkpoint_dir.rglob("*")):
            if path.is_file():
                archive.write(path, arcname=f"{checkpoint_dir.name}/{path.relative_to(checkpoint_dir)}")
    log("Checkpoint archive: {}", archive_path)
    return archive_path


def load_market_data() -> tuple[dict[str, pd.DataFrame], Path]:
    archive = load_dataset_archive()
    extracted = safe_extract_zip(archive, CFG.data_dir)
    data_dir = find_data_directory(extracted)
    raw: dict[str, pd.DataFrame] = {}
    for timeframe in ("M1", "M5", "M15", "M30", "H1", "H4"):
        path = find_timeframe_file(data_dir, timeframe)
        if path is None:
            raise FileNotFoundError(f"Missing {timeframe}.csv in {data_dir}")
        raw[timeframe] = load_ohlcv(path)
        log("Loaded {}: {:,} rows", timeframe, len(raw[timeframe]))
    return raw, data_dir


def run() -> dict[str, Any]:
    checkpoint_dir = ensure_checkpoint_dir()
    random.seed(CFG.seed)
    np.random.seed(CFG.seed)

    raw, data_dir = load_market_data()
    m1 = raw["M1"]
    if CFG.max_rows and len(m1) > CFG.max_rows:
        m1 = m1.tail(CFG.max_rows).reset_index(drop=True)
        log("Using the last {:,} M1 rows because MAX_ROWS is set", len(m1), level="WARN")

    atr = compute_atr(m1, CFG.atr_period)
    candidates = make_candidates(CFG.seed)
    outcome_cache: dict[tuple[float, int], dict[str, np.ndarray]] = {}
    signal_cache: dict[int, np.ndarray] = {}
    history: list[dict[str, Any]] = []
    evaluated: set[int] = set()
    best_result: Optional[dict[str, Any]] = None
    best_candidate: Optional[Candidate] = None
    milestones_saved: set[int] = set()

    n = len(m1)
    train_end = int(n * 0.60)
    validation_end = int(n * 0.80)
    purge = max(CFG.default_horizon, 60)
    validation_start = min(n, train_end + purge)
    test_start = min(n, validation_end + purge)

    supervisor: Optional[Supervisor]
    try:
        supervisor = Supervisor()
    except Exception as exc:
        log("Supervisor unavailable: {}. Continuing deterministically.", exc, level="WARN")
        supervisor = None

    log(
        "CRT optimization rows={} | train={} | validation={} | test={}",
        n,
        train_end,
        validation_end - validation_start,
        n - test_start,
    )

    for trial in range(CFG.max_trials):
        remaining = [candidate for candidate in candidates if candidate.candidate_id not in evaluated]
        if not remaining:
            break

        selected_id: Optional[int] = None
        supervisor_reply = ""
        if supervisor is not None and history:
            selected_id, supervisor_reply = supervisor.choose(remaining, history)
        if selected_id is None:
            # Deterministic fallback: evaluate candidates in the seeded order.
            selected_id = remaining[0].candidate_id
        candidate = next(c for c in candidates if c.candidate_id == selected_id)
        evaluated.add(candidate.candidate_id)

        key = (candidate.sl_atr_mult, candidate.horizon)
        if key not in outcome_cache:
            outcome_cache[key] = build_outcomes(
                m1,
                CFG.atr_period,
                candidate.sl_atr_mult,
                candidate.horizon,
            )
        outcomes = outcome_cache[key]

        if candidate.candidate_id not in signal_cache:
            signal_cache[candidate.candidate_id] = generate_crt_signals(
                m1,
                raw[candidate.reference_timeframe],
                atr,
                candidate,
            )
        signals = signal_cache[candidate.candidate_id]

        validation = metrics_for_slice(
            signals,
            outcomes,
            validation_start,
            validation_end,
            CFG.min_trades,
            CFG.min_coverage,
        )
        test = metrics_for_slice(
            signals,
            outcomes,
            test_start,
            n,
            CFG.min_trades,
            CFG.min_coverage,
        )
        result = {
            "trial": trial + 1,
            "candidate_id": candidate.candidate_id,
            "candidate": asdict(candidate),
            "validation": validation,
            "test_preview": test,
            "supervisor_reply": supervisor_reply,
        }
        history.append(result)
        save_json(checkpoint_dir / "trials.json", history)

        log(
            "Trial {} candidate={} validation accuracy={:.2%} coverage={:.2%} "
            "trades={} expectancy={:.3f}R test accuracy={:.2%}",
            trial + 1,
            candidate.candidate_id,
            validation["trade_accuracy"],
            validation["coverage"],
            validation["trades"],
            validation["expectancy_R"],
            test["trade_accuracy"],
        )

        if best_result is None or objective(validation) > objective(best_result["validation"]):
            best_result = result
            best_candidate = candidate
            save_checkpoint(checkpoint_dir, candidate, result)

        achieved = validation["eligible"] and validation["trade_accuracy"] >= CFG.target_accuracy
        for milestone in (20, 30, 80):
            if (
                milestone not in milestones_saved
                and validation["eligible"]
                and validation["trade_accuracy"] >= milestone / 100.0
            ):
                save_checkpoint(checkpoint_dir, candidate, result, milestone=milestone)
                milestones_saved.add(milestone)

        if achieved:
            log(
                "Target validation accuracy reached: {:.2%}. "
                "No further optimization trials will run.",
                validation["trade_accuracy"],
            )
            break