#!/usr/bin/env python3
"""
Analyze final portfolio balance under inflation-adjusted withdrawals.

$1M initial, constant annuity (real 2026$) withdrawn each year over 30 years.
60/40 stocks/bonds, 30-year rolling horizon, augmented by rolling and inverting.
Histogram shows amount left at end of 30 years.

Reuses Owl's rates data and roll/reverse logic. No Owl code is modified.

Run from project root:
  python scripts/cumulative_returns_analysis.py [--annuity 50000] [--initial 1000000] [--show]
"""

import argparse
import sys
import os

# Allow importing owlplanner from project
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))

import numpy as np  # noqa: E402
import pandas as pd  # noqa: E402
from itertools import product  # noqa: E402

# Owl rates: Stocks, Bonds, Inflation
from owlplanner.rates import (  # noqa: E402
    SP500,
    BondsBaa,
    Inflation,
    FROM,
    TO,
    apply_rate_sequence_transform,
)

# Plotting: matplotlib (Owl's standard) for histograms + fits
import matplotlib  # noqa: E402
from scipy import optimize  # noqa: E402
from scipy.stats import lognorm  # noqa: E402

# Parameters
HORIZON_YEARS = 30
STOCK_RATIO = 0.60
BOND_RATIO = 0.40
DEFAULT_INITIAL = 1_000_000
DEFAULT_ANNUITY = 50_000
OUTPUT_CSV = "cumulative_returns_results.csv"
OUTPUT_PNG = "cumulative_returns_histogram.png"


def simulate_final_balance(stock_pct, bond_pct, inflation_pct, initial, annuity_real):
    """
    Simulate 30-year portfolio with inflation-adjusted withdrawals.
    End of each year: grow by (1+r), then withdraw annuity_real * cum_inflation.
    Returns final balance (amount left at end of year 29).
    Rates in percent; annuity_real in dollars (constant real purchasing power).
    """
    stock_dec = np.asarray(stock_pct, dtype=float) / 100.0
    bond_dec = np.asarray(bond_pct, dtype=float) / 100.0
    inf_dec = np.asarray(inflation_pct, dtype=float) / 100.0
    portfolio_rates = STOCK_RATIO * stock_dec + BOND_RATIO * bond_dec

    balance = float(initial)
    cum_inflation = 1.0
    for n in range(HORIZON_YEARS):
        balance = balance * (1.0 + portfolio_rates[n])
        withdrawal = annuity_real * cum_inflation
        balance = max(0.0, balance - withdrawal)
        cum_inflation *= 1.0 + inf_dec[n]
    return balance


