"""
MLP-based activation score strategy.

Replaces the single-layer weighted sum of strategy_activation_scores.py with a
small multi-layer perceptron over the pre-normalised feature columns (55 for
BTC, 50 for alts; see FEATURE_COLS / the artifact's feature_cols):

    score = 1000 * tanh( W3 . tanh( W2 . tanh(W1.x + b1) + b2 ) + b3 )

Everything downstream of the score — crossunder entry/exit, confirmation,
trailing stop, regime/MVRV gates, metrics — is identical to the existing
strategy (imported or mirrored verbatim), so results stay directly comparable
and the TV-parity workflow carries over.

Weights are NOT optimizer params: they live in a JSON artifact
(strategies/params/mlp/mlp_weights_{ASSET}_{TF}.json) produced by
tools/train_mlp.py and referenced via the `mlp_weights_file` param.
The same weights are embedded into strategy_mlp_scores.pine as array literals
by tools/generate_pine_mlp_presets.py, with tanh implemented in Pine by the
identical clamped closed form (Pine v6 has no math.tanh):

    tanh(x) = x >= 20 ?  1.0
            : x <= -20 ? -1.0
            : (exp(2x) - 1) / (exp(2x) + 1)

Pilot limitations (documented intentionally):
  - `i_div_window` is unsupported — the MLP consumes raw single-bar RSID flags,
    matching _prepare_features()'s window=1 behaviour.
"""
import json
import os
import sys

import numpy as np
import pandas as pd

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from strategies.strategy_activation_scores import (  # noqa: E402
    _apply_trailing_stop,
    _prepare_features,
    calculate_metrics,  # re-exported: validate_strategy / optimize_strategy look it up on this module
)
from strategies.library_activation_scores import calculate_positions  # noqa: E402

__all__ = [
    "FEATURE_COLS", "FEATURE_WEIGHT_PARAMS", "ACTIVATION_NAME",
    "_tanh", "mlp_forward",
    "save_mlp_artifact", "load_mlp_artifact",
    "generate_signals", "calculate_metrics",
]

# Identifier recorded in artifacts so Python/Pine provenance is checkable.
ACTIVATION_NAME = "tanh_closed_form_clamp20"

