power_laws.py

#!/usr/bin/env python3
"""
powerlaw_lab.py (updated for the NEW spreadsheet format)

New default input format (from make_synthetic_input.py):
  unit_id, outcome

This version:
- Defaults to value_col="outcome" and unit_col="unit_id"
- Auto-detects CSV vs Excel (.xlsx/.xls)
- Lets you override columns if you want

Dependencies:
  pip install pandas numpy matplotlib
Optional:
  pip install openpyxl   # for .xlsx
"""

from __future__ import annotations

import argparse
import json
import math
import os
from dataclasses import dataclass, asdict
from typing import Dict, Optional, Tuple

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


# ----------------------------
# Utility math
# ----------------------------

def _norm_cdf(z: float) -> float:
    return 0.5 * (1.0 + math.erf(z / math.sqrt(2.0)))


def gini(x: np.ndarray) -> float:
    """Gini coefficient for non-negative values."""
    x = np.asarray(x, dtype=float)
    x = x[np.isfinite(x)]
    x = x[x >= 0]
    if x.size == 0:
        return float("nan")
    if np.allclose(x, 0):
        return 0.0
    x_sorted = np.sort(x)
    n = x_sorted.size
    cumx = np.cumsum(x_sorted)
    sumx = cumx[-1]
    return 1.0 - (2.0 * np.sum(cumx) / (n * sumx)) + (1.0 / n)


def top_share(x: np.ndarray, top_frac: float) -> float:
    """Share of total outcome contributed by the top fraction of units."""
    x = np.asarray(x, dtype=float)
    x = x[np.isfinite(x)]
    x = x[x >= 0]
    if x.size == 0:
        return float("nan")
    if top_frac <= 0:
        return 0.0
    if top_frac >= 1:
        return 1.0
    x_sorted = np.sort(x)[::-1]
    k = max(1, int(math.ceil(top_frac * x_sorted.size)))
    return float(np.sum(x_sorted[:k]) / max(np.sum(x_sorted), 1e-18))


def herfindahl(x: np.ndarray) -> float:
    """Herfindahl index on shares (0..1). Higher means more concentrated."""
    x = np.asarray(x, dtype=float)
    x = x[np.isfinite(x)]
    x = x[x > 0]
    if x.size == 0:
        return float("nan")
    s = x / np.sum(x)
    return float(np.sum(s * s))


# ----------------------------
# Tail models (x >= xmin)
# ----------------------------

@dataclass
class FitResult:
    model: str
    xmin: float
    n_tail: int
    params: Dict[str, float]
    loglik: float


def fit_powerlaw_tail(x: np.ndarray, xmin: float) -> FitResult:
    """
    Continuous power law for x >= xmin.
    MLE: alpha = 1 + n / sum(log(x/xmin))
    """
    xt = x[x >= xmin]
    xt = xt[np.isfinite(xt)]
    n = xt.size
    if n < 5:
        return FitResult("powerlaw", xmin, n, {"alpha": float("nan")}, float("nan"))

    logs = np.log(xt / xmin)
    denom = np.sum(logs)
    if denom <= 0:
        return FitResult("powerlaw", xmin, n, {"alpha": float("nan")}, float("nan"))

    alpha = 1.0 + n / denom
    ll = n * math.log((alpha - 1.0) / xmin) - alpha * np.sum(np.log(xt / xmin))
    return FitResult("powerlaw", xmin, n, {"alpha": float(alpha)}, float(ll))


def fit_lognormal_tail(x: np.ndarray, xmin: float) -> FitResult:
    """
    Lognormal on x>=xmin (not truncated). Practical approximation for comparison.
    MLE on log(x): mu = mean(log x), sigma = std(log x)
    """
    xt = x[x >= xmin]
    xt = xt[np.isfinite(xt)]
    n = xt.size
    if n < 5:
        return FitResult("lognormal", xmin, n, {"mu": float("nan"), "sigma": float("nan")}, float("nan"))

    lx = np.log(xt)
    mu = float(np.mean(lx))
    sigma = float(np.std(lx, ddof=0))
    if sigma <= 0:
        return FitResult("lognormal", xmin, n, {"mu": mu, "sigma": sigma}, float("nan"))

    ll = float(np.sum(
        -np.log(xt * sigma * math.sqrt(2.0 * math.pi)) - ((lx - mu) ** 2) / (2.0 * sigma ** 2)
    ))
    return FitResult("lognormal", xmin, n, {"mu": mu, "sigma": sigma}, ll)


