"""Portfolio construction from model scores."""

from __future__ import annotations

import pandas as pd


def top_k_equal_weight(scores: pd.DataFrame, k: int = 30) -> pd.DataFrame:
    """
    Build equal-weight long-only portfolio from cross-sectional scores.
    Input: MultiIndex (instrument, datetime) with score column or Series.
    """
    if isinstance(scores, pd.Series):
        scores = scores.to_frame("score")

    weights = []
    for dt, group in scores.groupby(level="datetime"):
        top = group.nlargest(k, "score")
        w = pd.Series(1.0 / len(top), index=top.index)
        weights.append(w)
    return pd.concat(weights).to_frame("weight")


def long_short_quantile(scores: pd.DataFrame, n_groups: int = 5) -> pd.DataFrame:
    weights = []
    for dt, group in scores.groupby(level="datetime"):
        group = group.copy()
        group["group"] = pd.qcut(group["score"].rank(method="first"), n_groups, labels=False)
        long = group[group["group"] == n_groups - 1]
        short = group[group["group"] == 0]
        w = pd.Series(0.0, index=group.index)
        if len(long):
            w.loc[long.index] = 0.5 / len(long)
        if len(short):
            w.loc[short.index] = -0.5 / len(short)
        weights.append(w.to_frame("weight"))
    return pd.concat(weights)