# The 49 CSV column names in the exact order produced by _prepare_features()
# (strategy_activation_scores.py cols 0-48). Single source of truth for the
# MLP input order; the artifact stores a copy and tests assert they match.
# NOTE: this is the feature-matrix order, which differs from config.WEIGHT_COLS
# order at positions 12-14 (params-JSON key order vs WEIGHT_COLS grouping).
FEATURE_COLS = [
    "stoch_norm",               # col 0
    "macd_pred_norm",           # col 1
    "osc_norm",                 # col 2
    "totalvol_norm",            # col 3  GL on-chain transfer volume percentrank
    "m3_momentum_norm",         # col 4
    "m2_tiny_norm",             # col 5
    "newaddr_norm",             # col 6  GL new addresses percentrank (demand/adoption)
    "stoch_div_norm",           # col 7
    "vwap_div_norm",            # col 8
    "stoch_peak_norm",          # col 9  (also the exit-confirmation gate)
    "sendaddr_norm",            # col 10 GL sending addresses percentrank (spending pressure)
    "m3_div_norm",              # col 11
    "bearish_engulfing_score",  # col 12
    "m2_nooff_norm",            # col 13
    "m2_div_norm",              # col 14
    "bullish_hammer_score",     # col 15
    "bullish_engulfing_score",  # col 16
    "shooting_star_score",      # col 17
    "btc_spx_corr_30",          # col 18
    "dxy_roc_norm",             # col 19
    "vix_pctrank_inv",          # col 20
    "btc_dom_roc_sign",         # col 21
    "us10y_roc_inv_sign",       # col 22
    "spy_above_200ema",         # col 23
    "gold_roc_pctrank",         # col 24
    "mvrv_zscore_value",        # col 25
    "mvrv_zscore_cont",         # col 26
    "nupl_norm",                # col 27
    "fed_net_liq_sign",         # col 28
    "gc_position",              # col 29
    "us2y_roc_inv_sign",        # col 30
    "yield_curve_sign",         # col 31
    "sopr_norm",                # col 32  SOPR percentrank: on-chain spending profit/loss state
    "rsid_reg_bull_norm",       # col 33
    "rsid_reg_bear_norm",       # col 34
    "rsid_hid_bull_norm",       # col 35
    "rsid_hid_bear_norm",       # col 36
    "rsid_rt_bull_norm",        # col 37
    "rsid_rt_bear_norm",        # col 38
    "rsid_slow_bull_norm",      # col 39
    "rsid_slow_bear_norm",      # col 40
    "rsid_delayed_peak_norm",   # col 41
    "rsid_delayed_dip_norm",    # col 42
    "oi_roc_norm",              # col 43
    "usdt_d_norm",              # col 44
    "basis_norm",               # col 45
    "cvd_norm",                 # col 46  OBV ROC percentrank: buy/sell pressure imbalance
    "btc_gold_norm",            # col 47
    "rsi_subtf_norm",           # col 48
    "bb_pct_b_norm",            # col 49  Bollinger Band %B: -1=at lower band, +1=at upper band
    "cvd_norm",                 # col 50  OBV ROC percentrank: buy/sell pressure imbalance
    "rvol_norm",                # col 51  Realized vol vs historical: volatility regime
    "active1y_norm",            # col 52  % supply inactive 1+yr: long-term holder concentration
    "hr_norm",                  # col 53  Hash rate ROC: miner commitment signal
    "sopr_norm",                # col 54  SOPR percentrank: recent spending profit/loss state
]

# i_w_* perceptron weight param corresponding to each FEATURE_COLS index
# (params-JSON key order, mirrored from the _prepare_features docstring).
# Used by tests (linear-degeneracy), trainer warm-starts, and codegen.
FEATURE_WEIGHT_PARAMS = [
    "i_w_stoch", "i_w_macd_pred", "i_w_osc", "i_w_totalvol",
    "i_w_m3_momentum", "i_w_m2_tiny", "i_w_newaddr", "i_w_stoch_div_osc",
    "i_w_vwap_div_osc", "i_w_stoch_peaking", "i_w_sendaddr",
    "i_w_m3_div_osc", "i_w_bearish_engulfing", "i_w_m2_div_osc_noOffset",
    "i_w_m2_div_osc", "i_w_bullish_hammer", "i_w_bullish_engulfing",
    "i_w_shooting_star", "i_w_btc_spx_corr", "i_w_dxy", "i_w_vix",
    "i_w_btc_dom", "i_w_us10y", "i_w_spy", "i_w_gold", "i_w_mvrv",
    "i_w_mvrv_cont", "i_w_nupl", "i_w_fed_net_liq", "i_w_gc_position",
    "i_w_us2y", "i_w_yield_curve", "i_w_qqq_spy_ratio",
    "i_w_rsid_reg_bull", "i_w_rsid_reg_bear",
    "i_w_rsid_hid_bull", "i_w_rsid_hid_bear",
    "i_w_rsid_rt_bull", "i_w_rsid_rt_bear",
    "i_w_rsid_slow_bull", "i_w_rsid_slow_bear",
    "i_w_rsid_delayed_peak", "i_w_rsid_delayed_dip",
    "i_w_oi_roc", "i_w_usdt_d", "i_w_basis",
    "i_w_fear_greed", "i_w_btc_gold", "i_w_rsi_subtf",
    "i_w_bb_pct_b",
]