def fit_exponential_tail(x: np.ndarray, xmin: float) -> FitResult:
    """
    Shifted exponential on y = x - xmin, for x>=xmin.
    MLE: lambda = 1/mean(y)
    """
    xt = x[x >= xmin]
    xt = xt[np.isfinite(xt)]
    n = xt.size
    if n < 5:
        return FitResult("exponential", xmin, n, {"lambda": float("nan")}, float("nan"))

    y = xt - xmin
    y = y[y >= 0]
    mean_y = float(np.mean(y))
    if mean_y <= 0:
        return FitResult("exponential", xmin, n, {"lambda": float("nan")}, float("nan"))

    lam = 1.0 / mean_y
    ll = float(n * math.log(lam) - lam * np.sum(y))
    return FitResult("exponential", xmin, n, {"lambda": float(lam)}, float(ll))


def vuong_test_loglik_diff(logpdf_diff: np.ndarray) -> Tuple[float, float]:
    """
    Vuong-style z-test for non-nested model comparison.
    Input: per-sample log-likelihood difference (model A - model B).
    Returns: (z, p_two_sided)
    """
    d = logpdf_diff[np.isfinite(logpdf_diff)]
    n = d.size
    if n < 10:
        return float("nan"), float("nan")
    m = float(np.mean(d))
    s = float(np.std(d, ddof=1))
    if s <= 0:
        return float("nan"), float("nan")
    z = math.sqrt(n) * m / s
    p = 2.0 * (1.0 - _norm_cdf(abs(z)))
    return float(z), float(p)


# ----------------------------
# Diagnostics
# ----------------------------

def ks_distance_powerlaw(x: np.ndarray, xmin: float, alpha: float) -> float:
    """
    KS distance between empirical CDF of tail and fitted powerlaw CDF.
    For x>=xmin, CDF: 1 - (x/xmin)^(1-alpha)
    """
    xt = np.sort(x[x >= xmin])
    n = xt.size
    if n < 5 or not np.isfinite(alpha) or alpha <= 1:
        return float("inf")

    emp_cdf = np.arange(1, n + 1) / n
    model_cdf = 1.0 - (xt / xmin) ** (1.0 - alpha)
    return float(np.max(np.abs(emp_cdf - model_cdf)))


def choose_xmin_auto(x: np.ndarray) -> float:
    """
    Auto-select xmin by minimizing KS distance over candidate cutoffs.
    Practical: scan candidate xmins from mid to high quantiles.
    """
    x = x[np.isfinite(x)]
    x = x[x > 0]
    if x.size < 50:
        return float(np.quantile(x, 0.5))

    qs = np.linspace(0.50, 0.90, 25)
    candidates = np.unique(np.quantile(x, qs))
    best = (float("inf"), None)

    for xmin in candidates:
        pl = fit_powerlaw_tail(x, float(xmin))
        if pl.n_tail < 30 or not np.isfinite(pl.params.get("alpha", float("nan"))):
            continue
        ks = ks_distance_powerlaw(x, float(xmin), pl.params["alpha"])
        if ks < best[0]:
            best = (ks, float(xmin))

    if best[1] is None:
        return float(np.quantile(x, 0.7))
    return float(best[1])


def ccdf_plot(x: np.ndarray, outpath: str, title: str) -> None:
    x = x[np.isfinite(x)]
    x = x[x > 0]
    if x.size == 0:
        return

    xs = np.sort(x)
    n = xs.size
    ccdf = 1.0 - (np.arange(1, n + 1) / n)

    plt.figure()
    plt.loglog(xs, ccdf, marker=".", linestyle="none")
    plt.xlabel("Outcome size (x)")
    plt.ylabel("P(X ≥ x)  (CCDF)")
    plt.title(title)
    plt.tight_layout()
    plt.savefig(outpath, dpi=180)
    plt.close()


def histogram_logx_plot(x: np.ndarray, outpath: str, title: str) -> None:
    x = x[np.isfinite(x)]
    x = x[x > 0]
    if x.size == 0:
        return

    lx = np.log(x)
    plt.figure()
    plt.hist(lx, bins=40)
    plt.xlabel("log(x)")
    plt.ylabel("count")
    plt.title(title)
    plt.tight_layout()
    plt.savefig(outpath, dpi=180)
    plt.close()


# ----------------------------
# Input loading (CSV or Excel)
# ----------------------------

