"""Reproducible fiscal/Fund stress tests. All invented parameters are scenarios.
Run: python3 simulations/run_simulations.py (requires NumPy).
Dollar values are trillions; rates are decimals; year 0 is normalized, not actual.
"""
from dataclasses import dataclass, asdict, replace
from pathlib import Path
import csv, itertools, json, math
import numpy as np

ROOT=Path(__file__).resolve().parent
OUT=ROOT/'results';OUT.mkdir(exist_ok=True)
SEED=20260927
START_GDP=32.0;START_DEBT=32.0

@dataclass
class Case:
    name:str='Integrated central'
    real:float=.03
    inflation:float=.04
    premium:float=0.0
    new_yield:float|None=None
    initial_rate:float=.033
    repricing:float=.20
    primary_start:float=-.026
    primary_target:float=0.0
    adjustment_years:int=10
    receipts:float=.15
    fund_return:float=.07
    fee:float=.0025
    withdrawal_rate:float=.03
    withdrawal_start:int=21
    mode:str='fund'
    ai_ramp:int=0
    revenue_delay:int=0
    revenue_end:int=999
    debt_target:float|None=None

def simulate(case,years=50,shock=None):
    """Exact annual cash-flow identity; primary surplus excludes resource/Fund flows.
    Interest uses the opening debt stock. Fiscal flows scale with current GDP.
    Receipts/contributions occur at year end. Payout is a fraction of opening Fund,
    capped at the assets available after returns and before new contributions.
    """
    d=START_DEBT;y=START_GDP;f=0.;price=1.;rate=case.initial_rate;total_w=0.
    rows=[dict(year=0,gdp=y,debt=d,fund=f,debt_ratio=100.,net_ratio=100.,fund_ratio=0.,price=1.,real_fund=0.,interest=0.,rate=rate,primary=0.,primary_pct=case.primary_start,receipts=0.,contribution=0.,withdrawal=0.,cumulative_withdrawals=0.,real_growth=0.,inflation=0.,fund_return=0.,debt_growth=0.)]
    for t in range(1,years+1):
        g=case.real if not case.ai_ramp else .018+(case.real-.018)*min(t/case.ai_ramp,1)
        pi=case.inflation;p=case.primary_start+(case.primary_target-case.primary_start)*min(t/case.adjustment_years,1)
        n=case.receipts if case.revenue_delay<t<=case.revenue_end else 0.
        gross_return=case.fund_return
        new_rate=case.new_yield if case.new_yield is not None else pi+.025+case.premium
        if shock:
            overrides=shock(t)
            g=overrides.get('real',g);pi=overrides.get('inflation',pi);p+=overrides.get('primary_delta',0)
            gross_return=overrides.get('fund_return',gross_return);n*=overrides.get('receipts_multiplier',1)
            new_rate=overrides.get('new_yield',new_rate)
        old_d=d;old_f=f
        y*= (1+g)*(1+pi);price*=1+pi
        rate=(1-case.repricing)*rate+case.repricing*new_rate
        interest=rate*old_d
        c=n if case.mode in ('fund','borrow') else 0.
        if case.mode in ('none','borrow'):n=0.
        if case.mode=='none':c=0.
        net_return=gross_return-case.fee
        available=old_f*(1+net_return)
        w=min(case.withdrawal_rate*old_f,available) if t>=case.withdrawal_start else 0.
        primary=p*y
        if case.debt_target is not None:
            primary=interest-case.debt_target*old_d-n+c-w;p=primary/y
        d=old_d+interest-primary-n+c-w
        f=available+c-w
        assert d>0 and f>=-1e-12, 'Outside positive-debt domain; extend model before interpreting'
        assert abs((d-f)-(old_d-old_f+interest-primary-n-old_f*net_return))<1e-8
        total_w+=w
        rows.append(dict(year=t,gdp=y,debt=d,fund=f,debt_ratio=100*d/y,net_ratio=100*(d-f)/y,fund_ratio=100*f/y,price=price,real_fund=f/price,interest=interest,rate=rate,primary=primary,primary_pct=p,receipts=n,contribution=c,withdrawal=w,cumulative_withdrawals=total_w,real_growth=g,inflation=pi,fund_return=net_return,debt_growth=d/old_d-1))
    return rows

def brief(rows):
    return {str(t):{k:rows[t][k] for k in ('debt_ratio','net_ratio','fund','real_fund','debt','gdp','withdrawal','primary_pct','interest','rate','cumulative_withdrawals')} for t in (10,20,30,50) if t<len(rows)}