def _tanh(x):
    """Clamped closed-form tanh, formula-identical to the Pine implementation.

    Inside (-20, 20): (exp(2x) - 1) / (exp(2x) + 1). At |x| >= 20 returns
    exactly +/-1.0 (Pine and Python clamp at the same point so neither side
    can diverge in the saturated region).
    """
    x = np.asarray(x, dtype=np.float64)
    out = np.empty_like(x)
    hi = x >= 20.0
    lo = x <= -20.0
    mid = ~(hi | lo)
    out[hi] = 1.0
    out[lo] = -1.0
    e2x = np.exp(2.0 * x[mid])
    out[mid] = (e2x - 1.0) / (e2x + 1.0)
    return out


def mlp_forward(X, layers):
    """Forward pass: X (T, F) float64, layers = [(W, b), ...] with W (n_out, n_in).

    Every layer (including the output layer) applies _tanh; the final scalar is
    scaled by 1000 so scores live in [-1000, +1000] like the existing
    activation-score convention.
    """
    a = np.asarray(X, dtype=np.float64)
    for W, b in layers:
        with np.errstate(divide="ignore", over="ignore", invalid="ignore"):
            z = a @ np.asarray(W, dtype=np.float64).T + np.asarray(b, dtype=np.float64)
        a = _tanh(z)
    return a[:, 0] * 1000.0


# ---------------------------------------------------------------------------
# Weights artifact I/O
# ---------------------------------------------------------------------------

def save_mlp_artifact(path, layers, asset=None, timeframe=None, training=None,
                      feature_cols=None):
    """Writes the weights artifact JSON. Floats serialised via repr (full
    precision) by json, so load() round-trips bit-exact float64."""
    feature_cols = list(feature_cols) if feature_cols is not None else list(FEATURE_COLS)
    arch = [int(np.asarray(layers[0][0]).shape[1])] + [
        int(np.asarray(W).shape[0]) for W, _ in layers
    ]
    payload = {
        "version": 1,
        "asset": asset,
        "timeframe": timeframe,
        "arch": arch,
        "activation": ACTIVATION_NAME,
        "feature_cols": feature_cols,
        "layers": [
            {"W": np.asarray(W, dtype=np.float64).tolist(),
             "b": np.asarray(b, dtype=np.float64).tolist()}
            for W, b in layers
        ],
        "training": training or {},
    }
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    with open(path, "w") as f:
        json.dump(payload, f)


_ARTIFACT_CACHE = {}


def load_mlp_artifact(path):
    """Loads (and caches by path+mtime) a weights artifact.

    Returns {"layers": [(W, b), ...] float64, "arch": [...],
             "feature_cols": [...], "meta": {...}}.

    Note: artifacts trained on the 49-feature set (pre-bb_pct_b_norm) are still
    loadable via load_mlp_artifact() and mlp_results_table; they carry their own
    feature_cols list and will fall back to 0 for any missing column.
    """
    abspath = os.path.abspath(path)
    mtime = os.path.getmtime(abspath)
    cached = _ARTIFACT_CACHE.get(abspath)
    if cached is not None and cached[0] == mtime:
        return cached[1]

    with open(abspath) as f:
        payload = json.load(f)
    layers = [
        (np.array(layer["W"], dtype=np.float64), np.array(layer["b"], dtype=np.float64))
        for layer in payload["layers"]
    ]
    art = {
        "layers": layers,
        "arch": payload["arch"],
        "feature_cols": payload["feature_cols"],
        "meta": {
            "version": payload.get("version"),
            "asset": payload.get("asset"),
            "timeframe": payload.get("timeframe"),
            "activation": payload.get("activation"),
            "training": payload.get("training", {}),
        },
    }
    _ARTIFACT_CACHE[abspath] = (mtime, art)
    return art


# ---------------------------------------------------------------------------
# Signal generation
# ---------------------------------------------------------------------------

