"""
lock_params.py — Automatically apply sign-stability locks to params JSON files.

Reads the sweep DB, identifies params that meet the "Lock NEGATIVE/POSITIVE (strong)"
criteria (N ≥ 200 top-quartile rows, ≤ 5% or ≥ 95% positive), then tightens the
corresponding params JSON range to exclude the disallowed sign.

Lock rules:
  NEGATIVE: set stop = 0.0  (np.arange exclusive, so only negative values sampled)
  POSITIVE: set start = step (first value is +step, only positive values sampled)

Only applies when evidence comes from N ≥ 200 rows for that specific asset/TF.
With --cross-asset: if ALL assets with sufficient data agree on a lock, also apply
the lock to assets/TFs that are still SPARSE (< 200 rows).

Usage:
    python3 tools/lock_params.py                    # dry run — show what would change
    python3 tools/lock_params.py --apply            # write changes to params JSONs
    python3 tools/lock_params.py --cross-asset      # broadcast confirmed locks to sparse assets
    python3 tools/lock_params.py --apply --cross-asset
    python3 tools/lock_params.py --asset COINBASE_BTCUSD --timeframe 1D
"""

import sqlite3
import json
import os
import sys
import glob
import argparse
from datetime import datetime

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import RESULTS_DIR

DB_FILE       = "results/sweep_database.db"
PARAMS_GLOB   = "strategies/params/params_strategy_activation_scores_*.json"
DB_SCHEMA_VER = 3          # must match auto_optimize_loop.py

MIN_N         = 200        # minimum top-quartile rows to issue a lock
MAX_PCT_NEG   = 5.0        # ≤ 5% positive → Lock NEGATIVE
MIN_PCT_POS   = 95.0       # ≥ 95% positive → Lock POSITIVE

WEIGHT_COLS = [
    "i_w_stoch", "i_w_macd_pred", "i_w_osc", "i_w_macd_bullish",
    "i_w_m3_momentum", "i_w_m2_tiny", "i_w_rsid_osc",
    "i_w_stoch_div_osc", "i_w_vwap_div_osc", "i_w_stoch_peaking",
    "i_w_stoch_bottoming", "i_w_m3_div_osc", "i_w_m2_div_osc",
    "i_w_m2_div_osc_noOffset", "i_w_bearish_engulfing",
    "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",
    # RSI divergence signals
    "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",
]


def load_all_params():
    """Return {(asset, tf): (path, dict)} for every params JSON found."""
    result = {}
    import re
    for jf in sorted(glob.glob(PARAMS_GLOB)):
        fname = os.path.basename(jf)
        stem = fname.replace("params_strategy_activation_scores_", "").replace(".json", "")
        parts = stem.rsplit("_", 1)
        if len(parts) == 2 and re.match(r"^\d+[HDMhdm]$", parts[1]):
            asset, tf = parts[0], parts[1]
            try:
                with open(jf) as f:
                    result[(asset, tf)] = (jf, json.load(f))
            except Exception as e:
                print(f"  [WARN] Could not load {jf}: {e}")
    return result


def compute_locks(conn, asset_filter=None, tf_filter=None):
    """
    Query the DB and return per-(asset, tf) lock decisions.

    Returns:
        locks_specific: {(asset, tf): {param: direction}}
          direction is "negative" or "positive"
        n_rows: {(asset, tf): int}  — top-quartile row count
    """
    where = f"WHERE schema_ver={DB_SCHEMA_VER} AND pnl_dd_percentile >= 75"
    params = []
    if asset_filter:
        where += " AND asset=?"; params.append(asset_filter)
    if tf_filter:
        where += " AND timeframe=?"; params.append(tf_filter)

    import pandas as pd
    df = pd.read_sql(f"SELECT * FROM sweep_results {where}", conn,
                     params=params if params else None)

    if df.empty:
        return {}, {}

    locks_specific = {}
    n_rows = {}

    for (asset, tf), grp in df.groupby(["asset", "timeframe"]):
        n = len(grp)
        n_rows[(asset, tf)] = n
        if n < MIN_N:
            continue
        combo_locks = {}
        for col in WEIGHT_COLS:
            if col not in grp.columns:
                continue
            data = grp[col].dropna()
            if len(data) < MIN_N:
                continue
            pct_pos = 100.0 * (data > 0).sum() / len(data)
            if pct_pos <= MAX_PCT_NEG:
                combo_locks[col] = "negative"
            elif pct_pos >= MIN_PCT_POS:
                combo_locks[col] = "positive"
        if combo_locks:
            locks_specific[(asset, tf)] = combo_locks

    return locks_specific, n_rows


