#!/usr/bin/env python3
"""
UnifiedMarginTracker — daily combined-margin estimate for the Best Ideas portfolio.

Aggregates margin estimates from each sleeve, scaled by fund_weights × fund_scale.
Reports combined utilisation and applies a dynamic scale cap when projected margin
would exceed ``margin_cap_frac`` (default 0.85) of the running NAV.

Sleeve margin sources
---------------------
* ``spy_theta``     — actual margin from ``evaluate_theta_margin``
                      column: ``margin_total_usd``
* ``vxx_regime``    — VXX regime stack daily CSV
                      column: ``margin_total_usd`` (added 2026-06)
* ``vxx_long_call`` — long OTM calls, paid premium → zero margin required
* ``macro_aw``      — from per-trade ``margin_reserved_usd`` (added 2026-06)
                      reconstructed as daily open-position margin
* ``tactical_aw``   — Tactical AW daily CSV column ``margin_usd`` (added 2026-06)
* ``tsmom``         — TSMOM daily CSV column ``margin_usd`` (added 2026-06)
* ``orb_zarattini``  — Zarattini 5m ORB daily CSV column ``margin_usd`` (added 2026-07)
"""

from __future__ import annotations

import warnings
from pathlib import Path

import numpy as np
import pandas as pd


# ---------------------------------------------------------------------------
# Column names in each sleeve's daily CSV that hold the margin series
# ---------------------------------------------------------------------------
_MARGIN_COL: dict[str, str] = {
    "spy_theta":    "margin_total_usd",
    "vxx_regime":   "margin_total_usd",
    "tactical_aw":  "margin_usd",
    "tsmom":        "margin_usd",
    "orb_zarattini": "margin_usd",
    "vol_edge":     "margin_usd",
    "ls_equity":    "margin_usd",
    "vxx_long_call": None,   # long options — zero margin
    "macro_aw":     None,    # computed from trades (see below)
    "equity_dip":   None,    # long equity — margin approximated at 15%
    "sector_momentum": None, # long equity — margin approximated at 15%
    "ride_rockets": None,    # long equity momentum blend — margin approximated at 15%
    "ma_slope_topn": None,     # long equity — margin approximated at 15%
    "ma_slope_inverse": None,  # inverse ETF — margin approximated at 15%
    "ma_slope_intraday": None, # long equity day-trade — margin approximated at 15%
}


def _load_margin_col(
    path: Path,
    col: str,
    idx: pd.DatetimeIndex,
    fallback_zero: bool = True,
) -> pd.Series:
    """Load a single margin column from a CSV and reindex to ``idx``."""
    if not path.is_file():
        if fallback_zero:
            return pd.Series(0.0, index=idx)
        raise FileNotFoundError(f"Margin CSV not found: {path}")
    df = pd.read_csv(path, parse_dates=["date"])
    df["date"] = pd.to_datetime(df["date"]).dt.normalize()
    df = df.set_index("date")
    if col not in df.columns:
        if fallback_zero:
            warnings.warn(
                f"Column '{col}' not in {path.name}; margin set to 0 for this sleeve. "
                "Re-run the sleeve backtest to populate this column.",
                stacklevel=3,
            )
            return pd.Series(0.0, index=idx)
        raise KeyError(f"Column '{col}' not in {path}")
    return df[col].reindex(idx).ffill().fillna(0.0).astype(float)


def _macro_aw_daily_margin(
    trades_path: Path,
    idx: pd.DatetimeIndex,
) -> pd.Series:
    """Reconstruct per-day open Macro AW margin from the trade log.

    For each business day, sum ``margin_reserved_usd`` for all trades where
    ``entry_date ≤ day < exit_date`` (position still open).
    """
    if not trades_path.is_file():
        return pd.Series(0.0, index=idx)
    trades = pd.read_csv(trades_path, parse_dates=["entry_date", "exit_date"])
    trades["entry_date"] = pd.to_datetime(trades["entry_date"]).dt.normalize()
    trades["exit_date"] = pd.to_datetime(trades["exit_date"]).dt.normalize()
    if "margin_reserved_usd" not in trades.columns:
        warnings.warn(
            "macro_aw trades CSV has no 'margin_reserved_usd' column. "
            "Re-run macro_aw_options_portfolio.py to add it.",
            stacklevel=3,
        )
        return pd.Series(0.0, index=idx)

    margin_by_day = pd.Series(0.0, index=idx)
    for _, row in trades.iterrows():
        m = float(row.get("margin_reserved_usd", 0.0) or 0.0)
        if m == 0.0:
            continue
        entry = pd.Timestamp(row["entry_date"])
        exit_ = pd.Timestamp(row["exit_date"])
        if pd.isna(entry) or pd.isna(exit_):
            continue
        entry = entry.normalize()
        exit_ = exit_.normalize()
        # Hold during [entry, exit) — exit day we close, margin released.
        days_open = idx[(idx >= entry) & (idx < exit_)]
        margin_by_day.loc[days_open] += m
    return margin_by_day


