# %%
# CTF Model Template
# ============================================================
# Usage: Copy this file, rename it (e.g. models/my-model.py),
#        and replace all TODO sections with your implementation.
#
# CTF Compliance rules (auto-checked by eapctf.ctf.validate):
#   - main() signature must be exactly: main(chars, features, daily_ret)
#   - Return DataFrame with columns: id, eom, w (no NaN, no duplicates)
#   - No network calls (socket, requests, urllib, http, etc.)
#   - No eval/exec/os.system/subprocess
#   - Set random seed for determinism
#   - No look-ahead bias: weights at eom use only data available at eom
#   - Extra packages (beyond pandas/numpy/scipy/sklearn/polars/joblib):
#     list them in requirements.txt next to this file
# ============================================================

import numpy as np
import pandas as pd


# ---------------------------------------------------------------------------
# Reproducibility
# ---------------------------------------------------------------------------

# TODO: Keep this. CTF requires deterministic output (Rule 4).
RANDOM_SEED = 42
np.random.seed(RANDOM_SEED)


# ---------------------------------------------------------------------------
# Hyperparameters
# ---------------------------------------------------------------------------

# TODO: Adjust these for your model.
WINDOW_LENGTH = 120   # months of training data per portfolio date (rolling)
N_FACTORS = 5         # example: number of factors, latent dims, etc.


# ---------------------------------------------------------------------------
# Step 1: Data preparation
# ---------------------------------------------------------------------------


def prepare_data(chars: pd.DataFrame, features: list[str]) -> pd.DataFrame:
    """Apply feature transformations to the characteristics DataFrame.

    Default implementation: rank-normalizes each feature to [-0.5, 0.5]
    via ECDF within each eom cross-section, matching the IPCA benchmark.
    Replace or extend for your model.

    Parameters
    ----------
    chars : pd.DataFrame
        Full characteristics panel (all firms, all dates).
    features : list[str]
        Feature column names to transform.

    Returns
    -------
    pd.DataFrame
        Same shape as input, with transformed feature columns.
    """
    chars = chars.copy()

    for feat in features:
        # ECDF rank within each cross-section
        is_zero = chars[feat] == 0
        chars[feat] = chars.groupby("eom")[feat].transform(
            lambda x: x.rank(method="max", pct=True)
        )
        chars.loc[is_zero, feat] = 0          # preserve exact zeros
        chars[feat] = chars[feat].fillna(0.5) # impute missing as median

    chars[features] -= 0.5  # center to [-0.5, 0.5]
    return chars


# ---------------------------------------------------------------------------
# Step 2: Model training
# ---------------------------------------------------------------------------


def train_model(
    X_train: pd.DataFrame,
    y_train: pd.Series,
):
    """Fit a model on one rolling training window.

    Parameters
    ----------
    X_train : pd.DataFrame
        Feature matrix for the training window.
        Shape: (n_obs, n_features). Index: (id, eom).
    y_train : pd.Series
        Target: 1-month-ahead excess return (ret_exc_lead1m).
        Same index as X_train.

    Returns
    -------
    model
        Fitted model object. Can be any object with a predict() method,
        or a dict of parameters, etc. — match what predict_returns() expects.
    """
    # TODO: Replace with your model.
    # Examples:
    #   from sklearn.linear_model import Ridge
    #   model = Ridge(alpha=1.0).fit(X_train, y_train)
    #
    #   from sklearn.ensemble import RandomForestRegressor
    #   model = RandomForestRegressor(n_estimators=100, random_state=RANDOM_SEED)
    #   model.fit(X_train, y_train)
    #
    # Default: OLS via pseudo-inverse (no sklearn dependency)
    X = X_train.values
    y = y_train.values
    beta = np.linalg.lstsq(X, y, rcond=None)[0]
    return beta


# ---------------------------------------------------------------------------
# Step 3: Prediction
# ---------------------------------------------------------------------------


def predict_returns(
    model,
    X_test: np.ndarray,
) -> np.ndarray:
    """Generate expected return predictions for portfolio construction.

    Parameters
    ----------
    model
        Output of train_model().
    X_test : np.ndarray
        Feature matrix at the portfolio formation date.
        Shape: (n_firms, n_features).

    Returns
    -------
    np.ndarray
        Predicted expected returns. Shape: (n_firms,).
    """
    # TODO: Match to train_model() output.
    # Default: linear prediction from OLS beta
    return X_test @ model


# ---------------------------------------------------------------------------
# Step 4: Portfolio weight construction
# ---------------------------------------------------------------------------


def construct_weights(
    predictions: np.ndarray,
    ids: list,
    eom: pd.Timestamp,
) -> pd.DataFrame:
    """Convert expected return predictions into portfolio weights.

    Parameters
    ----------
    predictions : np.ndarray
        Expected returns per firm. Shape: (n_firms,).
    ids : list
        Firm identifiers, aligned with predictions.
    eom : pd.Timestamp
        Portfolio formation date.

    Returns
    -------
    pd.DataFrame
        Columns: id, eom, w.  Ready to append to results.
    """
    # TODO: Replace with your weighting scheme.
    # Examples:
    #   MVO (IPCA style):
    #     inv_cov = np.linalg.inv(cov_matrix)
    #     w = inv_cov @ predictions
    #
    #   Rank-weighted long-short:
    #     ranks = predictions.argsort().argsort().astype(float)
    #     w = ranks - ranks.mean()
    #     w /= np.abs(w).sum()  # normalize to sum of abs weights = 1
    #
    # Default: z-score weighting (long-short, mean-zero)
    mu = predictions.mean()
    sigma = predictions.std()
    w = (predictions - mu) / (sigma + 1e-8)

    return pd.DataFrame({"id": ids, "eom": eom, "w": w})


