"""Core analyses for the SBA acquisition-loan book.
Input: raw/7a_all_asof_260630.pkl (from 01_load.py) and raw/prime_DPRIME.csv
(FRED series DPRIME, Bank Prime Loan Rate, https://fred.stlouisfed.org/series/DPRIME).
Output: tables/*.csv. Every table carries an n column.
Definitions (also in FINDINGS.md):
  disbursed     = LoanStatus not in {CANCLD, COMMIT}
  chargeoff rate (count) = CHGOFF / disbursed
  dollar loss rate = sum(GrossChargeOffAmount) / sum(GrossApproval), disbursed loans
  standard term = RevolverStatus 'N' and ProcessingMethod not an Express variant
  matured cohort = ApprovalFY 2010-2015 (primary); stress cohort = FY2005-2008
"""
import os, numpy as np, pandas as pd
HERE = os.path.dirname(os.path.abspath(__file__))
RAW = os.path.join(HERE, '..', 'raw'); OUT = os.path.join(HERE, '..', 'tables')
os.makedirs(OUT, exist_ok=True)
ASOF = pd.Timestamp('2026-06-30')

d = pd.read_pickle(os.path.join(RAW, '7a_all_asof_260630.pkl'))
d['FY'] = d['ApprovalFY'].astype(int)
d['disbursed'] = ~d['LoanStatus'].isin(['CANCLD', 'COMMIT'])
d['co'] = d['LoanStatus'].eq('CHGOFF')
d['express'] = d['ProcessingMethod'].fillna('').str.contains('Express', case=False)
d['revolver'] = d['RevolverStatus'].eq('Y')
d['std_term'] = (~d['revolver']) & (~d['express'])
d['coo'] = d['BusinessAge'].eq('Change of Ownership')
d['franchise'] = d['FranchiseCode'].notna()
d['co_amt'] = d['GrossChargeOffAmount'].fillna(0).where(d['co'], 0)
d['months_to_co'] = ((d['ChargeOffDate'] - d['ApprovalDate']).dt.days / 30.4375)
d['guar_pct'] = d['SBAGuaranteedApproval'] / d['GrossApproval']
bands = [0, 150_000, 350_000, 500_000, 1_000_000, 2_000_000, 3_500_000, 5_000_001]
labels = ['<=150k', '150k-350k', '350k-500k', '500k-1M', '1M-2M', '2M-3.5M', '3.5M-5M']
d['size_band'] = pd.cut(d['GrossApproval'], bands, labels=labels, right=True)
d['term_band'] = pd.cut(d['TermInMonths'], [-1, 59, 83, 119, 120, 239, 1000],
                        labels=['<5y', '5-7y', '7-<10y', '10y exactly', '10-<20y', '20y+'])

# prime rate on approval date
pr = pd.read_csv(os.path.join(RAW, 'prime_DPRIME.csv'), parse_dates=['observation_date'])
pr = pr.rename(columns={'observation_date': 'date', 'DPRIME': 'prime'})
pr['prime'] = pd.to_numeric(pr['prime'], errors='coerce'); pr = pr.dropna().sort_values('date')
tmp = d[['ApprovalDate']].reset_index().dropna().sort_values('ApprovalDate')
tmp = pd.merge_asof(tmp, pr, left_on='ApprovalDate', right_on='date', direction='backward')
d['prime'] = tmp.set_index('index')['prime']
d['spread'] = d['InitialInterestRate'] - d['prime']

def save(df, name):
    df.to_csv(os.path.join(OUT, name), index=False)
    print(f'wrote {name} rows={len(df)}')

def q(s, p): return s.quantile(p)

def wilson(k, n, z=1.96):
    if n == 0: return (np.nan, np.nan)
    p = k / n; den = 1 + z * z / n; c = (p + z * z / (2 * n)) / den; h = z * np.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / den
    return (round(100 * (c - h), 2), round(100 * (c + h), 2))

