#!/usr/bin/env python3
"""
R-CITY.rent-cagr-regression — what predicts city-level 5y ZORI CAGR?

Loads:
  - projects/cities/data/data.json (250 cities w/ ACS features, lat/lon, pop)
  - /tmp/cities_zori.csv (Zillow City ZORI sfrcondomfr, monthly, cached)
  - /tmp/cities_zhvi_all.csv (Zillow City ZHVI all-homes, monthly, cached)
  - /tmp/2023_Gaz_zcta_national.txt (Census 2023 Gazetteer — ZCTA centroids)
  - projects/zipdata/data/zip_master.json (31k ZCTAs w/ ACS/IRS/BLS features)

Pre-registered in projects/city-rent-research/RESEARCH.md.
OLS + partial-correlation matrix. GBM planned but sklearn unavailable in this
env so gbm = null (consistent with bah-research methodology).

NO FABRICATION. Cities without a real 2021-04 AND 2026-04 ZORI endpoint are
dropped from the target. ZHVI / ACS features with missing values get cohort-
mean imputation (and a missingness dummy is created but not in the published
feature set, per the BAH study's v1 convention).
"""

import csv
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, '..', '..'))

CITIES_JSON = os.path.join(HEDGE_ROOT, 'projects', 'cities', 'data', 'data.json')
ZIP_MASTER  = os.path.join(HEDGE_ROOT, 'projects', 'zipdata', 'data', 'zip_master.json')
ZORI_CSV    = '/tmp/cities_zori.csv'
ZHVI_CSV    = '/tmp/cities_zhvi_all.csv'
ZCTA_GAZ    = '/tmp/2023_Gaz_zcta_national.txt'

OUT = os.path.join(ROOT, 'data', 'regression_results.json')

# 5-year window endpoints in the Zillow monthly column headers
ZORI_START = '2021-04-30'
ZORI_END   = '2026-04-30'
N_YEARS    = 5

SOURCES = [
    'Zillow Research — City ZORI (sfrcondomfr smoothed, CC-BY) — target + asset-price control',
    'Zillow Research — City ZHVI all-homes (sfrcondo blend, CC-BY) — H1 asset-price control',
    'U.S. Census Population Estimates Program SUB-EST2024 — population (public domain)',
    'U.S. Census ACS 5-yr 2023 B19013 — median household income (public domain)',
    'U.S. Census ACS 5-yr 2023 B01002 — median age (public domain)',
    'U.S. Census 2023 Gazetteer Places — city lat/lon (public domain)',
    'U.S. Census 2023 Gazetteer ZCTA — ZCTA centroid for nearest-join (public domain)',
    'U.S. Census ACS 5-yr 2023 B15003 — bachelor_plus_pct (joined via nearest ZCTA)',
    'U.S. Census ACS 5-yr 2023 B23025 — unemployment_pct_acs (joined via nearest ZCTA)',
    'IRS SOI 2022 county-to-county migration — irs_inflow_2022 / irs_outflow_2022 (joined via nearest ZCTA)',
]


# ---------- Geographic helper ----------

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))


# ---------- Zillow per-city CAGR ----------

def load_zillow_series(path, label):
    """
    Returns {(state_abbrev, normalized_city_name): {YYYY-MM-DD: value}}.
    Header columns include 'RegionName', 'State', and many date columns.
    """
    out = {}
    with open(path, encoding='utf-8') as f:
        reader = csv.reader(f)
        header = next(reader)
        # Find name + state cols and date cols
        try:
            i_name  = header.index('RegionName')
            i_state = header.index('State')
        except ValueError:
            raise RuntimeError(f'{label}: header missing RegionName/State')
        # Date columns are anything like YYYY-MM-DD
        date_cols = [(i, c) for i, c in enumerate(header) if len(c) == 10 and c[4] == '-' and c[7] == '-']

        for row in reader:
            if len(row) < len(header):
                # Some rows have city names containing commas + quoted metros (we saw "New York-Newark...")
                # csv.reader handles standard quoting; if it doesn't, skip the malformed row.
                continue
            name = row[i_name].strip()
            state = row[i_state].strip().upper()
            if not name or not state:
                continue
            key = (state, normalize_name(name))
            series = {}
            for i, dc in date_cols:
                v = row[i].strip()
                if v == '' or v.lower() == 'nan':
                    continue
                try:
                    series[dc] = float(v)
                except ValueError:
                    continue
            out[key] = series
    print(f'[zillow] {label}: {len(out):,} city series', file=sys.stderr)
    return out


def normalize_name(s):
    """
    Match the join logic in projects/cities/scripts/pull_cities_dscr.py:
    lowercase, strip suffixes (city|town|village|borough|township|CDP|balance|pt.),
    strip punctuation, st.->saint.
    """
    s = s.lower().strip()
    s = s.replace('st.', 'saint').replace('st ', 'saint ')
    for suf in (' city', ' town', ' village', ' borough', ' township',
                ' cdp', ' balance', ' pt.', ' (balance)'):
        if s.endswith(suf):
            s = s[: -len(suf)]
    # Strip remaining punctuation except hyphen and space
    out = []
    for ch in s:
        if ch.isalnum() or ch in '- ':
            out.append(ch)
    return ''.join(out).strip()