def build_sleeve_margin_map(
    idx: pd.DatetimeIndex,
    *,
    active_sleeves: set[str],
    # sleeve daily CSV paths
    lit_mtm_daily: Path,
    vxx_regime_daily: Path,
    tsmom_daily: Path,
    tactical_aw_daily: Path,
    macro_aw_trades_path: Path,
    equity_dip_daily: Path | None = None,
    vol_edge_daily: Path | None = None,
    ls_equity_daily: Path | None = None,
    orb_zarattini_daily: Path | None = None,
    # scaling info
    fund_weights: dict[str, float],
    fund_scale: float,
    capital: float,
) -> dict[str, pd.Series]:
    """Return ``{sleeve_key: daily_margin_usd_at_scale_1}`` for every active sleeve.

    The series are at scale=1 (unlevered); multiply by ``fund_scale`` to get
    the margin consumed by the levered position.
    """
    result: dict[str, pd.Series] = {}

    for sleeve in active_sleeves:
        col = _MARGIN_COL.get(sleeve)

        if sleeve == "spy_theta":
            result[sleeve] = _load_margin_col(lit_mtm_daily, "margin_total_usd", idx)

        elif sleeve == "vxx_regime":
            result[sleeve] = _load_margin_col(vxx_regime_daily, "margin_total_usd", idx)

        elif sleeve == "tsmom":
            result[sleeve] = _load_margin_col(tsmom_daily, "margin_usd", idx)

        elif sleeve == "orb_zarattini":
            if orb_zarattini_daily is not None:
                result[sleeve] = _load_margin_col(orb_zarattini_daily, "margin_usd", idx)
            else:
                result[sleeve] = pd.Series(0.0, index=idx)

        elif sleeve == "tactical_aw":
            result[sleeve] = _load_margin_col(tactical_aw_daily, "margin_usd", idx)

        elif sleeve == "vol_edge":
            if vol_edge_daily is not None:
                result[sleeve] = _load_margin_col(vol_edge_daily, "margin_usd", idx)
            else:
                result[sleeve] = pd.Series(0.0, index=idx)

        elif sleeve == "ls_equity":
            if ls_equity_daily is not None:
                result[sleeve] = _load_margin_col(ls_equity_daily, "margin_usd", idx)
            else:
                result[sleeve] = pd.Series(0.0, index=idx)

        elif sleeve == "vxx_long_call":
            # Long premium, no margin required.
            result[sleeve] = pd.Series(0.0, index=idx)

        elif sleeve == "macro_aw":
            result[sleeve] = _macro_aw_daily_margin(macro_aw_trades_path, idx)

        elif sleeve in (
            "equity_dip", "sector_momentum", "johansen_etf", "ride_rockets",
            "ma_slope_topn", "ma_slope_inverse", "ma_slope_intraday",
            "qs_actionable_etf",
        ):
            # Long equity / ETF / day-trade: PM margin ~15% of position value.
            # Approximate from sleeve weight × capital × 0.15.
            w = fund_weights.get(sleeve, 0.0)
            base_notional = w * capital  # sleeve's capital fraction at NAV=$capital
            result[sleeve] = pd.Series(base_notional * 0.15, index=idx)

        else:
            result[sleeve] = pd.Series(0.0, index=idx)

    return result


def build_combined_margin(
    idx: pd.DatetimeIndex,
    sleeve_margin_map: dict[str, pd.Series],
    *,
    fund_weights: dict[str, float],
    fund_scale: float,
) -> pd.DataFrame:
    """Combine per-sleeve margins into a summary DataFrame.

    Each sleeve's standalone margin (computed at $100k notional) is scaled by:
    ``fund_weight × fund_scale`` because:
    * ``fund_weight`` is the fraction of NAV allocated to the sleeve, and
    * ``fund_scale`` is the overall leverage multiplier applied to the combined book.

    Columns returned:
    * ``margin_<sleeve>_usd``   — sleeve margin scaled by fund_weight × fund_scale
    * ``margin_combined_usd``   — total across all sleeves (at fund_weight × fund_scale)
    * ``margin_scale1_usd``     — total at scale=1 weighted by fund_weights (for gating)
    """
    out = pd.DataFrame(index=idx)
    combined_scale1 = pd.Series(0.0, index=idx)
    for sleeve, margin_s1 in sleeve_margin_map.items():
        w = float(fund_weights.get(sleeve, 0.0))
        # Scale by sleeve weight (fraction of NAV) and overall leverage.
        margin_weighted_s1 = margin_s1 * w   # at scale=1
        col = f"margin_{sleeve}_usd"
        out[col] = margin_weighted_s1 * fund_scale
        combined_scale1 = combined_scale1 + margin_weighted_s1

    out["margin_combined_usd"] = combined_scale1 * fund_scale
    out["margin_scale1_usd"] = combined_scale1
    return out


def apply_margin_cap(
    port_ret_unscaled: pd.Series,
    fund_scale: float,
    margin_scale1: pd.Series,
    nav: pd.Series,
    *,
    margin_cap_frac: float = 0.85,
) -> tuple[pd.Series, pd.Series]:
    """Apply a per-day dynamic scale cap so margin ≤ ``margin_cap_frac`` × prev_NAV.

    Parameters
    ----------
    port_ret_unscaled : daily portfolio return at scale=1 (normalized weights, no leverage)
    fund_scale : requested leverage multiplier
    margin_scale1 : combined margin at scale=1 per day
    nav : prior-day NAV (or starting capital on day 0)
    margin_cap_frac : max allowed margin / NAV (default 0.85 = 85%)

    Returns
    -------
    effective_ret : levered daily return, capped when margin would breach
    effective_scale : the scale actually applied each day (≤ fund_scale)
    """
    max_margin = nav * margin_cap_frac
    # Max scale allowed by margin cap
    # Avoid division by zero when margin_scale1 == 0 (no positions).
    safe = margin_scale1.replace(0.0, np.nan)
    max_scale_from_margin = max_margin / safe
    max_scale_from_margin = max_scale_from_margin.fillna(fund_scale)

    effective_scale = np.minimum(fund_scale, max_scale_from_margin.clip(lower=0.0))
    effective_ret = port_ret_unscaled * effective_scale
    return effective_ret, effective_scale