#!/usr/bin/env python3
# ─────────────────────────────────────────────────────────────────────────────
# compute.py — "If the US split into 4 quadrants and went to war, who's winning?"
#
# Assigns each contiguous state (+DC) to a quadrant (A=NW, B=NE, C=SW, D=SE),
# applies sub-state metro corrections to the population column, aggregates, and
# runs descriptive concentration statistics. Emits CSVs + a JSON payload the
# D3 page consumes. Pure stdlib.
#
# Data vintages (all fetched from Wikipedia, internally consistent):
#   • Population : Census Bureau July 1, 2025 estimates
#   • GDP        : 2025 nominal, US$ millions
#   • Land area  : Census land-area (sq mi)
#   • Metros     : Census July 1, 2025 MSA / metro-division estimates
# ─────────────────────────────────────────────────────────────────────────────
from __future__ import annotations

import csv
import json
from pathlib import Path

OUT = Path(__file__).parent

# ── Section 1 · Raw per-state data ───────────────────────────────────────────
# (population_2025, gdp_2025_millions, land_area_sqmi). AK/HI excluded: the meme
# map shows only the contiguous 48. DC included (it's on the map, in quadrant B).
STATES: dict[str, tuple[int, float, int]] = {
    'AL': (5_193_088, 341_154, 50_645),
    'AZ': (7_623_818, 598_189, 113_594),
    'AR': (3_114_791, 198_422, 52_035),
    'CA': (39_355_309, 4_251_000, 155_779),
    'CO': (6_012_561, 584_324, 103_642),
    'CT': (3_688_496, 376_455, 4_842),
    'DE': (1_059_952, 117_218, 1_949),
    'DC': (693_645, 192_618, 61),
    'FL': (23_462_518, 1_835_000, 53_625),
    'GA': (11_302_748, 924_829, 57_513),
    'ID': (2_029_733, 135_553, 82_643),
    'IL': (12_719_141, 1_202_000, 55_519),
    'IN': (6_973_333, 545_234, 35_826),
    'IA': (3_238_387, 277_110, 55_857),
    'KS': (2_977_220, 241_378, 81_759),
    'KY': (4_606_864, 306_897, 39_486),
    'LA': (4_618_189, 340_080, 43_204),
    'ME': (1_414_874, 102_844, 30_843),
    'MD': (6_265_347, 568_140, 9_707),
    'MA': (7_154_084, 820_105, 7_800),
    'MI': (10_127_884, 730_068, 56_539),
    'MN': (5_830_405, 531_465, 79_627),
    'MS': (2_954_160, 165_069, 46_923),
    'MO': (6_270_541, 468_470, 68_742),
    'MT': (1_144_694, 82_358, 145_546),
    'NE': (2_018_006, 198_073, 76_824),
    'NV': (3_282_188, 281_454, 109_781),
    'NH': (1_415_342, 125_523, 8_953),
    'NJ': (9_548_215, 887_175, 7_354),
    'NM': (2_125_498, 152_779, 121_298),
    'NY': (20_002_427, 2_468_000, 47_126),
    'NC': (11_197_968, 893_763, 48_618),
    'ND': (799_358, 81_883, 69_001),
    'OH': (11_900_510, 966_780, 40_861),
    'OK': (4_123_288, 274_421, 68_595),
    'OR': (4_273_586, 342_850, 95_988),
    'PA': (13_059_432, 1_056_000, 44_743),
    'RI': (1_114_521, 83_956, 1_034),
    'SC': (5_570_274, 378_831, 30_061),
    'SD': (935_094, 80_650, 75_811),
    'TN': (7_315_076, 589_818, 41_235),
    'TX': (31_709_821, 2_904_000, 261_232),
    'UT': (3_538_904, 315_973, 82_170),
    'VT': (644_663, 48_350, 9_217),
    'VA': (8_880_107, 798_448, 39_490),
    'WA': (8_001_020, 894_990, 66_456),
    'WV': (1_766_147, 109_277, 24_038),
    'WI': (5_972_787, 473_037, 54_158),
    'WY': (588_753, 52_622, 97_093),
}