CASES=[
 Case('Article target: enforce 2% debt growth',debt_target=.02),
 Case('Integrated central: primary balance by year 10'),
 Case('Stronger discipline: 1% primary surplus by year 10',primary_target=.01),
 Case('No fiscal correction: 2.6% primary deficit',primary_target=-.026),
 Case('Weaker AI: 1.8% real growth',real=.018),
 Case('Higher financing cost: new yield 8.5%',premium=.02),
 Case('2% inflation; yields fall with inflation',inflation=.02),
 Case('Resource shortfall: $15B annual net receipts',receipts=.015),
 Case('Fund financed by borrowing',mode='borrow'),
 Case('AI gains arrive gradually over 10 years',ai_ramp=10),
 Case('Resource receipts delayed 10 years',revenue_delay=10),
 Case('Finite receipts: contributions stop after year 20',revenue_end=20),
 Case('Adverse combination',real=.018,primary_target=-.015,premium=.02,receipts=.015,fund_return=.04),
]
paths={c.name:simulate(c) for c in CASES}

# Recession near the beginning of withdrawals; fixed stress, not a forecast.
def recession(t):
    if t==21:return {'real':-.03,'inflation':.05,'primary_delta':-.03,'fund_return':-.30,'receipts_multiplier':.5,'new_yield':.085}
    if t==22:return {'real':.005,'inflation':.04,'primary_delta':-.015,'fund_return':.02,'receipts_multiplier':.75,'new_yield':.075}
    return {}
paths['Recession and portfolio crash in years 21–22']=simulate(Case(),shock=recession)

# Fair allocation comparison: identical additional receipts and macro/fiscal paths.
allocation={label:simulate(replace(Case(),name=label,mode=mode)) for label,mode in [('No new receipts','none'),('Invest receipts in Fund','fund'),('Use receipts to reduce borrowing','paydown'),('Borrow to capitalize Fund','borrow')]}
accumulation={}
for r in (.04,.06,.07,.08):
    accumulation[str(r)]=simulate(Case(f'{r:.0%} accumulation',fund_return=r,fee=0.,withdrawal_start=999))

# A designed sensitivity grid. Its case frequencies have no probability meaning.
grid=[]
for g,pi,p,prem,n,r in itertools.product([.018,.025,.03,.035],[.02,.04,.05],[-.026,-.015,0.,.01],[0.,.015,.03],[.015,.15,.2],[.04,.07,.08]):
    c=Case(real=g,inflation=pi,primary_target=p,premium=prem,receipts=n,fund_return=r)
    a=simulate(c,30)[-1]
    grid.append(dict(real=g,inflation=pi,primary_target=p,new_yield=pi+.025+prem,receipts=n,gross_fund_return=r,**{k:a[k] for k in ('debt_ratio','net_ratio','fund')}))
heatmap={'yields':[.04,.045,.05,.055,.06,.065,.07,.075,.08,.085,.09],'primary_targets':[-.026,-.015,-.01,-.005,0.,.005,.01,.015,.02]}
heatmap['values']=[[simulate(Case(new_yield=y,primary_target=p),30)[-1]['debt_ratio'] for y in heatmap['yields']] for p in heatmap['primary_targets']]

