"""
Optimized Strategy: Walk-Forward Factor Selection with Trailing Meta-Selection
HW 2 Implementation

CTF Admin Modifications (2026-06-29):
--------------------------------------
1. Narrowed the broad `except Exception` in select_cluster_sharpe() to
   `except (ValueError, RuntimeError, np.linalg.LinAlgError)`.
   Reason: the security scanner rejects a catch-all `except Exception` that
   does not re-raise. The handler is a graceful fallback to Sharpe selection
   when the hierarchical clustering (squareform/linkage/fcluster) fails, so
   re-raising is not appropriate. The listed exception types cover the errors
   those routines actually raise, preserving the original fallback behavior.
"""

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
# ---------------------------------------------------------------------------
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
META_WINDOW_YEARS = 20 # Trailing window for meta-selection score calculation


# ===========================================================================
# 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.
    """
    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


# ===========================================================================
# 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 _sortino_series(tc):
    mean = tc.mean()
    downside = tc.clip(upper=0)
    stds = downside.std().replace(0, np.nan)
    return ((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_sortino(tc, valid, n_factors):
    """Top-N factors by Sortino ratio over the lookback slice."""
    return list(_sortino_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.
    """
    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 (ValueError, RuntimeError, np.linalg.LinAlgError):
        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,
    "Sortino": select_sortino
}


# ===========================================================================
# 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_inv_vol(tc, factors):
    """Weights proportional to inverse volatility of the factors."""
    tcf = tc[factors].fillna(0)
    stds = tcf.std().replace(0, np.nan)
    inv_vol = (1.0 / stds).fillna(0)
    return inv_vol.values / inv_vol.sum() if inv_vol.sum() > 1e-10 else np.full(len(factors), 1.0 / len(factors))


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,
    "InvVol": weight_inv_vol,
    "MeanVar": weight_mean_var
}


# ===========================================================================
# The candidate menu. Keep it SYMMETRIC (a full cross).
# ===========================================================================
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
]


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 over a trailing 20-year window.
    """
    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:
            # Trailing 20y history window to adapt to regimes
            start_date = rp - pd.DateOffset(months=1) - pd.DateOffset(years=META_WINDOW_YEARS) + pd.DateOffset(months=1)
            hist_window = hist.loc[start_date:]
            
            score = {}
            for ci in eligible:
                s = hist_window[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"
        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.
    """
    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: 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
    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")
    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"]))

    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)")

    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}")

    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}")