"""Factor neutralization (industry / market cap) via cross-sectional regression."""

from __future__ import annotations

import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression


def neutralize_factor(
    factor: pd.Series,
    exposures: pd.DataFrame,
    date_level: str = "datetime",
) -> pd.Series:
    """
    Neutralize factor against exposure columns (e.g. industry dummies, log market cap).
    Expects MultiIndex (instrument, datetime) or a panel with date column.
    """
    if isinstance(factor.index, pd.MultiIndex):
        df = factor.to_frame("factor").join(exposures, how="left")
        neutralized = []
        for dt, group in df.groupby(level=date_level):
            y = group["factor"].values
            x = group[exposures.columns].fillna(0).values
            if len(y) < x.shape[1] + 2:
                neutralized.append(group["factor"])
                continue
            reg = LinearRegression().fit(x, y)
            resid = y - reg.predict(x)
            neutralized.append(pd.Series(resid, index=group.index))
        return pd.concat(neutralized).sort_index()

    raise NotImplementedError("neutralize_factor currently supports MultiIndex panels only")


def zscore_by_date(factor: pd.Series, date_level: str = "datetime") -> pd.Series:
    def _z(x):
        return (x - x.mean()) / (x.std() + 1e-8)

    return factor.groupby(level=date_level, group_keys=False).apply(_z)