"""Fund-family classification.

Tags each symbol with its fund family (Vanguard, Fidelity, iShares/BlackRock,
T. Rowe Price, ...) from the Yahoo `longName`/`shortName` in the per-symbol
chart JSON, plus ticker-pattern fallbacks for the big families whose names
are sometimes missing.

Run:  python3 -m families   -> writes .cache/families.csv (sym, family, name)
"""

from __future__ import annotations

import csv
import json
import re
from collections import Counter
from pathlib import Path

ROOT = Path.home() / "prog/fin/stocks"
OUT = Path(__file__).parent / ".cache" / "families.csv"

# Ordered keyword rules: first match wins. Specific brands before generics.
RULES: list[tuple[str, str]] = [
    ("Vanguard",        r"vanguard"),
    ("Fidelity",        r"fidelity|fd[ae]{1,2}[uiv]?$|fdsi|fdls|fdsm"),
    ("iShares/BlackRock", r"\bishares\b"),
    ("BlackRock",       r"blackrock"),
    ("State Street/SPDR", r"state street|\bspdr\b|spts|sx[ae]x$|sfx[ae]$"),
    ("T. Rowe Price",   r"t\.? ?rowe price|^\btr[pd]|tlc[ae]d$|tltlx$|thdax$|trn[ax]$"),
    ("Schwab",          r"schwab"),
    ("MassMutual",      r"massmutual|mass mutual"),
    ("Franklin",        r"franklin"),
    ("Templeton",       r"templeton"),
    ("J.P. Morgan",     r"j\.? ?p\.? ?morgan|jpmorgan|jps[ae]x$|jp[ae]x$"),
    ("Invesco",         r"invesco"),
    ("American Century", r"american century"),
    ("Lord Abbett",     r"lord abbett"),
    ("Victory Capital", r"victory capital|victory shares|victory pioneers|victory pioneer|victory rs|victory fund"),
    ("DFA/Dimensional", r"dimensional|\bdfa\b|df[ae]x$|dgbex$|d[ij]a[ei]x?$"),
    ("Morgan Stanley",  r"morgan stanley"),
    ("Voya",            r"\bvoya\b"),
    ("Ned Davis",       r"ned davis"),
    ("Neuberger Berman", r"neuberger"),
    ("Dodge & Cox",     r"dodge ?cox"),
    ("Janus Henderson", r"janus"),
    ("PIMCO",           r"pimco"),
    ("Northern Trust",  r"northern trust"),
    ("Wells Fargo",     r"wells fargo"),
    ("Barclays",        r"barclays"),
    ("Carillon",        r"carillon"),
    ("Legg Mason",      r"legg mason"),
    ("Allianz",         r"allianz"),
    ("American Funds/Alliance", r"american funds|american fds"),
    ("American Beacon", r"american beacon|\bab (small|high|equity|income|bond)"),
    ("BNY Mellon",      r"bny |mellon"),
    ("Goldman Sachs AM", r"goldman sachs"),
    ("Nuveen",          r"nuveen"),
    ("Western Asset",   r"western asset"),
    ("Loomis Sayles",   r"loomis sayles"),
    ("MFS/Merrill",     r"\bmfs\b"),
    ("Principal",       r"principal "),
    ("PGIM",            r"pgim"),
    ("Hartford",        r"hartford"),
    ("ClearBridge",     r"clearbridge"),
    ("Nationwide",      r"nationwide"),
    ("abrdn",           r"abrdn"),
    ("John Hancock",    r"john hancock|jhanco|\bjh[a-z]{2}x$"),
    ("Virtus",          r"virtus"),
    ("SEI",             r"\bsei\b"),
    ("Empower",         r"empower"),
    ("Avantis",         r"avantis"),
    ("Allspring",       r"allspring"),
    ("Baird",           r"baird"),
    ("Nomura",          r"nomura"),
    ("Baron",           r"baron "),
    ("Artisan",         r"artisan"),
    ("Columbia",        r"columbia "),
    ("Calvert",         r"calvert"),
    ("First Eagle",     r"first eagle"),
    ("AQR",             r"\baqr\b"),
    ("Catalyst",        r"catalyst"),
    ("Touchstone",      r"touchstone"),
    ("NexPoint",        r"nexpoint"),
    ("Affiliated",      r"affiliated managers"),
    ("Putnam",          r"putnam"),
    ("Federated",       r"federated"),
    ("Schroders",       r"schroder"),
    ("Alger",           r"alger "),
    ("Oakmark",         r"oakmark"),
    ("DoubleLine",      r"doubleline"),
    ("Diamond Hill",    r"diamond hill"),
    ("First Trust",     r"first trust|firsttrust"),
    ("Transamerica",    r"transamerica"),
    ("Brookfield",      r"brookfield"),
    ("DWS",             r"\bdws\b"),
    ("Natixis",         r"natixis"),
    ("Barings/MML",     r"mml |barings"),
    ("Redmont",         r"redmont"),
    ("Parametric",      r"parametric"),
    ("Brown Advisory",  r"brown advisory"),
    ("Harbor",          r"harbor "),
    ("Madison",         r"madison "),
    ("Toews",           r"toews"),
    ("Conestoga",       r"conestoga"),
    ("Turner",          r"turner funds"),
    ("Cambiar",         r"cambiar"),
    ("Copeland",        r"copeland"),
    ("Pear Tree",       r"pear tree"),
    ("Iron/Unified",    r"\biron\b.*fund|unified series"),
    ("E-Valuator",      r"e-valuator"),
    ("SilverPepper",    r"silverpepper"),
    ("Global X",        r"global x"),
    ("Eaton Vance",     r"eaton vance"),
    ("William Blair",   r"william blair"),
    ("Xtrackers/DWS",   r"xtrackers"),
    ("DWS",             r"\bdws\b"),
    ("WisdomTree",      r"wisdomtree"),
    ("Innovator",       r"innovator "),
    ("iPath/UBS",       r"ipath"),
    ("UBS",             r"\bubs\b"),
    ("Cohen & Steers",  r"cohen & steers|cohen"),
    ("Royce",           r"royce"),
    ("Pacer",           r"pacer "),
    ("Capital Group",   r"capital group"),
    ("FlexShares",      r"flexshares"),
    ("Alps",            r"alps "),
    ("Segall Bryant",   r"segall bryant"),
    ("Guinness Atkinson", r"guinness"),
    ("Calamos",         r"calamos"),
    ("Glenmede",        r"glenmede"),
    ("KraneShares",     r"kraneshares"),
    ("ProShares",       r"proshares"),
    ("VanEck",          r"vaneck|van eck"),
    ("Guggenheim",      r"guggenheim"),
    ("Royce",           r"royce"),
    ("Harding Loevner", r"harding"),
    ("Burnham",         r"burnham"),
    ("Leuthold",        r"leuthold"),
    ("LSV Global",      r"lsv "),
    ("Cambria",         r"cambria"),
    ("QRAFT",           r"qraft"),
    ("TrueShares",      r"trueshares"),
    ("NorthSquare",     r"north ?square"),
    ("AdvisorShares",   r"advisors?hares"),
    ("GQ Partners",     r"gqg partners|gq partners"),
    ("Brandywine",      r"brandywine"),
    ("RBB",             r"rbb "),
    ("Strategy Shares", r"strategy shares"),
    ("Overlay Shares",  r"overlay shares"),
    ("Robo Global",     r"robo global"),
    ("Henry Schwebel",  r"schwebel|amg gw"),
    ("HC Fund Mgmt",    r"hcm "),
    ("Investment Mgrs Series", r"investment managers series"),
    ("United States (ETF)", r"^united states |united states (corporate|government|bond)"),
]
_RULES = [(name, re.compile(pat, re.I)) for name, pat in RULES]

