"""VOLT sufficient statistic (Beraja & Talamas 2026), computed from BDS establishment-age tables.

VOLT_tau = (y_tau * l_tau)^(1 - nu), with
  y_tau = mean employment of establishments older than tau / mean employment of all establishments
  l_tau = expected lifespan of establishments older than tau / expected lifespan of all establishments
Lifespans come from a synthetic cohort with a constant exit hazard within each BDS age bin (paper Section 4.1).
"""
import csv
from dataclasses import dataclass

NU = 0.75


@dataclass
class Bin:
    lo: int            # first age in the bin
    hi: int | None     # last age in the bin, None for the open-ended bin
    estabs: float
    emp: float
    exits: float | None  # None when BDS does not report exits (age 0) or suppresses them


def _num(x):
    try:
        return float(x)
    except ValueError:
        return None


# BDS age-bin labels -> (first age, last age). The left-censored bin starts at year - 1976.
def bin_bounds(label, year):
    lc = year - 1976
    fine = {"a) 0": (0, 0), "b) 1": (1, 1), "c) 2": (2, 2), "d) 3": (3, 3), "e) 4": (4, 4),
            "f) 5": (5, 5), "g) 6 to 10": (6, 10), "h) 11 to 15": (11, 15), "i) 16 to 20": (16, 20),
            "j) 21 to 25": (21, 25), "k) 26+": (26, lc - 1), "l) Left Censored": (lc, None)}
    coarse = {"a) 0": (0, 0), "b) 1 to 5": (1, 5), "c) 6 to 10": (6, 10), "d) 11+": (11, lc - 1),
              "e) Left Censored": (lc, None)}
    return fine.get(label) or coarse[label]


def hazard(b, bds=True):
    # BDS does not report exits in the age-0 bin; the paper sets that hazard to 0 (footnote 15).
    # Firm data observes first-period exits, so bds=False keeps them.
    if bds and b.lo == 0 and b.hi == 0:
        return 0.0
    if b.exits is None:
        raise ValueError(f"exits not reported for ages {b.lo}-{b.hi}")
    return b.exits / b.estabs


def survival_stats(bins, tau, bds=True):
    """Return (lbar, lbar_mature, nbar, nbar_mature) for maturity threshold tau."""
    bins = sorted(bins, key=lambda b: b.lo)
    if tau + 1 not in {b.lo for b in bins}:
        raise ValueError(f"tau={tau} is not supported: mature ages must start at a BDS bin boundary "
                         f"(allowed: {sorted(b.lo - 1 for b in bins if b.lo > 0)})")
    last = bins[-1]
    tK = last.lo
    dK = hazard(last, bds)
    if dK <= 0:
        raise ValueError("open-ended bin has zero exits: lifespan is not finite")

    def delta(k):
        for b in bins:
            if b.lo <= k and (b.hi is None or k <= b.hi):
                return hazard(b, bds)
        raise ValueError(f"no bin covers age {k}")

    P = [1.0]
    for k in range(tK):
        P.append(P[-1] * (1 - delta(k)))
    tail = P[tK] / dK
    lbar = sum(P[:tK]) + tail
    lbar_m = tau + 1 + (sum(P[tau + 1:tK]) + tail) / P[tau + 1]

    est = sum(b.estabs for b in bins)
    emp = sum(b.emp for b in bins)
    mature = [b for b in bins if b.lo > tau]
    nbar = emp / est
    nbar_m = sum(b.emp for b in mature) / sum(b.estabs for b in mature)
    return lbar, lbar_m, nbar, nbar_m


def volt(bins, tau, nu=NU):
    lbar, lbar_m, nbar, nbar_m = survival_stats(bins, tau)
    y = nbar_m / nbar
    l = lbar_m / lbar
    V = (y * l) ** (1 - nu)
    ly, ll = (1 - nu) * _log(y), (1 - nu) * _log(l)
    return {
        "tau": tau, "nu": nu, "VOLT": V, "y": y, "l": l,
        "lifespan_all": lbar, "lifespan_mature": lbar_m,
        "emp_per_estab_all": nbar, "emp_per_estab_mature": nbar_m,
        "size_effect": y ** (1 - nu), "lifespan_effect": l ** (1 - nu),
        # Paper's Table 1 shares: (effect - 1) / (VOLT - 1)
        "lifespan_share": (l ** (1 - nu) - 1) / (V - 1) if V != 1 else None,
        "size_share": (y ** (1 - nu) - 1) / (V - 1) if V != 1 else None,
        # Log shares add to one (Figure 5 decomposition)
        "lifespan_log_share": ll / (ll + ly) if (ll + ly) else None,
    }


def _log(x):
    import math
    return math.log(x)


def load_bins(path, year, key_col=None, key=None, age_col="eage"):
    out = []
    with open(path) as f:
        for r in csv.DictReader(f):
            if int(r["year"]) != year:
                continue
            if key_col and r[key_col] != key:
                continue
            lo, hi = bin_bounds(r[age_col], year)
            out.append(Bin(lo, hi, _num(r["estabs"]), _num(r["emp"]), _num(r["estabs_exit"])))
    return out


def industries(path, year):
    with open(path) as f:
        return sorted({r["vcnaics3"] for r in csv.DictReader(f) if int(r["year"]) == year})