STATE_NAMES = {
    'AL': 'Alabama',
    'AZ': 'Arizona',
    'AR': 'Arkansas',
    'CA': 'California',
    'CO': 'Colorado',
    'CT': 'Connecticut',
    'DE': 'Delaware',
    'DC': 'District of Columbia',
    'FL': 'Florida',
    'GA': 'Georgia',
    'ID': 'Idaho',
    'IL': 'Illinois',
    'IN': 'Indiana',
    'IA': 'Iowa',
    'KS': 'Kansas',
    'KY': 'Kentucky',
    'LA': 'Louisiana',
    'ME': 'Maine',
    'MD': 'Maryland',
    'MA': 'Massachusetts',
    'MI': 'Michigan',
    'MN': 'Minnesota',
    'MS': 'Mississippi',
    'MO': 'Missouri',
    'MT': 'Montana',
    'NE': 'Nebraska',
    'NV': 'Nevada',
    'NH': 'New Hampshire',
    'NJ': 'New Jersey',
    'NM': 'New Mexico',
    'NY': 'New York',
    'NC': 'North Carolina',
    'ND': 'North Dakota',
    'OH': 'Ohio',
    'OK': 'Oklahoma',
    'OR': 'Oregon',
    'PA': 'Pennsylvania',
    'RI': 'Rhode Island',
    'SC': 'South Carolina',
    'SD': 'South Dakota',
    'TN': 'Tennessee',
    'TX': 'Texas',
    'UT': 'Utah',
    'VT': 'Vermont',
    'VA': 'Virginia',
    'WA': 'Washington',
    'WV': 'West Virginia',
    'WI': 'Wisconsin',
    'WY': 'Wyoming',
}

# ── Section 2 · State → quadrant assignment ──────────────────────────────────
# Rule: assign each state to the quadrant holding the majority of its LANDMASS
# as drawn in the image (not real latitude, not population centroid). The two
# red lines fall near the KS/NE/MO/CO junction, so a handful of states are
# genuinely bisected — those are the SWING set, flagged and stress-tested.
QUADRANT = {
    # A — Northwest (interior west + northern plains)
    'WA': 'A',
    'OR': 'A',
    'ID': 'A',
    'MT': 'A',
    'WY': 'A',
    'NV': 'A',
    'UT': 'A',
    'CO': 'A',
    'ND': 'A',
    'SD': 'A',
    'NE': 'A',
    # B — Northeast (Midwest + Northeast + DC)
    'MN': 'B',
    'IA': 'B',
    'WI': 'B',
    'MI': 'B',
    'IL': 'B',
    'IN': 'B',
    'OH': 'B',
    'PA': 'B',
    'NY': 'B',
    'NJ': 'B',
    'CT': 'B',
    'RI': 'B',
    'MA': 'B',
    'VT': 'B',
    'NH': 'B',
    'ME': 'B',
    'MD': 'B',
    'DE': 'B',
    'DC': 'B',
    # C — Southwest (desert SW + southern plains + Texas)
    'CA': 'C',
    'AZ': 'C',
    'NM': 'C',
    'KS': 'C',
    'OK': 'C',
    'TX': 'C',
    # D — Southeast (the South + lower Midwest edge)
    'MO': 'D',
    'AR': 'D',
    'LA': 'D',
    'MS': 'D',
    'AL': 'D',
    'GA': 'D',
    'FL': 'D',
    'SC': 'D',
    'NC': 'D',
    'TN': 'D',
    'KY': 'D',
    'WV': 'D',
    'VA': 'D',
}

# Swing states: bisected by a line; base call + the side they'd flip to.
SWING = {'CO': 'C', 'KS': 'A', 'MO': 'B', 'WV': 'B', 'OK': 'C'}  # value = flip-to

QUAD_NAME = {'A': 'Northwest', 'B': 'Northeast', 'C': 'Southwest', 'D': 'Southeast'}
QUAD_COLOR = {'A': '#2a78d6', 'B': '#1baf7a', 'C': '#eda100', 'D': '#008300'}

# ── Section 3 · Sub-state metro corrections (population only) ─────────────────
# The user's methodology: whole states up top, then move metros that sit in a
# different quadrant than the rest of their state. Population is the corrected
# metric; GDP & land area stay at whole-state resolution (stated simplification).
#
# Each: (label, people, from_quad, to_quad, tier). Tier 1 = base case;
# Tier 2 = plausible-but-debatable, shown only in sensitivity.
CORRECTIONS = [
    ('SF–Oakland–Fremont MSA (Bay Area ex-South Bay)', 4_630_041, 'C', 'A', 1),
    ('Las Vegas–Henderson MSA', 2_407_226, 'A', 'C', 1),
    ('Northern Virginia (Arlington/Fairfax/Loudoun/PWC)', 2_340_000, 'D', 'B', 1),
    ('Houston–Pasadena–The Woodlands MSA', 7_904_627, 'C', 'D', 1),
    # Tier 2 — sensitivity only
    ('Austin–Round Rock–San Marcos MSA', 2_620_945, 'C', 'D', 2),
    ('Dallas–Plano–Irving metro division', 5_641_795, 'C', 'D', 2),
]