def monte_carlo(n=20000,years=30,primary_target=0.,mode='fund'):
    """Synthetic risk experiment; distributions are declared assumptions, not fitted.
    Reinitializing the seed gives paired shocks for each fiscal/allocation policy.
    """
    rng=np.random.default_rng(SEED)
    d=np.full(n,START_DEBT);y=np.full(n,START_GDP);f=np.zeros(n);price=np.ones(n);rate=np.full(n,.033)
    long_g=.03+rng.normal(0,.005,n);long_pi=.04+rng.normal(0,.005,n)
    zg=rng.normal(size=n);zp=rng.normal(size=n)
    debt_series=[d/y*100];net_series=[d/y*100];fund_series=[f.copy()]
    first_year_payout=None
    for t in range(1,years+1):
        z=rng.normal(size=(5,n));zg=.35*zg+math.sqrt(1-.35**2)*z[0];zp=.55*zp+math.sqrt(1-.55**2)*z[1]
        cycle=.015*zg
        g=np.clip(long_g+cycle,-.08,.10);pi=np.clip(long_pi+.008*zp,0.,.10)
        z_return=.35*z[0]-.20*z[1]+math.sqrt(1-.35**2-.2**2)*z[2]
        gross_return=np.exp(math.log(1.07)-.5*.12**2+.12*z_return)-1
        net_return=gross_return-.0025
        z_resources=.4*z[0]+math.sqrt(1-.4**2)*z[3]
        receipts=.15*np.exp(-.5*.35**2+.35*z_resources)
        # 100 percentage points of extra debt/GDP adds 1 percentage point to yield.
        risk=.01*np.maximum(d/y-1,0)
        new_rate=np.maximum(0,pi+.025+.005*z[4]+risk)
        rate=.8*rate+.2*new_rate
        old_d=d.copy();old_f=f.copy();y*=(1+g)*(1+pi);price*=1+pi
        primary_pct=-.026+(primary_target+.026)*min(t/10,1)+.4*cycle
        primary=primary_pct*y;interest=rate*old_d
        contribution=receipts if mode=='fund' else np.zeros(n)
        w=np.minimum(.03*old_f,old_f*(1+net_return)) if t>=21 else np.zeros(n)
        d=old_d+interest-primary-receipts+contribution-w
        f=old_f*(1+net_return)+contribution-w
        assert np.all(d>0) and np.all(f>=0)
        assert np.max(np.abs((d-f)-(old_d-old_f+interest-primary-receipts-old_f*net_return)))<1e-7
        if t==21:first_year_payout=w.copy()
        debt_series.append(d/y*100);net_series.append((d-f)/y*100);fund_series.append(f.copy())
    a=d/y*100;b=(d-f)/y*100
    q=[.05,.25,.5,.75,.95]
    report={'paths':n,'seed':SEED,'primary_target':primary_target,'allocation':mode,'debt_ratio_quantiles':dict(zip(map(str,q),np.quantile(a,q).tolist())),'net_ratio_quantiles':dict(zip(map(str,q),np.quantile(b,q).tolist())),'fund_quantiles':dict(zip(map(str,q),np.quantile(f,q).tolist())),'real_fund_quantiles':dict(zip(map(str,q),np.quantile(f/price,q).tolist())),'fraction_debt_below_start':float(np.mean(a<100)),'fraction_debt_below_50':float(np.mean(a<50)),'fraction_debt_below_25':float(np.mean(a<25)),'fraction_debt_above_150':float(np.mean(a>150)),'debt_bands':np.quantile(np.array(debt_series),[.05,.5,.95],axis=1).tolist(),'net_bands':np.quantile(np.array(net_series),[.05,.5,.95],axis=1).tolist(),'fund_bands':np.quantile(np.array(fund_series),[.05,.5,.95],axis=1).tolist()}
    return report,{'debt_ratio':a,'net_ratio':b,'fund':f,'real_fund':f/price}

mc,terminal=monte_carlo()
mc_pay,pay_terminal=monte_carlo(mode='paydown')
mc_weak,weak_terminal=monte_carlo(primary_target=-.015)
mc['paired_fund_minus_paydown_net_ratio_quantiles']=dict(zip(['.05','.5','.95'],np.quantile(terminal['net_ratio']-pay_terminal['net_ratio'],[.05,.5,.95]).tolist()))
mc['fraction_fund_lower_net_liabilities_than_paydown']=float(np.mean(terminal['net_ratio']<pay_terminal['net_ratio']))

def verify():
    checks=[]
    target=paths[CASES[0].name]
    for t in (10,20,30):assert abs(target[t]['debt_ratio']-100*(1.02/1.0712)**t)<1e-9
    checks.append('Published 2% debt-growth identity reproduced independently')
    for r,rows in accumulation.items():
        r=float(r)
        for t in (10,30,50):assert abs(rows[t]['fund']-.15*((1+r)**t-1)/r)<1e-8
    checks.append('Fund balances match closed-form end-year annuity')
    c=Case(primary_start=0,primary_target=0,receipts=0,initial_rate=.05,new_yield=.05,mode='none')
    rows=simulate(c)
    assert abs(rows[30]['debt']-START_DEBT*1.05**30)<1e-9
    checks.append('Zero-primary-balance debt matches independent compound-interest formula')
    # Equal financing and return rates: investing cash and avoiding debt have equal net position.
    c=Case(primary_start=0,primary_target=0,initial_rate=.05,new_yield=.05,fund_return=.05,fee=0,withdrawal_start=999)
    f=simulate(c);p=simulate(replace(c,mode='paydown'))
    assert max(abs((a['debt']-a['fund'])-b['debt']) for a,b in zip(f,p))<1e-8
    checks.append('Equal-rate allocation counterfactual has identical net liabilities')
    for name,rows in paths.items():
        for a in rows[1:]:
            assert abs(a['net_ratio']-(a['debt_ratio']-a['fund_ratio']))<1e-10
    checks.append('Debt, assets, withdrawals and net liabilities satisfy every annual cash-flow identity')
    assert len(grid)==1296
    checks.append('1,296 deterministic sensitivity combinations evaluated')
    return checks

