"""
Generate Pine Script preset block for strategy_mlp_scores.pine from trained MLP artifacts.

Reads strategies/params/mlp/mlp_weights_COINBASE_BTCUSD_{TF}.json and
optionally results/winners/optimization_winner_strategy_mlp_scores_COINBASE_BTCUSD_{TF}.csv
(for post-sweep thresholds; falls back to artifact's recommended_thresholds if missing).

Rewrites the auto-generated sentinel block inside strategy_mlp_scores.pine in-place.

Usage:
    python3 tools/generate_pine_mlp_presets.py [--dry-run]

After running, paste strategy_mlp_scores.pine into TradingView to compile-test.
"""

import argparse
import csv
import hashlib
import json
import math
import os
import re
import sys
from datetime import datetime

# ── Config ────────────────────────────────────────────────────────────────────

PINE_FILE    = "strategies/strategy_mlp_scores.pine"
WEIGHTS_DIR  = "strategies/params/mlp"
WINNERS_DIR  = "results/winners"

SENTINEL_START = "// ── AUTO-GENERATED MLP PRESETS"
SENTINEL_END   = "// ── END AUTO-GENERATED"

# Pilot scope: BTC only, all 5 TFs
_ASSETS = ["COINBASE_BTCUSD"]
_TFS    = ["4H", "6H", "8H", "12H", "1D"]

_FALLBACK_PRESET = "COINBASE_BTCUSD 12H"

# Presets whose arch is valid but whose trained weights are degenerate (too few
# bars, high regime variance).  They are aliased to another preset for BOTH
# weight loading AND thresholds so the chart still gets sensible behaviour.
# Key = degenerate preset, Value = preset to borrow from.
_PRESET_WEIGHT_ALIASES = {
    "COINBASE_BTCUSD 1D": "COINBASE_BTCUSD 12H",
}

# Threshold params emitted as Pine switch blocks (order determines Pine block order)
_THRESHOLD_PARAMS = [
    "i_long_entry_activation_threshold",
    "i_long_exit_activation_threshold",
    "i_long_exit_activation_confirmation_threshold",
    "i_use_long_exit_confirmation",
    "i_use_long_entry_confirmation",
    "i_trailing_stop_threshold",
    "i_regime_window",
    "i_regime_entry_min_score",
    "i_mvrv_suppress_bear",
    "i_exit_score_window",
    "i_entry_score_window",
]

_BOOL_PARAMS = {"i_use_long_entry_confirmation", "i_use_long_exit_confirmation", "i_mvrv_suppress_bear"}
_INT_PARAMS  = {"i_regime_window", "i_exit_score_window", "i_entry_score_window"}

# Metric cols that appear in winner CSVs but are not params
_METRIC_COLS = {
    "Total P&L %", "Max Drawdown %", "Sharpe Ratio", "Sortino Ratio",
    "Total Trades", "Calmar Ratio", "P&L/DD Ratio", "% In Market",
    "composite_score", "Composite", "verified", "GPU_Score", "CPU_Score",
    "_pnl_dd_percentile", "subperiod_consistent", "WFO_Score", "WFO_Fold_Calmars",
}


# ── Formatting helpers ────────────────────────────────────────────────────────

def fmt_float9(v: float) -> str:
    """Format float to 7 significant figures for Pine weight literals.
    Matches float32 training precision; ~3 fewer digits/weight vs prior 10g format.
    """
    s = f"{v:.7g}"
    if "." not in s and "e" not in s:
        s += ".0"
    return s


def fmt_threshold(name: str, v) -> str:
    """Format threshold param value as Pine literal (4-sig float, int, or bool)."""
    if name in _BOOL_PARAMS:
        if isinstance(v, str):
            return "true" if v.strip().lower() in ("true", "1", "yes") else "false"
        return "true" if float(v) > 0.5 else "false"
    if name in _INT_PARAMS:
        return str(int(round(float(v))))
    s = f"{float(v):.4g}"
    if "." not in s and "e" not in s:
        s += ".0"
    return s


def pine_type(name: str) -> str:
    if name in _BOOL_PARAMS:
        return "bool"
    if name in _INT_PARAMS:
        return "int"
    return "float"


def preset_key(asset: str, tf: str) -> str:
    return f"{asset} {tf}"


