import pandas as pd

from datamazing._conform import _concat
from datamazing.pandas.transformations import grouping
from datamazing.pandas.transformations.grouping import Grouper
from datamazing.pandas.transformations.intervals import merging


def _intervals_to_breakpoints(
    df: pd.DataFrame,
    by: list[str],
    interval: tuple[str, str],
    output_breakpoint: str = "breakpoint",
) -> pd.DataFrame:
    # Get all breakpoints (left and right) for each
    # group defined by "by" columns.
    if len(by) > 0:
        df = df.set_index(by)

    df = df[list(interval)]

    df.columns.name = "_interval"
    df = (
        df.stack()
        .to_frame(name=output_breakpoint)
        .reset_index(level=by)
        .reset_index(drop=True)
        .drop_duplicates()
    )
    return df


def _breakpoints_to_intervals(
    df: pd.DataFrame,
    by: list[str],
    breakpoint: str,
    output_interval: tuple[str, str] = ("start", "end"),
) -> pd.DataFrame:
    # Convert breakpoints to consecutive intervals
    # for each group defined by "by" columns.
    if df.empty:
        # return empty dataframe with correct columns
        columns = by + list(output_interval)
        dtypes = {}
        for col in by:
            dtypes[col] = df[col].dtype
        dtypes[output_interval[0]] = df[breakpoint].dtype
        dtypes[output_interval[1]] = df[breakpoint].dtype
        return pd.DataFrame(columns=columns).astype(dtypes)

    if len(by) == 0:
        s = df[breakpoint].sort_values()
        shifted_s = s.shift(-1)
    else:
        s = df.set_index(by)[breakpoint].sort_values()
        shifted_s = s.groupby(s.index).shift(-1)

    return (
        pd.concat(
            [
                s.to_frame(name=output_interval[0]),
                shifted_s.to_frame(name=output_interval[1]),
            ],
            axis=1,
        )
        .dropna()
        .reset_index(level=by)
        .reset_index(drop=True)
    )


def group_interval(
    df: pd.DataFrame,
    interval: tuple[str, str],
    by: list[str] | None = None,
) -> Grouper:
    """Group DataFrame by intervals.
    This groups by all overlapping interval parts.
    More formally, it will produce the same result, as if exploding
    the intervals, and then doing a regular group-by.

    Args:
        df (pd.DataFrame): DataFrame
        interval (tuple[str, str]): Interval columns (start, end).
        by (list[str]): Group keys. Defaults to None.
    """
    # to make an interval group by, we use the method described in the documentation
    # 'docs/interval-operations.md'.
    if by is None:
        by = []

    # first we get all breakspoints per group and convert these to
    # consecutive intervals
    breakpoints = _intervals_to_breakpoints(
        df,
        by,
        interval,
        output_breakpoint="_breakpoint",
    )

    partitions = _breakpoints_to_intervals(
        breakpoints,
        by,
        breakpoint="_breakpoint",
        output_interval=interval,
    )

    # partition the original dataframe
    partitioned_df = merging.merge_interval_interval(
        left=df,
        right=partitions,
        on=by,
        interval=interval,
        how="inner",
    )

    # do a regular group by
    return grouping.group(
        df=partitioned_df,
        by=_concat(by, interval),
    )