def _apply_trailing_stop_with_dz(entry_raw_arr, exit_raw_arr, exit_raw_fallback_arr,
                                   score_exit_arr, exit_threshold,
                                   close_arr, high_arr, trail_stop_pct,
                                   time_ok_arr=None):
    """Like _apply_trailing_stop but with dead-zone guard for exit smoothing.

    When a trade enters while score_exit[entry_bar] < exit_threshold (i.e. the
    SMA is already below the exit threshold — the "dead-zone"), the smoothed
    crossunder can never fire for that trade because there is no transition from
    above to below the threshold.  In that case we switch to the raw-score
    crossunder (exit_raw_fallback_arr) for the duration of that trade only.

    time_ok_arr (optional bool array): see _apply_trailing_stop — Pine gates
    strategy.close on timeCondition, so entries/exits only fire where
    time_ok_arr[t] is True; a position open at window end is held open.
    """
    n = len(close_arr)
    in_pos     = np.zeros(n, dtype=np.int32)
    exec_entry = np.zeros(n, dtype=bool)
    exec_exit  = np.zeros(n, dtype=bool)

    position   = False
    trade_high = 0.0
    fill_bar   = -1
    use_raw    = False

    for t in range(n):
        allowed = True if time_ok_arr is None else bool(time_ok_arr[t])
        if position:
            if t != fill_bar and high_arr[t] > trade_high:
                trade_high = high_arr[t]
            stop_price = trade_high * (1.0 - trail_stop_pct)
            trail_hit  = close_arr[t] <= stop_price

            exit_sig = exit_raw_fallback_arr[t] if use_raw else exit_raw_arr[t]
            if (trail_hit or exit_sig) and allowed:
                position   = False
                exec_exit[t] = True
                trade_high = 0.0
                fill_bar   = -1
                use_raw    = False
        else:
            if entry_raw_arr[t] and allowed:
                position   = True
                exec_entry[t] = True
                trade_high = close_arr[t]
                fill_bar   = t + 1
                use_raw    = score_exit_arr[t] < exit_threshold

        in_pos[t] = 1 if position else 0

    return in_pos, exec_entry, exec_exit


def _calculate_positions_with_dz(entry_raw_arr, exit_raw_arr, exit_raw_fallback_arr,
                                   score_exit_arr, exit_threshold):
    """Like calculate_positions but with dead-zone guard (no trailing stop)."""
    n = len(entry_raw_arr)
    in_pos     = np.zeros(n, dtype=np.int32)
    exec_entry = np.zeros(n, dtype=bool)
    exec_exit  = np.zeros(n, dtype=bool)

    position = False
    use_raw  = False

    for t in range(n):
        if position:
            exit_sig = exit_raw_fallback_arr[t] if use_raw else exit_raw_arr[t]
            if exit_sig:
                position = False
                exec_exit[t] = True
                use_raw = False
        else:
            if entry_raw_arr[t]:
                position = True
                exec_entry[t] = True
                use_raw = score_exit_arr[t] < exit_threshold

        in_pos[t] = 1 if position else 0

    return in_pos, exec_entry, exec_exit


