"""Maine Political Data Project - Partisan Score & Unenrolled Lean, v1.
Implements site-design/scoring-methodology-v1.md.
"""
import re, json
import numpy as np, pandas as pd
from scipy.optimize import nnls

import os
REG = os.environ.get('SCORES_REG', '/home/claude/w/reg61')            # registration v6.1 (out/ folder)
RES = os.environ.get('SCORES_RES', '/home/claude/w/v35')               # election results v3.5 (units_election_additions.csv)
# v2.1: SCORES_INACTIVE = comma list of snapshot dates whose INACTIVE registrants are added to the active ones.
# Default '' = active only (the v2 input). The v2.1 2022 runs use SCORES_INACTIVE=2022-11-08 (audit E9: the Nov 2022 SOS
# file moved ~216,000 voters to inactive status unevenly by party and town). See fix_v2.1_ei_2022.sh.
INACTIVE_SNAPS = [s for s in os.environ.get('SCORES_INACTIVE', '').split(',') if s]
OFFICES = ['President', 'Governor', 'U.S. Senate', 'U.S. House']
NONMAJOR_MAX = 0.15
PRIOR_M = None           # set in main after calibration (pseudo unenrolled registrants)
MAJOR = ['DEM', 'REP', 'UNE']
MINOR_COLS = ['GRN', 'LIB', 'NLB', 'CON', 'AME', 'ALL', 'PEO', 'SOC', 'YES', 'IND', 'REF']
CYCLES = list(range(1990, 2026, 2))

units = pd.read_csv(f'{REG}/units.csv')
uname = dict(zip(units.unit_id, units.canonical_name))
ucty = dict(zip(units.unit_id, units.county_abbr))
snaps = pd.read_csv(f'{REG}/snapshots.csv')
rw = pd.read_csv(f'{REG}/registration_town_wide.csv', low_memory=False)
if INACTIVE_SNAPS:
    _ia = pd.read_csv(f'{REG}/inactive_registration_town_wide.csv', low_memory=False)
    _ia = _ia[_ia.snapshot_date.isin(INACTIVE_SNAPS)]
    _num = MAJOR + MINOR_COLS + ['total']
    _a = rw[rw.snapshot_date.isin(INACTIVE_SNAPS)].set_index(['snapshot_date', 'unit_id'])
    _sum = _a[_num].fillna(0).add(_ia.set_index(['snapshot_date', 'unit_id'])[_num].fillna(0), fill_value=0).reindex(_a.index)
    _a[_num] = _sum
    rw = pd.concat([rw[~rw.snapshot_date.isin(INACTIVE_SNAPS)], _a.reset_index()], ignore_index=True)
rw['MIN'] = rw[MINOR_COLS].fillna(0).sum(axis=1)
for c in MAJOR: rw[c] = rw[c].fillna(0)

def snapshot_for(year):
    s = snaps[(snaps.snapshot_date.str[:4] == str(year)) & ~snaps.parties_reported.str.startswith('none')]
    gen = s[s.election.str.contains('General')]
    if len(gen): return gen.snapshot_date.iloc[-1], 'November general snapshot'
    pre = s[s.snapshot_date <= f'{year}-11-30']
    return pre.snapshot_date.iloc[-1], f'no enrollment in the Nov {year} file; used {pre.election.iloc[-1]} snapshot'

# ---------------- results ----------------
gen = pd.read_pickle(os.environ.get('SCORES_GEN', 'gen_all.pkl'))   # 1990-2024 general top-ticket rows of results_town_long.csv (prep_v21.py)
gen = gen[gen.row_type == 'town']
gen['year'] = gen.election_date.str[:4].astype(int)
# declared write-ins never count as the D or R nominee
gen.loc[gen.writein.astype(str).str.lower()=='true', 'cand_party'] = 'WRITEIN'

def race_table():
    races = []
    for cid, d in gen.groupby('contest_id'):
        cand = d[d.col_role.isin(['candidate', 'others', 'others_detail'])]
        cand = cand[cand.col_role != 'others_detail']
        tot = cand.votes.sum()
        dv = cand[(cand.col_role == 'candidate') & (cand.cand_party == 'DEM')].votes.sum()
        rv = cand[(cand.col_role == 'candidate') & (cand.cand_party == 'REP')].votes.sum()
        nm = 1 - (dv + rv) / tot if tot else np.nan
        dn = cand[(cand.col_role == 'candidate') & (cand.cand_party == 'DEM')].candidate_display.unique()
        rn = cand[(cand.col_role == 'candidate') & (cand.cand_party == 'REP')].candidate_display.unique()
        oth = cand[(cand.col_role == 'candidate') & ~cand.cand_party.isin(['DEM', 'REP'])].groupby('candidate_display').votes.sum().sort_values(ascending=False)
        top_oth = f'{oth.index[0]} {oth.iloc[0]/tot:.1%}' if len(oth) else ''
        ok = dv > 0 and rv > 0 and nm <= NONMAJOR_MAX
        why = '' if ok else ('no Democratic or no Republican candidate' if not (dv > 0 and rv > 0) else f'non-major candidates {nm:.1%} > 15%')
        races.append(dict(contest_id=cid, year=int(d.year.iloc[0]), office=d.office.iloc[0], district=d.district.iloc[0],
                          dem=', '.join(dn), rep=', '.join(rn), dem_votes=dv, rep_votes=rv, nonmajor_share=round(nm, 4),
                          largest_other=top_oth, qualifies=ok, exclusion_reason=why))
    return pd.DataFrame(races)

def race_votes(cid):
    d = gen[(gen.contest_id == cid) & (gen.col_role == 'candidate')]
    v = d.pivot_table(index='unit_id', columns='cand_party', values='votes', aggfunc='sum').fillna(0)
    out = pd.DataFrame({'D': v.get('DEM', 0), 'R': v.get('REP', 0)})
    return out

