import pandas as pd
import numpy as np
import yfinance as yf

# =========================================================
# CONFIG
# =========================================================
START_DATE = "2010-01-01"
END_DATE   = "2025-12-31"

SIP_START_DATE = "2018-04-01"
SIP_END_DATE   = "2025-10-31"   # set None for till last data

# Monthly salary SIP budget
MONTHLY_SIP = 10_000

# Tickers
NIFTY = "NIFTYBEES.NS"
GOLD  = "GOLDBEES.NS"

# Optimization range: N trading days between deployments
N_MIN = 2
N_MAX = 30  # inclusive

# Salary credit timing:
# "MS" = Month Start, "ME" = Month End
SALARY_CREDIT = "MS"

# Valuation tilt params (unchanged)
SMA_MONTHS   = 36
TILT_MULT    = 1.5
MAX_TILT     = 0.25
BASE_EQUITY  = 0.60
MIN_EQUITY   = 0.20
MAX_EQUITY_W = 0.90

# Daily approximation
TRADING_DAYS_PER_MONTH = 21
SMA_DAYS = SMA_MONTHS * TRADING_DAYS_PER_MONTH

# Crash/sideways definitions
CRASH_THRESHOLD = -0.15
SIDEWAYS_CAGR   = 0.06
ROLL_MONTHS_SIDEWAYS = 36

# =========================================================
# HELPERS
# =========================================================
SIP_START_DATE = pd.to_datetime(SIP_START_DATE)
SIP_END_DATE   = pd.to_datetime(SIP_END_DATE) if SIP_END_DATE else None

def in_sip_window(date):
    if date < SIP_START_DATE:
        return False
    if SIP_END_DATE and date > SIP_END_DATE:
        return False
    return True

def load_price(ticker):
    df = yf.download(
        ticker,
        start=START_DATE,
        end=END_DATE,
        auto_adjust=True,
        progress=False
    )
    if isinstance(df.columns, pd.MultiIndex):
        df.columns = df.columns.get_level_values(0)
    return df["Close"]

def xirr(cashflows):
    dates = np.array([cf[0] for cf in cashflows], dtype="datetime64[D]")
    amounts = np.array([cf[1] for cf in cashflows], dtype=float)

    def npv(rate):
        years = (dates - dates[0]).astype(int) / 365.25
        return np.sum(amounts / ((1 + rate) ** years))

    low, high = -0.99, 5.0
    for _ in range(250):
        mid = (low + high) / 2
        val = npv(mid)
        if abs(val) < 1e-6:
            return mid
        if val > 0:
            low = mid
        else:
            high = mid
    return mid

def sip_max_drawdown(value_series: pd.Series):
    peak = value_series.cummax()
    dd = value_series / peak - 1
    trough = dd.idxmin()
    peak_date = value_series.loc[:trough].idxmax()
    return peak_date, trough, float(dd.min())

