"""Cluster funds by RETURN DRIVER.

Input: fundlab/factor_results.json (per-fund loading vectors from the
39-sleeve full OLS). Distance: squared Euclidean on the full-window
loading vector (missing sleeves -> 0). Agglomerative, AVERAGE linkage
(Lance-Armstrong update), pure numpy (no scipy in the venv).

Output: fundlab/cluster_tree.json
  syms   - fund symbols (rows of the loading matrix)
  merges - [(a, b, dist), ...] in merge order, a/b ORIGINAL fund
           indices, b absorbed into a
  n      - number of funds

Cutting the tree at any k is O(n) (parent array over the first n-k
merges), so the app can offer a k slider with instant re-grouping.
Label a cluster by its top |median loading| sleeves; all-small ->
"idiosyncratic / alpha".

Run: python -m fundlab.cluster     (builds the tree, ~1-2 min)
"""
from __future__ import annotations

import json
import time
from pathlib import Path

import numpy as np

HERE = Path(__file__).parent
RESULTS = HERE / "factor_results.json"
TREE = HERE / "cluster_tree.json"
KMEANS = HERE / "cluster_kmeans.json"
from fundlab.factors import AXES, DRIVERS, cash_yield  # noqa: E402


def log(msg: str) -> None:
    print(time.strftime("%H:%M:%S"), msg, flush=True)


def loading_matrix() -> tuple[list[str], np.ndarray]:
    d = json.loads(RESULTS.read_text())
    yld = cash_yield()
    syms, rows = [], []
    for sym, v in d.items():
        # full window preferred; funds without 250 complete-case rows
        # there fall back to the 5y window
        f = v.get("full")
        if not (isinstance(f, dict) and "betas" in f):
            f = v.get("rec5")
        if not (isinstance(f, dict) and "betas" in f):
            continue
        row = np.array([f["betas"].get(s, 0.0) or 0.0 for s in DRIVERS],
                       dtype=float)
        # virtual CASH axis: the fund's intercept (annualized alpha) is
        # its yield capture - for a cash fund that IS the exposure,
        # because OLS puts a near-constant yield level in the intercept
        # rather than in the (small) shv/bil slope. Normalize by the
        # recent cash yield so 1.0 == "earns the cash rate".
        alpha = f.get("alpha_ann")
        row = np.append(row, (alpha / yld if alpha is not None and yld else 0.0))
        # winsorize at +/-4: everything beyond that is an OLS noise fit
        # (R2 ~ 0.01, |t| < 2 - e.g. fyhtx shv +63.8), never a real
        # exposure (no open-end fund runs 400% of a sleeve)
        row = np.clip(row, -4.0, 4.0)
        syms.append(sym)
        rows.append(row)
    return syms, np.vstack(rows)


def hclust(V: np.ndarray) -> list[list]:
    """Average-linkage agglomerative. Returns [(a, b, dist), ...] with
    ORIGINAL row indices, b absorbed into a, merge order."""
    n = len(V)
    ss = np.sum(V * V, axis=1)
    D = ss[:, None] + ss[None, :] - 2 * (V @ V.T)
    np.clip(D, 0, None)
    np.fill_diagonal(D, np.inf)
    order = np.arange(n)
    size = np.ones(n)
    merges: list[list] = []
    for step in range(n - 1):
        iu, ju = np.unravel_index(np.argmin(D), D.shape)
        i, j = int(iu), int(ju)
        if i > j:
            i, j = j, i
        a, b = int(order[i]), int(order[j])
        merges.append([a, b, float(D[i, j])])
        # Lance-Armstrong: new distance from i to every survivor k
        size_ij = size[i] + size[j]
        D[i, :] = (size[i] * D[i, :] + size[j] * D[j, :]) / size_ij
        D[:, i] = D[i, :]
        size[i] = size_ij
        D = np.delete(D, j, axis=0)
        D = np.delete(D, j, axis=1)
        order = np.delete(order, j)
        np.fill_diagonal(D, np.inf)
        if step % 500 == 499:
            log(f"  hclust: {step + 1}/{n - 1} merges")
    return merges


def cut(merges: list, n: int, k: int) -> dict[int, int]:
    """Cluster id (0..k-1 by size order) for each original index."""
    parent = list(range(n))

    def find(x: int) -> int:
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x
    for a, _b, _d in merges[: n - k]:
        a, b = int(a), int(_b)
        parent[find(b)] = find(a)
    roots = [find(i) for i in range(n)]
    # map roots to dense ids, largest cluster -> id 0
    import collections
    cnt = collections.Counter(roots)
    rank = {r: i for i, r in enumerate(sorted(cnt, key=cnt.get, reverse=True))}
    return {i: rank[r] for i, r in enumerate(roots)}


def cluster_means(syms: list[str], V: np.ndarray,
                  assign: dict[int, int]) -> list[dict]:
    """Per-cluster median loading profile (for labels)."""
    k = max(assign.values()) + 1
    out = []
    for c in range(k):
        idx = [i for i, s in assign.items() if s == c]
        med = np.median(V[idx], axis=0)
        top = np.argsort(-np.abs(med))[:4]
        prof = {AXES[t]: float(med[t]) for t in top
                if abs(med[t]) >= 0.05}
        out.append({"n": len(idx), "top": prof,
                    "syms": [syms[i] for i in idx]})
    return out


