"""
Strategy Template: Walk-Forward Factor Selection with
Meta-Selection

WHAT THIS DOES
==============
A long-short US equity strategy in four stages:

  Stage 1  For each of ~400 stock characteristics, build a monthly
           "factor portfolio" return: weight each stock by its
           cross-sectional rank (rank - 0.5), long the high end, short the
           low end.
  Stage 2  Define a MENU of candidate strategy configurations (lookback
           window x number of factors x selection rule x weighting rule)
           and simulate each one WALK-FORWARD: at each rebalance it picks
           factors using only past data, then earns the realized returns.
           Every config therefore has its own out-of-sample track record.
  Stage 3  META-SELECTION: at each rebalance, rank configs by the Sharpe
           ratio of their own out-of-sample history so far, and blend the
           top K. The strategy's METHODOLOGY is chosen by past data -- not
           by you peeking at the test period. This is the whole point.
  Stage 4  Convert the blended configs' factor weights into stock-level
           weights (exact replication), output (id, eom, w).

WHY THE META LAYER MATTERS
==========================
If you backtest 36 configurations, pick the best one over 1990-2023, and
report its 1990-2023 Sharpe, that number is biased: you used the answer to
choose the model. Here the config live at month t is chosen from
performance up to t-1 only, so the post-1990 numbers are honest.

RUN IT
======
    python 04-StrategyTemplate.py --data C:/data/jkp

Expects ctff_chars.parquet / ctff_features.parquet in --data
(see 02-DownloadJKPdata.py). Prints Sharpe ratios by period.

YOUR ASSIGNMENT
===============
Search for "EXTENSION POINT" below. Each is a self-contained place to add
your own idea WITHOUT breaking temporal integrity. Suggested order:
  1. Add a weighting scheme (easiest -- e.g. inverse-volatility).
  2. Add a selection rule (e.g. Sortino, or Sharpe shrunk toward zero).
  3. Change the meta-criterion (e.g. trailing 10y Sharpe instead of
     expanding; or penalize configs with high turnover).
  4. Add a factor-construction variant (e.g. decile spread portfolios).
Re-run after each change and write down what happened in 2004-2013 and
2014-2023 -- the honest decades -- not just the full period.
"""

import argparse
import time

import numpy as np
import pandas as pd
from scipy.cluster.hierarchy import linkage, fcluster
from scipy.spatial.distance import squareform

# ---------------------------------------------------------------------------
# Hyperparameters. Fix these BEFORE looking at test-period results, and
# resist the urge to tune them afterwards -- that re-introduces the bias
# this design removes.
# ---------------------------------------------------------------------------
REBAL_MONTHS = 6      # rebalance / meta-selection frequency (months)
META_TOP_K = 3        # how many configs the meta layer blends
META_MIN_OBS = 120    # months of config history before meta-selection starts
MIN_OBS = 60          # months of data required for a factor to be usable
MIN_STOCKS = 10       # min stocks per side of a factor portfolio
MIN_HISTORY = 24      # months of history before the first rebalance
SHRINKAGE = 0.3       # covariance shrinkage in mean-variance weighting
N_CLUSTERS_MULT = 3   # clusters = n_factors * this, in ClusterSharpe


