#!/usr/bin/env python3
"""Count how many independent signals a factor library really holds, then merge duplicates.

Standard library only. Modes:

* ``--self-test`` (default): three hidden sources, eleven noisy copies. Checks that
  clustering recovers the three groups and that merging keeps the diversification.
* ``--demo``: rebuilds the article's eleven momentum / value / size signals from the
  Ken French data bundled in ``linearmodels==7.0`` (the bundle ends in March 2017).
* ``--article-check``: recomputes the article tables from
  ``factor_dedup_article.csv``, the series frozen from that bundle. This is what
  CI runs, so a later linearmodels release cannot silently move the numbers.
* ``--csv FILE``: a wide CSV of returns (first column = date, one column per signal).
* ``--french FILE``: a CSV downloaded from the Ken French data library (values in
  percent; only the first, monthly block is read).

The effective number assumes equal weights and equal quality:

    N_eff = N / (1 + (N - 1) * avg_pairwise_correlation)

which is the same as N^2 / sum(correlation matrix). Each signal is scaled to the same
volatility before combining; the scaling uses the full sample, so treat the Sharpe
ratios as a description of the correlation structure, not as a tradable backtest.
"""

from __future__ import annotations

import argparse
import csv
import math
import random
import sys
from pathlib import Path


# ---------------------------------------------------------------- statistics


def mean(xs):
    return sum(xs) / len(xs)


def std(xs):
    m = mean(xs)
    return math.sqrt(sum((x - m) ** 2 for x in xs) / (len(xs) - 1))


def corr(a, b):
    ma, mb = mean(a), mean(b)
    num = sum((x - ma) * (y - mb) for x, y in zip(a, b))
    den = math.sqrt(sum((x - ma) ** 2 for x in a) * sum((y - mb) ** 2 for y in b))
    return num / den


def corr_matrix(series):
    n = len(series)
    c = [[1.0] * n for _ in range(n)]
    for i in range(n):
        for j in range(i + 1, n):
            c[i][j] = c[j][i] = corr(series[i], series[j])
    return c


def avg_pairwise(c):
    n = len(c)
    if n < 2:
        return 0.0
    return sum(c[i][j] for i in range(n) for j in range(i + 1, n)) / (n * (n - 1) / 2)


def effective_n(c):
    n = len(c)
    return n / (1 + (n - 1) * avg_pairwise(c))


def sharpe(xs, periods_per_year):
    return mean(xs) / std(xs) * math.sqrt(periods_per_year)


def scaled(xs):
    s = std(xs)
    return [x / s for x in xs]


def equal_weight(series):
    return [mean(row) for row in zip(*series)]


def strip(series, market):
    """Remove each series' OLS beta to the market column."""
    mm = mean(market)
    var = sum((x - mm) ** 2 for x in market)
    out = []
    for xs in series:
        mx = mean(xs)
        beta = sum((x - mx) * (y - mm) for x, y in zip(xs, market)) / var
        out.append([x - beta * y for x, y in zip(xs, market)])
    return out


# ---------------------------------------------------------------- clustering


def cluster(c, threshold):
    """Average-linkage clustering on correlation.

    Repeatedly joins the two groups with the highest average cross-correlation and
    stops when no pair of groups is correlated above ``threshold``.
    """
    groups = [[i] for i in range(len(c))]
    while len(groups) > 1:
        best, pair = None, None
        for a in range(len(groups)):
            for b in range(a + 1, len(groups)):
                r = mean([c[i][j] for i in groups[a] for j in groups[b]])
                if best is None or r > best:
                    best, pair = r, (a, b)
        if best < threshold:
            break
        a, b = pair
        groups[a] = groups[a] + groups[b]
        del groups[b]
    return groups


def merge(series, groups):
    """One signal per group: the equal-weight average of its vol-scaled members."""
    return [scaled(equal_weight([scaled(series[i]) for i in g])) for g in groups]


# ---------------------------------------------------------------- report


