from __future__ import annotations

import pandas as pd

from psx_signal.features.market_regime import build_market_regime
from psx_signal.features.price import add_price_features
from psx_signal.features.relative_strength import add_relative_strength
from psx_signal.features.technical import add_technical_features
from psx_signal.features.volume import add_volume_features


def build_features(ohlcv: pd.DataFrame, benchmark: pd.DataFrame) -> pd.DataFrame:
    source = ohlcv.sort_values(["symbol", "date"]).reset_index(drop=True).copy()
    analytical = source.copy()
    # Corporate-action-consistent series drive indicators; raw OHLC remains the execution series.
    for column in ("open", "high", "low", "close"):
        adjusted = f"adjusted_{column}"
        if adjusted in analytical:
            analytical[column] = analytical[adjusted].combine_first(analytical[column])

    def per_symbol(group: pd.DataFrame) -> pd.DataFrame:
        return add_volume_features(add_technical_features(add_price_features(group)))

    features = (
        analytical
        .groupby("symbol", group_keys=False, sort=False)
        .apply(per_symbol, include_groups=False)
        .reset_index(drop=True)
    )
    # pandas excludes grouping column with include_groups=False; restore from aligned source.
    if "symbol" not in features:
        features["symbol"] = analytical["symbol"].to_numpy()
    features = add_relative_strength(features, benchmark)
    regime = build_market_regime(benchmark, features)
    features = features.merge(regime, on="date", how="left", validate="many_to_one")
    raw = source[["symbol", "date", "open", "high", "low", "close"]].rename(
        columns={column: f"raw_{column}" for column in ("open", "high", "low", "close")}
    )
    features = features.merge(raw, on=["symbol", "date"], how="left", validate="one_to_one")
    for column in ("open", "high", "low", "close"):
        features[column] = features.pop(f"raw_{column}")
    return features.sort_values(["date", "symbol"])