def _sha8(path: str) -> str:
    try:
        with open(path, "rb") as f:
            return hashlib.sha256(f.read()).hexdigest()[:8]
    except OSError:
        return "????????"


# ── Load artifacts and thresholds ────────────────────────────────────────────

def load_presets() -> dict:
    """
    Returns {preset_key: {"thresholds": {...}, "layers": [(W,b),...], "arch": [...],
                           "artifact_path": str, "artifact_sha8": str}} for found presets.
    """
    import numpy as np
    from strategies.strategy_mlp_scores import load_mlp_artifact

    presets = {}

    # Winner CSVs for post-sweep thresholds
    winner_thresholds = {}
    winner_pattern = re.compile(
        r"optimization_winner_strategy_mlp_scores_([A-Z0-9_]+)_([A-Z0-9HD]+)\.csv$"
    )
    if os.path.isdir(WINNERS_DIR):
        for fname in sorted(os.listdir(WINNERS_DIR)):
            m = winner_pattern.match(fname)
            if not m:
                continue
            asset, tf = m.group(1), m.group(2)
            if asset not in _ASSETS or tf not in _TFS:
                continue
            fpath = os.path.join(WINNERS_DIR, fname)
            with open(fpath, newline="") as f:
                rows = list(csv.DictReader(f))
            if not rows:
                continue
            row = {k: v for k, v in rows[0].items() if k not in _METRIC_COLS}
            winner_thresholds[preset_key(asset, tf)] = row

    for asset in _ASSETS:
        for tf in _TFS:
            key = preset_key(asset, tf)

            # Determine artifact path: winner CSV's mlp_weights_file column takes
            # priority over the legacy per-TF default filename.
            artifact_path = os.path.join(WEIGHTS_DIR, f"mlp_weights_{asset}_{tf}.json")
            if key in winner_thresholds:
                csv_weights = winner_thresholds[key].get("mlp_weights_file", "").strip()
                if csv_weights and os.path.exists(csv_weights):
                    artifact_path = csv_weights
                    # This preset has a dedicated trained artifact — clear any alias
                    # that was set up when it was borrowing another TF's weights.
                    _PRESET_WEIGHT_ALIASES.pop(key, None)

            if not os.path.exists(artifact_path):
                print(f"  WARNING: no artifact for {key} ({artifact_path}) — skipping",
                      file=sys.stderr)
                continue
            try:
                art = load_mlp_artifact(artifact_path)
            except Exception as e:
                print(f"  WARNING: cannot load {artifact_path}: {e} — skipping",
                      file=sys.stderr)
                continue

            # Thresholds: winner CSV > artifact recommended_thresholds > zeros
            thr = {}
            rec = art["meta"].get("training", {}).get("recommended_thresholds", {})
            if key in winner_thresholds:
                # Filter to only the threshold params we emit
                w_row = winner_thresholds[key]
                thr = {p: w_row.get(p, rec.get(p, 0.0)) for p in _THRESHOLD_PARAMS}
                print(f"  {key}: using winner CSV thresholds + artifact {os.path.basename(artifact_path)}",
                      file=sys.stderr)
            elif rec:
                thr = {p: rec.get(p, 0.0) for p in _THRESHOLD_PARAMS}
                print(f"  {key}: using artifact recommended_thresholds", file=sys.stderr)
            else:
                thr = {p: 0.0 for p in _THRESHOLD_PARAMS}
                print(f"  {key}: no thresholds found — using zeros", file=sys.stderr)

            presets[key] = {
                "thresholds": thr,
                "layers": art["layers"],
                "arch": art["arch"],
                "artifact_path": artifact_path,
                "artifact_sha8": _sha8(artifact_path),
            }

    for asset in _ASSETS:
        for tf in _TFS:
            if preset_key(asset, tf) not in presets:
                print(f"  WARNING: preset {preset_key(asset, tf)} will use fallback defaults",
                      file=sys.stderr)

    return presets


# ── Weight array emission helpers ─────────────────────────────────────────────

def _row_literal(values) -> str:
    """Emit a list of floats as Pine array.from() call."""
    return "array.from(" + ", ".join(fmt_float9(float(v)) for v in values) + ")"