def describe(names, series, periods_per_year):
    c = corr_matrix(series)
    avg_single = mean([sharpe(s, periods_per_year) for s in series])
    combined = sharpe(equal_weight([scaled(s) for s in series]), periods_per_year)
    return {
        "n": len(series),
        "avg_corr": avg_pairwise(c),
        "n_eff": effective_n(c),
        "avg_single_sharpe": avg_single,
        "sqrt_n_estimate": avg_single * math.sqrt(len(series)),
        "combined_sharpe": combined,
        "corr": c,
    }


def report(names, series, periods_per_year, threshold):
    before = describe(names, series, periods_per_year)
    groups = cluster(before["corr"], threshold)
    merged = merge(series, groups)
    after = describe([f"group{k}" for k in range(len(groups))], merged, periods_per_year)

    print(f"signals: {before['n']}   average pairwise correlation: {before['avg_corr']:+.2f}")
    print(f"effective number:            {before['n_eff']:.2f}")
    print(f"average single Sharpe:       {before['avg_single_sharpe']:.2f}")
    print(f"sqrt(N) estimate:            {before['sqrt_n_estimate']:.2f}")
    print(f"equal-weight actual Sharpe:  {before['combined_sharpe']:.2f}")
    print(f"\nclusters at correlation > {threshold}:")
    for g in groups:
        print("  - " + ", ".join(names[i] for i in g))
    print(f"\nafter merge: {after['n']} signals, average correlation {after['avg_corr']:+.2f}")
    print(f"effective number:            {after['n_eff']:.2f}")
    print(f"equal-weight actual Sharpe:  {after['combined_sharpe']:.2f}")
    kept = after["combined_sharpe"] / before["combined_sharpe"]
    print(f"Sharpe kept after merging:   {kept:.0%}")
    return before, groups, after


# ---------------------------------------------------------------- loaders


def load_csv(path, percent):
    with open(path, newline="") as f:
        rows = list(csv.reader(f))
    names = [h.strip() for h in rows[0][1:]]
    cols = [[] for _ in names]
    for row in rows[1:]:
        if len(row) < len(names) + 1 or not row[0].strip():
            continue
        try:
            values = [float(v) for v in row[1 : len(names) + 1]]
        except ValueError:
            continue
        for k, v in enumerate(values):
            cols[k].append(v / 100 if percent else v)
    return names, cols


def load_french(path):
    """Read the first data block of a Ken French CSV (monthly YYYYMM rows, percent)."""
    with open(path, newline="", encoding="latin-1") as f:
        rows = list(csv.reader(f))
    names, cols, started = None, None, False
    for row in rows:
        cells = [x.strip() for x in row]
        if names is None:
            if len(cells) > 1 and cells[0] == "" and all(cells[1:]):
                names = cells[1:]
                cols = [[] for _ in names]
            continue
        if cells and cells[0].isdigit() and len(cells[0]) in (6, 8):
            started = True
            values = [float(v) for v in cells[1 : len(names) + 1]]
            if any(v <= -99.99 for v in values):  # French's missing-value marker
                continue
            for k, v in enumerate(values):
                cols[k].append(v / 100)
        elif started:
            break
    if names is None:
        sys.exit(f"{path}: no header row found")
    return names, cols


def load_demo():
    try:
        from linearmodels.datasets import french
    except ImportError:
        sys.exit("--demo needs the Ken French bundle from linearmodels==7.0: pip install 'linearmodels==7.0'")
    d = french.load()

    def col(name):
        return [float(v) for v in d[name].values]

    def spread(long, short):
        return [a - b for a, b in zip(col(long), col(short))]

    signals = {
        "Mom": col("Mom"),
        "Mom_small": spread("S1M5", "S1M1"),
        "Mom_mid": spread("S3M5", "S3M1"),
        "Mom_large": spread("S5M5", "S5M1"),
        "HML": col("HML"),
        "Value_small": spread("S1V5", "S1V1"),
        "Value_mid": spread("S3V5", "S3V1"),
        "Value_large": spread("S5V5", "S5V1"),
        "SMB": col("SMB"),
        "Size_midvalue": spread("S1V3", "S5V3"),
        "Size_midmom": spread("S1M3", "S5M3"),
    }
    dates = [str(v)[:7] for v in d["dates"].values]
    rf = col("RF")
    industries = {
        name: [x - r for x, r in zip(col(name), rf)]
        for name in ("NoDur", "Durbl", "Manuf", "Enrgy", "Chems", "BusEq",
                     "Telcm", "Utils", "Shops", "Hlth", "Money", "Other")
    }
    return list(signals), list(signals.values()), dates, industries, col("MktRF")


