# engine/strategy.py
import numpy as np


def allocation(sp_price, gd_price, sp_sma, gd_sma,
               base=0.6, tilt_mult=1.5,
               max_tilt=0.25, min_eq=0.2, max_eq=0.9):
    """Golden Nifty: SMA-momentum tilt between NiftyBees and GoldBees."""
    rd = (sp_price / sp_sma) - (gd_price / gd_sma)
    tilt = np.clip(-rd * tilt_mult, -max_tilt, max_tilt)
    eq_w = np.clip(base * (1 + tilt), min_eq, max_eq)
    return eq_w, 1 - eq_w


def alpha_shield_allocation(midcap_price, midcap_sma60):
    """
    Alpha Shield: Dynamic 70/30 Midcap+Gold based on 60-month SMA valuation.

    When midcap is expensive (price >> 5yr SMA) → reduce midcap, increase gold.
    When midcap is cheap (price << 5yr SMA) → increase midcap aggressively.

    Formula: midcap% = clip(70% - (price/sma60 - 1) × 60%, 40%, 92%)
    Backtested XIRR: ~16.9% p.a. over 12+ years (vs 15.6% static 70/30).
    """
    ratio = midcap_price / midcap_sma60
    eq_w = float(np.clip(0.70 - (ratio - 1.0) * 0.60, 0.40, 0.92))
    return eq_w, 1 - eq_w


# Strategy registry: maps strategy_name → engine configuration
STRATEGY_REGISTRY = {
    "golden_nifty": {
        "equity_symbol": "NIFTYBEES.NS",
        "gold_symbol":   "GOLDBEES.NS",
        "sma_months":    36,
        "allocation_fn": "golden_nifty",
    },
    "alpha_shield": {
        "equity_symbol": "JUNIORBEES.NS",
        "gold_symbol":   "GOLDBEES.NS",
        "sma_months":    60,
        "allocation_fn": "alpha_shield",
    },
}

DEFAULT_STRATEGY = "golden_nifty"


def get_strategy_config(strategy_name: str) -> dict:
    return STRATEGY_REGISTRY.get(strategy_name) or STRATEGY_REGISTRY[DEFAULT_STRATEGY]


def compute_weights(strategy_name: str, equity_price: float, gold_price: float,
                    equity_hist, gold_hist, sma_months: int):
    """Dispatch allocation to the correct strategy function."""
    if strategy_name == "alpha_shield":
        sma60 = equity_hist.rolling(sma_months).mean().iloc[-1]
        return alpha_shield_allocation(equity_price, sma60)
    # default: golden_nifty
    eq_sma = equity_hist.rolling(sma_months).mean().iloc[-1]
    gd_sma = gold_hist.rolling(sma_months).mean().iloc[-1]
    return allocation(equity_price, gold_price, eq_sma, gd_sma)