def _score_to_signals(df, params, stoch_peak_arr, close_arr, high_arr):
    """Threshold/regime/crossunder/position block — mirrors
    strategy_activation_scores.generate_signals lines 332-421 verbatim (that
    block is inline in the original and cannot be imported without modifying
    it). df must already contain 'activation_score'.
    """
    entry_threshold     = params.get('i_long_entry_activation_threshold',              106.0)
    exit_threshold      = params.get('i_long_exit_activation_threshold',               140.5)
    exit_conf_threshold = params.get('i_long_exit_activation_confirmation_threshold',  32.4125)
    use_exit_conf       = params.get('i_use_long_exit_confirmation',                   1.0)
    use_entry_conf      = params.get('i_use_long_entry_confirmation',                  False)

    score = df['activation_score']

    # --- Exit score smoothing (window=1 is identity; smooths exit crossunder only) ---
    exit_window = max(1, int(params.get('i_exit_score_window', 1)))
    score_exit = score.rolling(exit_window, min_periods=1).mean() if exit_window > 1 else score

    # --- Entry score smoothing (window=1 is identity; smooths entry crossunder only) ---
    entry_window = max(1, int(params.get('i_entry_score_window', 1)))
    score_entry = score.rolling(entry_window, min_periods=1).mean() if entry_window > 1 else score

    # --- Regime filter (entries only; window=0 disables) ---
    regime_window = int(params.get('i_regime_window', 0))
    regime_min    = float(params.get('i_regime_entry_min_score', -1000.0))
    if regime_window > 0:
        regime_ok = score.rolling(window=regime_window, min_periods=1).mean() > regime_min
    else:
        regime_ok = pd.Series(True, index=df.index)

    # --- MVRV bear-regime suppression ---
    suppress_bear = bool(params.get('i_mvrv_suppress_bear', False))
    if suppress_bear and 'mvrv_regime' in df.columns:
        regime_ok = regime_ok & (df['mvrv_regime'] >= 0)

    # --- Exit signal (2-bar crossunder on smoothed score + confirmation) ---
    stoch_peak = pd.Series(stoch_peak_arr != 0, index=df.index)
    if use_exit_conf:
        exit_raw = (
            (score_exit.shift(2) >= exit_threshold) &
            (score_exit.shift(1) < exit_threshold) &
            (score_exit < exit_conf_threshold)
        )
        # Dead-zone fallback: same crossunder logic but on raw score
        exit_raw_fallback = (
            (score.shift(2) >= exit_threshold) &
            (score.shift(1) < exit_threshold) &
            (score < exit_conf_threshold)
        ) if exit_window > 1 else exit_raw
    else:
        # Non-confirmation path gates on stoch_is_peaking (peak_norm == -1 when peaking)
        exit_raw = (
            (score_exit.shift(2) >= exit_threshold) &
            (score_exit.shift(1) < exit_threshold) &
            stoch_peak
        )
        exit_raw_fallback = (
            (score.shift(2) >= exit_threshold) &
            (score.shift(1) < exit_threshold) &
            stoch_peak
        ) if exit_window > 1 else exit_raw

    # --- Entry signal (uses smoothed score; momentum check stays on raw score) ---
    if use_entry_conf:
        entry_raw = (
            (score_entry.shift(2) >= entry_threshold) &
            (score_entry.shift(1) < entry_threshold) &
            (score > score.shift(1)) &
            ~exit_raw.shift(1, fill_value=False)
        )
    else:
        entry_raw = (score_entry.shift(1) >= entry_threshold) & (score_entry < entry_threshold)

    entry_raw = entry_raw & regime_ok
    time_start = params.get('_pine_time_start')
    time_end = params.get('_pine_time_end')
    # time_ok_arr mirrors Pine's `timeCondition`. It gates entry_raw/exit_raw
    # (below) AND is passed to the trailing-stop replay so the trail/score close
    # also respects the window — a position open at window end is held open, like
    # Pine (which cannot fire strategy.close after endDate). None when no window
    # is set (optimizer path) → trail behaves as before.
    time_ok_arr = None
    if (time_start is not None or time_end is not None) and 'time' in df.columns:
        time = pd.to_datetime(df['time'])
        time_ok = pd.Series(True, index=df.index)
        if time_start is not None:
            time_ok = time_ok & (time >= pd.Timestamp(time_start))
        if time_end is not None:
            time_ok = time_ok & (time <= pd.Timestamp(time_end))
        entry_raw = entry_raw & time_ok
        exit_raw = exit_raw & time_ok
        exit_raw_fallback = exit_raw_fallback & time_ok
        time_ok_arr = time_ok.to_numpy(dtype=bool)

    # --- Executed positions (entry only if flat, exit only if long) ---
    trail_stop_pct = float(params.get('i_trailing_stop_threshold', 0.0)) / 100.0
    entry_arr    = entry_raw.fillna(False).values.astype(bool)
    exit_arr     = exit_raw.fillna(False).values.astype(bool)
    fallback_arr = exit_raw_fallback.fillna(False).values.astype(bool)
    score_exit_arr = score_exit.values

    if exit_window > 1:
        if trail_stop_pct > 0.0:
            in_pos, exec_entry, exec_exit = _apply_trailing_stop_with_dz(
                entry_arr, exit_arr, fallback_arr,
                score_exit_arr, exit_threshold,
                close_arr, high_arr, trail_stop_pct,
                time_ok_arr=time_ok_arr,
            )
        else:
            in_pos, exec_entry, exec_exit = _calculate_positions_with_dz(
                entry_arr, exit_arr, fallback_arr,
                score_exit_arr, exit_threshold,
            )
    elif trail_stop_pct > 0.0:
        in_pos, exec_entry, exec_exit = _apply_trailing_stop(
            entry_arr, exit_arr, close_arr, high_arr, trail_stop_pct,
            time_ok_arr=time_ok_arr,
        )
    else:
        in_pos, exec_entry, exec_exit = calculate_positions(entry_arr, exit_arr)

    df['position']      = in_pos.astype(int)
    df['execute_entry'] = exec_entry
    df['execute_exit']  = exec_exit
    return df