# ---------------------------------------------------------------- article lock

ARTICLE_CSV = Path(__file__).with_name("factor_dedup_article.csv")
SIGNAL_NAMES = [
    "Mom", "Mom_small", "Mom_mid", "Mom_large",
    "HML", "Value_small", "Value_mid", "Value_large",
    "SMB", "Size_midvalue", "Size_midmom",
]
INDUSTRY_NAMES = [
    "NoDur", "Durbl", "Manuf", "Enrgy", "Chems", "BusEq",
    "Telcm", "Utils", "Shops", "Hlth", "Money", "Other",
]
# Printed to two decimals, matching the factor post. Frozen from linearmodels 7.0.
ARTICLE = {
    "span": ("1949-01", "1983-01", "1983-02", "2017-03"),
    "avg_corr": "0.13",
    "n_eff": "4.81",
    "avg_single_sharpe": "0.46",
    "sqrt_n_estimate": "1.51",
    "combined_sharpe": "1.00",
    "groups": (
        ("HML", "Value_large", "Value_mid"),
        ("Mom", "Mom_large", "Mom_mid", "Mom_small"),
        ("SMB", "Size_midmom", "Size_midvalue"),
        ("Value_small",),
    ),
    "merged_n_eff": "4.10",
    "merged_sharpe": "1.02",
    "half1": ("0.13", "4.75", "1.33", "1.22"),
    "half2": ("0.12", "4.95", "0.77", "0.90"),
    "by_count": ("2.62", "0.87"),
    "by_source": ("3.99", "0.90"),
    "ind_raw": ("0.65", "1.47"),
    "ind_resid": ("-0.01", "12.83"),
    "value_small_large": "0.25",
    "value_small_hml": "0.64",
    "value_small_mid": "0.60",
    "value_group_avg": "0.50",
    "mom_within": "0.80",
}


def fmt2(x):
    return f"{x:.2f}"


def require(label, got, want):
    if got != want:
        sys.exit(f"article check failed for {label}: got {got}, post says {want}")


def pairwise(names, series, a, b):
    c = corr_matrix(series)
    return c[names.index(a)][names.index(b)]