SLEEVE_NAME = {
    "qqq": "US growth", "ivv": "US large blend", "iwm": "US small",
    "efa": "intl developed", "vwo": "EM equity", "vnq": "REIT",
    "shv": "cash/ultra-short", "shy": "short T 1-3y",
    "ief": "T 7-10y", "tlt": "long T", "agg": "core bond (AGG)",
    "vblix": "vol (pure)", "vweax": "HY corporate", "vmbix": "MBS",
    "lqd": "IG corporate", "hyg": "high yield", "pff": "preferreds",
    "emb": "EM debt", "tip": "TIPS/inflation", "vtv": "US value",
    "djp": "gas", "gsg": "oil", "gld": "gold", "fxe": "EUR",
    "fxy": "JPY", "xlk": "tech", "xlf": "financials", "xle": "energy",
    "xlv": "healthcare", "xlp": "staples", "xlu": "utilities",
    "xly": "consumer disc", "xlb": "materials", "dbmf": "CTA/futures",
    "dbb": "commodities",
    "cash": "cash (yield)",
}

# the CASH axis is a yield LEVEL (alpha/cash-yield, 1.0 = earns the cash
# rate), not a return SENSITIVITY like the betas. In raw Euclidean space a
# pure cash fund (0,...,0,~0.5) sits inside the cloud of low-exposure
# balanced/target-date funds and k-means swallows it. Emphasize the axis
# for DISTANCE only (labels still use the raw values).
CASH_EMPHASIS = 2.0


def emphasized(V: np.ndarray, ci: int) -> np.ndarray:
    W = V.copy()
    W[:, ci] *= CASH_EMPHASIS
    return W


def label(prof: dict) -> str:
    # a real driver shows up as >=0.3 on a median loading; below that the
    # cluster is balanced mixes with no single dominant driver. The CASH
    # axis is a level, not a beta: in a mixed low-exposure cluster the
    # median cash is diluted, so it gets a slightly lower threshold and
    # must clearly beat the runner-up to count as THE driver.
    if not prof:
        return "no dominant driver (balanced/idio)"
    names = sorted(prof, key=lambda s: -abs(prof[s]))
    top_s, top_v = names[0], prof[names[0]]
    runner = abs(prof[names[1]]) if len(names) > 1 else 0.0
    if top_s == "cash":
        if top_v >= 0.25 and top_v >= 1.5 * runner:
            return " + ".join(f"{SLEEVE_NAME.get(s, s)} {prof[s]:+.2f}"
                              for s in names[:3])
        return "no dominant driver (balanced/idio)"
    if abs(top_v) < 0.30:
        return "no dominant driver (balanced/idio)"
    return " + ".join(f"{SLEEVE_NAME.get(s, s)} {prof[s]:+.2f}"
                      for s in names[:3])


def kmeans(V: np.ndarray, k: int, seed: int = 0,
            iters: int = 100) -> np.ndarray:
    """k-means++ init, deterministic seed. Returns labels (len n)."""
    rng = np.random.default_rng(seed)
    n = len(V)
    # k-means++ init
    C = np.empty((k, V.shape[1]))
    C[0] = V[rng.integers(n)]
    d2 = ((V - C[0]) ** 2).sum(1)
    for j in range(1, k):
        probs = d2 / d2.sum()
        C[j] = V[rng.choice(n, p=probs)]
        d2 = np.minimum(d2, ((V - C[j]) ** 2).sum(1))
    for _ in range(iters):
        dist = ((V[:, None, :] - C[None, :, :]) ** 2).sum(2)   # n x k
        lab = dist.argmin(1)
        newC = C.copy()
        for j in range(k):
            m = lab == j
            if m.any():
                newC[j] = V[m].mean(0)
        if np.allclose(newC, C):
            break
        C = newC
    return dist.argmin(1)


def run_kmeans(k: int = 30) -> dict:
    syms, V = loading_matrix()
    lab = kmeans(emphasized(V, AXES.index("cash")), k)
    out: dict[int, dict] = {}
    for c in range(k):
        idx = [i for i, l in enumerate(lab) if l == c]
        if not idx:
            continue
        med = np.median(V[idx], axis=0)
        top = np.argsort(-np.abs(med))[:4]
        prof = {AXES[t]: float(med[t]) for t in top
                if abs(med[t]) >= 0.08}
        out[c] = {"n": len(idx), "top": prof, "label": label(prof),
                  "syms": [syms[i] for i in idx]}
    return {"k": k, "clusters": out}


def run() -> None:
    t0 = time.time()
    syms, V = loading_matrix()
    log(f"cluster: {len(syms)} funds x {V.shape[1]} drivers")
    merges = hclust(V)
    TREE.write_text(json.dumps({"syms": syms, "merges": merges,
                                "n": len(syms), "sleeves": AXES}))
    log(f"hierarchical tree written in {time.time() - t0:.0f}s -> {TREE.name}")
    # NOTE: fixed-k cuts of the tree peel small pure-sleeve clusters off a
    # giant blend blob (most funds are multi-sleeve mixes) - the useful
    # grouping is k-means on the loading vectors instead.
    res = run_kmeans(30)
    KMEANS.write_text(json.dumps(res))
    log(f"k-means k=30 written in {time.time() - t0:.0f}s -> {KMEANS.name}")
    for c in sorted(res["clusters"], key=lambda c: -res["clusters"][c]["n"]):
        v = res["clusters"][c]
        log(f"  n={v['n']:4d}  {v['label']}")


if __name__ == "__main__":
    run()