# ===========================================================================
# Stage 1: factor portfolio returns
# ===========================================================================
def compute_factor_ls_returns(chars, feature_names, min_stocks, min_obs):
    """
    Linear-rank long-short return for every feature, every month.

    weight_i = pct_rank(feature_i) - 0.5          (mean ~ 0)
    LS_t     = sum_i(weight_i * ret_i) / sum_i(|weight_i|)

    CONVENTION (memorize this): the LS value stored at month t uses
    ret_exc_lead1m, i.e. it is the return earned over month t+1. It is
    KNOWN only at the end of month t+1. All later stages must respect this.
    """
    print("  Stage 1: factor LS returns...")
    t0 = time.time()

    valid_features = [f for f in feature_names if f in chars.columns]
    df = chars[["eom", "ret_exc_lead1m"] + valid_features]
    df = df[df["ret_exc_lead1m"].notna()]

    months = sorted(df["eom"].unique())
    ls_matrix = np.full((len(months), len(valid_features)), np.nan)

    for m_idx, eom in enumerate(months):
        month_data = df[df["eom"] == eom]
        if len(month_data) < min_stocks * 2:
            continue

        ret = month_data["ret_exc_lead1m"].values
        ranks = month_data[valid_features].rank(pct=True).values
        w = np.nan_to_num(ranks, nan=0.5) - 0.5
        not_nan = ~np.isnan(month_data[valid_features].values)
        w = w * not_nan

        num = (ret[:, None] * w).sum(axis=0)
        den = np.abs(w).sum(axis=0)
        ok = (not_nan.sum(axis=0) >= min_stocks * 2) & (den > 1e-10)
        row = np.full(len(valid_features), np.nan)
        row[ok] = num[ok] / den[ok]
        ls_matrix[m_idx] = row

    ls_df = pd.DataFrame(ls_matrix, index=pd.to_datetime(months),
                         columns=valid_features).sort_index()
    ls_df = ls_df[ls_df.columns[ls_df.notna().sum() >= min_obs]]
    print(f"    {ls_df.shape[0]} months x {ls_df.shape[1]} factors "
          f"({time.time()-t0:.0f}s)")
    return ls_df

    # ----------------------------------------------------------------------
    # EXTENSION POINT 1: factor construction.
    # Ideas: decile spread (top 10% minus bottom 10%, equal-weighted);
    # z-score weights instead of rank weights; industry-neutral ranks
    # (rank within `sic` groups); volatility-scaled factor returns.
    # Add a parallel function and a new menu dimension "ls_type".
    # ----------------------------------------------------------------------


# ===========================================================================
# Selection rules: history -> ordered list of factor names
# ===========================================================================
def _sharpe_series(tc):
    stds = tc.std().replace(0, np.nan)
    return ((tc.mean() / stds) * np.sqrt(12)).fillna(0)


def select_sharpe(tc, valid, n_factors):
    """Top-N factors by annualized Sharpe over the lookback slice."""
    return list(_sharpe_series(tc[valid]).nlargest(
        min(n_factors, len(valid))).index)


def select_cluster_sharpe(tc, valid, n_factors):
    """
    Diversified top-N: cluster factors by return correlation, then take the
    best-Sharpe factor from each of the best N clusters. Prevents picking
    five flavours of the same momentum factor.
    """
    tcv = tc[valid]
    if len(tcv) < 12 or len(valid) < n_factors:
        return select_sharpe(tc, valid, n_factors)

    corr = tcv.corr().fillna(0).clip(-1, 1)
    dist = (1.0 - corr.abs()).values.copy()
    np.fill_diagonal(dist, 0)
    dist = np.clip((dist + dist.T) / 2, 0, None)
    try:
        condensed = np.nan_to_num(squareform(dist, checks=False),
                                  nan=1.0, posinf=1.0, neginf=0.0)
        Z = linkage(condensed, method="average")
        n_clusters = min(max(n_factors * N_CLUSTERS_MULT, 20), len(valid))
        labels = fcluster(Z, t=n_clusters, criterion="maxclust")
    except Exception:
        return select_sharpe(tc, valid, n_factors)

    scores = _sharpe_series(tcv)
    cdf = pd.DataFrame({"factor": valid, "cluster": labels,
                        "score": scores.fillna(-999).values})
    best = cdf.loc[cdf.groupby("cluster")["score"].idxmax()]
    best = best[best["score"] > -998].sort_values("score", ascending=False)
    return best.head(n_factors)["factor"].tolist()


SELECT_FN = {
    "Sharpe": select_sharpe,
    "ClusterSharpe": select_cluster_sharpe,
    # ----------------------------------------------------------------------
    # EXTENSION POINT 2: add a selection rule and list it in the menu below.
    # Ideas: Sortino ratio; Sharpe shrunk toward the cross-factor mean
    # (penalizes short lucky streaks); |Sharpe| with sign-flipping (short
    # the factor if its Sharpe is negative); selection on t-statistics.
    # Signature: fn(tc, valid, n_factors) -> list of factor names, where
    # tc is the lookback slice of LS returns (NaN already filled with 0).
    # ----------------------------------------------------------------------
}


# ===========================================================================
# Weighting rules: history + chosen factors -> weight vector (sums to 1)
# ===========================================================================
def weight_ew(tc, factors):
    return np.full(len(factors), 1.0 / len(factors))


def weight_rank_sharpe(tc, factors):
    """Best factor by Sharpe gets the largest weight (rank weights)."""
    order = _sharpe_series(tc[factors]).rank(method="first")
    w = order.values.astype(float)
    return w / w.sum()