# ── Section 4 · Aggregation helpers ──────────────────────────────────────────
def base_totals() -> dict[str, dict[str, float]]:
    agg = {q: {'pop': 0, 'gdp': 0.0, 'area': 0, 'n': 0} for q in 'ABCD'}
    for st, (pop, gdp, area) in STATES.items():
        q = QUADRANT[st]
        agg[q]['pop'] += pop
        agg[q]['gdp'] += gdp
        agg[q]['area'] += area
        agg[q]['n'] += 1
    return agg


def apply_corrections(pop_by_quad: dict[str, int], tier_max: int) -> dict[str, int]:
    out = dict(pop_by_quad)
    for _label, people, frm, to, tier in CORRECTIONS:
        if tier <= tier_max:
            out[frm] -= people
            out[to] += people
    return out


def hhi(shares: list[float]) -> float:
    """Herfindahl index over member shares within a quadrant (0..1)."""
    return sum(s * s for s in shares)


def coef_var(values: list[float]) -> float:
    n = len(values)
    if n < 2:
        return 0.0
    mean = sum(values) / n
    if mean == 0:
        return 0.0
    var = sum((v - mean) ** 2 for v in values) / n
    return (var**0.5) / mean


# ── Section 5 · Build everything ─────────────────────────────────────────────
def main() -> None:
    agg = base_totals()

    pop_state_level = {q: agg[q]['pop'] for q in 'ABCD'}
    pop_corrected = apply_corrections(pop_state_level, tier_max=1)
    pop_tier2 = apply_corrections(pop_state_level, tier_max=2)

    # Swing flip: move every swing state to its alternate side, recompute pop.
    pop_swing = {q: agg[q]['pop'] for q in 'ABCD'}
    for st, flip_to in SWING.items():
        cur = QUADRANT[st]
        pop_swing[cur] -= STATES[st][0]
        pop_swing[flip_to] += STATES[st][0]
    pop_swing_corrected = apply_corrections(pop_swing, tier_max=1)

    nat_pop = sum(pop_state_level.values())
    nat_gdp = sum(agg[q]['gdp'] for q in 'ABCD')
    nat_area = sum(agg[q]['area'] for q in 'ABCD')

    # Per-quadrant intra-concentration (on CORRECTED-basis membership pop, but
    # HHI/CoV computed on whole-state pops within each quadrant's base members).
    members = {q: [st for st in STATES if QUADRANT[st] == q] for q in 'ABCD'}
    intra = {}
    for q in 'ABCD':
        pops = sorted((STATES[st][0] for st in members[q]), reverse=True)
        tot = sum(pops)
        shares = [p / tot for p in pops]
        top_state = max(members[q], key=lambda s: STATES[s][0])
        intra[q] = {
            'hhi': hhi(shares),
            'cov': coef_var([float(p) for p in pops]),
            'top_state': top_state,
            'top_state_share': STATES[top_state][0] / tot,
            'n_members': len(pops),
        }

    # ---- CSV 1: per-state ----
    with (OUT / 'quadrant_states.csv').open('w', newline='') as f:
        w = csv.writer(f)
        w.writerow(
            [
                'abbr',
                'state',
                'quadrant',
                'quadrant_name',
                'population_2025',
                'gdp_2025_musd',
                'land_area_sqmi',
                'is_swing',
                'swing_flip_to',
            ]
        )
        for st in sorted(STATES, key=lambda s: (QUADRANT[s], -STATES[s][0])):
            pop, gdp, area = STATES[st]
            w.writerow(
                [
                    st,
                    STATE_NAMES[st],
                    QUADRANT[st],
                    QUAD_NAME[QUADRANT[st]],
                    pop,
                    gdp,
                    area,
                    'yes' if st in SWING else 'no',
                    SWING.get(st, ''),
                ]
            )

    # ---- CSV 2: corrections ----
    with (OUT / 'quadrant_corrections.csv').open('w', newline='') as f:
        w = csv.writer(f)
        w.writerow(['metro', 'population', 'from_quadrant', 'to_quadrant', 'tier'])
        for label, people, frm, to, tier in CORRECTIONS:
            w.writerow([label, people, frm, to, tier])

    # ---- CSV 3: quadrant summary ----
    with (OUT / 'quadrant_summary.csv').open('w', newline='') as f:
        w = csv.writer(f)
        w.writerow(
            [
                'quadrant',
                'name',
                'n_states',
                'pop_state_level',
                'pop_corrected',
                'pop_sensitivity_tier2',
                'pop_share_corrected_pct',
                'gdp_2025_musd',
                'gdp_share_pct',
                'land_area_sqmi',
                'area_share_pct',
                'pop_density_per_sqmi',
                'gdp_per_capita_usd',
                'intra_hhi',
                'intra_cov',
                'top_state',
                'top_state_share_pct',
            ]
        )
        for q in 'ABCD':
            pc = pop_corrected[q]
            w.writerow(
                [
                    q,
                    QUAD_NAME[q],
                    agg[q]['n'],
                    pop_state_level[q],
                    pc,
                    pop_tier2[q],
                    round(100 * pc / nat_pop, 2),
                    round(agg[q]['gdp'], 0),
                    round(100 * agg[q]['gdp'] / nat_gdp, 2),
                    agg[q]['area'],
                    round(100 * agg[q]['area'] / nat_area, 2),
                    # density & per-capita use the WHOLE-STATE pop basis so
                    # numerator/denominator stay consistent (metro pop moves
                    # but metro GDP/land do not — mixing bases is an artifact).
                    round(pop_state_level[q] / agg[q]['area'], 1),
                    round(agg[q]['gdp'] * 1_000_000 / pop_state_level[q], 0),
                    round(intra[q]['hhi'], 4),
                    round(intra[q]['cov'], 3),
                    intra[q]['top_state'],
                    round(100 * intra[q]['top_state_share'], 1),
                ]
            )

    # ---- JSON payload for the D3 page ----
    payload = {
        'meta': {
            'national_pop': nat_pop,
            'national_gdp_musd': nat_gdp,
            'national_area_sqmi': nat_area,
            'pop_vintage': 'Census July 1, 2025 estimate',
            'gdp_vintage': '2025 nominal',
        },
        'quadrant_name': QUAD_NAME,
        'quadrant_color': QUAD_COLOR,
        'quadrants': [],
        'corrections': [
            {'metro': l, 'pop': p, 'from': frm, 'to': to, 'tier': t}
            for (l, p, frm, to, t) in CORRECTIONS
        ],
        'scenarios': {
            'state_level': pop_state_level,
            'corrected': pop_corrected,
            'tier2': pop_tier2,
            'swing_flip': pop_swing_corrected,
        },
        'states': [
            {
                'abbr': st,
                'name': STATE_NAMES[st],
                'q': QUADRANT[st],
                'pop': STATES[st][0],
                'gdp': STATES[st][1],
                'area': STATES[st][2],
                'swing': st in SWING,
            }
            for st in STATES
        ],
    }
    for q in 'ABCD':
        pc = pop_corrected[q]
        payload['quadrants'].append(
            {
                'id': q,
                'name': QUAD_NAME[q],
                'color': QUAD_COLOR[q],
                'n_states': agg[q]['n'],
                'pop_state_level': pop_state_level[q],
                'pop': pc,
                'pop_share': round(100 * pc / nat_pop, 2),
                'gdp': round(agg[q]['gdp'], 0),
                'gdp_share': round(100 * agg[q]['gdp'] / nat_gdp, 2),
                'area': agg[q]['area'],
                'area_share': round(100 * agg[q]['area'] / nat_area, 2),
                'density': round(pop_state_level[q] / agg[q]['area'], 1),
                'gdp_per_capita': round(agg[q]['gdp'] * 1_000_000 / pop_state_level[q], 0),
                'hhi': round(intra[q]['hhi'], 4),
                'cov': round(intra[q]['cov'], 3),
                'top_state': STATE_NAMES[intra[q]['top_state']],
                'top_state_share': round(100 * intra[q]['top_state_share'], 1),
            }
        )

    (OUT / 'quadrant_data.json').write_text(json.dumps(payload, indent=2))

    # ---- console checksum + headline ----
    print(f'National pop (48+DC): {nat_pop:,}  [checksum ~333M expected]')
    print(f'National GDP (musd) : {nat_gdp:,.0f}')
    print()
    hdr = f'{"Q":<2}{"name":<11}{"n":>3}{"state-lvl":>13}{"corrected":>13}{"share%":>8}{"GDP$M":>12}{"gdp%":>7}{"HHI":>7}{"top state":>16}'
    print(hdr)
    for q in 'ABCD':
        p = payload['quadrants'][-4 + 'ABCD'.index(q)]
    for qd in payload['quadrants']:
        print(
            f'{qd["id"]:<2}{qd["name"]:<11}{qd["n_states"]:>3}'
            f'{qd["pop_state_level"]:>13,}{qd["pop"]:>13,}{qd["pop_share"]:>7}%'
            f'{qd["gdp"]:>12,.0f}{qd["gdp_share"]:>6}%{qd["hhi"]:>7.3f}'
            f'{qd["top_state"]:>16} ({qd["top_state_share"]}%)'
        )
    print()
    print(
        'Corrected ranking by population:',
        ' > '.join(f'{q}' for q in sorted('ABCD', key=lambda q: -pop_corrected[q])),
    )
    print('Swing-flip scenario pop:', {q: f'{pop_swing_corrected[q]:,}' for q in 'ABCD'})


if __name__ == '__main__':
    main()