"""
Synthetic control from scratch: weights, fit, lift and a placebo test.

Companion to https://vexpower.com/mmm/synthetic-control-method.

Usage
    python synthetic-control.py                      # seeded synthetic example
    python synthetic-control.py my_data.csv 40       # your own wide CSV

Your CSV: one row per period, a column called `treated` with the treated
unit's outcome, and one column per donor unit (any names). Optional columns
`week`, `period` and `synthetic` are ignored. The number after the file is the
count of pre-period rows (the campaign starts on the next row). The CSV you
can download from the widget on the page has exactly this shape.

Needs numpy and pandas. No covariates, no intercept: this is the plain
constrained least squares version (weights >= 0, summing to 1). For real
work, look at augmented synthetic control (the augsynth R package, used by
GeoLift) or CausalImpact.
"""
import sys

import numpy as np
import pandas as pd


def make_example(seed: int = 7, n: int = 52, t0: int = 40, lift: float = 0.08, noise: float = 0.02) -> tuple[pd.DataFrame, int]:
    r = np.random.default_rng(seed)
    t = np.arange(n)
    trend, season = t / n, np.sin(2 * np.pi * (t + 6) / n)
    cyc = np.zeros(n)
    for i in range(1, n):
        cyc[i] = 0.85 * cyc[i - 1] + r.normal()
    cyc /= np.abs(cyc).max()
    names = ["north", "south", "east", "west", "coast", "valley", "metro", "lakes"]
    signals = {}
    for name in names:
        lvl, g, s, c = r.uniform(60, 160), r.uniform(-0.3, 0.6), r.uniform(0, 1), r.uniform(-1, 1)
        signals[name] = lvl * (1 + g * trend + 0.15 * s * season + 0.06 * c * cyc)
    w = np.zeros(len(names))
    pick = r.choice(len(names), 3, replace=False)
    w[pick] = r.uniform(0.3, 1.3, 3)
    w /= w.sum()
    base = sum(w[j] * signals[nm] for j, nm in enumerate(names))
    treated = base * (1 + r.normal(0, noise, n)) + np.where(t >= t0, lift * base, 0)
    df = pd.DataFrame({"week": t + 1, "treated": treated})
    for nm in names:
        df[nm] = signals[nm] * (1 + r.normal(0, noise, n))
    print("True donor weights used to build the treated unit:", {k: round(float(v), 3) for k, v in zip(names, w)})
    print(f"True lift: {lift:.0%} of the untreated level, {np.mean(lift * base[t0:]):.2f} per week")
    return df, t0


def project_simplex(v: np.ndarray) -> np.ndarray:
    """Euclidean projection onto {w >= 0, sum(w) = 1} (Duchi et al. 2008)."""
    u = np.sort(v)[::-1]
    css = np.cumsum(u)
    rho = np.nonzero(u - (css - 1) / np.arange(1, len(v) + 1) > 0)[0][-1]
    theta = (css[rho] - 1) / (rho + 1)
    return np.maximum(v - theta, 0)


def fit_weights(Y0: np.ndarray, y: np.ndarray, iters: int = 20000) -> np.ndarray:
    """min ||y - Y0 w||^2  s.t. w >= 0, sum(w) = 1.  Y0 is T x J.

    Accelerated projected gradient (FISTA), the same algorithm as the widget.
    Donor series are usually highly collinear, which is why it needs many
    iterations; a general-purpose solver such as scipy's SLSQP can stop early
    on problems like this, so check the pre-period fit whatever you use.
    """
    scale = np.sqrt(np.mean(y ** 2))
    X, yy = Y0 / scale, y / scale
    G, Xy = X.T @ X, X.T @ yy
    step = 1 / (2 * np.linalg.eigvalsh(G).max())
    w = np.full(X.shape[1], 1 / X.shape[1])
    z, t = w.copy(), 1.0
    for _ in range(iters):
        w_next = project_simplex(z - step * 2 * (G @ z - Xy))
        t_next = (1 + np.sqrt(1 + 4 * t * t)) / 2
        z = w_next + ((t - 1) / t_next) * (w_next - w)
        w, t = w_next, t_next
    return w


def run(y: np.ndarray, donors: np.ndarray, t0: int):
    w = fit_weights(donors[:t0], y[:t0])
    synth = donors @ w
    gap = y - synth
    pre_rmspe = np.sqrt(np.mean(gap[:t0] ** 2))
    post_rmspe = np.sqrt(np.mean(gap[t0:] ** 2))
    return w, synth, gap, post_rmspe / pre_rmspe


if len(sys.argv) > 1:
    df = pd.read_csv(sys.argv[1])
    t0 = int(sys.argv[2])
else:
    df, t0 = make_example()

donor_cols = [c for c in df.columns if c not in ("week", "period", "treated", "synthetic")]
y = df["treated"].to_numpy(float)
D = df[donor_cols].to_numpy(float)

w, synth, gap, ratio = run(y, D, t0)
print("\nFitted weights:")
for c, v in sorted(zip(donor_cols, w), key=lambda x: -x[1]):
    print(f"  {c:>10}: {v:.3f}")
mape = np.mean(np.abs(gap[:t0]) / y[:t0])
lift = gap[t0:].mean()
print(f"\nPre-period MAPE: {mape:.2%}")
print(f"Estimated lift: {lift:.2f} per period ({lift / synth[t0:].mean():.1%} of the synthetic counterfactual)")
print(f"Total over {len(y) - t0} post periods: {gap[t0:].sum():.1f}")

# Equal-weight, level-shifted average of donors: what a simple DiD would use.
avg = D.mean(axis=1)
did_cf = avg + (y[:t0] - avg[:t0]).mean()
print(f"Equal-weight DiD-style estimate: {(y[t0:] - did_cf[t0:]).mean():.2f} per period")

# Placebo in space: pretend each donor was treated, fit it from the others.
ratios = []
for j, c in enumerate(donor_cols):
    others = np.delete(D, j, axis=1)
    ratios.append(run(D[:, j], others, t0)[3])
rank = 1 + sum(r >= ratio for r in ratios)
print(f"\nPlacebo test: treated post/pre RMSPE ratio {ratio:.2f} ranks {rank} of {len(ratios) + 1}"
      f" (permutation p-value {rank / (len(ratios) + 1):.2f})")