"""TV & anime baby-name tracker: analysis of SSA national baby-name data vs StashGrid /characters/ names.
Usage: python analyze.py <names.zip> <characters_db.json> <franchises.json> <outdir>
Reproducible: reads only the official SSA names.zip (yobYYYY.txt files: Name,Sex,Count; names with >=5 births per sex/year).
"""
import sys, zipfile, re, json, csv, os, statistics
from collections import defaultdict

zpath, dbpath, frpath, out = sys.argv[1:5]
os.makedirs(out, exist_ok=True)

# ---- load SSA ----
counts = defaultdict(dict)          # (name, 'F'/'M') -> {year: count}
totals = defaultdict(int)           # (year, sex) -> total births in file
ranks = {}                          # (name, sex, year) -> rank
with zipfile.ZipFile(zpath) as z:
    years = sorted(int(m.group(1)) for n in z.namelist() if (m := re.match(r'yob(\d{4})\.txt$', n)))
    for y in years:
        rows = []
        for line in z.read(f'yob{y}.txt').decode('utf-8').splitlines():
            n, s, c = line.split(','); c = int(c)
            counts[(n, s)][y] = c; totals[(y, s)] += c; rows.append((n, s, c))
        for s in 'FM':
            ordered = sorted((r for r in rows if r[1] == s), key=lambda r: -r[2])
            for i, r in enumerate(ordered, 1): ranks[(r[0], s, y)] = i
FIRST, LAST = years[0], years[-1]

# ---- character given names ----
db = json.load(open(dbpath)); fr = json.load(open(frpath))
TITLES = {'mr', 'mrs', 'ms', 'miss', 'dr', 'doctor', 'professor', 'lord', 'lady', 'king', 'queen', 'prince', 'princess',
          'sir', 'the', 'captain', 'agent', 'uncle', 'aunt', 'general', 'count', 'baron'}
# surname-first or nickname cases: given name to test (checked by hand against /characters/ descriptions)
OVERRIDE = json.load(open(os.path.join(os.path.dirname(frpath), 'name_overrides.json'))) if os.path.exists(os.path.join(os.path.dirname(frpath), 'name_overrides.json')) else {}

import unicodedata
def given(full):
    if full in OVERRIDE: return OVERRIDE[full]
    full = unicodedata.normalize('NFKD', full).encode('ascii', 'ignore').decode()
    toks = [t.strip('.,"\'()') for t in full.split()]
    toks = [t for t in toks if t and t.lower() not in TITLES]
    return toks[0] if toks else None

def mean(xs): return sum(xs) / len(xs) if xs else 0.0

DEBUT = json.load(open(os.path.join(os.path.dirname(frpath), 'char_debut.json')))  # char_debut overrides
rows = []
for fkey, chars in db.items():
    for ch in chars:
        a = DEBUT.get(ch['name'], fr[fkey]['anchor'])
        g = given(ch['name'])
        if not g or not re.fullmatch(r"[A-Za-z]+", g): continue
        g = g[0].upper() + g[1:].lower()
        sexes = {'female': ['F'], 'male': ['M']}.get(ch['sex'], ['F', 'M'])
        for s in sexes:
            ser = counts.get((g, s), {})
            if not ser:
                rows.append(dict(franchise=fkey, character=ch['name'], given=g, sex=s, anchor=a, in_ssa=False)); continue
            fy = min(ser); py = max(ser, key=lambda y: ser[y])
            before = mean([ser.get(y, 0) for y in range(a - 3, a)])
            after_years = [y for y in range(a + 1, a + 4) if y <= LAST]
            after = mean([ser.get(y, 0) for y in after_years])
            pre_total = sum(c for y, c in ser.items() if y < a)
            rows.append(dict(franchise=fkey, character=ch['name'], given=g, sex=s, anchor=a, in_ssa=True,
                             first_year=fy, peak_year=py, peak=ser[py], before_avg=round(before, 1), after_avg=round(after, 1),
                             lift=round(after / before, 2) if before else None, latest=ser.get(LAST, 0),
                             latest_rank=ranks.get((g, s, LAST)), pre_anchor_total=pre_total,
                             new_since_anchor=fy >= a, since_anchor_total=sum(c for y, c in ser.items() if y >= a)))

with open(os.path.join(out, 'results.csv'), 'w', newline='') as f:
    keys = ['franchise', 'character', 'given', 'sex', 'anchor', 'in_ssa', 'first_year', 'peak_year', 'peak', 'before_avg', 'after_avg',
            'lift', 'latest', 'latest_rank', 'pre_anchor_total', 'new_since_anchor', 'since_anchor_total']
    w = csv.DictWriter(f, fieldnames=keys); w.writeheader(); [w.writerow(r) for r in rows]

# compact series for charts / a future interactive tracker: anchor-10 .. LAST
series = {}
for r in rows:
    if not r['in_ssa']: continue
    k = f"{r['given']}|{r['sex']}"
    if k in series: continue
    lo = max(FIRST, r['anchor'] - 15)
    series[k] = {str(y): counts[(r['given'], r['sex'])].get(y, 0) for y in range(lo, LAST + 1)}
json.dump({'source': 'SSA national baby names (names.zip), names with 5+ births', 'years': [FIRST, LAST],
           'rows': rows, 'series': series}, open(os.path.join(out, 'tracker.json'), 'w'))
print('years', FIRST, LAST, '| characters', sum(len(v) for v in db.values()), '| rows', len(rows),
      '| in SSA', sum(1 for r in rows if r['in_ssa']))