def loss_stats(g):
    disb = g[g['disbursed']]
    n = len(disb)
    return pd.Series({
        'n_disbursed': n,
        'n_chargeoff': int(disb['co'].sum()),
        'chargeoff_rate_pct': round(100 * disb['co'].mean(), 2) if n else np.nan,
        'ci95_low_pct': wilson(int(disb['co'].sum()), n)[0], 'ci95_high_pct': wilson(int(disb['co'].sum()), n)[1],
        'dollar_loss_rate_pct': round(100 * disb['co_amt'].sum() / disb['GrossApproval'].sum(), 2) if n else np.nan,
        'still_active_exempt_pct': round(100 * disb['LoanStatus'].eq('EXEMPT').mean(), 2) if n else np.nan,
        'approved_usd_m': round(disb['GrossApproval'].sum() / 1e6, 1),
    })

# ---------- T01 volume and size by FY ----------
rows = []
for fy, g in d.groupby('FY'):
    a = g['GrossApproval']
    rows.append({'FY': fy, 'n_loans': len(g), 'approved_usd_bn': round(a.sum() / 1e9, 2),
                 'mean_usd': round(a.mean()), 'p25_usd': q(a, .25), 'median_usd': q(a, .5), 'p75_usd': q(a, .75), 'p90_usd': q(a, .9),
                 'share_le_150k_pct': round(100 * (a <= 150_000).mean(), 1), 'share_gt_1m_pct': round(100 * (a > 1_000_000).mean(), 1),
                 'share_express_pct': round(100 * g['express'].mean(), 1), 'share_revolver_pct': round(100 * g['revolver'].mean(), 1),
                 'n_std_term': int(g['std_term'].sum()), 'median_std_term_usd': q(g.loc[g['std_term'], 'GrossApproval'], .5),
                 'partial_year': fy == 2026})
save(pd.DataFrame(rows), 'T01_volume_size_by_fy.csv')

# ---------- T02 rates vs prime ----------
r = d[(d['FY'] >= 2010) & d['InitialInterestRate'].between(0.5, 30) & d['prime'].notna()]
rows = []
for (fy, fv), g in r.groupby(['FY', 'FixedorVariableInterestInd']):
    rows.append({'FY': fy, 'fixed_or_variable': fv, 'n': len(g), 'median_rate': q(g['InitialInterestRate'], .5),
                 'p10_rate': q(g['InitialInterestRate'], .1), 'p90_rate': q(g['InitialInterestRate'], .9),
                 'mean_prime_at_approval': round(g['prime'].mean(), 2), 'median_spread_over_prime': round(q(g['spread'], .5), 2),
                 'p90_spread_over_prime': round(q(g['spread'], .9), 2)})
save(pd.DataFrame(rows), 'T02_rates_vs_prime_by_fy.csv')
rows = []
rr = r[(r['FY'] >= 2024) & (r['FixedorVariableInterestInd'] == 'V')]
for (sb, ex), g in rr.groupby(['size_band', 'express'], observed=True):
    rows.append({'size_band': sb, 'express': ex, 'n': len(g), 'median_spread': round(q(g['spread'], .5), 2), 'p75_spread': round(q(g['spread'], .75), 2),
                 'p90_spread': round(q(g['spread'], .9), 2), 'max_spread': round(g['spread'].max(), 2), 'median_rate': q(g['InitialInterestRate'], .5)})
save(pd.DataFrame(rows), 'T02b_variable_spread_by_size_FY2024_2026.csv')
# monthly series for chart: median rate and prime
m = r[r['FixedorVariableInterestInd'] == 'V'].copy(); m['month'] = m['ApprovalDate'].dt.to_period('M').astype(str)
mm = m.groupby('month').agg(n=('spread', 'size'), median_rate=('InitialInterestRate', 'median'), prime=('prime', 'median'), median_spread=('spread', 'median')).reset_index()
save(mm, 'T02c_monthly_variable_rate_vs_prime.csv')