def emit_var_declarations(arch, indent="") -> list:
    """Emit 'var float[]' declarations for all weight/bias arrays (must appear before any if-block that assigns to them)."""
    lines = []
    lines.append(f"{indent}// MLP weight arrays (populated on barstate.isfirst from the preset block below)")
    lines.append(f"{indent}// Architecture: {' → '.join(str(d) for d in arch)}")
    for i in range(len(arch) - 1):
        n_in, n_out = arch[i], arch[i + 1]
        lines.append(f"{indent}var float[] _w{i+1} = array.new_float({n_out * n_in}, 0.0)  // W{i+1} [{n_out}×{n_in}]")
        lines.append(f"{indent}var float[] _b{i+1} = array.new_float({n_out}, 0.0)")
    lines.append("")
    return lines


def emit_weight_fn(arr, fn_name, indent="    ") -> list:
    """
    Emit a Pine named function that builds ONE weight matrix or bias vector and returns it.
    Moving data into functions keeps the main script body small (Pine's main-body size limit
    is separate from per-function limits).

    The function has no arguments and returns the array — caller does:
        _w1 := fn_name()
    """
    import numpy as np
    arr = np.asarray(arr, dtype=np.float64)
    lines = [f"{fn_name}() =>"]

    if arr.ndim == 1 or (arr.ndim == 2 and arr.shape[0] == 1):
        # Bias vector or single-row W matrix: fits in one array.from
        row = arr.ravel()
        lines.append(f"{indent}{_row_literal(row)}")
    else:
        # Multi-row weight matrix: build row-by-row using local _tw/_tr accumulators.
        # Each function has its own local scope so no cross-function naming conflicts.
        n_out = arr.shape[0]
        lines.append(f"{indent}float[] _tw = {_row_literal(arr[0])}")
        lines.append(f"{indent}float[] _tr = {_row_literal(arr[1])}")
        lines.append(f"{indent}array.concat(_tw, _tr)")
        for row_i in range(2, n_out):
            lines.append(f"{indent}_tr := {_row_literal(arr[row_i])}")
            lines.append(f"{indent}array.concat(_tw, _tr)")
        lines.append(f"{indent}_tw")

    return lines


def emit_weight_init(layers, arch, indent="        ") -> list:
    """
    Return list of Pine lines that initialize _w1/_b1/.._b3 from the given layers.

    Uses row-by-row array.from() calls with array.concat() to build each W matrix,
    avoiding any single very-large array.from() call.

    _tw/_tr are declared with 'float[]' only on first use; subsequent layers reuse
    them with ':=' to avoid duplicate-declaration errors in Pine v6.
    """
    import numpy as np

    lines = []
    tw_declared = False
    tr_declared = False

    for li, (W, b) in enumerate(layers):
        W = np.asarray(W, dtype=np.float64)    # shape [n_out, n_in]
        b = np.asarray(b, dtype=np.float64)    # shape [n_out]
        n_out, n_in = W.shape
        w_var = f"_w{li+1}"
        b_var = f"_b{li+1}"

        if n_out == 1:
            # Single row — can fit in one array.from
            lines.append(f"{indent}{w_var} := {_row_literal(W[0])}")
        else:
            # Multiple rows: build using _tw/_tr accumulators
            if tw_declared:
                lines.append(f"{indent}_tw := {_row_literal(W[0])}")
            else:
                lines.append(f"{indent}float[] _tw = {_row_literal(W[0])}")
                tw_declared = True

            if n_out > 1:
                if tr_declared:
                    lines.append(f"{indent}_tr := {_row_literal(W[1])}")
                else:
                    lines.append(f"{indent}float[] _tr = {_row_literal(W[1])}")
                    tr_declared = True
                lines.append(f"{indent}array.concat(_tw, _tr)")
                for row_i in range(2, n_out):
                    lines.append(f"{indent}_tr := {_row_literal(W[row_i])}")
                    lines.append(f"{indent}array.concat(_tw, _tr)")
            lines.append(f"{indent}{w_var} := _tw")

        # Bias vector
        lines.append(f"{indent}{b_var} := {_row_literal(b)}")

    return lines


# ── Generate Pine block ───────────────────────────────────────────────────────