def check_article_numbers(names, series, dates, industries, market):
    """Fail if the computed tables drift from the published post."""
    before = describe(names, series, 12)
    groups = cluster(before["corr"], 0.5)
    merged_series = merge(series, groups)
    after = describe([f"group{k}" for k in range(len(groups))], merged_series, 12)
    require("average correlation", fmt2(before["avg_corr"]), ARTICLE["avg_corr"])
    require("effective number", fmt2(before["n_eff"]), ARTICLE["n_eff"])
    require("average single Sharpe", fmt2(before["avg_single_sharpe"]), ARTICLE["avg_single_sharpe"])
    require("sqrt(N) estimate", fmt2(before["sqrt_n_estimate"]), ARTICLE["sqrt_n_estimate"])
    require("equal-weight Sharpe", fmt2(before["combined_sharpe"]), ARTICLE["combined_sharpe"])
    got_groups = tuple(sorted(tuple(sorted(names[i] for i in g)) for g in groups))
    require("clusters", got_groups, ARTICLE["groups"])
    require("merged effective number", fmt2(after["n_eff"]), ARTICLE["merged_n_eff"])
    require("merged Sharpe", fmt2(after["combined_sharpe"]), ARTICLE["merged_sharpe"])

    half = len(dates) // 2
    require("sample span", (dates[0], dates[half - 1], dates[half], dates[-1]), ARTICLE["span"])
    for label, lo, hi, want in (
        ("first half", 0, half, ARTICLE["half1"]),
        ("second half", half, len(dates), ARTICLE["half2"]),
    ):
        part = [s[lo:hi] for s in series]
        c = corr_matrix(part)
        everything = sharpe(equal_weight([scaled(s) for s in part]), 12)
        merged_sharpe = sharpe(equal_weight(merge(part, groups)), 12)
        require(label, (fmt2(avg_pairwise(c)), fmt2(effective_n(c)), fmt2(everything), fmt2(merged_sharpe)), want)

    pick = [names.index(n) for n in ("Mom", "Mom_small", "Mom_mid", "Mom_large", "HML", "SMB")]
    lib = [series[i] for i in pick]
    c = corr_matrix(lib)
    by_count = sharpe(equal_weight([scaled(s) for s in lib]), 12)
    by_source = merge(lib, [[0, 1, 2, 3], [4], [5]])
    require("count weighting", (fmt2(effective_n(c)), fmt2(by_count)), ARTICLE["by_count"])
    require(
        "source weighting",
        (fmt2(effective_n(corr_matrix(by_source))), fmt2(sharpe(equal_weight(by_source), 12))),
        ARTICLE["by_source"],
    )

    raw = corr_matrix(industries)
    resid = corr_matrix(strip(industries, market))
    require("industry raw", (fmt2(avg_pairwise(raw)), fmt2(effective_n(raw))), ARTICLE["ind_raw"])
    require("industry residual", (fmt2(avg_pairwise(resid)), fmt2(effective_n(resid))), ARTICLE["ind_resid"])
    require("small vs large value", fmt2(pairwise(names, series, "Value_small", "Value_large")), ARTICLE["value_small_large"])
    require("small value vs HML", fmt2(pairwise(names, series, "Value_small", "HML")), ARTICLE["value_small_hml"])
    require("small vs mid value", fmt2(pairwise(names, series, "Value_small", "Value_mid")), ARTICLE["value_small_mid"])
    group_avg = mean([
        pairwise(names, series, "Value_small", other)
        for other in ("HML", "Value_mid", "Value_large")
    ])
    require("small value vs merged value group", fmt2(group_avg), ARTICLE["value_group_avg"])
    moms = ["Mom", "Mom_small", "Mom_mid", "Mom_large"]
    mom_pairs = [
        pairwise(names, series, a, b)
        for i, a in enumerate(moms) for b in moms[i + 1 :]
    ]
    require("momentum within-group correlation", fmt2(mean(mom_pairs)), ARTICLE["mom_within"])


def article_check():
    if not ARTICLE_CSV.is_file():
        sys.exit(f"missing {ARTICLE_CSV.name}; it is the linearmodels 7.0 series behind the post")
    with ARTICLE_CSV.open(newline="") as f:
        rows = list(csv.reader(f))
    names = [h.strip() for h in rows[0][1:]]
    wanted = SIGNAL_NAMES + INDUSTRY_NAMES + ["MktRF"]
    if names != wanted:
        sys.exit(f"{ARTICLE_CSV.name} columns do not match the frozen layout")
    dates, cols = [], [[] for _ in names]
    for row in rows[1:]:
        if not row or not row[0].strip():
            continue
        dates.append(row[0].strip())
        for k, value in enumerate(row[1 : len(names) + 1]):
            cols[k].append(float(value))
    n_sig = len(SIGNAL_NAMES)
    n_ind = len(INDUSTRY_NAMES)
    check_article_numbers(
        names[:n_sig],
        cols[:n_sig],
        dates,
        cols[n_sig : n_sig + n_ind],
        cols[n_sig + n_ind],
    )
    print("article-check passed")


# ---------------------------------------------------------------- self-test


def self_test():
    rng = random.Random(7)
    months = 600
    sources = [[rng.gauss(0.004, 0.03) for _ in range(months)] for _ in range(3)]
    copies = [4, 4, 3]
    names, series, truth = [], [], []
    for s, k in enumerate(copies):
        for j in range(k):
            names.append(f"src{s}_v{j}")
            series.append([x + rng.gauss(0, 0.015) for x in sources[s]])
            truth.append(s)

    c = corr_matrix(series)
    n = len(series)
    direct = n * n / sum(sum(row) for row in c)
    assert abs(direct - effective_n(c)) < 1e-9, "N_eff formula and matrix sum disagree"

    before, groups, after = report(names, series, 12, 0.5)
    recovered = sorted(sorted({truth[i] for i in g}) for g in groups)
    assert recovered == [[0], [1], [2]], f"clusters do not match the three sources: {groups}"
    assert 2.5 < before["n_eff"] < 3.5, before["n_eff"]
    assert before["sqrt_n_estimate"] > 1.4 * before["combined_sharpe"], "sqrt(N) should overstate"
    assert after["combined_sharpe"] > 0.95 * before["combined_sharpe"], "merging lost too much"
    print("\nself-test passed")


