"""CTF model: one-factor residual targets and beta-aware portfolio books.

This is a small, distinct experiment motivated by the public Residual Books
submission. A stock's training label is its next-month return after removing
the monthly intercept and exposure to a trailing-daily market beta. Portfolio
directions hedge the same beta before position caps. Development is pre-1990;
earlier models' test results have already been observed.
"""

import numpy as np
import pandas as pd
import lightgbm as lgb
from sklearn.covariance import LedoitWolf
from sklearn.linear_model import Ridge


TRAIN_YEARS = 5
LOOKBACK_DAYS = 756
MIN_DAILY_OBS = 250
BETA_LOOKBACK_DAYS = 252
MIN_BETA_OBS = 200
N_SIDE = 125
MIN_NAMES = 50
POSITION_CAP = 0.02
RIDGE_ALPHA = 100.0
LGBM_PARAMS = {
    "objective": "regression",
    "n_estimators": 500,
    "learning_rate": 0.03,
    "num_leaves": 63,
    "min_child_samples": 1000,
    "random_state": 42,
    "n_jobs": 4,
    "deterministic": True,
    "force_col_wise": True,
    "verbosity": -1,
}


def is_test(values: pd.Series) -> pd.Series:
    return values.astype("string").str.strip().str.lower().isin(
        ["1", "1.0", "true", "t"]
    )


def unit_gross_dollar_neutral(scores: np.ndarray) -> np.ndarray:
    centered = scores - scores.mean()
    gross = np.abs(centered).sum()
    return centered / gross if gross > 0 else np.zeros_like(centered)


def capped_weights(weights: np.ndarray) -> np.ndarray:
    """Keep long and short books at 0.5 gross each, if each is feasible."""
    def allocate(side: np.ndarray) -> np.ndarray:
        active = side > 0
        if active.sum() * POSITION_CAP < 0.5:
            return side
        result = np.zeros_like(side)
        remaining = 0.5
        while active.any() and remaining > 0:
            positions = np.flatnonzero(active)
            proposed = side[active] * remaining / side[active].sum()
            over = proposed > POSITION_CAP
            if not over.any():
                result[positions] = proposed
                break
            result[positions[over]] = POSITION_CAP
            active[positions[over]] = False
            remaining -= POSITION_CAP * int(over.sum())
        return result

    return allocate(np.maximum(weights, 0)) - allocate(np.maximum(-weights, 0))


def month_zscore(values: pd.Series, dates: pd.Series) -> np.ndarray:
    centered = values - values.groupby(dates).transform("mean")
    scale = values.groupby(dates).transform("std").replace(0, np.nan)
    return np.arctan((centered / scale).fillna(0).to_numpy(dtype=np.float32))


def standardize(values: np.ndarray) -> np.ndarray:
    deviation = values.std()
    if not np.isfinite(deviation) or deviation < 1e-12:
        return np.zeros_like(values, dtype=np.float64)
    return (values - values.mean()) / deviation


