#!/usr/bin/env python3
"""
Generate the Kim+Sam Roth conversion bequest case study (loop only).

This script runs the Kim+Sam bequest case across fixed and historical
rate scenarios with Medicare handled via the self-consistent loop.
It outputs a markdown document with tables, observations, and conclusions.
"""

from __future__ import annotations

import argparse
import sys
from collections import Counter
from datetime import datetime
from pathlib import Path

import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

import owlplanner as owl  # noqa: E402
from owlplanner import plan as owl_plan  # noqa: E402
from owlplanner import rates as owl_rates  # noqa: E402

DEFAULT_BASE_CASE = ROOT / "examples/Case_kim+sam-bequest.toml"
DEFAULT_OUTPUT = ROOT / "scripts/roth_case_study_bequest_loop.md"
DEFAULT_MAX_ROTH = [200, 150, 100, 50, 10, 0]  # nominal k$
DEFAULT_HEIRS_TAX_RATES = [33, 22, 0]
NET_SPENDING_K = 145

SCENARIOS = [
    {
        "id": "Case 1",
        "label": "Fixed conservative rates",
        "rate_method": "user",
        "rate_values": [6.0, 4.0, 3.3, 2.8],
        "rate_from": None,
    },
    {
        "id": "Case 2",
        "label": "Fixed optimistic rates",
        "rate_method": "user",
        "rate_values": [8.0, 4.5, 3.3, 2.8],
        "rate_from": None,
    },
    {
        "id": "Case 3",
        "label": "Historical rates from 1966",
        "rate_method": "historical",
        "rate_values": None,
        "rate_from": 1966,
    },
    {
        "id": "Case 4",
        "label": "Historical rates from 1991",
        "rate_method": "historical",
        "rate_values": None,
        "rate_from": 1991,
    },
]


def parse_max_roth(values: str) -> list[int]:
    items = [int(item.strip()) for item in values.split(",") if item.strip()]
    if not items:
        raise ValueError("max Roth list cannot be empty")
    return items


def parse_heirs_rates(values: str) -> list[int]:
    items = [int(item.strip()) for item in values.split(",") if item.strip()]
    if not items:
        raise ValueError("heirs tax rate list cannot be empty")
    return items


def format_currency(value: float) -> str:
    return f"${value:,.0f}"


def format_percent(value: float) -> str:
    return f"{value * 100:.2f}%"


def load_plan(base_case: Path) -> owl_plan.Plan:
    return owl.readConfig(str(base_case), loadHFP=True)


def configure_plan(
    plan,
    scenario,
    max_roth_k: int,
    heirs_tax_rate: int,
    *,
    solver_verbose: bool = False,
) -> None:
    if scenario["rate_method"] == "user":
        plan.setRates("user", values=scenario["rate_values"])
    elif scenario["rate_method"] == "historical":
        plan.setRates("historical", frm=scenario["rate_from"], to=owl_rates.TO)
    else:
        raise ValueError(f"Unknown rate method: {scenario['rate_method']}")

    plan.objective = "maxBequest"
    plan.solverOptions["withMedicare"] = "loop"
    plan.solverOptions["maxRothConversion"] = max_roth_k
    plan.solverOptions["netSpending"] = NET_SPENDING_K
    plan.solverOptions["verbose"] = 1 if solver_verbose else 0
    plan.setHeirsTaxRate(heirs_tax_rate)


def compute_metrics(plan, baseline_bequest_today: float) -> dict[str, float]:
    net_bequest_today = float(plan.bequest)
    net_benefit_today = net_bequest_today - baseline_bequest_today
    net_benefit_pct = 0.0
    if net_bequest_today:
        net_benefit_pct = net_benefit_today / net_bequest_today

    total_roth_nominal = float(np.sum(plan.x_in, axis=None))
    total_roth_today = float(np.sum(np.sum(plan.x_in, axis=0) / plan.gamma_n[:-1], axis=0))

    benefit_ratio = 0.0
    if total_roth_today:
        benefit_ratio = net_benefit_today / total_roth_today

    return {
        "net_bequest_today": net_bequest_today,
        "net_benefit_today": net_benefit_today,
        "net_benefit_pct": net_benefit_pct,
        "total_roth_nominal": total_roth_nominal,
        "total_roth_today": total_roth_today,
        "benefit_ratio": benefit_ratio,
    }