# ---------------------------------------------------------------- main


def main():
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    g = p.add_mutually_exclusive_group()
    g.add_argument("--self-test", action="store_true")
    g.add_argument("--article-check", action="store_true")
    g.add_argument("--demo", action="store_true")
    g.add_argument("--csv")
    g.add_argument("--french")
    p.add_argument("--percent", action="store_true", help="CSV values are in percent")
    p.add_argument("--strip", help="remove each signal's beta to this column first")
    p.add_argument("--threshold", type=float, default=0.5)
    p.add_argument("--periods-per-year", type=int, default=12)
    a = p.parse_args()

    if a.demo:
        names, series, dates, industries, market = load_demo()
        ppy = a.periods_per_year
        print(f"Ken French data via linearmodels, {dates[0]} to {dates[-1]}, monthly, gross of costs\n")
        print("== 1. eleven momentum / value / size signals")
        _, groups, _ = report(names, series, ppy, a.threshold)

        print("\n== 2. same clusters, each half of the sample")
        half = len(dates) // 2
        for label, lo, hi in (("first half", 0, half), ("second half", half, len(dates))):
            part = [s[lo:hi] for s in series]
            c = corr_matrix(part)
            everything = sharpe(equal_weight([scaled(s) for s in part]), ppy)
            merged = sharpe(equal_weight(merge(part, groups)), ppy)
            print(
                f"{label} {dates[lo]}..{dates[hi - 1]}: avg corr {avg_pairwise(c):+.2f}, "
                f"N_eff {effective_n(c):.2f}, Sharpe all {everything:.2f} vs merged {merged:.2f}"
            )

        print("\n== 3. a library with four momentum variants, HML and SMB")
        pick = [names.index(n) for n in ("Mom", "Mom_small", "Mom_mid", "Mom_large", "HML", "SMB")]
        lib = [series[i] for i in pick]
        c = corr_matrix(lib)
        by_count = sharpe(equal_weight([scaled(s) for s in lib]), ppy)
        by_source = merge(lib, [[0, 1, 2, 3], [4], [5]])
        c_src = corr_matrix(by_source)
        print(f"equal weight per signal: N_eff {effective_n(c):.2f}, Sharpe {by_count:.2f}")
        print(f"merge momentum first:    N_eff {effective_n(c_src):.2f}, "
              f"Sharpe {sharpe(equal_weight(by_source), ppy):.2f}")

        print("\n== 4. twelve industry portfolios, before and after removing the market")
        ind = list(industries.values())
        raw = corr_matrix(ind)
        res = corr_matrix(strip(ind, market))
        print(f"raw excess returns:   avg corr {avg_pairwise(raw):+.2f}, N_eff {effective_n(raw):.2f}")
        print(f"market beta removed:  avg corr {avg_pairwise(res):+.2f}, N_eff {effective_n(res):.2f}")
        check_article_numbers(names, series, dates, list(industries.values()), market)
        print("\ndemo matches the frozen article numbers")
        return

    if a.article_check:
        article_check()
        return

    if a.csv or a.french:
        names, series = load_csv(a.csv, a.percent) if a.csv else load_french(a.french)
        if a.strip:
            if a.strip not in names:
                sys.exit(f"--strip column {a.strip!r} not in {names}")
            k = names.index(a.strip)
            market = series[k]
            names = names[:k] + names[k + 1 :]
            series = strip(series[:k] + series[k + 1 :], market)
        report(names, series, a.periods_per_year, a.threshold)
        return

    self_test()


if __name__ == "__main__":
    main()