# Ticker-pattern fallbacks (used when the name is empty or unmatched).
TICKER_RULES: list[tuple[str, str]] = [
    ("Vanguard",        r"^vt[bcdfghjkmnpqrsuw][a-z]x?$"),
    ("Fidelity",        r"^fd[a-z]{2,4}$"),
    ("iShares/BlackRock", r"^i[a-z]{3}$"),
    ("State Street/SPDR", r"^s[pfx][a-z]{2}$|^st[a-z]{2,3}x?$"),
    ("T. Rowe Price",   r"^t[lm][a-z]x?$"),
    ("MassMutual",      r"^m[a-z]{4}$"),
    ("Franklin",        r"^f[ae][a-z]{2,3}x?$"),
    ("DFA/Dimensional", r"^d[fgij][ae]x?$|^df[a-z]x?$"),
    ("Schwab",          r"^sw[a-z]{2,3}$"),
    ("Vanguard",        r"^vb[a-z]x?$"),
    ("Global X",        r"^[a-z]{3}x$"),
    ("WisdomTree",      r"^w[a-z]{3}$"),
    ("Xtrackers/DWS",   r"^x[a-z]{3}$"),
    ("Pacer",           r"^p[a-z]{2}x$"),
]
_TICKER_RULES = [(name, re.compile(pat, re.I)) for name, pat in TICKER_RULES]

UNKNOWN = "(unknown)"


def classify(sym: str, name: str) -> str:
    n = (name or "").strip()
    if n:
        for fam, rx in _RULES:
            if rx.search(n):
                return fam
    for fam, rx in _TICKER_RULES:
        if rx.fullmatch(sym):
            return fam
    return UNKNOWN


def load_names(root: Path = ROOT) -> dict[str, str]:
    out: dict[str, str] = {}
    for f in root.glob("*.json"):
        sym = f.name.rsplit(".", 1)[0]
        try:
            m = json.loads(f.read_bytes())["chart"]["result"][0]["meta"]
            out[sym] = (m.get("longName") or m.get("shortName") or "").strip()
        except Exception:
            out[sym] = ""
    return out


def build(root: Path = ROOT, out: Path = OUT) -> Counter:
    names = load_names(root)
    out.parent.mkdir(parents=True, exist_ok=True)
    with out.open("w", newline="") as fh:
        w = csv.writer(fh)
        w.writerow(["sym", "family", "name"])
        for sym in sorted(names):
            w.writerow([sym, classify(sym, names[sym]), names[sym]])
    return Counter(classify(s, n) for s, n in names.items())


if __name__ == "__main__":
    c = build()
    total = sum(c.values())
    print(f"{len(c)} families, {total} symbols")
    for k, v in c.most_common(40):
        print(f"{v:6d}  {k}")
    print(f"\nunknown: {c[UNKNOWN]} ({c[UNKNOWN]/total:.1%})")