def cagr(v0, v1, n):
    if v0 is None or v1 is None or v0 <= 0 or v1 <= 0 or n <= 0:
        return None
    return (v1 / v0) ** (1.0 / n) - 1.0


# ---------- ZCTA centroid + nearest-neighbor ----------

def load_zcta_centroids(path):
    """Returns list of (zcta, lat, lon) tuples."""
    out = []
    with open(path, encoding='utf-8') as f:
        header = f.readline().split('\t')
        header = [h.strip() for h in header]
        i_zcta = header.index('GEOID')
        i_lat  = header.index('INTPTLAT')
        i_lon  = header.index('INTPTLONG')
        for line in f:
            parts = line.rstrip('\n').split('\t')
            if len(parts) <= max(i_zcta, i_lat, i_lon):
                continue
            try:
                z = parts[i_zcta].strip().zfill(5)
                lat = float(parts[i_lat].strip())
                lon = float(parts[i_lon].strip())
            except ValueError:
                continue
            out.append((z, lat, lon))
    print(f'[zcta] loaded {len(out):,} ZCTA centroids', file=sys.stderr)
    return out


def load_zip_master(path):
    """Returns {zcta: rec}."""
    with open(path) as f:
        d = json.load(f)
    by_zcta = {}
    for z in d.get('zips', []):
        by_zcta[z['zcta']] = z
    print(f'[zipmaster] loaded {len(by_zcta):,} ZCTA records', file=sys.stderr)
    return by_zcta


def nearest_zcta(city_lat, city_lon, zcta_centroids):
    """Brute-force nearest by haversine. n_zcta ≈ 33k × n_city = 250 → 8M ops — fine."""
    if city_lat is None or city_lon is None:
        return None, None
    best_z = None; best_d = float('inf')
    for z, lat, lon in zcta_centroids:
        d = haversine_km(city_lat, city_lon, lat, lon)
        if d < best_d:
            best_d = d; best_z = z
    return best_z, best_d


# ---------- Per-city feature compute ----------

def compute_city_features(city, zori, zhvi, zcta_centroids, zip_master):
    """Compute one city's full feature row. Called concurrently."""
    name = city['name']
    state = city['state']
    key = (state, normalize_name(name))

    zori_series = zori.get(key) or {}
    zhvi_series = zhvi.get(key) or {}

    zori_5y_cagr = cagr(zori_series.get(ZORI_START), zori_series.get(ZORI_END), N_YEARS)
    zhvi_5y_cagr = cagr(zhvi_series.get(ZORI_START), zhvi_series.get(ZORI_END), N_YEARS)

    # Nearest ZCTA join
    z, dist_km = nearest_zcta(city.get('lat'), city.get('lon'), zcta_centroids)
    zrec = zip_master.get(z) if z else None

    bachelor_plus_pct = zrec.get('bachelor_plus_pct') if zrec else None
    unemployment_pct = zrec.get('unemployment_pct_acs') if zrec else None
    inflow  = zrec.get('irs_inflow_2022') if zrec else None
    outflow = zrec.get('irs_outflow_2022') if zrec else None
    net_mig = None
    if inflow is not None and outflow is not None:
        net_mig = float(inflow) - float(outflow)

    return {
        'city_id': city['city_id'],
        'name': name,
        'state': state,
        'lat': city.get('lat'),
        'lon': city.get('lon'),
        'population': city.get('population'),
        'median_income': city.get('median_income'),
        'median_age': city.get('median_age'),
        'nearest_zcta': z,
        'nearest_zcta_dist_km': dist_km,
        'bachelor_plus_pct': bachelor_plus_pct,
        'unemployment_pct_acs': unemployment_pct,
        'irs_inflow_2022': inflow,
        'irs_outflow_2022': outflow,
        'irs_net_migration_2022': net_mig,
        'zori_5y_cagr': zori_5y_cagr,
        'zhvi_5y_cagr': zhvi_5y_cagr,
    }


def build_panel(cities, zori, zhvi, zcta_centroids, zip_master):
    print(f'[panel] computing features for {len(cities)} cities (parallel)...',
          file=sys.stderr)
    rows = []
    with ThreadPoolExecutor(max_workers=8) as ex:
        futures = [
            ex.submit(compute_city_features, c, zori, zhvi, zcta_centroids, zip_master)
            for c in cities
        ]
        for f in futures:
            rows.append(f.result())
    df = pd.DataFrame(rows)
    print(f'[panel] {len(df)} rows built', file=sys.stderr)
    return df


# ---------- Feature engineering ----------