def apply_lock_to_param(param_dict, col, direction):
    """
    Modify param_dict[col] in-place to enforce the sign lock.

    NEGATIVE: set stop = 0.0  (arange exclusive → only negatives sampled)
    POSITIVE: set start = step (first sampled value is +step)

    Returns (changed: bool, description: str)
    """
    if col not in param_dict:
        return False, "param not found in JSON"

    entry = param_dict[col]

    # Values-only entries (discrete list) — check if already effectively locked
    if "values" in entry and "start" not in entry:
        vals = entry["values"]
        if direction == "negative" and all(v <= 0 for v in vals):
            return False, f"already locked (values={vals}, all non-positive)"
        if direction == "positive" and all(v >= 0 for v in vals):
            return False, f"already locked (values={vals}, all non-negative)"
        return False, f"values-only param with mixed signs {vals} — manual review needed"

    if "start" not in entry or "stop" not in entry:
        return False, "no start/stop keys"

    current_start = float(entry["start"])
    current_stop  = float(entry["stop"])
    step          = float(entry.get("step", 5.0))

    if direction == "negative":
        # Already fully negative?
        if current_stop <= 0.0:
            return False, f"already locked (stop={current_stop})"
        new_stop = 0.0
        # Sanity: must leave at least a few values
        if current_start >= new_stop:
            return False, f"would produce empty range (start={current_start} ≥ stop=0)"
        entry["stop"] = new_stop
        return True, f"stop {current_stop} → {new_stop}"

    elif direction == "positive":
        # Already fully positive?
        if current_start >= step:
            return False, f"already locked (start={current_start})"
        new_start = step
        if new_start >= current_stop:
            return False, f"would produce empty range (start={new_start} ≥ stop={current_stop})"
        entry["start"] = new_start
        return True, f"start {current_start} → {new_start}"

    return False, "unknown direction"


def main():
    parser = argparse.ArgumentParser(
        description="Apply sign-stability locks to params JSON files.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument("--apply",        action="store_true",
                        help="Write changes to params JSONs (default: dry run)")
    parser.add_argument("--cross-asset",  action="store_true",
                        help="Broadcast locks confirmed by all data-rich assets to sparse ones too")
    parser.add_argument("--asset",        help="Restrict analysis to one asset")
    parser.add_argument("--timeframe",    help="Restrict analysis to one timeframe")
    args = parser.parse_args()

    if not os.path.exists(DB_FILE):
        print(f"DB not found: {DB_FILE}")
        sys.exit(1)

    conn = sqlite3.connect(DB_FILE)
    all_params = load_all_params()

    print(f"\n{'='*65}")
    print(f"  lock_params.py  —  {'DRY RUN' if not args.apply else 'APPLYING CHANGES'}")
    print(f"  DB schema ver : {DB_SCHEMA_VER}")
    print(f"  Threshold     : N ≥ {MIN_N} top-quartile rows, ≤{MAX_PCT_NEG}% or ≥{MIN_PCT_POS}% positive")
    if args.cross_asset:
        print(f"  Cross-asset   : ON — locks confirmed by all data-rich assets broadcast to sparse")
    print(f"{'='*65}\n")

    locks_specific, n_rows = compute_locks(conn, args.asset, args.timeframe)
    conn.close()

    if not locks_specific:
        print("No lock-eligible combos found (need ≥ 200 top-quartile rows per combo).")
        return

    # ── Cross-asset inference ──────────────────────────────────────────────
    # Find params where EVERY data-rich (asset, tf) agrees on the same lock.
    cross_locks = {}
    if args.cross_asset:
        rich_combos = [(a, t) for (a, t), lk in locks_specific.items()]
        if len(rich_combos) >= 2:
            # Intersection of lock decisions across all rich combos
            shared = None
            for combo, lk in locks_specific.items():
                if shared is None:
                    shared = dict(lk)
                else:
                    # Keep only params where direction agrees
                    shared = {p: d for p, d in shared.items()
                              if p in lk and lk[p] == d}
            if shared:
                cross_locks = shared
                print(f"  Cross-asset locks (agreed by all {len(rich_combos)} rich combos):")
                for p, d in sorted(cross_locks.items()):
                    print(f"    {p:<40} → lock {d.upper()}")
                print()

    # ── Build full change plan ─────────────────────────────────────────────
    # For each params JSON file, determine what changes to make.
    # Priority: specific lock > cross-asset lock
    total_changes = 0
    total_skipped = 0

    for (asset, tf), (path, param_dict) in sorted(all_params.items()):
        # Gather applicable locks for this combo
        combo_locks = {}
        source_label = {}

        # Specific locks (evidence from this exact asset/TF)
        if (asset, tf) in locks_specific:
            for p, d in locks_specific[(asset, tf)].items():
                combo_locks[p] = d
                source_label[p] = f"(N={n_rows.get((asset,tf),0)} rows)"

        # Cross-asset locks for combos without their own data
        if args.cross_asset and (asset, tf) not in locks_specific:
            for p, d in cross_locks.items():
                if p not in combo_locks:
                    combo_locks[p] = d
                    source_label[p] = "(cross-asset)"

        if not combo_locks:
            continue

        print(f"  {asset}-{tf}  [{path}]")
        file_changed = False

        for param, direction in sorted(combo_locks.items()):
            changed, desc = apply_lock_to_param(param_dict, param, direction)
            src = source_label.get(param, "")
            if changed:
                verb = "APPLY " if args.apply else "WOULD "
                print(f"    {verb}Lock {direction.upper():<8} {param:<40} {desc}  {src}")
                total_changes += 1
                file_changed = True
            else:
                print(f"    SKIP  Lock {direction.upper():<8} {param:<40} {desc}  {src}")
                total_skipped += 1

        if file_changed and args.apply:
            with open(path, "w") as f:
                json.dump(param_dict, f, indent=2)
                f.write("\n")
            print(f"    → Written: {path}")

        print()

    print(f"{'='*65}")
    if args.apply:
        print(f"  Applied {total_changes} lock(s).  Skipped {total_skipped} (already locked or no change).")
    else:
        print(f"  DRY RUN: {total_changes} lock(s) would be applied.  {total_skipped} already OK.")
        print(f"  Re-run with --apply to write changes.")
    print(f"{'='*65}\n")


if __name__ == "__main__":
    main()