def daily_matrix(
    daily_ret: pd.DataFrame, ids: np.ndarray
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    selected = daily_ret.loc[daily_ret["id"].isin(ids), ["id", "date", "ret_exc"]]
    if selected.empty:
        return np.array([], dtype="datetime64[ns]"), np.array([]), np.empty((0, 0))
    dates = pd.to_datetime(selected["date"]).to_numpy(dtype="datetime64[ns]")
    names, name_index = np.unique(selected["id"].to_numpy(dtype=np.int64), return_inverse=True)
    days, day_index = np.unique(dates, return_inverse=True)
    matrix = np.full((len(days), len(names)), np.nan, dtype=np.float32)
    matrix[day_index, name_index] = selected["ret_exc"].to_numpy(dtype=np.float32)
    return days, names, matrix


def trailing_betas(
    ids: np.ndarray,
    eom: pd.Timestamp,
    days: np.ndarray,
    daily_ids: np.ndarray,
    returns: np.ndarray,
    market: np.ndarray,
) -> np.ndarray:
    """Estimate stock betas with only the prior 252 daily observations."""
    beta = np.zeros(len(ids), dtype=np.float64)
    if not len(days) or not len(daily_ids):
        return beta
    end = np.searchsorted(days, np.datetime64(eom, "ns"), side="right")
    start = max(0, end - BETA_LOOKBACK_DAYS)
    if end - start < MIN_BETA_OBS:
        return beta

    positions = np.searchsorted(daily_ids, ids)
    matched = positions < len(daily_ids)
    matched[matched] = daily_ids[positions[matched]] == ids[matched]
    if not matched.any():
        return beta
    window = returns[start:end, positions[matched]].astype(np.float64, copy=False)
    market_window = market[start:end, None]
    observed = np.isfinite(window) & np.isfinite(market_window)
    count = observed.sum(axis=0)
    denominator_count = np.maximum(count, 1)
    stock_mean = np.where(observed, window, 0.0).sum(axis=0) / denominator_count
    market_mean = np.where(observed, market_window, 0.0).sum(axis=0) / denominator_count
    stock_centered = np.where(observed, window - stock_mean, 0.0)
    market_centered = np.where(observed, market_window - market_mean, 0.0)
    denominator = np.square(market_centered).sum(axis=0)
    valid = (count >= MIN_BETA_OBS) & (denominator > 1e-12)
    estimates = np.zeros(len(count), dtype=np.float64)
    estimates[valid] = (
        (stock_centered[:, valid] * market_centered[:, valid]).sum(axis=0)
        / denominator[valid]
    )
    beta[matched] = np.clip(estimates, -5.0, 5.0)
    return beta


def residual_targets(
    train_data: pd.DataFrame,
    days: np.ndarray,
    daily_ids: np.ndarray,
    returns: np.ndarray,
    market: np.ndarray,
) -> np.ndarray:
    """Remove monthly intercept and trailing-beta exposure from known labels."""
    labels = np.zeros(len(train_data), dtype=np.float32)
    for eom, index in train_data.groupby("eom", sort=True).indices.items():
        group = train_data.iloc[index]
        realized = group["ret_exc_lead1m"].to_numpy(dtype=np.float64)
        beta = trailing_betas(
            group["id"].to_numpy(dtype=np.int64),
            pd.Timestamp(eom),
            days,
            daily_ids,
            returns,
            market,
        )
        centered_beta = beta - beta.mean()
        centered_return = realized - realized.mean()
        beta_variance = float(centered_beta @ centered_beta)
        slope = (
            float(centered_beta @ centered_return) / beta_variance
            if beta_variance > 1e-12
            else 0.0
        )
        residual = centered_return - slope * centered_beta
        scale = residual.std(ddof=1)
        if np.isfinite(scale) and scale > 1e-12:
            labels[index] = np.arctan(residual / scale).astype(np.float32)
    return labels


def covariance_weights(
    scores: np.ndarray,
    ids: np.ndarray,
    betas: np.ndarray,
    eom: pd.Timestamp,
    days: np.ndarray,
    daily_ids: np.ndarray,
    returns: np.ndarray,
) -> np.ndarray:
    """Hedge dollar and estimated market beta before applying position caps."""
    order = np.argsort(scores, kind="stable")
    chosen = np.unique(np.r_[order[:N_SIDE], order[-N_SIDE:]])
    weights = np.zeros(len(ids), dtype=np.float64)
    if len(chosen) < MIN_NAMES or len(days) == 0:
        return capped_weights(unit_gross_dollar_neutral(scores))

    end = np.searchsorted(days, np.datetime64(eom, "ns"), side="right")
    start = max(0, end - LOOKBACK_DAYS)
    if end - start < MIN_DAILY_OBS:
        return capped_weights(unit_gross_dollar_neutral(scores))

    positions = np.searchsorted(daily_ids, ids[chosen])
    in_daily = positions < len(daily_ids)
    matched = np.zeros(len(chosen), dtype=bool)
    matched[in_daily] = daily_ids[positions[in_daily]] == ids[chosen[in_daily]]
    chosen = chosen[matched]
    positions = positions[matched]
    if len(chosen) < MIN_NAMES:
        return capped_weights(unit_gross_dollar_neutral(scores))

    window = returns[start:end, positions]
    enough = np.isfinite(window).sum(axis=0) >= MIN_DAILY_OBS
    chosen = chosen[enough]
    window = window[:, enough]
    if len(chosen) < MIN_NAMES:
        return capped_weights(unit_gross_dollar_neutral(scores))

    window = np.where(np.isfinite(window), window, 0).astype(np.float64)
    window -= window.mean(axis=0, keepdims=True)
    covariance = LedoitWolf(assume_centered=True).fit(window).covariance_
    ridge = float(np.trace(covariance) / len(chosen))
    covariance.flat[::len(chosen) + 1] += ridge
    controls = np.ones((len(chosen), 1), dtype=np.float64)
    chosen_beta = betas[chosen]
    if np.std(chosen_beta) > 1e-6:
        controls = np.column_stack([controls, chosen_beta])
    directions = np.linalg.solve(covariance, np.column_stack([scores[chosen], controls]))
    alpha, hedges = directions[:, 0], directions[:, 1:]
    hedge_coefficients = np.linalg.pinv(controls.T @ hedges) @ (controls.T @ alpha)
    direction = alpha - hedges @ hedge_coefficients
    weights[chosen] = capped_weights(unit_gross_dollar_neutral(direction))
    if not np.isfinite(weights).all():
        raise ValueError(f"Nonfinite weights for {eom}")
    return weights


def fit_year(
    panel: pd.DataFrame,
    feature_names: list[str],
    daily_ret: pd.DataFrame,
    year: int,
) -> tuple[pd.DataFrame, int]:
    panel = panel.sort_values(["eom", "id"], kind="stable").reset_index(drop=True)
    panel["eom"] = pd.to_datetime(panel["eom"])
    train = (panel["eom"].dt.year < year) & panel["ret_exc_lead1m"].notna()
    forecast = panel["eom"].dt.year == year
    if not forecast.any():
        raise ValueError(f"No rows to forecast in {year}")
    if panel.loc[train, "eom"].nunique() < 36:
        output = panel.loc[forecast, ["id", "eom", "ctff_test"]].copy()
        ranked = panel.loc[forecast].groupby("eom")[feature_names].rank(pct=True).fillna(0.5)
        output["score"] = ranked.mean(axis=1).to_numpy(dtype=np.float64)
        output["w"] = output.groupby("eom")["score"].transform(
            lambda s: capped_weights(unit_gross_dollar_neutral(s.to_numpy()))
        )
        return output[["id", "eom", "ctff_test", "w"]], int(train.sum())

    train_data = panel.loc[train].reset_index(drop=True)
    forecast_data = panel.loc[forecast, ["id", "eom", "ctff_test"]].reset_index(drop=True)
    days, daily_ids, returns = daily_matrix(daily_ret, panel["id"].unique())
    market_by_day = daily_ret.groupby("date", sort=True)["ret_exc"].mean()
    market_by_day.index = pd.to_datetime(market_by_day.index)
    market = market_by_day.reindex(pd.to_datetime(days)).to_numpy(dtype=np.float64)
    raw_train = train_data[feature_names].to_numpy(dtype=np.float32, copy=True)
    raw_forecast = panel.loc[forecast, feature_names].to_numpy(dtype=np.float32, copy=True)
    raw_train[~np.isfinite(raw_train)] = np.nan
    raw_forecast[~np.isfinite(raw_forecast)] = np.nan
    target = residual_targets(train_data, days, daily_ids, returns, market)
    boosted = lgb.LGBMRegressor(**LGBM_PARAMS)
    boosted.fit(raw_train, target)
    boosted_scores = boosted.predict(raw_forecast)
    del raw_train, raw_forecast, boosted

    ranked_train = train_data.groupby("eom")[feature_names].rank(pct=True).fillna(0.5)
    ranked_forecast = panel.loc[forecast].groupby("eom")[feature_names].rank(pct=True).fillna(0.5)
    ridge = Ridge(alpha=RIDGE_ALPHA)
    ridge.fit(ranked_train.to_numpy(dtype=np.float32) - 0.5, target)
    ridge_scores = ridge.predict(ranked_forecast.to_numpy(dtype=np.float32) - 0.5)
    del ranked_train, ranked_forecast, ridge

    weights = np.zeros(len(forecast_data), dtype=np.float64)
    for eom, index in forecast_data.groupby("eom", sort=True).indices.items():
        score = 0.5 * standardize(boosted_scores[index]) + 0.5 * standardize(
            ridge_scores[index]
        )
        month_ids = forecast_data.loc[index, "id"].to_numpy(dtype=np.int64)
        beta = trailing_betas(
            month_ids, pd.Timestamp(eom), days, daily_ids, returns, market
        )
        weights[index] = covariance_weights(
            score,
            month_ids,
            beta,
            pd.Timestamp(eom),
            days,
            daily_ids,
            returns,
        )
    forecast_data["w"] = weights
    return forecast_data, int(train.sum())


def main(
    chars: pd.DataFrame, features: pd.DataFrame, daily_ret: pd.DataFrame
) -> pd.DataFrame:
    """Return columns id, eom, w for the CTF test stock-month observations."""
    feature_names = sorted(features["features"].astype(str).tolist())
    required = {"id", "eom", "ret_exc_lead1m", "ctff_test", *feature_names}
    if required - set(chars.columns):
        raise ValueError("Missing required characteristic columns")
    if {"id", "date", "ret_exc"} - set(daily_ret.columns):
        raise ValueError("Missing required daily return columns")

    dates = pd.to_datetime(chars["eom"])
    test = is_test(chars["ctff_test"])
    years = sorted(dates.loc[test].dt.year.unique())
    if not years:
        years = [int(dates.max().year)]
    daily = daily_ret.copy()
    daily["date"] = pd.to_datetime(daily["date"])
    daily = daily.sort_values("date", kind="stable").reset_index(drop=True)
    daily_dates = daily["date"].to_numpy(dtype="datetime64[ns]")
    output = []
    for year in years:
        in_window = dates.ge(pd.Timestamp(int(year) - TRAIN_YEARS, 1, 1)) & dates.lt(
            pd.Timestamp(int(year) + 1, 1, 1)
        )
        panel = chars.loc[
            in_window, ["id", "eom", "ret_exc_lead1m", "ctff_test", *feature_names]
        ].copy()
        daily_start = pd.Timestamp(int(year) - TRAIN_YEARS - 1, 1, 1)
        daily_end = pd.Timestamp(int(year) + 1, 1, 1)
        lo = np.searchsorted(daily_dates, np.datetime64(daily_start), side="left")
        hi = np.searchsorted(daily_dates, np.datetime64(daily_end), side="left")
        result, training_rows = fit_year(panel, feature_names, daily.iloc[lo:hi], int(year))
        if test.any():
            result = result.loc[is_test(result["ctff_test"])]
        else:
            result = result.loc[result["eom"] == result["eom"].max()]
        output.append(result[["id", "eom", "w"]])
        print(f"{year}: trained on {training_rows:,} rows, output {len(result):,} rows", flush=True)

    portfolio = pd.concat(output, ignore_index=True)
    portfolio = portfolio.sort_values(["eom", "id"], kind="stable").reset_index(drop=True)
    if portfolio.empty or portfolio.isna().any().any() or portfolio.duplicated(["id", "eom"]).any():
        raise ValueError("Invalid CTF portfolio")
    portfolio["id"] = portfolio["id"].astype("int64")
    portfolio["w"] = portfolio["w"].astype("float64")
    return portfolio[["id", "eom", "w"]]