# ---------------------------------------------------------------------------
# Main rolling-window loop
# ---------------------------------------------------------------------------


def run_model(
    chars: pd.DataFrame,
    features: list[str],
    window_length: int = WINDOW_LENGTH,
) -> pd.DataFrame:
    """Run rolling-window train/predict loop over all portfolio dates.

    For each date in the test period (ctff_test == True):
      1. Assemble training window: [date - window_length months, date)
      2. Fit model on training data
      3. Predict expected returns at formation date
      4. Construct and store portfolio weights

    Parameters
    ----------
    chars : pd.DataFrame
        Prepared characteristics panel (after prepare_data()).
    features : list[str]
        Feature column names.
    window_length : int
        Number of months in the rolling training window.

    Returns
    -------
    pd.DataFrame
        Columns: id, eom, w — one row per firm per portfolio date.
    """
    pf_dates = chars.loc[chars["ctff_test"], "eom_ret"].sort_values().unique()
    results = []

    for d in pf_dates:
        # --- Training window ---
        # Use data strictly before d (no look-ahead)
        window_start = d + pd.DateOffset(days=1) - pd.DateOffset(months=window_length) - pd.DateOffset(days=1)
        train_mask = (chars["eom_ret"] < d) & (chars["eom_ret"] >= window_start)
        train = chars.loc[train_mask, ["id", "eom", "ret_exc_lead1m", *features]]
        train = train.set_index(["id", "eom"])

        X_train = train.drop(columns=["ret_exc_lead1m"])
        y_train = train["ret_exc_lead1m"]

        if len(X_train) < len(features) + 1:
            # Not enough data to fit; skip this date
            continue

        # --- Fit model ---
        model = train_model(X_train, y_train)

        # --- Predict at formation date ---
        test_slice = chars.loc[chars["eom_ret"] == d, ["id", *features]]
        X_test = test_slice[features].values
        ids = test_slice["id"].tolist()

        predictions = predict_returns(model, X_test)

        # --- Construct weights ---
        w_df = construct_weights(predictions, ids, eom=d)
        results.append(w_df)

    if not results:
        return pd.DataFrame(columns=["id", "eom", "w"])

    return pd.concat(results, ignore_index=True)


# ---------------------------------------------------------------------------
# CTF entry point (DO NOT rename or change signature)
# ---------------------------------------------------------------------------


def main(
    chars: pd.DataFrame,
    features: pd.DataFrame,  # noqa: ARG001 — feature name list, not matrix
    daily_ret: pd.DataFrame,  # noqa: ARG001 — unused in most models
) -> pd.DataFrame:
    """CTF submission entry point.

    Called by eapctf.ctf.run_local() with the three standard CTF datasets.

    Parameters
    ----------
    chars : pd.DataFrame
        Firm characteristics panel. Key columns:
          id, eom, eom_ret, ret_exc_lead1m, ctff_test, + 402 feature cols
    features : pd.DataFrame
        Single-column DataFrame with column 'features' listing all
        characteristic names present in chars.
    daily_ret : pd.DataFrame
        Daily excess returns: id, date, ret_exc.
        Most models do not use this for weight construction.

    Returns
    -------
    pd.DataFrame
        Columns: id (int), eom (date string YYYY-MM-DD), w (float).
        One row per firm per portfolio formation date.
        Only test-period dates (ctff_test == True) should appear.
    """
    # Convert date columns
    chars["eom"] = pd.to_datetime(chars["eom"])
    chars["eom_ret"] = pd.to_datetime(chars["eom_ret"])

    # Extract feature names
    feature_list = features["features"].tolist()

    # TODO: Optionally filter features to a subset.
    # Example (JKP 153 only):
    #   from eapctf.utils.constants import JKP_153
    #   feature_list = [f for f in JKP_153 if f in feature_list]

    # Prepare features
    chars = prepare_data(chars, feature_list)

    # Run rolling-window model
    weights = run_model(chars, feature_list, window_length=WINDOW_LENGTH)

    # Ensure eom is formatted as YYYY-MM-DD string (CTF requirement)
    weights["eom"] = pd.to_datetime(weights["eom"]).dt.strftime("%Y-%m-%d")

    return weights


# ---------------------------------------------------------------------------
# Local testing (not used by CTF infrastructure)
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import sys
    sys.path.insert(0, ".")

    print("Loading CTF data...")
    _chars = pd.read_parquet("data/ctf/ctff_chars.parquet")
    _features = pd.read_parquet("data/ctf/ctff_features.parquet")
    _daily_ret = pd.read_parquet("data/ctf/ctff_daily_ret.parquet")

    print("Running model...")
    _weights = main(_chars, _features, _daily_ret)
    print(f"Weights shape: {_weights.shape}")
    print(_weights.head())

    _weights.to_parquet("data/ctf/template_weights.parquet", index=False)
    print("Saved to data/ctf/template_weights.parquet")