# ---------------- voting-unit mapping (registration unit -> results unit, per year) ----------------
notes_by_year = {}
for y, d in gen.groupby('year'):
    n = d[d.unit_note.notna()][['unit_id', 'unit_name', 'unit_note', 'sos_label']].drop_duplicates()
    notes_by_year[y] = n

def stem(name):
    s = re.sub(r'\b(Twp|Plt|Township|Plantation)\b\.?', '', name).strip()
    return s

# v2.3 (audit 2026-10-14, S1): a host line is accepted only if it is in the SAME COUNTY as the registration unit, and a unit
# that has its own result line in the other district's U.S. House race of the same year is never mapped into this race.
# Before v2.3 the bare place name was matched across counties, so Washington (Knox, CD 1) was pulled into the CD 2 race
# (2014-2022) through Washington Twp (Franklin) lines, and Unity (Waldo, CD 2) into the 2024 CD 1 race through Albion's
# "Albion/Unity Twp" line (Unity Twp is the Kennebec township).
_cty_lookup = dict(ucty)
try:
    _adds0 = pd.read_csv(f'{RES}/units_election_additions.csv')
    _cty_lookup.update({u: c for u, c in zip(_adds0.unit_id, _adds0.ucty) if u not in _cty_lookup})
except Exception:
    pass
_own_house = {}
for (_y, _cid), _d in gen[gen.office == 'U.S. House'].groupby(['year', 'contest_id']):
    _own_house[(_y, _cid)] = set(_d.unit_id)
def _own_line_elsewhere(u, year, res_units):
    """True if this is a U.S. House race and unit u has its own line in the OTHER district's House race of this year."""
    hs = [us for (y, _c), us in _own_house.items() if y == year]
    if not any(res_units <= us for us in hs): return False          # not a U.S. House race: rule does not apply
    return any(u in us and not res_units <= us for us in hs)

def build_map(year, reg_units, res_units):
    """Return dict reg_unit -> vote_unit, plus list of dropped reg units."""
    m, dropped = {}, []
    res_units = set(res_units)
    adds = pd.read_csv(f'{RES}/units_election_additions.csv')
    comb = dict(zip(adds.unit_id, adds.canonical_name))
    for u in reg_units:
        if u in res_units: m[u] = u; continue
        if _own_line_elsewhere(u, year, res_units): dropped.append(u); continue
        nm = stem(uname.get(u, u)); cty = ucty.get(u)
        host = None
        # combined result units named after components (e.g. 'Guilford / Kingsbury Plt'), same county only
        for cu in res_units:
            if cu in comb and _cty_lookup.get(cu) == cty and any(stem(p.strip()).lower() == nm.lower() for p in comb[cu].split('/')):
                host = cu; break
        if host is None:
            n = notes_by_year.get(year, pd.DataFrame(columns=['unit_id', 'unit_note', 'sos_label']))
            for _, r in n.iterrows():
                txt = f"{r.unit_note} {r.sos_label}"
                if 'includes' in str(r.unit_note) and re.search(r'\b' + re.escape(nm) + r'\b', txt, re.I) and r.unit_id in res_units \
                        and _cty_lookup.get(r.unit_id) == cty:
                    host = r.unit_id; break
        if host: m[u] = host
        else: dropped.append(u)
    return m, dropped

# ---------------- model ----------------
def fit(X, y, w):
    sw = np.sqrt(w)
    coef, _ = nnls(X * sw[:, None], y * sw)
    return coef

def run_cycle(year, races):
    snap, snap_note = snapshot_for(year)
    R = rw[rw.snapshot_date == snap].set_index('unit_id')[['DEM', 'REP', 'UNE', 'MIN', 'total']]
    rows, coefs, maps = [], [], {}
    for _, rc in races[(races.year == year) & races.qualifies].iterrows():
        V = race_votes(rc.contest_id)
        mp, dropped = build_map(year, R.index, V.index)
        maps[rc.contest_id] = (mp, dropped)
        Rv = R.loc[list(mp)].assign(vu=[mp[u] for u in mp]).groupby('vu')[['DEM', 'REP', 'UNE', 'MIN']].sum()
        J = Rv.join(V, how='inner')
        J = J[(J[['DEM', 'REP', 'UNE', 'MIN']].sum(axis=1) > 0)]
        tot = J[['DEM', 'REP', 'UNE', 'MIN']].sum(axis=1).values
        X = (J[['DEM', 'REP', 'MIN', 'UNE']].values / tot[:, None])
        cD = fit(X, J.D.values / tot, tot); cR = fit(X, J.R.values / tot, tot)
        coefs.append(dict(contest_id=rc.contest_id, year=year, office=rc.office, district=rc.district, snapshot=snap,
                          n_units=len(J), dem_reg_D=cD[0], rep_reg_D=cD[1], min_reg_D=cD[2], une_reg_D=cD[3],
                          dem_reg_R=cR[0], rep_reg_R=cR[1], min_reg_R=cR[2], une_reg_R=cR[3]))
        J['ED'] = cD[0] * J.DEM + cD[1] * J.REP; J['ER'] = cR[0] * J.DEM + cR[1] * J.REP
        J['MD'] = cD[2] * J.MIN; J['MR'] = cR[2] * J.MIN
        J['UD'] = J.D - J.ED - J.MD; J['UR'] = J.R - J.ER - J.MR
        J['contest_id'] = rc.contest_id; J['office'] = rc.office; J['year'] = year
        rows.append(J.reset_index().rename(columns={'vu': 'vote_unit', 'index': 'vote_unit'}))
    return (pd.concat(rows) if rows else pd.DataFrame()), pd.DataFrame(coefs), maps, snap, snap_note