# ---------- T03 terms ----------
rows = []
# CHGOFF rows excluded: SBA overwrites TermInMonths on charged-off loans (see FINDINGS trap)
for fy, g in d[d['std_term'] & ~d['co']].groupby('FY'):
    t = g['TermInMonths']
    rows.append({'FY': fy, 'n_std_term': len(g), 'median_term_months': q(t, .5), 'share_120m_pct': round(100 * (t == 120).mean(), 1),
                 'share_ge_240m_pct': round(100 * (t >= 240).mean(), 1), 'share_lt_120m_pct': round(100 * (t < 120).mean(), 1)})
save(pd.DataFrame(rows), 'T03_terms_by_fy_std_term_excl_chargeoffs.csv')

# ---------- T04 charge-off by vintage ----------
rows = []
for fy, g in d.groupby('FY'):
    s = loss_stats(g); s['FY'] = fy
    s2 = loss_stats(g[g['std_term']]); s['std_term_n'] = s2['n_disbursed']; s['std_term_chargeoff_rate_pct'] = s2['chargeoff_rate_pct']; s['std_term_dollar_loss_pct'] = s2['dollar_loss_rate_pct']
    s3 = loss_stats(g[g['express']]); s['express_n'] = s3['n_disbursed']; s['express_chargeoff_rate_pct'] = s3['chargeoff_rate_pct']
    s['censored_warning'] = 'right-censored, do not use as loss rate' if fy >= 2016 else ''
    rows.append(s)
t4 = pd.DataFrame(rows); t4 = t4[['FY'] + [c for c in t4.columns if c != 'FY']]
save(t4, 'T04_chargeoff_by_vintage.csv')

MAT = d[d['FY'].between(2010, 2015)]
STRESS = d[d['FY'].between(2005, 2008)]

def by_dim(frame, col, min_n=0, label=None):
    out = frame.groupby(col, observed=True).apply(loss_stats, include_groups=False).reset_index()
    out = out[out['n_disbursed'] >= min_n]
    return out

# ---------- T05 by size / term / processing / business age / collateral / franchise, matured cohort ----------
dims = [('size_band', 'size band'), ('express', 'Express (any)'), ('revolver', 'revolver'),
        ('BusinessAge', 'business age'), ('CollateralInd', 'collateral flag'), ('franchise', 'franchise'), ('SoldSecMrktInd', 'sold secondary market'),
        ('BusinessType', 'business type')]
parts = []
for col, lab in dims:
    for cname, frame in [('FY2010-2015 all', MAT), ('FY2010-2015 std term', MAT[MAT['std_term']]), ('FY2005-2008 all', STRESS), ('FY2005-2008 std term', STRESS[STRESS['std_term']])]:
        o = by_dim(frame, col); o.insert(0, 'cohort', cname); o.insert(0, 'dimension', lab); o = o.rename(columns={col: 'value'}); o['value'] = o['value'].astype(str)
        parts.append(o)
save(pd.concat(parts, ignore_index=True), 'T05_chargeoff_by_dimension_matured.csv')

# ---------- T06 by industry ----------
SECT = {'11': 'Agriculture, forestry, fishing', '21': 'Mining, oil and gas', '22': 'Utilities', '23': 'Construction', '31': 'Manufacturing', '32': 'Manufacturing', '33': 'Manufacturing',
        '42': 'Wholesale trade', '44': 'Retail trade', '45': 'Retail trade', '48': 'Transportation and warehousing', '49': 'Transportation and warehousing', '51': 'Information',
        '52': 'Finance and insurance', '53': 'Real estate, rental and leasing', '54': 'Professional, scientific, technical services', '55': 'Management of companies',
        '56': 'Administrative, support, waste management', '61': 'Educational services', '62': 'Health care and social assistance', '71': 'Arts, entertainment, recreation',
        '72': 'Accommodation and food services', '81': 'Other services (except public administration)', '92': 'Public administration'}