def worst_rolling_xirr(ledger_df: pd.DataFrame, periods: int):
    dates = ledger_df.index
    worst = None
    if len(dates) < periods + 2:
        return None

    for i in range(len(dates) - periods):
        start = dates[i]
        end = dates[i + periods]
        window = ledger_df.loc[start:end]
        if len(window) < max(6, periods // 2):
            continue

        cashflows = [(d, -float(row["Cash_Added"])) for d, row in window.iterrows()]
        cashflows.append((end, float(window["Portfolio_Value"].iloc[-1])))

        try:
            rx = xirr(cashflows)
            if worst is None or rx < worst["XIRR"]:
                worst = {"Start": start, "End": end, "XIRR": rx}
        except Exception:
            pass

    return worst

def rolling_cagr(series: pd.Series, periods: int, years: float):
    r = series / series.shift(periods)
    return (r ** (1 / years)) - 1

# =========================================================
# DATA
# =========================================================
prices = pd.DataFrame({
    "NIFTY": load_price(NIFTY),
    "GOLD":  load_price(GOLD)
}).dropna()

# business day frequency
prices = prices.resample("B").last().dropna()

# SMA / deviation on daily data
sma_nifty = prices["NIFTY"].rolling(SMA_DAYS).mean()
sma_gold  = prices["GOLD"].rolling(SMA_DAYS).mean()

dev_nifty = (prices["NIFTY"] / sma_nifty) - 1
dev_gold  = (prices["GOLD"] / sma_gold) - 1
rel_dev   = dev_nifty - dev_gold

# Restrict to SIP window once for speed
sip_prices = prices.loc[
    (prices.index >= SIP_START_DATE) &
    ((prices.index <= SIP_END_DATE) if SIP_END_DATE else True)
].copy()

sip_rel_dev = rel_dev.loc[sip_prices.index]

# =========================================================
# SALARY CREDIT DATES (MONTH START/END) ALIGNED TO TRADING DAYS
# =========================================================
def get_salary_dates(trading_index: pd.DatetimeIndex, credit: str):
    if credit not in ("MS", "ME"):
        raise ValueError("SALARY_CREDIT must be 'MS' or 'ME'.")

    # calendar month anchors
    start = trading_index.min().to_period("M").to_timestamp()
    end   = trading_index.max().to_period("M").to_timestamp(how="end")

    # choose all months in range
    months = pd.date_range(start=start, end=end, freq="MS")
    if credit == "ME":
        months = months + pd.offsets.MonthEnd(0)

    # map each month anchor to nearest trading day in index
    out = []
    for d in months:
        if d < trading_index.min() or d > trading_index.max():
            continue
        loc = trading_index.get_indexer([d], method="nearest")[0]
        td = trading_index[loc]
        if in_sip_window(td):
            out.append(td)

    return pd.DatetimeIndex(sorted(set(out)))

SALARY_DATES = get_salary_dates(sip_prices.index, SALARY_CREDIT)

# =========================================================
# CORE ENGINE: MONTHLY SALARY -> DEPLOY IN EQUAL PARTS ON EVERY Nth TRADING DAY
# =========================================================
def build_deployment_schedule(trading_index: pd.DatetimeIndex, salary_dates: pd.DatetimeIndex, n: int):
    """
    For each month: credit MONTHLY_SIP on salary date, then deploy it in equal chunks
    every n trading days until next salary date (exclusive).
    """
    deploy_map = {}  # date -> cash_to_invest

    for i, s_date in enumerate(salary_dates):
        if i < len(salary_dates) - 1:
            next_salary = salary_dates[i + 1]
            month_window = trading_index[(trading_index >= s_date) & (trading_index < next_salary)]
        else:
            month_window = trading_index[trading_index >= s_date]

        if len(month_window) == 0:
            continue

        # deployment dates: every nth trading day starting at salary date
        deploy_dates = month_window[::n]
        if len(deploy_dates) == 0:
            continue

        per_order = MONTHLY_SIP / len(deploy_dates)

        for d in deploy_dates:
            deploy_map[d] = deploy_map.get(d, 0.0) + per_order

    deploy_dates = pd.DatetimeIndex(sorted(deploy_map.keys()))
    return deploy_dates, deploy_map

def run_strategy_for_n(sip_prices: pd.DataFrame, sip_rel_dev: pd.Series, salary_dates: pd.DatetimeIndex, n: int):
    deploy_dates, deploy_map = build_deployment_schedule(sip_prices.index, salary_dates, n)

    nifty_units = 0.0
    gold_units  = 0.0
    total_invested = 0.0
    prev_value = None

    rows = []

    for date in deploy_dates:
        cash = float(deploy_map[date])
        total_invested += cash

        nifty_price = float(sip_prices.loc[date, "NIFTY"])
        gold_price  = float(sip_prices.loc[date, "GOLD"])
        rd = float(sip_rel_dev.loc[date]) if not pd.isna(sip_rel_dev.loc[date]) else np.nan

        tilt = 0.0
        if not np.isnan(rd):
            tilt = np.clip(-rd * TILT_MULT, -MAX_TILT, MAX_TILT)

        equity_w = BASE_EQUITY * (1 + tilt)
        equity_w = min(max(equity_w, MIN_EQUITY), MAX_EQUITY_W)
        gold_w   = 1 - equity_w

        nifty_buy = cash * equity_w
        gold_buy  = cash * gold_w

        nifty_units += nifty_buy / nifty_price
        gold_units  += gold_buy  / gold_price

        nifty_val = nifty_units * nifty_price
        gold_val  = gold_units  * gold_price
        port_val  = nifty_val + gold_val

        unrealized = port_val - total_invested
        if prev_value is None:
            period_pnl = 0.0
        else:
            period_pnl = port_val - prev_value - cash
        prev_value = port_val

        rows.append({
            "Date": date,
            "Cash_Added": cash,
            "Total_Invested": total_invested,
            "Equity_Weight": equity_w,
            "Gold_Weight": gold_w,
            "Rel_Deviation": rd,
            "Portfolio_Value": port_val,
            "Period_PnL": period_pnl,
            "Unrealized_PnL": unrealized
        })

    ledger = pd.DataFrame(rows).set_index("Date")
    if ledger.empty:
        return None

    # XIRR
    cashflows = [(d, -float(r["Cash_Added"])) for d, r in ledger.iterrows()]
    cashflows.append((ledger.index[-1], float(ledger["Portfolio_Value"].iloc[-1])))
    strat_xirr = xirr(cashflows)

    # Drawdown
    peak_date, trough_date, mdd = sip_max_drawdown(ledger["Portfolio_Value"])

    # Worst rolling ~24M XIRR: convert 24 months -> periods using actual deploy frequency for this N
    # Estimate periods/month using actual average deployments per month
    months = ledger.index.to_period("M")
    dep_per_month = months.value_counts().mean()
    roll_periods = int(round(24 * dep_per_month))
    roll_periods = max(6, roll_periods)

    worst = worst_rolling_xirr(ledger, roll_periods)
    worst_xirr = np.nan if worst is None else float(worst["XIRR"])

    # Time underwater
    underwater = (ledger["Portfolio_Value"] < ledger["Total_Invested"]).mean()

    return {
        "N_days": n,
        "Deployments": int(len(ledger)),
        "XIRR": float(strat_xirr),
        "Max_Drawdown": float(mdd),
        "DD_Peak": peak_date.date(),
        "DD_Trough": trough_date.date(),
        "Worst_Rolling_24M_XIRR": worst_xirr,
        "Underwater_%": float(underwater * 100),
        "Avg_Deployments_Per_Month": float(dep_per_month),
        "Rolling_24M_Periods_Used": int(roll_periods)
    }, ledger

# =========================================================
# OPTIMIZATION RUN
# =========================================================
results = []
best_ledgers = {}  # optionally store some ledgers

for n in range(N_MIN, N_MAX + 1):
    out = run_strategy_for_n(sip_prices, sip_rel_dev, SALARY_DATES, n)
    if out is None:
        continue
    metrics, ledger = out
    results.append(metrics)

    # Store a few ledgers if you want to inspect later
    if n in (2, 5, 10, 15, 20, 30):
        best_ledgers[n] = ledger

res_df = pd.DataFrame(results)
if res_df.empty:
    raise RuntimeError("No results computed. Check date ranges and tickers.")

# Sort by worst rolling 24M XIRR first (most important for SIP survivability),
# then by Max Drawdown (less negative is better), then by overall XIRR.
res_df_sorted = res_df.sort_values(
    by=["Worst_Rolling_24M_XIRR", "Max_Drawdown", "XIRR"],
    ascending=[False, False, False]
)

pd.set_option("display.width", 200)
pd.set_option("display.max_columns", 50)

print("\n=== OPTIMIZATION SUMMARY ===")
print(f"Salary credit: {SALARY_CREDIT} | Monthly SIP: ₹{MONTHLY_SIP:,}")
print(f"N range: {N_MIN} → {N_MAX} (trading days)")
print("\nTop 15 configs (ranked by Worst Rolling ~24M XIRR, then DD, then XIRR):")
print(res_df_sorted.head(15).to_string(index=False))

print("\nWorst 10 configs:")
print(res_df_sorted.tail(10).to_string(index=False))

# Export results
out_file = "SIPXAR_N_Optimization.xlsx"
with pd.ExcelWriter(out_file, engine="xlsxwriter") as writer:
    res_df_sorted.to_excel(writer, sheet_name="Optimization_Ranked", index=False)
    res_df.sort_values("N_days").to_excel(writer, sheet_name="Optimization_By_N", index=False)

print(f"\nExcel exported: {out_file}")

# =========================================================
# OPTIONAL: SAVE LEDGERS FOR SELECTED N VALUES
# =========================================================
# Uncomment if you want ledgers saved as separate sheets
# with pd.ExcelWriter("SIPXAR_Selected_Ledgers.xlsx", engine="xlsxwriter") as writer:
#     for n, led in best_ledgers.items():
#         led.to_excel(writer, sheet_name=f"N_{n}")
# print("Selected ledgers exported: SIPXAR_Selected_Ledgers.xlsx")
