"""Combine multiple alphas into a single combined weight.

Future combination methods can be registered below.
"""

import logging
from typing import Callable

import pandas as pd

from pipeline.common.schema import COMBO_COLUMNS

logger = logging.getLogger(__name__)


def _equal_weight(alpha_dfs: list[pd.DataFrame]) -> pd.DataFrame:
    """Equal-weight combination: mean of all alpha weights per (symbol_id, date).

    If any alpha has NaN for a symbol/date, that alpha is skipped for that row.
    """
    # Stack all alphas with (symbol_id, date, alpha_name) as key
    combined = pd.concat(alpha_dfs, ignore_index=True)
    # Group by symbol_id + date, take mean of weights
    result = combined.groupby(["symbol_id", "date"])["weight"].mean().reset_index()
    return result


# Registry of combo methods — add new functions + register them here
COMBO_METHODS: dict[str, Callable] = {
    "equal_weight": _equal_weight,
}


def combine_alphas(
    alpha_paths: list[str],
    combo_name: str,
    method: str = "equal_weight",
) -> pd.DataFrame:
    """Load alphas from parquet, combine, and return combo weights.

    Args:
        alpha_paths: List of paths to alpha parquet files.
        combo_name: Name identifier for this combo.
        method: Combination method ('equal_weight').

    Returns:
        DataFrame with COMBO_COLUMNS.

    Raises:
        ValueError: If method is unknown or alpha grids don't align.
    """
    if method not in COMBO_METHODS:
        raise ValueError(
            f"Unknown combo method: {method}. Options: {list(COMBO_METHODS)}"
        )

    alpha_dfs = []
    for path in alpha_paths:
        df = pd.read_parquet(path)
        alpha_dfs.append(df)
        logger.info("Loaded alpha: %s (%d rows)", path, len(df))

    # Verify alignment: all alphas must share the same (symbol_id, date) pairs
    keys = [set(zip(df["symbol_id"], pd.to_datetime(df["date"]).astype(str))) for df in alpha_dfs]
    common = keys[0]
    for i, k in enumerate(keys[1:], 1):
        if k != common:
            logger.warning("Alpha %d has different (symbol_id, date) grid — intersection used", i)
            common = common.intersection(k)

    combine_fn = COMBO_METHODS[method]
    result = combine_fn(alpha_dfs)
    result["combo_name"] = combo_name
    result = result[COMBO_COLUMNS]
    result = result.sort_values(["symbol_id", "date"]).reset_index(drop=True)

    logger.info(
        "Combo '%s': %d symbols × %d dates, weight range [%.4f, %.4f]",
        combo_name,
        result["symbol_id"].nunique(),
        result["date"].nunique(),
        result["weight"].min(),
        result["weight"].max(),
    )
    return result