def generate_block(presets: dict) -> str:
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M")
    ordered_keys = [preset_key(a, tf) for a in _ASSETS for tf in _TFS]
    present_keys = [k for k in ordered_keys if k in presets]

    lines = []
    lines.append(f"{SENTINEL_START} — do not edit manually ──────────────────────────────────────────")
    lines.append(f"// Run tools/generate_pine_mlp_presets.py to regenerate from trained weight artifacts.")
    lines.append(f"// Last generated: {timestamp}  |  {len(present_keys)}/{len(ordered_keys)} presets found")
    lines.append("")

    # ── Preset auto-detection boilerplate ──
    lines.append("// ── Preset auto-detection from chart symbol + timeframe ─────────────────────")
    lines.append("_tf_secs  = timeframe.in_seconds(timeframe.period)")
    lines.append('_tf_label = _tf_secs == 86400 ? "1D" :')
    lines.append('             _tf_secs == 14400 ? "4H" :')
    lines.append('             _tf_secs == 21600 ? "6H" :')
    lines.append('             _tf_secs == 28800 ? "8H" :')
    lines.append('             _tf_secs == 43200 ? "12H" : ""')
    lines.append("")

    all_option_keys = [preset_key(a, tf) for a in _ASSETS for tf in _TFS]
    options_str = '["", ' + ", ".join(f'"{k}"' for k in all_option_keys) + "]"
    lines.append(f'i_preset_override = input.string("", "Force preset (blank = auto)",')
    lines.append(f'    options={options_str},')
    lines.append(f'    group=group_neural_activation_thresholds, display=display.none)')
    lines.append("")
    lines.append('_auto_key = syminfo.prefix + "_" + syminfo.ticker + " " + _tf_label')
    lines.append('_preset   = i_preset_override != "" ? i_preset_override : _auto_key')
    lines.append("// ─────────────────────────────────────────────────────────────────────────────")
    lines.append("")

    # ── Threshold switch blocks ──
    # Get fallback preset's thresholds (use first present key, or zeros)
    fallback_key = _FALLBACK_PRESET
    if fallback_key not in presets and present_keys:
        fallback_key = present_keys[0]
    fallback_thr = presets[fallback_key]["thresholds"] if fallback_key in presets else {p: 0.0 for p in _THRESHOLD_PARAMS}

    for param in _THRESHOLD_PARAMS:
        ptype = pine_type(param)
        fb_val = fmt_threshold(param, fallback_thr.get(param, 0.0))
        lines.append(f"{ptype} {param} = switch _preset")
        for key in ordered_keys:
            # Each preset uses its own winner-CSV thresholds (alias only affects weights).
            if key in presets:
                val = fmt_threshold(param, presets[key]["thresholds"].get(param, 0.0))
                weight_src = _PRESET_WEIGHT_ALIASES.get(key)
                alias_note = f"  // thresholds: own CSV, weights: {weight_src}" if weight_src else ""
                lines.append(f'    "{key}" => {val}{alias_note}')
        lines.append(f"    => {fb_val}  // fallback: {fallback_key}")
        lines.append("")

    # ── Var declarations for weight arrays (must precede the if barstate.isfirst block) ──
    ref_arch = presets[present_keys[0]]["arch"] if present_keys else [49, 16, 8, 1]

    # Guard: drop any preset whose arch doesn't match ref_arch — mismatched arch causes
    # Pine runtime "index N is out of bounds" because the forward pass loop bounds are
    # hardcoded to ref_arch layer sizes.
    arch_ok_keys = []
    for key in present_keys:
        if presets[key]["arch"] != ref_arch:
            print(
                f"  WARNING: skipping {key} — arch {presets[key]['arch']} != ref {ref_arch}. "
                f"Retrain with --hidden {' '.join(str(d) for d in ref_arch[1:-1])}",
                file=sys.stderr,
            )
        else:
            arch_ok_keys.append(key)

    lines.extend(emit_var_declarations(ref_arch))

    # ── Weight functions (one per W/b matrix per preset) ──
    # Moving weight data into named functions keeps the main script body small.
    # Pine enforces a main-body size limit separately from per-function limits.
    # Each function builds one array and returns it; the caller does _w1 := f_...().
    import numpy as np
    for key in ordered_keys:
        if key not in arch_ok_keys:
            continue
        if key in _PRESET_WEIGHT_ALIASES:
            continue  # aliased — no weight functions needed; loads via target's functions
        p = presets[key]
        safe = key.replace(" ", "_")
        arch_str = str(p["arch"])
        sha8 = p["artifact_sha8"]
        lines.append(f"// {key}  arch: {arch_str}  artifact: {sha8}")
        for li, (W, b) in enumerate(p["layers"]):
            lines.extend(emit_weight_fn(np.asarray(W), f"_fw_{safe}_w{li+1}"))
            lines.append("")
            lines.extend(emit_weight_fn(np.asarray(b), f"_fw_{safe}_b{li+1}"))
            lines.append("")

    # ── Call-site: tiny if-blocks that just invoke the functions ──
    # Build reverse alias map: target_key → [alias_keys...]
    alias_reverse: dict[str, list[str]] = {}
    for alias_key, target_key in _PRESET_WEIGHT_ALIASES.items():
        alias_reverse.setdefault(target_key, []).append(alias_key)

    for key in ordered_keys:
        if key not in arch_ok_keys:
            continue
        if key in _PRESET_WEIGHT_ALIASES:
            continue  # emitted as part of the target's block below
        safe = key.replace(" ", "_")
        n_layers = len(presets[key]["layers"])
        # Build condition: this key + any keys that alias to it
        extra_aliases = alias_reverse.get(key, [])
        cond_parts = [f'_preset == "{key}"'] + [f'_preset == "{ak}"' for ak in extra_aliases]
        cond = " or ".join(cond_parts)
        cond_str = f"({cond})" if len(cond_parts) > 1 else cond
        lines.append(f'if barstate.isfirst and {cond_str}')
        for li in range(n_layers):
            lines.append(f'    _w{li+1} := _fw_{safe}_w{li+1}()')
            lines.append(f'    _b{li+1} := _fw_{safe}_b{li+1}()')
        lines.append("")

    lines.append(f"{SENTINEL_END} ─────────────────────────────────────────────────────────────────────")

    return "\n".join(lines) + "\n"