d['sector'] = d['NaicsCode'].str[:2].map(SECT)
MAT = d[d['FY'].between(2010, 2015)]; STRESS = d[d['FY'].between(2005, 2008)]
a = by_dim(MAT[MAT['std_term']], 'sector'); a.insert(0, 'cohort', 'FY2010-2015 std term')
b = by_dim(MAT, 'sector'); b.insert(0, 'cohort', 'FY2010-2015 all')
c = by_dim(STRESS[STRESS['std_term']], 'sector'); c.insert(0, 'cohort', 'FY2005-2008 std term')
save(pd.concat([a, b, c]).sort_values(['cohort', 'chargeoff_rate_pct'], ascending=[True, False]), 'T06_chargeoff_by_sector_matured.csv')
d['naics6'] = d['NaicsCode'] + ' ' + d['NaicsDescription'].fillna('')
MAT = d[d['FY'].between(2010, 2015)]
n6 = by_dim(MAT[MAT['std_term']], 'naics6', min_n=300).sort_values('chargeoff_rate_pct', ascending=False)
n6.insert(0, 'cohort', 'FY2010-2015 std term, n>=300')
save(n6, 'T06b_chargeoff_by_naics6_matured_n300.csv')

# ---------- T07 franchise brands ----------
d['brand'] = d['FranchiseCode'] + ' ' + d['FranchiseName'].fillna('')
MAT = d[d['FY'].between(2010, 2015)]
fb = by_dim(MAT[MAT['franchise']], 'brand', min_n=100).sort_values('chargeoff_rate_pct', ascending=False)
fb.insert(0, 'cohort', 'FY2010-2015 all franchise loans, brand n>=100')
save(fb, 'T07_chargeoff_by_franchise_brand_matured_n100.csv')

# ---------- T08 lenders ----------
rows = []
for fy, g in d[d['FY'].between(2010, 2025)].groupby('FY'):
    cnt = g.groupby('BankName').size().sort_values(ascending=False); usd = g.groupby('BankName')['GrossApproval'].sum().sort_values(ascending=False)
    sh = usd / usd.sum()
    rows.append({'FY': fy, 'n_loans': len(g), 'n_lenders': g['BankName'].nunique(), 'top10_share_count_pct': round(100 * cnt.head(10).sum() / cnt.sum(), 1),
                 'top10_share_usd_pct': round(100 * usd.head(10).sum() / usd.sum(), 1), 'top100_share_usd_pct': round(100 * usd.head(100).sum() / usd.sum(), 1),
                 'hhi_usd': round(10000 * (sh ** 2).sum()), 'lenders_with_1_loan': int((cnt == 1).sum())})
save(pd.DataFrame(rows), 'T08_lender_concentration_by_fy.csv')
recent = d[d['FY'].between(2023, 2025)]
def lender_table(frame, label):
    g = frame.groupby('BankName')
    t = pd.DataFrame({'n_loans': g.size(), 'approved_usd_m': (g['GrossApproval'].sum() / 1e6).round(1), 'median_loan_usd': g['GrossApproval'].median(),
                      'express_share_pct': (100 * g['express'].mean()).round(1)}).sort_values('approved_usd_m', ascending=False).head(30).reset_index()
    t.insert(0, 'population', label); t.insert(1, 'rank_usd', range(1, len(t) + 1))
    return t
save(pd.concat([lender_table(recent, 'All 7(a) FY2023-2025'), lender_table(recent[recent['coo']], 'Change of Ownership flag FY2023-2025')]), 'T08b_top_lenders_FY2023_2025.csv')
# lender loss rates, matured, std term, n>=500 (BankName is CURRENT assignee: merger trap)
MAT = d[d['FY'].between(2010, 2015)]
ll = by_dim(MAT[MAT['std_term']], 'BankName', min_n=500).sort_values('chargeoff_rate_pct')
ll['median_loan_usd'] = ll['BankName'].map(MAT[MAT['std_term']].groupby('BankName')['GrossApproval'].median())
ll.insert(0, 'cohort', 'FY2010-2015 std term, lender n>=500; BankName = current assignee')
save(ll, 'T08c_lender_chargeoff_matured_n500.csv')

