#!/usr/bin/env python3
"""
Retirement trajectory projection for Dane.

Models portfolio growth from age 25 -> 65 under three post-PhD savings
scenarios, plus shows how negligible the one-time summer 401k contribution is.

All figures are in TODAY'S DOLLARS (real returns), so they're directly
comparable to current purchasing power. Edit the ASSUMPTIONS block and re-run.

    /tmp/retire-venv/bin/python retirement_projection.py
"""

import matplotlib
matplotlib.use("Agg")  # no display needed; save to file
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter

# ----------------------------------------------------------------------------
# ASSUMPTIONS  (edit these and re-run)
# ----------------------------------------------------------------------------
REAL_RETURN      = 0.07     # 7% real (inflation-adjusted) annual return
SWR              = 0.04     # 4% safe withdrawal rate
START_AGE        = 25
END_AGE          = 65
START_PORTFOLIO  = 48_000   # Roth IRA ($16k) + Schwab Stocks ($32k) today

# 529 -> Roth IRA pipeline (grandparent-funded). $24.9k total, $7k/yr cap.
PIPELINE = {25: 7_000, 26: 7_000, 27: 7_000, 28: 3_900}   # sums to $24,900

# Post-PhD own savings (401k + brokerage). PhD done end of 2027 => income 2028.
POST_PHD_START_AGE = 27     # 2028 = age 27
SCENARIOS = {
    "Floor — no post-PhD saving":  0,
    "Moderate — $20k/yr post-PhD": 20_000,
    "Aggressive — $35k/yr post-PhD": 35_000,
}

# The contested decision: one summer of 10% Roth 401k = ~$2,240 at age 25.
SUMMER_401K = 2_240

# ----------------------------------------------------------------------------
# MODEL
# ----------------------------------------------------------------------------
def project(post_phd_rate, summer_401k=0):
    """Year-by-year: grow existing balance, then add year-end contributions.
    Returns (ages, balances)."""
    ages, balances = [], []
    bal = START_PORTFOLIO + summer_401k  # summer contribution lands at age 25
    for age in range(START_AGE, END_AGE + 1):
        ages.append(age)
        balances.append(bal)
        # contribution for THIS year (added at year-end, after growth)
        contrib = PIPELINE.get(age, 0)
        if age >= POST_PHD_START_AGE:
            contrib += post_phd_rate
        bal = bal * (1 + REAL_RETURN) + contrib
    return ages, balances

# ----------------------------------------------------------------------------
# PLOT
# ----------------------------------------------------------------------------
fig, ax = plt.subplots(figsize=(12, 7.5))
colors = {"Floor — no post-PhD saving": "#9aa0a6",
          "Moderate — $20k/yr post-PhD": "#1a73e8",
          "Aggressive — $35k/yr post-PhD": "#188038"}

finals = {}
for label, rate in SCENARIOS.items():
    ages, bals = project(rate)
    finals[label] = bals[-1]
    ax.plot(ages, bals, label=label, linewidth=2.6, color=colors[label])
    # final value annotation
    final = bals[-1]
    income = final * SWR
    ax.annotate(f"  ${final/1e6:.2f}M\n  (${income/1e3:.0f}k/yr @ 4%)",
                xy=(65, final), xytext=(65.4, final),
                va="center", fontsize=10, fontweight="bold",
                color=colors[label])

# The summer 401k question, framed HONESTLY.
# $2,240 grows to the same gross amount in ANY account. The real question is
# Roth-401k vs taxable-brokerage (Dane's actual alternative — he'd save it, not
# spend it). The only delta is the tax treatment of the growth.
LTCG_RATE = 0.15
summer_gross = SUMMER_401K * (1 + REAL_RETURN) ** (END_AGE - START_AGE)  # grows either way
summer_gain  = summer_gross - SUMMER_401K
summer_taxable_net = summer_gross - summer_gain * LTCG_RATE   # taxable: pay LTCG on gain
roth_advantage = summer_gross - summer_taxable_net            # Roth keeps the tax

# vertical marker: PhD done / income starts
ax.axvline(POST_PHD_START_AGE, color="#d93025", linestyle="--", alpha=0.55, linewidth=1.4)
ax.annotate("PhD done /\nincome starts\n(~2028)", xy=(POST_PHD_START_AGE, ax.get_ylim()[1]*0.0),
            xytext=(POST_PHD_START_AGE + 0.3, finals["Aggressive — $35k/yr post-PhD"]*0.62),
            fontsize=9, color="#d93025")

# annotation box: the summer 401k decision, framed honestly
txt = (f"The contested decision (summer 10% Roth 401k, ~${SUMMER_401K:,}):\n"
       f"• Grows to ~${summer_gross/1e3:.0f}k by 65 in ANY account\n"
       f"• vs taxable brokerage (your real alt.), the Roth tax\n"
       f"   benefit is only ~${roth_advantage/1e3:.1f}k — the rest happens\n"
       f"   whether it's Roth or not.\n"
       f"• On a ${finals['Aggressive — $35k/yr post-PhD']/1e6:.1f}M trajectory: a rounding error.")
ax.text(0.03, 0.97, txt, transform=ax.transAxes, fontsize=9.5,
        va="top", ha="left",
        bbox=dict(boxstyle="round,pad=0.5", facecolor="#fef7e0", edgecolor="#f9ab00"))

# formatting
ax.set_title("Dane's Retirement Trajectory — today's dollars, 7% real return",
             fontsize=14, fontweight="bold", pad=14)
ax.set_xlabel("Age", fontsize=11)
ax.set_ylabel("Portfolio value (today's $)", fontsize=11)
ax.yaxis.set_major_formatter(FuncFormatter(lambda v, _: f"${v/1e6:.1f}M"))
ax.set_xlim(START_AGE, 69.5)
ax.set_ylim(bottom=0)
ax.grid(True, alpha=0.25)
ax.legend(loc="center left", bbox_to_anchor=(0.03, 0.55), fontsize=10, framealpha=0.9)

fig.tight_layout()
out = "/Users/danesabo/Documents/Finances/retirement_projection.png"
fig.savefig(out, dpi=140, bbox_inches="tight")

# ----------------------------------------------------------------------------
# TEXT SUMMARY
# ----------------------------------------------------------------------------
print("=" * 64)
print("RETIREMENT PROJECTION  (today's dollars, 7% real return)")
print("=" * 64)
print(f"Start: ${START_PORTFOLIO:,} at age {START_AGE}")
print(f"529->Roth pipeline: ${sum(PIPELINE.values()):,} over ages "
      f"{min(PIPELINE)}-{max(PIPELINE)}")
print(f"Post-PhD saving starts at age {POST_PHD_START_AGE}\n")
print(f"{'Scenario':<32}{'Age 65':>12}{'Income @4%':>14}")
print("-" * 58)
for label, final in finals.items():
    print(f"{label:<32}{'$'+format(final/1e6, '.2f')+'M':>12}"
          f"{'$'+format(final*SWR/1e3, '.0f')+'k/yr':>14}")
print("-" * 58)
print(f"\nSummer 10% Roth 401k (~${SUMMER_401K:,} one-time at age 25):")
print(f"  grows to ~${summer_gross:,.0f} by age 65 in ANY account")
print(f"  Roth-vs-taxable tax benefit only: ~${roth_advantage:,.0f}")
print(f"  (= {roth_advantage/finals['Aggressive — $35k/yr post-PhD']*100:.3f}% "
      f"of the aggressive total — a rounding error)")
print(f"\nChart saved: {out}")