# ── Rewrite Pine file ─────────────────────────────────────────────────────────

def rewrite_pine(block: str, dry_run: bool) -> None:
    if not os.path.exists(PINE_FILE):
        sys.exit(f"Pine file not found: {PINE_FILE}")
    with open(PINE_FILE) as f:
        content = f.read()

    start_idx = content.find(SENTINEL_START)
    end_idx   = content.find(SENTINEL_END)
    if start_idx == -1 or end_idx == -1:
        sys.exit(f"Sentinel markers not found in {PINE_FILE}. "
                 f"Expected '{SENTINEL_START}' and '{SENTINEL_END}'.")

    end_idx += len(SENTINEL_END)
    # Advance past any trailing dash characters in the end sentinel
    while end_idx < len(content) and content[end_idx] in " -─\t":
        end_idx += 1
    # Move back to include the last char of the end sentinel line
    # (find the newline after the end sentinel)
    nl = content.find("\n", content.find(SENTINEL_END))
    if nl != -1:
        end_idx = nl + 1
    else:
        end_idx = len(content)

    new_content = content[:start_idx] + block + content[end_idx:]

    if dry_run:
        print("── DRY RUN — would write the following block ──")
        print(block)
        print(f"── (Pine file not modified) ──")
    else:
        with open(PINE_FILE, "w") as f:
            f.write(new_content)
        print(f"✅ Rewrote {PINE_FILE}  ({len(present_keys_global)}/{len(all_keys_global)} presets)")
        print(f"   Next step: paste into TradingView and run:")
        print(f"       python3 tools/check_pine.py --mark-valid")


present_keys_global = []
all_keys_global = []


def main():
    global present_keys_global, all_keys_global

    parser = argparse.ArgumentParser(
        description="Regenerate MLP preset block in strategy_mlp_scores.pine")
    parser.add_argument("--dry-run", action="store_true",
                        help="Print generated block without modifying Pine file")
    args = parser.parse_args()

    # Ensure we can import project modules
    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

    print("Loading MLP artifacts...", file=sys.stderr)
    presets = load_presets()

    all_keys_global   = [preset_key(a, tf) for a in _ASSETS for tf in _TFS]
    present_keys_global = [k for k in all_keys_global if k in presets]

    if not presets:
        print("No artifacts found — nothing to generate.", file=sys.stderr)
        sys.exit(1)

    print(f"\nGenerating block for {len(present_keys_global)} preset(s)...", file=sys.stderr)
    block = generate_block(presets)
    rewrite_pine(block, args.dry_run)


if __name__ == "__main__":
    main()