def generate_signals(df, **params):
    """Generates entry/exit signals from the MLP activation score.

    Params:
      mlp_weights_file — path to the weights artifact JSON (required)
      i_long_entry_activation_threshold / i_long_exit_activation_threshold /
      i_long_exit_activation_confirmation_threshold — score thresholds
      i_use_long_entry_confirmation / i_use_long_exit_confirmation — flags
      i_trailing_stop_threshold — pct (0 disables)
      i_regime_window / i_regime_entry_min_score / i_mvrv_suppress_bear — gates
    """
    weights_file = params.get("mlp_weights_file")
    if not weights_file:
        raise ValueError("strategy_mlp_scores requires the 'mlp_weights_file' param")
    art = load_mlp_artifact(weights_file)
    art_cols = art["feature_cols"]
    if art_cols != FEATURE_COLS:
        # Allow artifacts trained on a strict subset of FEATURE_COLS (e.g. the old
        # 49-feature set during the transition to 50 features).  Missing columns
        # default to 0 in _prepare_features, which is safe for near-zero-weight cols.
        extra = [c for c in FEATURE_COLS if c not in art_cols]
        missing = [c for c in art_cols if c not in FEATURE_COLS]
        import warnings
        if missing:
            # Approach-B probe artifacts carry extra computed columns (e.g. in_long_position,
            # bars_held_norm) not in FEATURE_COLS. Allow them if they exist in the input df;
            # _prepare_features will read them directly. Fatal only if they'd silently be 0.
            missing_and_absent = [c for c in missing if c not in df.columns]
            if missing_and_absent:
                raise ValueError(
                    f"Artifact {weights_file} expects columns not in FEATURE_COLS or df: "
                    f"{missing_and_absent}"
                )
            warnings.warn(
                f"Artifact {weights_file} expects extra columns (probe features): {missing}. "
                "These will be read from the input df.",
                stacklevel=2,
            )
        if extra:
            warnings.warn(
                f"Artifact {weights_file} was trained on {len(art_cols)} features; "
                f"current FEATURE_COLS has {len(FEATURE_COLS)} (new: {extra}). "
                "New columns default to 0 — retrain for full performance.",
                stacklevel=2,
            )

    df = df.copy()
    df.columns = df.columns.str.lower().str.strip()

    X = _prepare_features(df, art_cols)
    df['activation_score'] = mlp_forward(X, art["layers"])

    stoch_peak_arr = X[:, 9]  # col 9 = stoch_peak_norm (exit gate)
    close_arr = df['close'].values.astype(np.float64)
    high_arr  = df['high'].values.astype(np.float64)

    return _score_to_signals(df, params, stoch_peak_arr, close_arr, high_arr)
