"""
apply_shap_locks.py — Lock <2% SHAP-importance params to {"values": [0]}
and unlock any KEEP_UNLOCKED params that were previously locked to [0].

Usage:
    python3 tools/apply_shap_locks.py          # dry run
    python3 tools/apply_shap_locks.py --apply  # write changes
"""

import json
import glob
import sys
import os

PARAMS_GLOB = "strategies/params/params_strategy_activation_scores_*.json"

# SHAP report rank 14+ (<2%) — lock to 0
LOCK_TO_ZERO = {
    "i_w_us2y", "i_w_m3_momentum", "i_w_rsid_rt_bear", "i_w_gold",
    "i_w_usdt_d", "i_w_m2_div_osc", "i_w_vwap_div_osc", "i_w_oi_roc",
    "i_w_rsid_hid_bull", "i_w_rsid_reg_bear", "i_w_osc", "i_w_gc_position",
    "i_w_btc_spx_corr", "i_w_stoch_div_osc", "i_w_m3_div_osc",
    "i_w_yield_curve", "i_w_mvrv_cont", "i_w_shooting_star", "i_w_dxy",
    "i_w_spy", "i_w_us10y", "i_w_vix", "i_w_stoch", "i_w_fed_net_liq",
    "i_w_rsid_slow_bull", "i_w_rsid_delayed_dip", "i_w_rsid_reg_bull",
    "i_w_mvrv", "i_w_rsid_rt_bull", "i_w_rsid_hid_bear",
    "i_w_fear_greed", "i_w_qqq_spy_ratio",
}

# User-specified keep-unlocked list (unlock if currently {"values": [0]})
KEEP_UNLOCKED = {
    "i_w_btc_dom", "i_w_bearish_engulfing", "i_w_bullish_hammer",
    "i_w_btc_gold", "i_w_nupl", "i_w_macd_pred", "i_w_bullish_engulfing",
    "i_w_rsi_subtf", "i_w_m2_tiny",
    # also STRONG above 2% from the report
    "i_w_rsid_delayed_peak", "i_w_rsid_slow_bear", "i_w_basis",
    "i_w_m2_div_osc_noOffset",
}

# Default range when unlocking a param that was {"values": [0]}
DEFAULT_RANGE = {"start": -100.0, "stop": 100.0, "step": 5.0}

apply = "--apply" in sys.argv

print(f"\n{'='*65}")
print(f"  apply_shap_locks.py  —  {'APPLYING CHANGES' if apply else 'DRY RUN'}")
print(f"  Locking {len(LOCK_TO_ZERO)} WEAK params to 0")
print(f"  Protecting/unlocking {len(KEEP_UNLOCKED)} STRONG params")
print(f"{'='*65}\n")

files = sorted(glob.glob(PARAMS_GLOB))
total_locked = 0
total_unlocked = 0
total_skipped = 0

for fpath in files:
    fname = os.path.basename(fpath)
    with open(fpath) as f:
        params = json.load(f)

    file_changes = []

    for key in list(params.keys()):
        entry = params[key]

        # Lock WEAK params to {"values": [0]}
        if key in LOCK_TO_ZERO:
            if entry == {"values": [0]}:
                total_skipped += 1
            else:
                file_changes.append(("LOCK  → 0", key, str(entry)))
                if apply:
                    params[key] = {"values": [0]}
                total_locked += 1

        # Unlock KEEP_UNLOCKED params if they're currently {"values": [0]}
        elif key in KEEP_UNLOCKED:
            if entry == {"values": [0]}:
                file_changes.append(("UNLOCK← range", key, str(entry)))
                if apply:
                    params[key] = dict(DEFAULT_RANGE)
                total_unlocked += 1

    if file_changes:
        print(f"  {fname}")
        for action, key, old in file_changes:
            verb = "  " if apply else "  WOULD "
            print(f"    {verb}{action:<16} {key:<40}  was: {old}")
        if apply:
            with open(fpath, "w") as f:
                json.dump(params, f, indent=2)
                f.write("\n")
            print(f"    → Written")
        print()

print(f"{'='*65}")
if apply:
    print(f"  Applied: {total_locked} locked, {total_unlocked} unlocked, {total_skipped} already-locked skipped.")
else:
    print(f"  DRY RUN: {total_locked} would lock, {total_unlocked} would unlock, {total_skipped} already OK.")
    print(f"  Re-run with --apply to write changes.")
print(f"{'='*65}\n")