# ---------- T09 geography ----------
g = d[d['FY'].between(2023, 2025)].groupby('ProjectState')
geo = pd.DataFrame({'n_loans_FY23_25': g.size(), 'approved_usd_m': (g['GrossApproval'].sum() / 1e6).round(1), 'median_loan_usd': g['GrossApproval'].median(),
                    'n_change_of_ownership': g['coo'].sum(), 'coo_share_pct': (100 * g['coo'].mean()).round(1)}).reset_index()
MAT = d[d['FY'].between(2010, 2015)]
geo = geo.merge(by_dim(MAT[MAT['std_term']], 'ProjectState')[['ProjectState', 'n_disbursed', 'chargeoff_rate_pct']].rename(columns={'n_disbursed': 'matured_std_n', 'chargeoff_rate_pct': 'matured_std_chargeoff_pct'}), on='ProjectState', how='left')
save(geo.sort_values('approved_usd_m', ascending=False), 'T09_geography_by_state.csv')

# ---------- T10 change of ownership flag ----------
c = d[d['FY'] >= 2019]
rows = []
for fy, g in c.groupby('FY'):
    k = g[g['coo']]
    rows.append({'FY': fy, 'n_all': len(g), 'n_coo': len(k), 'coo_share_count_pct': round(100 * len(k) / len(g), 1),
                 'coo_share_usd_pct': round(100 * k['GrossApproval'].sum() / g['GrossApproval'].sum(), 1), 'coo_usd_bn': round(k['GrossApproval'].sum() / 1e9, 2),
                 'coo_median_usd': k['GrossApproval'].median(), 'coo_p25_usd': q(k['GrossApproval'], .25), 'coo_p75_usd': q(k['GrossApproval'], .75),
                 'coo_share_gt_1m_pct': round(100 * (k['GrossApproval'] > 1e6).mean(), 1), 'coo_median_term_m': k['TermInMonths'].median(),
                 'coo_share_term_ge_240_pct': round(100 * (k['TermInMonths'] >= 240).mean(), 1), 'coo_franchise_pct': round(100 * k['franchise'].mean(), 1),
                 'coo_express_pct': round(100 * k['express'].mean(), 1), 'coo_median_rate': k['InitialInterestRate'].median(),
                 'coo_median_spread_variable': k.loc[k['FixedorVariableInterestInd'] == 'V', 'spread'].median(),
                 'coo_variable_pct': round(100 * (k['FixedorVariableInterestInd'] == 'V').mean(), 1),
                 'coo_median_jobs': k['JobsSupported'].median(), 'partial_year': fy == 2026})
save(pd.DataFrame(rows), 'T10_change_of_ownership_by_fy.csv')
cc = d[d['coo'] & d['FY'].between(2023, 2025)]
ind = cc.groupby('naics6').agg(n=('GrossApproval', 'size'), approved_usd_m=('GrossApproval', lambda s: round(s.sum() / 1e6, 1)), median_usd=('GrossApproval', 'median')).sort_values('n', ascending=False).head(40).reset_index()
save(ind, 'T10b_coo_top_industries_FY2023_2025.csv')
br = cc[cc['franchise']].groupby('brand').agg(n=('GrossApproval', 'size'), median_usd=('GrossApproval', 'median')).sort_values('n', ascending=False).head(30).reset_index()
save(br, 'T10c_coo_top_franchise_brands_FY2023_2025.csv')
sz = cc.groupby('size_band', observed=True).size().rename('n').reset_index(); sz['share_pct'] = (100 * sz['n'] / sz['n'].sum()).round(1)
save(sz, 'T10d_coo_size_distribution_FY2023_2025.csv')

