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

This script runs the Kim+Sam spending 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 rates as owl_rates  # noqa: E402

DEFAULT_BASE_CASE = ROOT / "examples/Case_kim+sam-spending.toml"
DEFAULT_OUTPUT = ROOT / "scripts/roth_case_study_loop.md"
DEFAULT_MAX_ROTH = [200, 150, 100, 50, 10, 0]  # nominal k$

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 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, *, 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 = "maxSpending"
    plan.solverOptions["withMedicare"] = "loop"
    plan.solverOptions["maxRothConversion"] = max_roth_k
    plan.solverOptions["verbose"] = 1 if solver_verbose else 0


def compute_metrics(plan, baseline_total_spending_today: float) -> dict[str, float]:
    net_spending_today = plan.g_n[0] / plan.gamma_n[0]
    total_net_spending_today = float(np.sum(plan.g_n / plan.gamma_n[:-1], axis=0))
    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))

    net_benefit_today = total_net_spending_today - baseline_total_spending_today
    net_benefit_pct = 0.0
    if total_net_spending_today:
        net_benefit_pct = net_benefit_today / total_net_spending_today

    benefit_ratio = 0.0
    if total_roth_today:
        benefit_ratio = net_benefit_today / total_roth_today

    return {
        "net_spending_today": net_spending_today,
        "total_net_spending_today": total_net_spending_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,
    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, 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']}")

    baseline_total_spending_today = float(
        np.sum(baseline_plan.g_n / baseline_plan.gamma_n[:-1], axis=0)
    )

    rows: list[dict[str, float]] = []
    for max_roth in max_roth_list:
        plan = load_plan(base_case)
        configure_plan(plan, scenario, max_roth, solver_verbose=solver_verbose)
        print(f"\n==== {scenario['id']} ======= {scenario['label']} ==== ${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']} max={max_roth}")

        metrics = compute_metrics(plan, baseline_total_spending_today)
        rows.append(
            {
                "scenario_id": scenario["id"],
                "scenario_label": scenario["label"],
                "max_roth_k": max_roth,
                "net_spending_today": metrics["net_spending_today"],
                "total_net_spending_today": metrics["total_net_spending_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 $)"] = df["net_spending_today"].map(format_currency)
    df["total net spending (today's $)"] = df["total_net_spending_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 $)",
            "total net spending (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 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 is {format_currency(best['net_benefit_today'])}, "
                f"ratio {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, best net benefits occur with maximum Roth caps "
                f"between {min(caps)}k and {max(caps)}k (mode {mode_cap}k)."
            )
        )
        conclusions.append(
            (
                "The best-case net spending uplift ranges from "
                f"{format_currency(min(benefit_values))} to {format_currency(max(benefit_values))} "
                "in today's dollars, while benefits remain on the order of ~1% of total spending."
            )
        )
        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 only a few cents of added spending."
            )
        )

    observations.extend(
        [
            (
                "Net spending benefits (in dollars) remain positive across all scenarios, but the "
                "magnitude is modest relative to total lifetime spending."
            ),
            (
                "Fixed-rate scenarios show similar net benefit percentages despite different "
                "absolute spending levels, suggesting conversions are driven more by tax mechanics "
                "than return levels."
            ),
            (
                "The benefit-to-conversion ratio is a separate metric; historical-rate sequences "
                "show slightly higher ratios, indicating that conversions during weaker return "
                "periods can be incrementally more efficient per dollar converted."
            ),
            (
                "Medicare and IRMAA premiums are accounted for via the self-consistent loop, so "
                "both the net benefit and the benefit-to-conversion ratio already 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],
    verbose: bool,
    scenarios: list[dict[str, object]],
    solver_verbose: bool,
) -> str:
    sections: list[str] = []
    sections.append("# Roth conversion case study (Kim + Sam)\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. The goal here is to quantify how the "
            "maximum allowed Roth conversions change available net spending under comparable "
            "market conditions, while keeping the rest of the plan assumptions fixed.\n\n"
            "All runs use the Medicare/IRMAA self-consistent loop option, so each market case "
            "appears once. Net spending is an annual value (today's dollars), totals are over the "
            "full plan, and net benefit is computed relative to the zero-conversion baseline. Two "
            "additional columns report total Roth conversions in today's dollars and the ratio of "
            "net spending benefit to total Roth conversions (both in today's dollars).\n\n"
        )
    )
    sections.append(f"Generated: {datetime.now().isoformat(sep=' ', timespec='seconds')}\n\n")

    all_rows: list[dict[str, float]] = []
    for scenario in scenarios:
        df, rows = run_table(
            base_case,
            scenario,
            max_roth_list,
            verbose=verbose,
            solver_verbose=solver_verbose,
        )
        all_rows.extend(rows)
        title = f"{scenario['id']}: {scenario['label']} (loop)"
        sections.append(render_section(title, 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")

    return "".join(sections)


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate Roth 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$.",
    )
    args = parser.parse_args()

    base_case = args.base_case
    output_path = args.output
    max_roth_list = args.max_roth
    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,
        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()
