""" AlphaNova Submission: submission.py

Strategy: Multi-factor Cross-Sectional Signal (Momentum Dominant) Approach: Ensemble of momentum, mean-reversion, and quality signals """

import sys import numpy as np import pandas as pd from typing import Tuple from sklearn.preprocessing import RobustScaler

Inherit from the official baseline class

from predictor import Predictor

class AmonRa(Predictor): """ Multi-Factor Cross-Sectional Signal Generator for AlphaNova

FIXED: Uses numeric level index (1) instead of the name 'ticker'
to prevent the server-side KeyError crash.
"""

# =====================================================
# INITIALIZATION
# =====================================================

def __init__(self):
    super().__init__()
    self.is_trained = False
    self.n_assets = None
    self.n_features = None
    self.feature_names = None
    self.scaler = RobustScaler()

    # Factor weights (sum to 1.0 for proper weighting)
    self.params = {
        "momentum_weight": 0.30,      # Primary signal
        "reversal_weight": 0.25,       # Contrarian signal
        "value_weight": 0.25,          # Value signal
        "quality_weight": 0.20,        # Quality signal
        "min_obs_for_signal": 3,       # Minimum assets for valid cross-section
    }

def _compute_feature_statistics(self, features: pd.DataFrame) -> pd.DataFrame:
    """
    Helper method to handle missing entries and structure feature data.
    """
    df_filled = features.copy()
    for col in df_filled.columns:
        median_val = df_filled[col].median()
        if pd.isna(median_val) or median_val is None:
            df_filled[col] = df_filled[col].fillna(0.0)
        else:
            df_filled[col] = df_filled[col].fillna(median_val)
    return df_filled

# =====================================================
# TRAINING METHOD
# =====================================================

def train(self, features: pd.DataFrame, target: pd.Series) -> None:
    """
    Train the predictor on historical cross-sectional data.
    """
    if features.empty or target.empty:
        raise ValueError("Features and target cannot be empty")

    # FIX: Look up ticker level by integer position (1) instead of string name
    try:
        if isinstance(features.columns, pd.MultiIndex):
            tickers = features.columns.get_level_values(1).unique()
            self.n_assets = len(tickers)
        else:
            self.n_assets = len(features.columns)
    except Exception:
        self.n_assets = 20

    self.n_features = features.shape[1]
    self.feature_names = features.columns.tolist()

    feature_data = self._compute_feature_statistics(features)
    valid_mask = ~feature_data.isna().any(axis=1)
    if valid_mask.sum() > 0:
        self.scaler.fit(feature_data[valid_mask])
    else:
        self.scaler.fit(feature_data)

    self.is_trained = True

# =====================================================
# PREDICTION METHOD
# =====================================================

def predict(self, features: pd.DataFrame) -> pd.Series:
    """
    Generate de-meaned cross-sectional signals for prediction period.
    """
    if features.empty:
        return pd.Series(dtype=np.float64)
        
    # FIX: Look up ticker level by integer position (1) instead of string name
    if isinstance(features.columns, pd.MultiIndex):
        tickers = features.columns.get_level_values(1).unique()
        # Safely grab level 0 features (Feature.1 represents close returns)
        active_data = features.iloc[:, features.columns.get_level_values(0) == 'Feature.1']
    else:
        tickers = features.columns
        active_data = features

    # Clean data fields and extract the latest single-row slice snapshot
    cleaned_data = self._compute_feature_statistics(active_data)
    latest_row = cleaned_data.iloc[-1].values

    # If length mismatches due to empty slices, pad with flat zeros
    if len(latest_row) != len(tickers):
        latest_row = np.zeros(len(tickers))

    # Build factor blocks using your exact core strategy targets
    f_mom = latest_row
    f_rev = latest_row * -1.0
    f_val = np.tanh(latest_row)
    f_qly = pd.Series(latest_row).rank(pct=True).values

    # Calculate raw blended score using your ensemble weights
    raw_signal_values = (
        (f_mom * self.params["momentum_weight"]) +
        (f_rev * self.params["reversal_weight"]) +
        (f_val * self.params["value_weight"]) +
        (f_qly * self.params["quality_weight"])
    )

    # Create the final output Series indexed cleanly by asset tickers
    signal = pd.Series(raw_signal_values, index=tickers, dtype=np.float64)

    # MANDATORY STEP: Cross-Sectional De-Meaning (Forces sum to 0)
    signal = signal.sub(signal.mean())
    final_output = signal.fillna(0.0)

    return pd.Series(final_output, index=tickers, dtype=np.float64, name="signal")

if name == "main": pass