def run_analysis(show_plot=False, initial=DEFAULT_INITIAL, annuity=DEFAULT_ANNUITY):
    """Run full analysis: augment data, simulate withdrawals, fit two Gaussians, plot."""
    if not show_plot:
        matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    stock = np.asarray(SP500, dtype=float)
    bond = np.asarray(BondsBaa, dtype=float)
    inflation = np.asarray(Inflation, dtype=float)
    n_data = len(stock)

    # Valid starting years: need 30 consecutive years
    y_start_min = FROM
    y_end_max = TO - HORIZON_YEARS + 1
    start_years = list(range(y_start_min, y_end_max + 1))

    # Augmented: (reverse, roll) in {0,1} x {0..29}, matching Owl runHistoricalRange
    reverse_roll_pairs = list(product([False, True], range(HORIZON_YEARS)))

    rows = []
    for year in start_years:
        i0 = year - FROM
        if i0 + HORIZON_YEARS > n_data:
            continue
        # Extract 30-year sequences: stocks, bonds, inflation
        tau_kn = np.stack([
            stock[i0 : i0 + HORIZON_YEARS],
            bond[i0 : i0 + HORIZON_YEARS],
            inflation[i0 : i0 + HORIZON_YEARS],
        ], axis=0)

        for reverse, roll in reverse_roll_pairs:
            tau_transformed = apply_rate_sequence_transform(tau_kn.copy(), reverse, roll)
            final_bal = simulate_final_balance(
                tau_transformed[0, :],
                tau_transformed[1, :],
                tau_transformed[2, :],
                initial,
                annuity,
            )
            rows.append({
                "start_year": year,
                "roll": roll,
                "inversion": 1 if reverse else 0,
                "final_balance": final_bal,
            })

    df = pd.DataFrame(rows)

    # Write results to CSV
    df.to_csv(OUTPUT_CSV, index=False)
    print(f"Wrote {len(df)} rows to {OUTPUT_CSV} (initial=${initial:,.0f}, annuity=${annuity:,.0f}/yr real)")

    # Two log-normal mixture fit (positive values only)
    x = df["final_balance"].values
    fit_result = fit_two_lognormals(x)

    # Plot
    fig, axes = plt.subplots(2, 1, figsize=(9, 10), sharex=True)

    # Top: histogram with fitted log-normals
    ax1 = axes[0]
    n_bins = min(60, max(20, len(x) // 30))
    counts, bin_edges, patches = ax1.hist(
        x, bins=n_bins, density=True, alpha=0.6, color="steelblue",
        edgecolor="white", linewidth=0.5, label="Data"
    )
    bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2
    x_min_plot = x[x > 0].min() if np.any(x > 0) else 1.0
    x_plot = np.linspace(max(0.1, x_min_plot * 0.5), x.max(), 500)

    if fit_result["success"]:
        p = fit_result["params"]
        ln1 = lognorm.pdf(x_plot, p["s1"], scale=np.exp(p["mu1"]))
        ln2 = lognorm.pdf(x_plot, p["s2"], scale=np.exp(p["mu2"]))
        mix1 = p["w1"] * ln1
        mix2 = (1 - p["w1"]) * ln2
        ax1.plot(x_plot, mix1, "b-", alpha=0.8, linewidth=1.5,
                 label=f"LogNorm 1 (μ={p['mu1']:.2f}, σ={p['s1']:.2f})")
        ax1.plot(x_plot, mix2, "g-", alpha=0.8, linewidth=1.5,
                 label=f"LogNorm 2 (μ={p['mu2']:.2f}, σ={p['s2']:.2f})")
        ax1.plot(x_plot, mix1 + mix2, "k--", alpha=0.9, linewidth=2, label="Mixture")
    else:
        ax1.text(0.5, 0.9, "Two log-normal fit did not converge", transform=ax1.transAxes, ha="center")

    ax1.set_ylabel("Density")
    ax1.set_title(
        f"Final Balance After 30 Years (60/40) – Augmented Historical\n"
        f"${initial/1e6:.1f}M initial, ${annuity/1e3:.0f}k/yr real · n={len(x)} scenarios"
    )
    ax1.legend(loc="upper right", fontsize=8)
    ax1.grid(True, alpha=0.3)

    # Bottom: residuals (data histogram - fit) if fit succeeded
    ax2 = axes[1]
    if fit_result["success"]:
        p = fit_result["params"]
        # Only evaluate fit at positive bin centers (log-normal undefined at 0)
        bc_pos = np.maximum(bin_centers, 0.01)
        fit_density = (
            p["w1"] * lognorm.pdf(bc_pos, p["s1"], scale=np.exp(p["mu1"]))
            + (1 - p["w1"]) * lognorm.pdf(bc_pos, p["s2"], scale=np.exp(p["mu2"]))
        )
        bin_width = bin_edges[1] - bin_edges[0]
        # Histogram density at bin centers (approximate)
        data_density = counts
        residuals = data_density - fit_density
        ax2.bar(bin_centers, residuals, width=bin_width * 0.9, alpha=0.6, color="orange", label="Residuals")
        ax2.axhline(0, color="black", linewidth=0.8)
        ax2.set_ylabel("Residual (data - fit)")
        ax2.set_title("Fit residuals")
    else:
        ax2.text(0.5, 0.5, "No fit to show residuals", transform=ax2.transAxes, ha="center")
    ax2.set_xlabel("Final balance ($)")
    ax2.legend()
    ax2.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(OUTPUT_PNG, dpi=150, bbox_inches="tight")
    print(f"Saved {OUTPUT_PNG}")
    if show_plot:
        plt.show()

    # Print fit summary
    if fit_result["success"]:
        p = fit_result["params"]
        med1, med2 = np.exp(p["mu1"]), np.exp(p["mu2"])
        print("\nTwo log-normal fit (μ, σ = params of log(X)):")
        print(f"  LogNorm 1: w={p['w1']:.3f}, μ={p['mu1']:.2f}, σ={p['s1']:.2f} → median=${med1:,.0f}")
        print(f"  LogNorm 2: w={1-p['w1']:.3f}, μ={p['mu2']:.2f}, σ={p['s2']:.2f} → median=${med2:,.0f}")
        print(f"  Negative log-likelihood: {fit_result['nll']:.2f}")
    else:
        print("\nTwo log-normal fit did not converge. Try different initial guesses or check data.")


def _neg_loglik_lognorm_mix(params, x):
    """Negative log-likelihood for two log-normal mixture. x must be strictly positive."""
    w1, mu1, log_s1, mu2, log_s2 = params
    w1 = max(1e-6, min(1 - 1e-6, w1))
    s1 = np.exp(log_s1)
    s2 = np.exp(log_s2)
    p = w1 * lognorm.pdf(x, s1, scale=np.exp(mu1)) + (1 - w1) * lognorm.pdf(x, s2, scale=np.exp(mu2))
    p = np.clip(p, 1e-300, None)
    return -np.sum(np.log(p))


def fit_two_lognormals(x):
    """
    Fit a mixture of two log-normals to data x.
    Uses only positive values (log-normal requires x > 0).
    Returns dict with success, params {w1, mu1, s1, mu2, s2}, nll.
    """
    x = np.asarray(x, dtype=float)
    x_pos = x[x > 0]
    n = len(x_pos)
    if n < 10:
        return {"success": False, "params": None, "nll": np.nan}

    # Initial guess from log-space: fit mixture of Gaussians to log(x)
    log_x = np.log(x_pos)
    threshold = np.percentile(log_x, 40)
    low = log_x[log_x <= threshold]
    high = log_x[log_x > threshold]
    if len(low) == 0:
        low = log_x[log_x <= np.median(log_x)]
    if len(high) == 0:
        high = log_x[log_x > np.median(log_x)]
    mu1_init = np.mean(low)
    mu2_init = np.mean(high)
    s1_init = max(np.std(low), 0.1)
    s2_init = max(np.std(high), 0.1)
    w1_init = len(low) / n

    log_x_min, log_x_max = log_x.min(), log_x.max()
    bounds = [
        (1e-6, 1 - 1e-6),
        (log_x_min - 2, log_x_max + 2),
        (-5, 3),
        (log_x_min - 2, log_x_max + 2),
        (-5, 3),
    ]
    x0 = [w1_init, mu1_init, np.log(max(s1_init, 0.01)), mu2_init, np.log(max(s2_init, 0.01))]

    try:
        res = optimize.minimize(
            _neg_loglik_lognorm_mix,
            x0,
            args=(x_pos,),
            method="L-BFGS-B",
            bounds=bounds,
            options={"maxiter": 2000},
        )
        if not res.success:
            return {"success": False, "params": None, "nll": res.fun}

        w1, mu1, log_s1, mu2, log_s2 = res.x
        return {
            "success": True,
            "params": {
                "w1": float(np.clip(w1, 1e-6, 1 - 1e-6)),
                "mu1": float(mu1),
                "s1": float(np.exp(log_s1)),
                "mu2": float(mu2),
                "s2": float(np.exp(log_s2)),
            },
            "nll": float(res.fun),
        }
    except Exception:
        return {"success": False, "params": None, "nll": np.nan}


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Analyze final balance after 30 years with inflation-adjusted withdrawals."
    )
    parser.add_argument(
        "--annuity", type=float, default=DEFAULT_ANNUITY,
        help=f"Annual withdrawal in real dollars (default: {DEFAULT_ANNUITY:,.0f})",
    )
    parser.add_argument(
        "--initial", type=float, default=DEFAULT_INITIAL,
        help=f"Initial portfolio amount (default: {DEFAULT_INITIAL:,.0f})",
    )
    parser.add_argument("--show", action="store_true", help="Show the plot interactively")
    args = parser.parse_args()
    run_analysis(show_plot=args.show, initial=args.initial, annuity=args.annuity)