def weight_mean_var(tc, factors):
    """w ~ Sigma^{-1} mu with diagonal shrinkage; long-only, sums to 1."""
    tcf = tc[factors].fillna(0)
    n = len(factors)
    if len(tcf) < max(12, n + 2):
        return np.full(n, 1.0 / n)
    mu, cov = tcf.mean().values, tcf.cov().values
    cov = (1 - SHRINKAGE) * cov + SHRINKAGE * (np.trace(cov) / n) * np.eye(n)
    try:
        w = np.linalg.inv(cov) @ mu
    except np.linalg.LinAlgError:
        return np.full(n, 1.0 / n)
    w = np.maximum(w, 0)
    return w / w.sum() if w.sum() > 1e-10 else np.full(n, 1.0 / n)


WEIGHT_FN = {
    "EW": weight_ew,
    "RankSharpe": weight_rank_sharpe,
    "MeanVar": weight_mean_var,
    # ----------------------------------------------------------------------
    # EXTENSION POINT 3: add a weighting scheme and list it in the menu.
    # Ideas: inverse volatility (w_k ~ 1/sigma_k); risk parity; equal
    # risk contribution; mean-variance with a long-short factor book.
    # Signature: fn(tc, factors) -> np.array of weights summing to 1.
    # ----------------------------------------------------------------------
}


# ===========================================================================
# The candidate menu. Keep it SYMMETRIC (a full cross), not cherry-picked.
# ===========================================================================
CONFIG_GRID = [
    {"window": w, "n_factors": n, "selection": s, "weighting": g}
    for w in (10, 20, None)            # lookback years; None = expanding
    for n in (5, 10)                   # number of factors held
    for s in SELECT_FN                 # selection rules registered above
    for g in WEIGHT_FN                 # weighting rules registered above
]
# ---------------------------------------------------------------------------
# EXTENSION POINT 4: menu design. Add dimensions (rebalance frequency,
# ls_type from Extension 1, ...) or values -- but beware: a bigger menu
# means more noise for the meta layer to sort through, and each config
# needs ~10 years of track record before it can be chosen. Quality over
# quantity.
# ---------------------------------------------------------------------------


def config_label(cfg):
    w = "all" if cfg["window"] is None else f"{cfg['window']}y"
    return f"{w}/{cfg['n_factors']}/{cfg['selection']}/{cfg['weighting']}"