# ---------- T11 timing of charge-offs (hazard), matured cohort ----------
rows = []
for cname, frame in [('FY2010-2015 std term', d[d['FY'].between(2010, 2015) & d['std_term'] & d['disbursed']]),
                     ('FY2005-2008 std term', d[d['FY'].between(2005, 2008) & d['std_term'] & d['disbursed']])]:
    n = len(frame); cos = frame.loc[frame['co'], 'months_to_co']
    for yr in range(1, 16):
        rows.append({'cohort': cname, 'n_disbursed': n, 'years_since_approval': yr, 'cum_chargeoff_pct_of_loans': round(100 * (cos <= 12 * yr).sum() / n, 2),
                     'cum_share_of_all_chargeoffs_pct': round(100 * (cos <= 12 * yr).sum() / len(cos), 1), 'n_chargeoffs_total': len(cos)})
save(pd.DataFrame(rows), 'T11_chargeoff_timing_curve.csv')

# ---------- T12 early loss at fixed horizon, by vintage and CoO flag (handles censoring) ----------
rows = []
for fy, g in d[d['FY'].between(2010, 2022) & d['disbursed'] & d['std_term']].groupby('FY'):
    for grp_name, gg in [('all std term', g), ('Change of Ownership flag', g[g['coo']]), ('Existing or >2y (FY2018+ category)', g[g['BusinessAge'] == 'Existing or more than 2 years old']),
                         ('Startup', g[g['BusinessAge'] == 'Startup, Loan Funds will Open Business'])]:
        if len(gg) == 0: continue
        h36 = gg['co'] & (gg['months_to_co'] <= 36); h48 = gg['co'] & (gg['months_to_co'] <= 48)
        rows.append({'FY': fy, 'group': grp_name, 'n_disbursed': len(gg), 'chargeoff_within_36m_pct': round(100 * h36.mean(), 2),
                     'chargeoff_within_48m_pct': round(100 * h48.mean(), 2) if fy <= 2021 else np.nan, 'median_loan_usd': gg['GrossApproval'].median()})
save(pd.DataFrame(rows), 'T12_early_chargeoff_fixed_horizon.csv')

# ---------- T13 CARES-era monthly volume ----------
w = d[d['ApprovalDate'].between('2018-10-01', '2023-09-30')].copy(); w['month'] = w['ApprovalDate'].dt.to_period('M').astype(str)
mw = w.groupby('month').agg(n_loans=('GrossApproval', 'size'), approved_usd_m=('GrossApproval', lambda s: round(s.sum() / 1e6, 1)), median_usd=('GrossApproval', 'median'),
                             n_coo=('coo', 'sum'), share_gt_1m_pct=('GrossApproval', lambda s: round(100 * (s > 1e6).mean(), 1)),
                             median_guar_pct=('guar_pct', lambda s: round(100 * s.median(), 1)), share_guar_ge_89_pct=('guar_pct', lambda s: round(100 * (s >= 0.895).mean(), 1))).reset_index()
save(mw, 'T13_monthly_volume_FY2019_2023.csv')

# ---------- T14 guaranty percentage by size, recent ----------
gp = d[d['FY'].between(2024, 2026)]
rows = []
for (sb, ex), g in gp.groupby(['size_band', 'express'], observed=True):
    rows.append({'size_band': sb, 'express': ex, 'n': len(g), 'median_guar_pct': round(100 * g['guar_pct'].median(), 1), 'p10_guar_pct': round(100 * q(g['guar_pct'], .1), 1),
                 'p90_guar_pct': round(100 * q(g['guar_pct'], .9), 1), 'share_sold_secondary_pct': round(100 * g['SoldSecMrktInd'].eq('Y').mean(), 1)})
save(pd.DataFrame(rows), 'T14_guaranty_pct_by_size_FY2024_2026.csv')
d.to_pickle(os.path.join(RAW, '7a_enriched.pkl'))
print('done')
