from typing import Literal

import pandas as pd

Label = Literal["right", "center"]
Closed = Literal["left", "right"]


class Roller:
    def __init__(
        self,
        df: pd.DataFrame,
        on: str,
        window_size: pd.Timedelta,
        closed: Closed,
        label: Label,
    ):
        if label not in ["right", "center"]:
            raise NotImplementedError(
                f"Unsupported label {label}."
                "Supported labels are 'right' and 'center'."
            )

        self.df = df
        self.on = on
        self.window_size = window_size
        self.closed = closed
        self.label = label

    def agg(self, method: str | dict):
        center = self.label == "center"

        df = (
            self.df.sort_values(by=self.on)
            .set_index(keys=self.on)
            .rolling(window=self.window_size, center=center, closed=self.closed)
            .aggregate(method, numeric_only=True)
            .reset_index()
        )

        return df


def rolling(
    df: pd.DataFrame,
    on: str,
    window_size: pd.Timedelta,
    closed: Closed = "right",
    label: Label = "right",
):
    """Aggregate data using rolling windows.

    Args:
        df (pd.DataFrame): DataFrame
        on (str): Column to roll on.
        window_size (pd.Timedelta): Size of the rolling window.
        closed (Closed, optional): Which side of the interval to include.
            Defaults to "right".
        label (Label, optional): Which side of the interval to use as label.
            Defaults to "right".
    """
    return Roller(df, on, window_size, closed, label)