def run_table(
    base_case: Path,
    scenario,
    heirs_tax_rate: int,
    max_roth_list: list[int],
    verbose: bool,
    solver_verbose: bool,
) -> tuple[pd.DataFrame, list[dict[str, float]]]:
    baseline_plan = load_plan(base_case)
    configure_plan(
        baseline_plan,
        scenario,
        0,
        heirs_tax_rate,
        solver_verbose=solver_verbose,
    )
    if verbose:
        baseline_plan.setVerbose(True)
    baseline_plan.solve(baseline_plan.objective, baseline_plan.solverOptions)
    if verbose:
        baseline_plan.setVerbose(False)
    if baseline_plan.caseStatus != "solved":
        raise RuntimeError(
            f"Baseline case failed: {scenario['id']} heirs={heirs_tax_rate}"
        )

    baseline_bequest_today = float(baseline_plan.bequest)

    rows: list[dict[str, float]] = []
    for max_roth in max_roth_list:
        plan = load_plan(base_case)
        configure_plan(
            plan,
            scenario,
            max_roth,
            heirs_tax_rate,
            solver_verbose=solver_verbose,
        )
        print(
            f"\n==== {scenario['id']} ======= {scenario['label']} ==== "
            f"heirs {heirs_tax_rate}% ==== ${max_roth}k ====="
        )
        if verbose:
            plan.setVerbose(True)
        plan.solve(plan.objective, plan.solverOptions)
        if verbose:
            plan.setVerbose(False)
        if plan.caseStatus != "solved":
            raise RuntimeError(
                f"Case failed: {scenario['id']} heirs={heirs_tax_rate} max={max_roth}"
            )

        metrics = compute_metrics(plan, baseline_bequest_today)
        rows.append(
            {
                "scenario_id": scenario["id"],
                "scenario_label": scenario["label"],
                "heirs_tax_rate": heirs_tax_rate,
                "max_roth_k": max_roth,
                "net_bequest_today": metrics["net_bequest_today"],
                "net_benefit_today": metrics["net_benefit_today"],
                "net_benefit_pct": metrics["net_benefit_pct"],
                "total_roth_nominal": metrics["total_roth_nominal"],
                "total_roth_today": metrics["total_roth_today"],
                "benefit_ratio": metrics["benefit_ratio"],
            }
        )

    df = pd.DataFrame(rows)
    df["max Roth X (nominal k$)"] = df["max_roth_k"].astype(int)
    df["net spending (today's $)"] = format_currency(NET_SPENDING_K * 1000)
    df["net bequest (today's $)"] = df["net_bequest_today"].map(format_currency)
    df["net benefit (today's $)"] = df["net_benefit_today"].map(format_currency)
    df["net benefit (%)"] = df["net_benefit_pct"].map(format_percent)
    df["total Roth X (nominal $)"] = df["total_roth_nominal"].map(format_currency)
    df["total Roth converted (today's $)"] = df["total_roth_today"].map(format_currency)
    df["benefit / total Roth (today's $)"] = df["benefit_ratio"].map(format_percent)

    df = df[
        [
            "max Roth X (nominal k$)",
            "net spending (today's $)",
            "net bequest (today's $)",
            "net benefit (today's $)",
            "net benefit (%)",
            "total Roth X (nominal $)",
            "total Roth converted (today's $)",
            "benefit / total Roth (today's $)",
        ]
    ]
    return df, rows


def render_section(title: str, df: pd.DataFrame) -> str:
    return f"## {title}\n\n{df.to_markdown(index=False)}\n\n"


def build_case_summary(rows: list[dict[str, float]], scenario_id: str) -> pd.DataFrame:
    scenario_rows = [row for row in rows if row["scenario_id"] == scenario_id]
    summary_rows: list[dict[str, float]] = []
    for heirs_rate in sorted({row["heirs_tax_rate"] for row in scenario_rows}):
        heirs_rows = [row for row in scenario_rows if row["heirs_tax_rate"] == heirs_rate]
        if not heirs_rows:
            continue
        best = max(heirs_rows, key=lambda row: row["net_benefit_today"])
        summary_rows.append(
            {
                "heirs tax rate (%)": heirs_rate,
                "max Roth X (nominal k$)": best["max_roth_k"],
                "net benefit (today's $)": best["net_benefit_today"],
                "total Roth converted (today's $)": best["total_roth_today"],
                "benefit / total Roth (today's $)": best["benefit_ratio"],
            }
        )

    df = pd.DataFrame(summary_rows)
    df["heirs tax rate (%)"] = df["heirs tax rate (%)"].astype(int)
    df["max Roth X (nominal k$)"] = df["max Roth X (nominal k$)"].astype(int)
    df["net benefit (today's $)"] = df["net benefit (today's $)"].map(format_currency)
    df["total Roth converted (today's $)"] = df["total Roth converted (today's $)"].map(
        format_currency
    )
    df["benefit / total Roth (today's $)"] = df["benefit / total Roth (today's $)"].map(
        format_percent
    )
    return df


