f/data.py

"""Data ingestion: Yahoo Finance CSV dumps -> parquet cache.

Layout expected in the data root (default ~/prog/fin/stocks):
    {SYM}-history.csv       Date,Open,High,Low,Close,Adj Close,Volume
    {SYM}-dividend.csv      Date,Dividends            (optional)
    {SYM}-capitalGain.csv   Date,Capital Gains        (optional)
    {SYM}.json              Yahoo chart API payload  (metadata only)

All panels are date x symbol DataFrames. The "adj" panel (Adj Close)
already bakes in dividends and capital gains, so it is the correct
input for pre-tax total returns. The dividend/capgain panels are used
only by the tax engine (to tax distributions).

First call builds a parquet cache (default: .cache/) in ~1 min for
4k symbols; subsequent calls load in well under a second.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path

import pandas as pd

DEFAULT_ROOT = Path("~/prog/fin/stocks").expanduser()
CACHE_DIR = Path(__file__).parent / ".cache"


@dataclass
class Bundle:
    adj: pd.DataFrame      # adjusted close (total-return adjusted)
    close: pd.DataFrame    # raw close
    div: pd.DataFrame      # dividend distributions per share
    capg: pd.DataFrame     # capital gain distributions per share
    names: dict            # symbol -> long name (best effort)


def _read_panel(root: Path, suffix: str, col: str) -> pd.DataFrame:
    cols = []
    syms = []
    for f in sorted(root.glob(f"*-{suffix}.csv")):
        sym = f.name[: -len(f"-{suffix}.csv")]
        df = pd.read_csv(f, index_col=0)
        df.index = pd.to_datetime(df.index, format="mixed", errors="coerce")
        df = df[~df.index.isna()]
        df = df[~df.index.duplicated(keep="last")].sort_index()
        if col not in df.columns:
            continue
        cols.append(df[col])
        syms.append(sym)
    if not cols:
        return pd.DataFrame()
    out = pd.concat(cols, axis=1)
    out.columns = syms
    return out.sort_index()


def _read_names(root: Path, max_files: int = 5000) -> dict:
    names = {}
    for f in root.glob("*.json"):
        try:
            meta = json.loads(f.read_text())["chart"]["result"][0]["meta"]
            names[f.stem] = meta.get("longName", f.stem)
        except Exception:
            names[f.stem] = f.stem
    return names


# a bundle is ~2 GB resident; memoize the last one so Streamlit re-runs
# (every widget change!) don't reload the parquet panels
_BUNDLE_CACHE: dict[tuple, Bundle] = {}


def load_bundle(root: Path = DEFAULT_ROOT, cache: Path = CACHE_DIR,
                rebuild: bool = False) -> Bundle:
    """Load the full dataset, using (or building) a parquet cache."""
    key = (str(root), rebuild)
    hit = _BUNDLE_CACHE.get(key)
    if hit is not None:
        return hit
    cache = Path(cache)
    cache.mkdir(parents=True, exist_ok=True)
    cache_files = [cache / f"panel_{k}.parquet" for k in ("adj", "close", "div", "capg")]

    if not rebuild and all(f.exists() for f in cache_files):
        adj = pd.read_parquet(cache_files[0])
        close = pd.read_parquet(cache_files[1])
        div = pd.read_parquet(cache_files[2])
        capg = pd.read_parquet(cache_files[3])
    else:
        adj = _read_panel(root, "history", "Adj Close")
        close = _read_panel(root, "history", "Close")
        div = _read_panel(root, "dividend", "Dividends").fillna(0.0)
        capg = _read_panel(root, "capitalGain", "Capital Gains").fillna(0.0)
        adj.to_parquet(cache_files[0])
        close.to_parquet(cache_files[1])
        div.to_parquet(cache_files[2])
        capg.to_parquet(cache_files[3])
    names_path = cache / "names.json"
    if names_path.exists():
        names = json.loads(names_path.read_text())
    else:
        names = _read_names(root)
        names_path.write_text(json.dumps(names))
    bundle = Bundle(adj, close, div, capg, names)
    # keep only the newest bundle: there is no room for two
    _BUNDLE_CACHE.clear()
    _BUNDLE_CACHE[key] = bundle
    return bundle


def search_symbols(bundle: Bundle, query: str, limit: int = 200) -> list[str]:
    """Case-insensitive search over symbols and names.

    Ranked: exact symbol, symbol prefix, name prefix, symbol substring,
    name substring.
    """
    q = query.lower()
    buckets: dict[int, list[str]] = {}
    for s in bundle.adj.columns:
        sl, nl = s.lower(), bundle.names.get(s, "").lower()
        if sl == q:
            buckets.setdefault(0, []).append(s)
        elif sl.startswith(q):
            buckets.setdefault(1, []).append(s)
        elif nl.startswith(q):
            buckets.setdefault(2, []).append(s)
        elif q in sl:
            buckets.setdefault(3, []).append(s)
        elif q in nl:
            buckets.setdefault(4, []).append(s)
    out: list[str] = []
    for k in sorted(buckets):
        out.extend(buckets[k])
    return out[:limit]