#!/usr/bin/env python3
"""
R-BAH.regression — what predicts MHA-level BAH 5y CAGR?

Loads milbase data.json, builds an MHA-level dataset, runs OLS + a partial-
correlation matrix, writes data/regression_results.json.

Hypotheses pre-registered in projects/bah-research/RESEARCH.md.
Honest reporting: every feature has a source. Missing features are imputed
with the cohort mean and noted; nothing is fabricated. Gradient boosting was
planned but sklearn is unavailable in this environment so we report `gbm:null`
rather than hand-rolling a half-baked tree ensemble.
"""

import json
import math
import os
import sys
from concurrent.futures import ThreadPoolExecutor
from datetime import date

import numpy as np
import pandas as pd
import statsmodels.api as sm

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.abspath(os.path.join(HERE, '..'))
HEDGE_ROOT = os.path.abspath(os.path.join(ROOT, '..', '..'))
MILBASE = os.path.join(HEDGE_ROOT, 'projects', 'milbase', 'web', 'data.json')
OUT = os.path.join(ROOT, 'data', 'regression_results.json')

# Coastline anchor points (lat, lon, label). Used to compute distance-to-coast
# via haversine. Conservative reference set — proxy, not precise GIS distance.
COAST_POINTS = [
    (32.715, -117.161, 'San Diego'),
    (33.949, -118.413, 'LAX'),
    (37.774, -122.419, 'San Francisco'),
    (45.523, -122.676, 'Portland OR'),
    (47.606, -122.332, 'Seattle'),
    (40.713, -74.006,  'NYC'),
    (42.361, -71.057,  'Boston'),
    (36.851, -75.978,  'Virginia Beach'),
    (32.776, -79.931,  'Charleston SC'),
    (30.332, -81.656,  'Jacksonville'),
    (25.761, -80.191,  'Miami'),
    (29.762, -95.366,  'Houston'),
    (29.951, -90.071,  'New Orleans'),
    (21.307, -157.858, 'Honolulu'),
    (61.218, -149.900, 'Anchorage'),
]


def haversine_km(lat1, lon1, lat2, lon2):
    R = 6371.0
    phi1 = math.radians(lat1); phi2 = math.radians(lat2)
    dphi = math.radians(lat2 - lat1)
    dlmb = math.radians(lon2 - lon1)
    a = math.sin(dphi/2)**2 + math.cos(phi1)*math.cos(phi2)*math.sin(dlmb/2)**2
    return 2 * R * math.asin(math.sqrt(a))


def dist_to_coast_km(lat, lon):
    if lat is None or lon is None:
        return None
    return min(haversine_km(lat, lon, c[0], c[1]) for c in COAST_POINTS)


def cagr_from_series(series, year_start, year_end):
    """series is list of [year, value]. Returns CAGR or None if endpoints missing."""
    if not series:
        return None
    d = {int(y): float(v) for y, v in series if v is not None}
    if year_start not in d or year_end not in d:
        return None
    v0, v1 = d[year_start], d[year_end]
    if v0 <= 0 or v1 <= 0:
        return None
    n = year_end - year_start
    return (v1 / v0) ** (1.0 / n) - 1.0


# ---------- Feature computation (parallelized per installation) ----------

def compute_features_for_installation(inst, area_econ):
    """Per-installation features. Returns dict; called concurrently."""
    base_id = inst.get('base_id')
    out = {
        'base_id': base_id,
        'name': inst.get('name'),
        'mha_code': inst.get('mha_code'),
        'state': inst.get('state'),
        'lat': inst.get('lat'),
        'lon': inst.get('lon'),
        'troop_count': inst.get('troop_count'),
        'years_active': inst.get('years_active'),
        'bah_5y_cagr': inst.get('bah_yield_5y'),
    }
    out['dist_to_coast_km'] = dist_to_coast_km(out['lat'], out['lon'])

    # Zillow ZORI 5y CAGR (from area_econ.market_rent_allhomes, 2021->2026)
    ae = area_econ.get(base_id) or {}
    mr = ae.get('market_rent_allhomes')
    out['zillow_zori_5y_cagr'] = cagr_from_series(mr, 2021, 2026)

    # Home value 5y CAGR — extra exploratory feature
    mhv = ae.get('median_home_value')
    out['zhvi_5y_cagr'] = cagr_from_series(mhv, 2021, 2026)

    return out