checks=verify()
sources=[
 {'label':'CBO February 2026 outlook','url':'https://www.cbo.gov/publication/62105','use':'Primary deficit 2.6% of GDP in 2026; average real growth 1.8% over 2027–2036. Dated federal projections, not calibration of this simulation.'},
 {'label':'CBO subsequent outlook updates','url':'https://www.cbo.gov/topics/economy/outlook-budget-and-economy','use':'Context for the dated baseline; August 20 tariff update is not incorporated into this normalized exercise.'},
 {'label':'Interior FY2025 resource disbursements','url':'https://www.doi.gov/node/64051','use':'$14.61B total disbursements; $5.01B to Treasury. Neither is a forecast of new investable revenue.'},
 {'label':'Federal Reserve longer-run strategy','url':'https://www.federalreserve.gov/monetarypolicy/monetary-policy-strategy-tools-and-communications-statement-on-longer-run-goals-monetary-policy-strategy-2025.htm','use':'2% PCE inflation objective; the simulation uses GDP-price inflation and does not presume a policy change.'}
]
results={'numpy_version':np.__version__,'date':'2026-09-27','seed':SEED,'units':'USD trillions; debt ratios in percent of GDP','initial_conditions':{'gdp':START_GDP,'debt':START_DEBT,'fund':0,'note':'Deliberate rounded scale and 100% ratio, not measured current balances.'},'parameters':[asdict(c) for c in CASES],'scenarios':{n:brief(v) for n,v in paths.items()},'paths':paths,'allocation':{n:brief(v) for n,v in allocation.items()},'allocation_paths':allocation,'accumulation':{r:brief(v) for r,v in accumulation.items()},'heatmap':heatmap,'monte_carlo':{'balance_by_year10':mc,'same_receipts_paydown':mc_pay,'persistent_1_5_deficit':mc_weak},'sources':sources,'checks':checks,'grid_cases':len(grid)}
(OUT/'results.json').write_text(json.dumps(results,indent=2,allow_nan=False))
with (OUT/'scenario_paths.csv').open('w',newline='') as f:
    writer=csv.DictWriter(f,fieldnames=['scenario']+list(next(iter(paths.values()))[0]));writer.writeheader()
    for name,rows in paths.items():writer.writerows(dict(scenario=name,**row) for row in rows)
with (OUT/'sensitivity_grid.csv').open('w',newline='') as f:
    writer=csv.DictWriter(f,fieldnames=list(grid[0]));writer.writeheader();writer.writerows(grid)
with (OUT/'monte_carlo_terminal.csv').open('w',newline='') as f:
    writer=csv.writer(f);writer.writerow(['path','fund_case_debt_pct','fund_case_net_pct','fund_T','fund_real_T','paydown_debt_pct','persistent_deficit_debt_pct'])
    for i in range(20000):writer.writerow([i+1,terminal['debt_ratio'][i],terminal['net_ratio'][i],terminal['fund'][i],terminal['real_fund'][i],pay_terminal['debt_ratio'][i],weak_terminal['debt_ratio'][i]])
print('SCENARIOS (year30 debt/GDP, net liabilities/GDP, Fund $T)')
for name,rows in paths.items():
    a=rows[30];print(f'{name}: {a["debt_ratio"]:.2f}%, {a["net_ratio"]:.2f}%, ${a["fund"]:.3f}T')
print('ALLOCATION',json.dumps({n:brief(v)['30'] for n,v in allocation.items()}))
print('MC',json.dumps({n:{k:v for k,v in data.items() if not k.endswith('bands')} for n,data in results['monte_carlo'].items()}))
print('REQUIRED_PRIMARY',json.dumps({'year1':paths[CASES[0].name][1],'peak_pct':max(paths[CASES[0].name][1:31],key=lambda x:x['primary_pct'])}))
print('CHECKS',len(checks),'passed')