def load_table(path: str, sheet: Optional[str] = None) -> pd.DataFrame:
    """
    Loads:
      - .csv
      - .xlsx / .xls (needs openpyxl/xlrd depending)
    """
    ext = os.path.splitext(path.lower())[1]
    if ext == ".csv":
        return pd.read_csv(path)
    if ext in {".xlsx", ".xls"}:
        return pd.read_excel(path, sheet_name=sheet if sheet else 0)
    raise ValueError(f"Unsupported file type '{ext}'. Use .csv or .xlsx/.xls")


def resolve_columns(df: pd.DataFrame, value_col: Optional[str], unit_col: Optional[str]) -> Tuple[str, Optional[str]]:
    """
    New default format:
      unit_id, outcome

    If not provided, we try to auto-pick:
      value_col: outcome (preferred), else first numeric column
      unit_col: unit_id if present
    """
    cols = list(df.columns)

    # unit_col default
    if unit_col is None and "unit_id" in cols:
        unit_col = "unit_id"
    elif unit_col is not None and unit_col not in cols:
        raise ValueError(f"unit_col '{unit_col}' not found. Available columns: {cols}")

    # value_col default
    if value_col is None:
        if "outcome" in cols:
            value_col = "outcome"
        else:
            # fallback: first numeric-looking column
            numeric_cols = [c for c in cols if pd.api.types.is_numeric_dtype(df[c])]
            if not numeric_cols:
                raise ValueError("No numeric columns found. Provide --value_col explicitly.")
            value_col = numeric_cols[0]
    else:
        if value_col not in cols:
            raise ValueError(f"value_col '{value_col}' not found. Available columns: {cols}")

    return value_col, unit_col


# ----------------------------
# Main analysis routine
# ----------------------------

@dataclass
class OutcomeSummary:
    n: int
    n_positive: int
    mean: float
    median: float
    p90: float
    p99: float
    gini: float
    herfindahl: float
    top_1pct_share: float
    top_5pct_share: float
    top_10pct_share: float
    top_10_units_share: float


def analyze_outcomes(
    df: pd.DataFrame,
    value_col: str,
    unit_col: Optional[str],
    xmin: Optional[float],
    auto_xmin: bool,
    outdir: str,
) -> Dict:
    os.makedirs(outdir, exist_ok=True)

    # Coerce values to numeric
    x_raw = pd.to_numeric(df[value_col], errors="coerce").to_numpy(dtype=float)
    x_pos = x_raw[np.isfinite(x_raw)]
    x_pos = x_pos[x_pos > 0]

    summ = OutcomeSummary(
        n=int(len(x_raw)),
        n_positive=int(len(x_pos)),
        mean=float(np.mean(x_pos)) if x_pos.size else float("nan"),
        median=float(np.median(x_pos)) if x_pos.size else float("nan"),
        p90=float(np.quantile(x_pos, 0.90)) if x_pos.size else float("nan"),
        p99=float(np.quantile(x_pos, 0.99)) if x_pos.size else float("nan"),
        gini=float(gini(x_pos)),
        herfindahl=float(herfindahl(x_pos)),
        top_1pct_share=float(top_share(x_pos, 0.01)),
        top_5pct_share=float(top_share(x_pos, 0.05)),
        top_10pct_share=float(top_share(x_pos, 0.10)),
        top_10_units_share=float(np.sum(np.sort(x_pos)[-10:]) / max(np.sum(x_pos), 1e-18)) if x_pos.size else float("nan"),
    )

    if xmin is None:
        xmin_use = choose_xmin_auto(x_pos) if auto_xmin else float(np.quantile(x_pos, 0.7)) if x_pos.size else float("nan")
    else:
        xmin_use = float(xmin)

    pl = fit_powerlaw_tail(x_pos, xmin_use)
    ln = fit_lognormal_tail(x_pos, xmin_use)
    ex = fit_exponential_tail(x_pos, xmin_use)

    verdict = "unknown"
    if np.isfinite(pl.loglik) and np.isfinite(ex.loglik):
        if pl.loglik - ex.loglik > 5.0:
            verdict = "heavy_tail_likely"

    if verdict == "heavy_tail_likely" and np.isfinite(ln.loglik) and np.isfinite(pl.loglik):
        if ln.loglik > pl.loglik + 5.0:
            verdict = "heavy_tail_likely_lognormal_like"
        elif pl.loglik > ln.loglik + 5.0:
            verdict = "powerlaw_tail_likely"
        else:
            verdict = "heavy_tail_likely_ambiguous"

    # Save plots
    ccdf_plot(x_pos, os.path.join(outdir, "ccdf_loglog.png"), "CCDF (log-log)")
    histogram_logx_plot(x_pos, os.path.join(outdir, "hist_logx.png"), "Histogram of log(x)")

    result = {
        "input": {
            "value_col": value_col,
            "unit_col": unit_col,
        },
        "summary": asdict(summ),
        "xmin": xmin_use,
        "fits": {
            "powerlaw": asdict(pl),
            "lognormal": asdict(ln),
            "exponential": asdict(ex),
        },
        "verdict": verdict,
        "notes": [
            "Default expected columns are unit_id and outcome.",
            "This tool checks for heavy tails and concentration.",
            "A 'powerlaw tail likely' result is suggestive, not a proof.",
            "Mixing multiple regimes can fake a tail. Segment your data if needed.",
        ],
    }

    with open(os.path.join(outdir, "summary.json"), "w", encoding="utf-8") as f:
        json.dump(result, f, indent=2)

    return result