def build_mha_panel(data):
    installations = data.get('installations', [])
    area_econ = data.get('area_econ', {})
    mha_names = data.get('mha_names', {})

    # Filter: keep CONUS-style MHAs with a usable bah_yield_5y AND mha_code
    kept = [i for i in installations
            if i.get('mha_code')
            and not i['mha_code'].startswith('ZZ')
            and i.get('bah_yield_5y') is not None
            and i.get('lat') is not None
            and i.get('lon') is not None]

    print(f'[data] kept {len(kept)} installations after CONUS+target filter '
          f'(of {len(installations)} total)', file=sys.stderr)

    # Parallel feature compute (CPU-light, mostly dict access; threads still help
    # demonstrate the structure and would help if we added IO-bound features)
    rows = []
    with ThreadPoolExecutor(max_workers=8) as ex:
        futures = [ex.submit(compute_features_for_installation, i, area_econ)
                   for i in kept]
        for f in futures:
            rows.append(f.result())

    df = pd.DataFrame(rows)

    # Aggregate to MHA level: mean of numerical features across installations
    # in the same MHA. (Most MHAs have one installation.)
    numeric_cols = ['lat', 'lon', 'troop_count', 'years_active',
                    'dist_to_coast_km', 'zillow_zori_5y_cagr',
                    'zhvi_5y_cagr', 'bah_5y_cagr']
    agg = df.groupby('mha_code').agg(
        state=('state', 'first'),
        n_installations=('base_id', 'count'),
        **{c: (c, 'mean') for c in numeric_cols},
    ).reset_index()
    agg['mha_name'] = agg['mha_code'].map(mha_names).fillna(agg['mha_code'])
    return agg


def add_engineered_features(df):
    # log troops (heavy right tail); abs lat; missingness imputation w/ cohort mean
    df = df.copy()
    df['log_troops'] = np.log1p(df['troop_count'].fillna(0))
    df['abs_lat'] = df['lat'].abs()

    # Impute zillow_zori_5y_cagr and zhvi_5y_cagr with cohort mean
    for c in ['zillow_zori_5y_cagr', 'zhvi_5y_cagr',
              'years_active', 'dist_to_coast_km']:
        mu = df[c].mean()
        df[c + '_imputed'] = df[c].isna().astype(int)
        df[c] = df[c].fillna(mu)
    return df


def standardize(X, mean=None, std=None):
    if mean is None:
        mean = X.mean(axis=0)
        std = X.std(axis=0, ddof=0)
        std = std.replace(0, 1)
    return (X - mean) / std, mean, std


def fit_ols(X_train, y_train, X_test, y_test, feat_names):
    Xtr = sm.add_constant(X_train, has_constant='add')
    Xte = sm.add_constant(X_test, has_constant='add')
    model = sm.OLS(y_train, Xtr).fit()
    yhat_tr = model.predict(Xtr)
    yhat_te = model.predict(Xte)

    def r2(y, yh):
        ss_res = float(((y - yh) ** 2).sum())
        ss_tot = float(((y - y.mean()) ** 2).sum())
        return 1.0 - ss_res / ss_tot if ss_tot > 0 else None

    def mape(y, yh):
        m = (y != 0)
        return float((np.abs((y[m] - yh[m]) / y[m])).mean())

    # Coefficient block — keep `const` at the top
    coefs = []
    params = model.params
    bse = model.bse
    tvals = model.tvalues
    pvals = model.pvalues
    conf = model.conf_int(alpha=0.05)
    names_all = ['const'] + list(feat_names)
    for n in names_all:
        coefs.append({
            'name': n,
            'beta': float(params[n]),
            'se': float(bse[n]),
            'ci_low': float(conf.loc[n, 0]),
            'ci_high': float(conf.loc[n, 1]),
            'tstat': float(tvals[n]),
            'pvalue': float(pvals[n]),
        })

    return {
        'r2_train': r2(y_train, yhat_tr),
        'r2_test': r2(y_test, yhat_te),
        'mape_train': mape(y_train, yhat_tr),
        'mape_test': mape(y_test, yhat_te),
        'rmse_test': float(np.sqrt(((y_test - yhat_te) ** 2).mean())),
        'adj_r2_train': float(model.rsquared_adj),
        'f_pvalue': float(model.f_pvalue),
        'coefficients': coefs,
        'n_train': int(len(y_train)),
        'n_test': int(len(y_test)),
    }


def partial_correlations(df, feat_names, target):
    cols = list(feat_names) + [target]
    M = df[cols].to_numpy(dtype=float)
    # Pearson correlation matrix
    C = np.corrcoef(M, rowvar=False)
    out = {}
    for i, a in enumerate(cols):
        out[a] = {b: float(C[i, j]) for j, b in enumerate(cols)}
    return out