def add_engineered_features(df):
    df = df.copy()
    df['log_population'] = np.log1p(df['population'].fillna(0))
    df['abs_lat'] = df['lat'].abs()
    # Cohort-mean imputation for non-target features
    for c in ['bachelor_plus_pct', 'unemployment_pct_acs',
              'irs_net_migration_2022', 'median_age',
              'median_income', 'zhvi_5y_cagr']:
        mu = df[c].mean()
        df[c + '_imputed'] = df[c].isna().astype(int)
        df[c] = df[c].fillna(mu)
    return df


# ---------- OLS + diagnostics ----------

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())

    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)
    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, partials, target='zori_5y_cagr'):
    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})'

    def row(feat, expected_sign):
        if feat not in coefs:
            return {'verdict': f'{feat} not in model'}
        c = coefs[feat]
        return {
            'expected_sign': expected_sign,
            'beta_standardized': c['beta'],
            'pvalue': c['pvalue'],
            'bivariate_r': partials.get(feat, {}).get(target),
            'verdict': verdict(c['beta'], c['pvalue'], expected_sign),
        }

    return {
        'H1_zori_lags_zhvi':       row('zhvi_5y_cagr',          '+'),
        'H2_education_premium':    row('bachelor_plus_pct',     '+'),
        'H3_net_migration':        row('irs_net_migration_2022','+'),
        'H4_unemployment_inverse': row('unemployment_pct_acs',  '-'),
        'H5_income_level':         row('median_income',         '+'),
        'H6_population_size':      row('log_population',        '+'),
    }


def main():
    print('[run] loading cities, ZORI, ZHVI, ZCTA centroids, zip_master...',
          file=sys.stderr)
    with open(CITIES_JSON) as f:
        cities_data = json.load(f)
    cities = cities_data['cities']
    print(f'[run] {len(cities)} cities loaded', file=sys.stderr)

    # Parallel IO-bound loads: ZORI + ZHVI + ZCTA + ZIP master concurrently
    with ThreadPoolExecutor(max_workers=4) as ex:
        f_zori = ex.submit(load_zillow_series, ZORI_CSV, 'ZORI')
        f_zhvi = ex.submit(load_zillow_series, ZHVI_CSV, 'ZHVI all-homes')
        f_gaz  = ex.submit(load_zcta_centroids, ZCTA_GAZ)
        f_zm   = ex.submit(load_zip_master, ZIP_MASTER)
        zori = f_zori.result()
        zhvi = f_zhvi.result()
        zcta_centroids = f_gaz.result()
        zip_master = f_zm.result()

    panel = build_panel(cities, zori, zhvi, zcta_centroids, zip_master)

    # Drop cities w/o a real ZORI 5y CAGR (NO IMPUTATION of the target)
    n_before = len(panel)
    panel = panel.dropna(subset=['zori_5y_cagr']).reset_index(drop=True)
    n_after = len(panel)
    print(f'[panel] dropped {n_before - n_after} rows missing zori_5y_cagr target '
          f'-> {n_after} cities for fit', file=sys.stderr)

    panel = add_engineered_features(panel)

    feat_names = [
        'zhvi_5y_cagr',
        'bachelor_plus_pct',
        'irs_net_migration_2022',
        'unemployment_pct_acs',
        'median_income',
        'log_population',
        'abs_lat',
        'lon',
        'median_age',
    ]
    target = 'zori_5y_cagr'

    # Final dropna across feature set (post-imputation this should be 0)
    keep = panel[feat_names + [target]].dropna()
    keep_idx = keep.index
    panel = panel.loc[keep_idx].reset_index(drop=True)
    print(f'[panel] {len(panel)} rows after final dropna', file=sys.stderr)

    # 80/20 split
    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)

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

    # GBM (sklearn-dependent; reports null if absent)
    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': {
            c + '_imputed_count': int(panel[c + '_imputed'].sum())
            for c in ['bachelor_plus_pct', 'unemployment_pct_acs',
                      'irs_net_migration_2022', 'median_age',
                      'median_income', 'zhvi_5y_cagr']
            if c + '_imputed' in panel.columns
        },
        'attrition': {
            'n_cities_input': len(cities),
            'n_after_zori_target_filter': int(n),
        },
        '_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',
            'window': f'{ZORI_START} → {ZORI_END} ({N_YEARS}-yr CAGR)',
            'sources': SOURCES,
            'limitations': [
                f'n={n} max — small-sample cross-section',
                'No spatial-cluster correction on standard errors',
                'Nearest-ZCTA centroid join is an approximation for large multi-ZCTA cities',
                'IRS migration is 2022-only (single-year), not a 5y cumulative net flow',
                'H5 risks reverse causation: high-rent cities select for higher-income earners',
                '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  n        = {n}", file=sys.stderr)
    print(f"  n_train  = {len(train_idx)}, n_test = {len(test_idx)}", file=sys.stderr)
    print(f"  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()