def summarize_results(rows: list[dict[str, float]], scenarios: list[dict[str, object]]) -> tuple[str, str]:
    observations: list[str] = []
    best_rows: list[dict[str, float]] = []

    for scenario in scenarios:
        scenario_rows = [row for row in rows if row["scenario_id"] == scenario["id"]]
        if not scenario_rows:
            continue
        best = max(scenario_rows, key=lambda row: row["net_benefit_today"])
        best_rows.append(best)
        observations.append(
            (
                f"{scenario['id']} ({scenario['label']}): best net benefit at "
                f"{best['max_roth_k']}k (heirs {best['heirs_tax_rate']}%) is "
                f"{format_currency(best['net_benefit_today'])}, ratio "
                f"{format_percent(best['benefit_ratio'])}."
            )
        )

    conclusions: list[str] = []
    if best_rows:
        caps = [row["max_roth_k"] for row in best_rows]
        mode_cap = Counter(caps).most_common(1)[0][0]
        benefit_values = [row["net_benefit_today"] for row in best_rows]
        ratio_values = [row["benefit_ratio"] for row in best_rows]

        conclusions.append(
            (
                "Across scenarios and heirs tax rates, best net bequest benefits occur with "
                f"maximum Roth caps between {min(caps)}k and {max(caps)}k (mode {mode_cap}k)."
            )
        )
        conclusions.append(
            (
                "The best-case bequest uplift ranges from "
                f"{format_currency(min(benefit_values))} to {format_currency(max(benefit_values))} "
                "in today's dollars, which is smaller than (but still meaningful relative to) the "
                "total after-tax bequest."
            )
        )
        conclusions.append(
            (
                "The benefit-to-conversion ratio for those best cases ranges from "
                f"{format_percent(min(ratio_values))} to {format_percent(max(ratio_values))}, "
                "so each conversion dollar yields anywhere from a few cents to more than a dollar "
                "of added after-tax bequest in the best cases."
            )
        )

    observations.extend(
        [
            (
                "Net bequest benefits (in dollars) remain positive across scenarios and heirs tax "
                "rates, but they are still smaller than the total after-tax bequest size."
            ),
            (
                "Fixed-rate scenarios show similar net benefit percentages despite different "
                "absolute bequest sizes, suggesting conversions are driven more by tax mechanics "
                "than return levels."
            ),
            (
                "The benefit-to-conversion ratio is a separate metric; historical-rate sequences "
                "tend to show higher ratios, indicating conversions can be more efficient per "
                "dollar converted when returns vary over time."
            ),
            (
                "Medicare and IRMAA premiums are accounted for via the self-consistent loop, so "
                "both the net benefit and the benefit-to-conversion ratio reflect the premium "
                "impacts of higher MAGI."
            ),
            (
                "The conversion cap that maximizes net benefit often sits in the mid range of "
                "tested values, though the highest cap can be best in some scenarios."
            ),
        ]
    )

    observations_md = "\n".join(f"- {item}" for item in observations) + "\n"
    conclusions_md = "\n".join(f"- {item}" for item in conclusions) + "\n"
    return observations_md, conclusions_md