# ----------------------------
# Bet scoring (unchanged)
# ----------------------------

BET_DIMENSIONS = [
    ("upside_ceiling", "Upside ceiling: could this be 10x to 100x your effort or capital?"),
    ("downside_cap", "Downside cap: is max loss bounded and survivable?"),
    ("leverage", "Leverage: does it scale without linear hours?"),
    ("feedback_speed", "Feedback speed: can you learn in days or weeks?"),
    ("distribution_edge", "Distribution edge: do you have a real path to reach buyers fast?"),
]


def prompt_score(prompt: str) -> int:
    while True:
        s = input(f"{prompt}  [0,1,2] : ").strip()
        if s in {"0", "1", "2"}:
            return int(s)
        if s.lower() in {"low", "l"}:
            return 0
        if s.lower() in {"med", "m", "mid"}:
            return 1
        if s.lower() in {"high", "h"}:
            return 2
        print("Enter 0, 1, 2 or low, med, high.")


def score_bets_interactive(outpath: str) -> None:
    rows = []
    print("Enter bets. Leave name empty to stop.\n")

    while True:
        name = input("Bet name: ").strip()
        if not name:
            break

        scores = {}
        total = 0
        for key, text in BET_DIMENSIONS:
            v = prompt_score(text)
            scores[key] = v
            total += v

        notes = input("One-line note (optional): ").strip()
        rows.append({"bet": name, **scores, "total_0_to_10": total, "note": notes})

        print(f"Total score: {total}/10\n")

    if not rows:
        print("No bets saved.")
        return

    pd.DataFrame(rows).sort_values("total_0_to_10", ascending=False).to_csv(outpath, index=False)
    print(f"Saved: {outpath}")


# ----------------------------
# CLI
# ----------------------------

def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description="Power law lab for outcomes and power-law bet scoring.")
    sub = p.add_subparsers(dest="cmd", required=True)

    a = sub.add_parser("analyze", help="Analyze outcome data from CSV or Excel.")
    a.add_argument("data", help="Path to data file (.csv, .xlsx, .xls).")
    a.add_argument("--sheet", default=None, help="Excel sheet name (optional).")
    a.add_argument("--value_col", default=None, help="Outcome column. Defaults to 'outcome' if present.")
    a.add_argument("--unit_col", default=None, help="Unit id column. Defaults to 'unit_id' if present.")
    a.add_argument("--xmin", type=float, default=None, help="Tail cutoff xmin. If omitted, you can use --auto_xmin.")
    a.add_argument("--auto_xmin", action="store_true", help="Auto choose xmin by KS-min scan.")
    a.add_argument("--outdir", default="powerlaw_results", help="Output folder for plots and summary.json.")

    b = sub.add_parser("score_bets", help="Interactive bet scoring to CSV.")
    b.add_argument("--out", default="bet_scores.csv", help="Output CSV path.")

    return p


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()

    if args.cmd == "analyze":
        df = load_table(args.data, sheet=args.sheet)
        value_col, unit_col = resolve_columns(df, args.value_col, args.unit_col)

        res = analyze_outcomes(
            df=df,
            value_col=value_col,
            unit_col=unit_col,
            xmin=args.xmin,
            auto_xmin=bool(args.auto_xmin),
            outdir=args.outdir,
        )

        print(json.dumps(res["summary"], indent=2))
        print(f"\nValue column: {res['input']['value_col']}")
        print(f"Unit column:  {res['input']['unit_col']}")
        print(f"Tail xmin used: {res['xmin']}")
        print(f"Verdict: {res['verdict']}")
        print(f"Outputs in: {args.outdir}")

    elif args.cmd == "score_bets":
        score_bets_interactive(args.out)


if __name__ == "__main__":
    main()