# ===========================================================================
# Stage 2: walk-forward simulation of every config
# ===========================================================================
def simulate_configs(ls_returns, configs, rebal_months, min_obs, min_history):
    """
    At each rebalance rp, every config selects factors/weights from LS rows
    <= rp - 1 month, then earns the realized LS returns until the next
    rebalance. Returns the per-config schedules and out-of-sample returns.
    """
    print("  Stage 2: walk-forward config simulation...")
    t0 = time.time()
    all_months = ls_returns.index.sort_values()
    rebal_points = list(all_months[::rebal_months])

    oos = pd.DataFrame(np.nan, index=all_months, columns=range(len(configs)))
    schedule = {}
    sel_cache = {}

    for i, rp in enumerate(rebal_points):
        sel_end = rp - pd.DateOffset(months=1)   # last KNOWN LS row
        ls_avail = ls_returns.loc[:sel_end]
        if len(ls_avail.dropna(how="all")) < min_history:
            continue

        next_rp = rebal_points[i + 1] if i + 1 < len(rebal_points) else None
        block = (all_months[(all_months >= rp) & (all_months < next_rp)]
                 if next_rp is not None else all_months[all_months >= rp])

        threshold = min(min_obs, max(6, len(ls_avail) // 2))
        valid = list(ls_avail.columns[ls_avail.notna().sum() >= threshold])
        if not valid:
            continue

        per_cfg = {}
        for ci, cfg in enumerate(configs):
            if cfg["window"] is None:
                tc = ls_avail
            else:
                start = (sel_end - pd.DateOffset(years=cfg["window"])
                         + pd.DateOffset(months=1))
                tc = ls_avail.loc[start:]
            tc = tc.fillna(0)
            if len(tc) < 12:
                continue

            key = (rp, cfg["window"], cfg["selection"], cfg["n_factors"])
            if key not in sel_cache:
                sel_cache[key] = SELECT_FN[cfg["selection"]](
                    tc, valid, cfg["n_factors"])
            factors = sel_cache[key]
            if not factors:
                continue

            weights = WEIGHT_FN[cfg["weighting"]](tc, factors)
            per_cfg[ci] = (factors, weights)
            oos.loc[block, ci] = (
                ls_returns.loc[block, factors].fillna(0).values @ weights)

        if per_cfg:
            schedule[rp] = per_cfg

    print(f"    {len(configs)} configs x {len(schedule)} rebalances "
          f"({time.time()-t0:.0f}s)")
    return rebal_points, schedule, oos


# ===========================================================================
# Stage 3: meta-selection
# ===========================================================================
def meta_select(rebal_points, schedule, oos, configs, top_k, meta_min_obs):
    """
    At each rebalance, blend the top_k configs by the annualized Sharpe of
    their OWN out-of-sample returns up to t-1. Until any config has
    meta_min_obs months of history, hold an equal blend of all configs.
    """
    print("  Stage 3: meta-selection...")
    chosen, log = {}, []
    for rp in rebal_points:
        if rp not in schedule:
            continue
        active = list(schedule[rp].keys())
        hist = oos.loc[:rp - pd.DateOffset(months=1)]
        counts = hist.notna().sum()

        eligible = [ci for ci in active if counts[ci] >= meta_min_obs]
        if eligible:
            score = {}
            for ci in eligible:
                s = hist[ci].dropna()
                sd = s.std()
                score[ci] = (s.mean() / sd) * np.sqrt(12) if sd > 0 else -np.inf
            top = sorted(score, key=score.get, reverse=True)[:top_k]
            mode = "meta"
            # --------------------------------------------------------------
            # EXTENSION POINT 5: the meta-criterion. This expanding Sharpe
            # is the simplest defensible choice. Ideas: trailing 10-year
            # Sharpe (adapts faster, noisier); minimum Sharpe across
            # sub-windows (rewards consistency); Sharpe net of an estimated
            # transaction-cost drag; performance-weighted blend of ALL
            # configs (softmax over scores) instead of a top-K cutoff.
            # Whatever you do, only use `hist` -- rows <= rp - 1 month.
            # --------------------------------------------------------------
        else:
            top, mode = active, "ew_prior"

        chosen[rp] = top
        log.append({"rebal": rp, "mode": mode,
                    "chosen": ";".join(config_label(configs[ci])
                                       for ci in top)})
    n_meta = sum(1 for r in log if r["mode"] == "meta")
    print(f"    {len(chosen)} rebalances ({n_meta} meta-selected)")
    return chosen, log


# ===========================================================================
# Stage 4: stock-level weights
# ===========================================================================
def compute_stock_weights(chars, schedule, chosen):
    """
    Blend the chosen configs equally and map factor weights to stocks:
      w_i = (1/K) * sum_cfg sum_k  fw_k * (rank_i - 0.5) / sum_j|rank_j - 0.5|
    Each config's stock portfolio exactly replicates its factor-level
    return, so Stage 3's statistics describe the portfolio actually held.
    """
    print("  Stage 4: stock-level weights...")
    t0 = time.time()
    rebal_dates = sorted(chosen.keys())
    if not rebal_dates:
        return pd.DataFrame(columns=["id", "eom", "w"])

    all_factors = sorted({f for rp in rebal_dates for ci in chosen[rp]
                          for f in schedule[rp][ci][0]})
    available = [f for f in all_factors if f in chars.columns]
    df_work = chars[["id", "eom"] + available]
    months = sorted(df_work["eom"].unique())

    results = []
    for eom in months:
        applicable = [d for d in rebal_dates if d <= eom]
        if not applicable:
            continue
        rp = applicable[-1]
        cfg_ids = chosen[rp]

        month_data = df_work[df_work["eom"] == eom]
        if len(month_data) < 50:
            continue

        needed = sorted({f for ci in cfg_ids for f in schedule[rp][ci][0]
                         if f in month_data.columns})
        if not needed:
            continue
        ranks_all = month_data[needed].rank(pct=True)

        stock_w = np.zeros(len(month_data))
        blend = 1.0 / len(cfg_ids)
        for ci in cfg_ids:
            factors, weights = schedule[rp][ci]
            for factor, fw in zip(factors, weights):
                if factor not in ranks_all.columns:
                    continue
                r = ranks_all[factor].values
                nan_mask = np.isnan(r)
                rw = np.where(nan_mask, 0.0, r - 0.5)
                denom = np.abs(rw).sum()
                if denom < 1e-10 or (~nan_mask).sum() < MIN_STOCKS * 2:
                    continue
                stock_w += blend * fw * rw / denom

        nz = np.abs(stock_w) > 1e-15
        if nz.any():
            results.append(pd.DataFrame({"id": month_data["id"].values[nz],
                                         "eom": eom, "w": stock_w[nz]}))

    output = pd.concat(results, ignore_index=True)
    print(f"    {len(output):,} rows, {output['eom'].nunique()} months "
          f"({time.time()-t0:.0f}s)")
    return output


# ===========================================================================
# Entry point (course interface): main(chars, features, daily_ret)
# ===========================================================================
def main(chars, features, daily_ret):
    feature_names = features["features"].tolist()
    chars = chars.copy()
    chars["eom"] = pd.to_datetime(chars["eom"])
    n_months = chars["eom"].nunique()
    print(f"Data: {len(chars):,} stock-months, {len(feature_names)} "
          f"features, {n_months} months")

    if n_months < 200:   # small validation datasets: loosen thresholds
        min_obs, min_stocks = max(6, n_months // 5), 5
        min_history, meta_min_obs = 12, max(24, n_months // 3)
    else:
        min_obs, min_stocks = MIN_OBS, MIN_STOCKS
        min_history, meta_min_obs = MIN_HISTORY, META_MIN_OBS

    ls = compute_factor_ls_returns(chars, feature_names, min_stocks, min_obs)
    rebal_points, schedule, oos = simulate_configs(
        ls, CONFIG_GRID, REBAL_MONTHS, min_obs, min_history)
    chosen, meta_log = meta_select(
        rebal_points, schedule, oos, CONFIG_GRID, META_TOP_K, meta_min_obs)

    main.meta_log = meta_log          # diagnostics for graders/notebooks
    main.config_oos = oos
    main.config_labels = [config_label(c) for c in CONFIG_GRID]

    output = compute_stock_weights(chars, schedule, chosen)
    assert set(output.columns) == {"id", "eom", "w"}
    assert output["w"].notna().all() and len(output) > 0
    return output


# ===========================================================================
# Stand-alone runner with evaluation
# ===========================================================================
if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", required=True,
                    help="directory containing ctff_*.parquet")
    ap.add_argument("--out", default=None,
                    help="path for the weights CSV "
                         "(default: <data>/weights.csv)")
    args = ap.parse_args()

    chars = pd.read_parquet(f"{args.data}/ctff_chars.parquet")
    features = pd.read_parquet(f"{args.data}/ctff_features.parquet")
    weights = main(chars, features,
                   pd.DataFrame(columns=["id", "date", "ret_exc"]))

    # Gross portfolio returns: weights at eom t earn ret_exc_lead1m of t
    chars["eom"] = pd.to_datetime(chars["eom"])
    weights["eom"] = pd.to_datetime(weights["eom"])
    merged = weights.merge(
        chars[["id", "eom", "ret_exc_lead1m"]].dropna(),
        on=["id", "eom"], how="inner")
    port = (merged["w"] * merged["ret_exc_lead1m"]).groupby(
        merged["eom"]).sum().sort_index()

    def sharpe(r):
        return r.mean() / r.std() * np.sqrt(12) if len(r) >= 12 else np.nan

    print("\n  SHARPE RATIOS BY PERIOD (gross)")
    for name, (a, b) in {
        "1960-1989": ("1960", "1989"), "1990-2003": ("1990", "2003"),
        "2004-2013": ("2004", "2013"), "2014-2023": ("2014", "2023"),
        "1990-2023": ("1990", "2023"),
    }.items():
        sub = port.loc[a:b]
        print(f"  {name}: {sharpe(sub):7.3f}  ({len(sub)} months)")

    # One-way turnover (per month, as fraction of gross book)
    piv = weights.pivot_table(index="eom", columns="id", values="w",
                              fill_value=0.0)
    to = piv.diff().abs().sum(axis=1) / 2
    print(f"\n  Avg monthly one-way turnover: {to.iloc[1:].mean():.3f}")
    # NOTE: returns above are GROSS. A simple net adjustment:
    # net_t = gross_t - 2 * turnover_t * cost_per_side (e.g. 0.0025).

    # Portfolio weights in JKP CTF submission format (id, eom YYYY-MM-DD, w).
    # main() guarantees no missing weights; graders score only ctff_test rows
    # but all rows are exported. See https://jkpfactors.com/ctf/rules .
    out_path = args.out or f"{args.data}/weights.csv"
    submission = weights[["id", "eom", "w"]].copy()
    submission["eom"] = submission["eom"].dt.strftime("%Y-%m-%d")
    submission.to_csv(out_path, index=False)
    print(f"  Wrote {len(submission):,} weight rows -> {out_path}")