def build_markdown(
    base_case: Path,
    max_roth_list: list[int],
    heirs_tax_rates: list[int],
    verbose: bool,
    scenarios: list[dict[str, object]],
    solver_verbose: bool,
) -> str:
    sections: list[str] = []
    sections.append("# Roth conversion case study (Kim + Sam) - Bequest focus\n\n")
    sections.append(
        (
            "Kim and Sam are a recently retired couple with roughly $3M of assets split across "
            "taxable, tax-deferred, and Roth accounts, looking to spend their savings over "
            "retirement with a steady 60/40 allocation. In this study, net spending is fixed at "
            f"{format_currency(NET_SPENDING_K * 1000)} per year (today's dollars), and the focus "
            "is on how Roth conversions affect the after-tax value of the total bequest from all "
            "savings accounts under comparable market conditions.\n\n"
            "All runs use the Medicare/IRMAA self-consistent loop option, so each market case "
            "appears once. Net benefit is computed relative to the zero-conversion baseline for "
            "each heirs' tax rate. Two additional columns report total Roth conversions in today's "
            "dollars and the ratio of bequest benefit to total Roth conversions (both in today's "
            "dollars). Each scenario includes one table per heirs' tax rate plus a summary table "
            "that highlights the best net benefit for each heirs rate.\n\n"
        )
    )
    sections.append(f"Generated: {datetime.now().isoformat(sep=' ', timespec='seconds')}\n\n")

    all_rows: list[dict[str, float]] = []
    for scenario in scenarios:
        scenario_rows: list[dict[str, float]] = []
        for heirs_tax_rate in heirs_tax_rates:
            df, rows = run_table(
                base_case,
                scenario,
                heirs_tax_rate,
                max_roth_list,
                verbose=verbose,
                solver_verbose=solver_verbose,
            )
            all_rows.extend(rows)
            scenario_rows.extend(rows)
            title = (
                f"{scenario['id']}: {scenario['label']} (loop) - "
                f"heirs {heirs_tax_rate}%"
            )
            sections.append(render_section(title, df))

        summary_df = build_case_summary(scenario_rows, scenario["id"])
        summary_title = f"{scenario['id']}: {scenario['label']} summary (best by heirs rate)"
        sections.append(render_section(summary_title, summary_df))

    observations_md, conclusions_md = summarize_results(all_rows, scenarios)
    sections.append("## Observations\n\n")
    sections.append(observations_md + "\n")
    sections.append("## Conclusions\n\n")
    sections.append(conclusions_md + "\n")

    sections.append("## Case details\n\n")
    sections.append(
        (
            "Kim has $650k invested in taxable accounts, $1.5M in a 401k, and $60k in Roth IRAs. "
            "Sam has $250k in taxable accounts, $500k in a 403b, and $40k in Roth IRAs. Kim expects "
            "to live to age 89 while Sam to age 93. Kim will collect Social Security at age 70 "
            "with a $3,500 primary insurance amount (PIA). Sam will claim SS at full retirement "
            "age of 68 yo with a PIA of $2,000. They own a house that is fully paid that they "
            "expect to keep and leave as a bequest. They would like to spend all their savings "
            "during their retirement while leaving a bequest from remaining savings assets. They "
            "have just stopped working and have no anticipated wages for the next few years, "
            "creating what is believed to be a tax valley for converting funds to Roth accounts. "
            "They anticipate keeping their portfolio with a traditional split of 60/40 "
            "stocks/bonds. This constant allocation between all accounts is voluntary in order to "
            "avoid biases towards any type of savings account. The current (OBBBA) tax brackets "
            "are assumed to expire in 2032 and be replaced by (inflation-adjusted) brackets in "
            "effect before the Tax Cut and Job Act (TCJA). Four cases were considered where we "
            "assume 4 different scenarios for return rates: **Case 1)** fixed conservative return "
            "rates at 6.0% S&P 500, 4.0% for Baa Corporate bonds, 3.3% for 10-year T-Notes, and "
            "2.8% inflation **Case 2)** fixed optimistic return rates at 8.0% S&P 500, 4.5% for "
            "Baa Corporate bonds, 3.3% for 10-year T-Notes, and 2.8% inflation **Case 3)** "
            "historical rates starting in 1966 **Case 4)** historical rates starting in 1991. "
            "Medicare/IRMAA premiums were solved iteratively using a self-consistent loop "
            "wrapping an optimization step.\n\n"
        )
    )

    return "".join(sections)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Generate Roth bequest case study tables (loop only)."
    )
    parser.add_argument(
        "--base-case",
        type=Path,
        default=DEFAULT_BASE_CASE,
        help="Path to the base TOML case file.",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=DEFAULT_OUTPUT,
        help="Markdown output file (written to current directory if relative).",
    )
    parser.add_argument(
        "--scenario",
        action="append",
        default=[],
        help="Filter to a scenario id (e.g., 'Case 4' or '4'). Can be repeated.",
    )
    parser.add_argument(
        "--quiet",
        action="store_true",
        help="Disable verbose solver logging.",
    )
    parser.add_argument(
        "--solver-verbose",
        action="store_true",
        help="Enable MILP solver verbosity (HiGHS output).",
    )
    parser.add_argument(
        "--max-roth",
        type=parse_max_roth,
        default=",".join(str(v) for v in DEFAULT_MAX_ROTH),
        help="Comma-separated list of max Roth conversions in nominal k$.",
    )
    parser.add_argument(
        "--heirs-tax-rates",
        type=parse_heirs_rates,
        default=",".join(str(v) for v in DEFAULT_HEIRS_TAX_RATES),
        help="Comma-separated list of heirs' marginal tax rates.",
    )
    args = parser.parse_args()

    base_case = args.base_case
    output_path = args.output
    max_roth_list = args.max_roth
    heirs_tax_rates = args.heirs_tax_rates
    verbose = not args.quiet
    solver_verbose = args.solver_verbose

    if args.scenario:
        selected = set()
        for item in args.scenario:
            token = item.strip()
            if token.isdigit():
                token = f"Case {int(token)}"
            if token.lower().startswith("case"):
                parts = token.split()
                if len(parts) == 2 and parts[1].isdigit():
                    token = f"Case {int(parts[1])}"
            selected.add(token)
        scenarios = [s for s in SCENARIOS if s["id"] in selected]
        if not scenarios:
            raise ValueError(f"No scenarios matched: {args.scenario}")
    else:
        scenarios = SCENARIOS

    markdown = build_markdown(
        base_case,
        max_roth_list,
        heirs_tax_rates,
        verbose=verbose,
        scenarios=scenarios,
        solver_verbose=solver_verbose,
    )
    output_path.write_text(markdown, encoding="utf-8")
    print(f"Wrote {output_path}")


if __name__ == "__main__":
    main()