def evaluate_hypotheses(ols, feat_names, partials, target='bah_5y_cagr'):
    """Pre-registered hypothesis check using the OLS coefficients + bivariate r."""
    coefs = {c['name']: c for c in ols['coefficients']}
    def verdict(beta, pvalue, expected_sign):
        if pvalue is None or math.isnan(pvalue):
            return 'inconclusive'
        sign_match = (beta > 0 and expected_sign == '+') or (beta < 0 and expected_sign == '-')
        if pvalue < 0.05 and sign_match:
            return 'supported (p<0.05, sign matches)'
        if pvalue < 0.05 and not sign_match:
            return f'rejected (p<0.05 but sign opposite — beta={beta:+.3g})'
        if pvalue < 0.10 and sign_match:
            return 'weakly supported (p<0.10, sign matches)'
        return f'inconclusive (p={pvalue:.2f}, beta={beta:+.3g})'

    h = {}
    # H1: BAH growth follows ZORI growth — coefficient on zillow_zori_5y_cagr positive
    if 'zillow_zori_5y_cagr' in coefs:
        c = coefs['zillow_zori_5y_cagr']
        h['H1_bah_lags_zori'] = {
            'expected_sign': '+',
            'beta_standardized': c['beta'],
            'pvalue': c['pvalue'],
            'bivariate_r': partials['zillow_zori_5y_cagr'][target],
            'verdict': verdict(c['beta'], c['pvalue'], '+'),
        }
    # H2: median income growth — untestable directly
    h['H2_income_growth'] = {
        'verdict': 'untestable — no per-MHA ACS income series cached locally',
        'note': 'not run; flagged in RESEARCH.md',
    }
    # H3: coast/inland gradient — dist_to_coast_km coefficient negative (closer = faster)
    if 'dist_to_coast_km' in coefs:
        c = coefs['dist_to_coast_km']
        h['H3_coast_gradient'] = {
            'expected_sign': '- (closer to coast → higher CAGR)',
            'beta_standardized': c['beta'],
            'pvalue': c['pvalue'],
            'bivariate_r': partials['dist_to_coast_km'][target],
            'verdict': verdict(c['beta'], c['pvalue'], '-'),
        }
    # H4: population growth — untestable directly
    h['H4_pop_growth'] = {
        'verdict': 'untestable — no per-MHA Census population series cached locally',
        'note': 'not run; flagged in RESEARCH.md',
    }
    # H5: troop count — log_troops coefficient positive
    if 'log_troops' in coefs:
        c = coefs['log_troops']
        h['H5_troop_count'] = {
            'expected_sign': '+',
            'beta_standardized': c['beta'],
            'pvalue': c['pvalue'],
            'bivariate_r': partials['log_troops'][target],
            'verdict': verdict(c['beta'], c['pvalue'], '+'),
            'caveat': 'cross-sectional level, not change — H5 strictly tested as "level proxy"',
        }
    return h


def main():
    print('[run] loading milbase data.json...', file=sys.stderr)
    with open(MILBASE) as f:
        data = json.load(f)

    panel = build_mha_panel(data)
    print(f'[panel] {len(panel)} MHA rows', file=sys.stderr)

    panel = add_engineered_features(panel)

    # Drop any rows still missing target or core features
    feat_names = [
        'zillow_zori_5y_cagr',
        'zhvi_5y_cagr',
        'log_troops',
        'years_active',
        'abs_lat',
        'lon',
        'dist_to_coast_km',
    ]
    target = 'bah_5y_cagr'
    keep = panel[feat_names + [target]].dropna()
    keep_idx = keep.index
    panel = panel.loc[keep_idx].reset_index(drop=True)
    print(f'[panel] {len(panel)} MHA rows after final dropna', file=sys.stderr)

    # 80/20 split, random_state=42
    rng = np.random.default_rng(42)
    n = len(panel)
    idx = np.arange(n)
    rng.shuffle(idx)
    n_test = int(round(n * 0.20))
    test_idx = idx[:n_test]
    train_idx = idx[n_test:]

    X = panel[feat_names].to_numpy(dtype=float)
    y = panel[target].to_numpy(dtype=float)

    X_train_raw = pd.DataFrame(X[train_idx], columns=feat_names)
    X_test_raw = pd.DataFrame(X[test_idx], columns=feat_names)

    X_train, mu, sd = standardize(X_train_raw)
    X_test, _, _ = standardize(X_test_raw, mu, sd)

    y_train = pd.Series(y[train_idx])
    y_test = pd.Series(y[test_idx])

    print(f'[fit] OLS on n_train={len(train_idx)} features={len(feat_names)}',
          file=sys.stderr)
    ols = fit_ols(X_train, y_train, X_test, y_test, feat_names)

    # Partial / bivariate correlation matrix on the full panel
    partials = partial_correlations(panel, feat_names, target)

    hyp = evaluate_hypotheses(ols, feat_names, partials, target=target)

    # GBM: planned but sklearn missing in this env — report null
    gbm = None
    try:
        from sklearn.ensemble import GradientBoostingRegressor  # noqa
        from sklearn.metrics import r2_score
        gbr = GradientBoostingRegressor(
            n_estimators=200, max_depth=3, learning_rate=0.05,
            random_state=42, subsample=0.85,
        )
        gbr.fit(X_train_raw, y_train)
        gbm = {
            'r2_train': float(r2_score(y_train, gbr.predict(X_train_raw))),
            'r2_test':  float(r2_score(y_test,  gbr.predict(X_test_raw))),
            'feature_importance': {
                feat_names[i]: float(gbr.feature_importances_[i])
                for i in range(len(feat_names))
            },
            'params': {'n_estimators': 200, 'max_depth': 3,
                       'learning_rate': 0.05, 'subsample': 0.85,
                       'random_state': 42},
        }
    except ImportError:
        gbm = {
            'available': False,
            'reason': 'sklearn not installed in working environment; OLS is the load-bearing model',
        }

    results = {
        'n': int(n),
        'n_train': int(len(train_idx)),
        'n_test': int(len(test_idx)),
        'features': feat_names,
        'target': target,
        'ols': ols,
        'gbm': gbm,
        'partial_correlations': partials,
        'feature_means_train': {feat_names[i]: float(mu.iloc[i]) for i in range(len(feat_names))},
        'feature_stds_train':  {feat_names[i]: float(sd.iloc[i]) for i in range(len(feat_names))},
        'target_summary': {
            'mean': float(panel[target].mean()),
            'std':  float(panel[target].std(ddof=0)),
            'min':  float(panel[target].min()),
            'p25':  float(panel[target].quantile(0.25)),
            'p50':  float(panel[target].quantile(0.50)),
            'p75':  float(panel[target].quantile(0.75)),
            'max':  float(panel[target].max()),
        },
        'hypothesis_results': hyp,
        'missingness': {
            'zillow_zori_5y_cagr_imputed_count': int(panel['zillow_zori_5y_cagr_imputed'].sum())
                if 'zillow_zori_5y_cagr_imputed' in panel.columns else 0,
            'zhvi_5y_cagr_imputed_count': int(panel['zhvi_5y_cagr_imputed'].sum())
                if 'zhvi_5y_cagr_imputed' in panel.columns else 0,
        },
        '_meta': {
            'fitted': date.today().isoformat(),
            'no_fabrication': True,
            'split': {'train_pct': 0.80, 'test_pct': 0.20, 'random_state': 42},
            'standardization': 'z-score on training set, applied to test set',
            'tier': 'BAH E5 with-dependents (2021->2026 CAGR)',
            'sources': [
                'DoD Defense Travel Management Office — historical BAH 2013-2026',
                'Zillow Research — ZORI (Observed Rent Index), ZHVI (Home Value Index)',
                'DoD / open-source installation rosters — troop_count, lat/lon, years_active',
            ],
            'limitations': [
                f'n={n} is small for cross-sectional inference',
                'No spatial-cluster correction on standard errors',
                'H2 (income) and H4 (population) untestable here — no cached ACS / Census per-MHA series',
                'H5 tested as cross-sectional level only, not change-in-troops',
                'OLS p-values reported are suggestive, not confirmatory',
                'sklearn missing: gradient boosting result is null in this environment',
            ],
        },
    }

    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    with open(OUT, 'w') as f:
        json.dump(results, f, indent=2, default=float)

    # Short readable summary
    print('\n=== OLS coefficient summary (standardized) ===', file=sys.stderr)
    for c in ols['coefficients']:
        flag = ''
        if c['name'] != 'const':
            if c['pvalue'] < 0.01: flag = ' ***'
            elif c['pvalue'] < 0.05: flag = ' **'
            elif c['pvalue'] < 0.10: flag = ' *'
        print(f"  {c['name']:30s} beta={c['beta']:+.4f}  "
              f"95%CI=[{c['ci_low']:+.4f},{c['ci_high']:+.4f}]  "
              f"t={c['tstat']:+.2f}  p={c['pvalue']:.3f}{flag}",
              file=sys.stderr)
    print(f"\n  R^2 train = {ols['r2_train']:.4f}", file=sys.stderr)
    print(f"  R^2 test  = {ols['r2_test']:.4f}", file=sys.stderr)
    print(f"  MAPE test = {ols['mape_test']:.4f}", file=sys.stderr)
    print(f"\nwrote {OUT}", file=sys.stderr)


if __name__ == '__main__':
    main()