from typing import Literal

import pandas as pd

from datamazing.pandas.transformations import basic, merging, quantities


def _intersect_intervals(
    df: pd.DataFrame,
    left_interval: tuple[str, str],
    right_interval: tuple[str, str],
    output_interval: tuple[str, str],
) -> pd.DataFrame:
    # When comparing two intervals A and B, interval B
    # will partition interval A into three parts:
    # - the left-cut (part of A left of B)
    # - the overlapping part (part of A overlapping B)
    # - the right-cut (part of A right of B)
    # This functions returns the overlaps of intervals
    df = df.copy()
    df[output_interval[0]] = df[[left_interval[0], right_interval[0]]].max(
        axis=1, skipna=False
    )
    df[output_interval[1]] = df[[left_interval[1], right_interval[1]]].min(
        axis=1, skipna=False
    )

    df = df.drop(
        columns=(
            set(left_interval)
            .union(set(right_interval))
            .difference(set(output_interval))
        )
    )

    is_valid = df[output_interval[0]] < df[output_interval[1]]

    df = df[is_valid]

    return df


def _left_cut_intervals(
    df: pd.DataFrame,
    left_interval: tuple[str, str],
    right_interval: tuple[str, str],
    output_interval: tuple[str, str],
) -> pd.DataFrame:
    # When comparing two intervals A and B, interval B
    # will partition interval A into three parts:
    # - the left-cut (part of A left of B)
    # - the overlapping part (part of A overlapping B)
    # - the right-cut (part of A right of B)
    # This functions returns the left-cuts of intervals

    df = df.copy()
    df[output_interval[0]], df[output_interval[1]] = df[left_interval[0]], df[
        [left_interval[1], right_interval[0]]
    ].min(axis=1, skipna=False)

    df = df.drop(
        columns=(
            set(left_interval)
            .union(set(right_interval))
            .difference(set(output_interval))
        )
    )

    is_valid = df[output_interval[0]] < df[output_interval[1]]

    df = df[is_valid]

    return df


def _right_cut_intervals(
    df: pd.DataFrame,
    left_interval: tuple[str, str],
    right_interval: tuple[str, str],
    output_interval: tuple[str, str],
) -> pd.DataFrame:
    # When comparing two intervals A and B, interval B
    # will partition interval A into three parts:
    # - the left-cut (part of A left of B)
    # - the overlapping part (part of A overlapping B)
    # - the right-cut (part of A right of B)
    # This functions returns the right-cuts of intervals
    df = df.copy()
    df[output_interval[0]], df[output_interval[1]] = (
        df[[left_interval[0], right_interval[1]]].max(axis=1, skipna=False),
        df[left_interval[1]],
    )

    df = df.drop(
        columns=(
            set(left_interval)
            .union(set(right_interval))
            .difference(set(output_interval))
        )
    )

    is_valid = df[output_interval[0]] < df[output_interval[1]]
    df = df[is_valid]
    return df


def _merge_interval_interval_inner(
    left: pd.DataFrame,
    right: pd.DataFrame,
    on: list[str],
    interval: tuple[str, str],
    suffixes: tuple[str, str],
) -> pd.DataFrame:
    # to make an inner interval join, we use the method described in the documentation
    # 'docs/interval-operations.md'.

    # do a regular inner join on the specified on-keys
    overlapping = merging.merge(
        left,
        right,
        on=on,
        how="inner",
        suffixes=suffixes,
    )

    left_interval = (interval[0] + suffixes[0], interval[1] + suffixes[0])
    right_interval = (interval[0] + suffixes[1], interval[1] + suffixes[1])

    # get the intersection of the left and right intervals
    # (disjoint intervals will be filtered out here)
    merged = _intersect_intervals(
        overlapping,
        left_interval,
        right_interval,
        output_interval=interval,
    )

    return merged


def _merge_interval_interval_left(
    left: pd.DataFrame,
    right: pd.DataFrame,
    on: list[str],
    interval: tuple[str, str],
    suffixes: tuple[str, str],
) -> pd.DataFrame:
    # to make a left interval join, we use the method described in the documentation
    # 'docs/interval-operations.md'.

    # do an inner interval join first
    overlapping = _merge_interval_interval_inner(
        left,
        right,
        on,
        interval,
        suffixes,
    )

    # get the bounding interval (i.e. min/max range) of all intervals in right
    # DataFrame per on-key combination
    match_stats = overlapping.groupby(on, dropna=False, as_index=False).agg(
        {interval[0]: "min", interval[1]: "max"}
    )
    stats_interval = (
        "_min_overlap_start",
        "_max_overlap_end",
    )
    match_stats = match_stats.rename(
        columns={
            interval[0]: stats_interval[0],
            interval[1]: stats_interval[1],
        }
    )

    left = merging.merge(left, match_stats, on=on, how="left")

    # get all left rows that has no match on the keys at all
    has_no_match = left[left[stats_interval[0]].isna() & left[stats_interval[1]].isna()]
    has_no_match = has_no_match.drop(columns=list(stats_interval))

    # get all left rows that has at least one match on the keys,
    # but with part of the interval not overlapping with any right intervals
    has_match_left_cuts = _left_cut_intervals(
        left,
        interval,
        stats_interval,
        output_interval=interval,
    )

    has_match_right_cuts = _right_cut_intervals(
        left,
        interval,
        stats_interval,
        output_interval=interval,
    )

    # combine all rows to be added to the inner-join result
    rest = basic.concat_by_name(
        [has_match_left_cuts, has_match_right_cuts, has_no_match]
    )

    # rename potentially colliding column names using the given suffixes
    for col in rest.columns:
        if col not in overlapping.columns:
            rest = rest.rename(columns={col: col + suffixes[0]})

    # append the inner-join result with the rest
    return basic.concat_by_name([overlapping, rest])


def merge_interval_interval(
    left: pd.DataFrame,
    right: pd.DataFrame,
    interval: tuple[str, str],
    on: list[str] | None = None,
    how: Literal["left", "inner"] = "inner",
    suffixes: tuple[str, str] = ("_x", "_y"),
):
    """Merge DataFrame as an interval-interval-join.
    This matches all rows where the specified on-keys are equal
    AND where the interval in the left and right DataFrame overlaps
    (for a left-join, it will also include all interval parts in the left DataFrame,
    for which there is no match in the right DataFrame).
    More formally, it will produce the same result, as if exploding the intervals
    of the left and right DataFrame, and then doing a regular join.

    Args:
        left (pd.DataFrame): Left DataFrame
        right (pd.DataFrame): Right DataFrame
        on (list[str]): Columns to do an equi-join on.
        interval (tuple[str, str]): Interval columns (start, end).
        how (JoinType, optional): Join type ("inner", "left"). Defaults to "inner".
        suffixes (tuple[str, str], optional): Suffixes to use for colliding column
            names. Defaults to ("_x", "_y").
    """
    if how == "inner":
        return _merge_interval_interval_inner(
            left,
            right,
            on,
            interval,
            suffixes,
        )
    elif how == "left":
        return _merge_interval_interval_left(
            left,
            right,
            on,
            interval,
            suffixes,
        )
    else:
        raise NotImplementedError(
            f"Join type '{how}' not implemented for interval-interval-join."
        )


def merge_point_interval(
    left: pd.DataFrame,
    right: pd.DataFrame,
    left_point: str,
    right_interval: tuple[str, str],
    on: list[str] | None = None,
    how: Literal["left", "inner"] = "inner",
    suffixes: tuple[str, str] = ("_x", "_y"),
):
    """Merge DataFrame as a point-interval-join.
    This matches all rows where the specified on-keys are equal
    AND where the point in the left DataFrame is contained within
    the interval in the right DataFrame.
    More formally, it will produce the same result, as if exploding the intervals
    of the right DataFrame, and then doing a regular join.

    Args:
        left (pd.DataFrame): Left DataFrame
        right (pd.DataFrame): Right DataFrame
        on (list[str]): Columns to do an equi-join on.
        left_point (str): Point column in left DataFrame.
        right_interval (tuple[str, str]): Interval columns (start, end) in right
            DataFrame.
        how (JoinType, optional): Join type ("inner", "left"). Defaults to "inner".
        suffixes (tuple[str, str], optional): Suffixes to use for colliding column
            names. Defaults to ("_x", "_y").
    """
    # we make a point-interval join by converting the point to a tiny interval
    # and doing an interval-interval join.
    left = left.copy()
    left[right_interval[0]] = left[left_point]
    left[right_interval[1]] = left[left_point] + quantities.get_epsilon(
        left[left_point].dtype
    )

    df = merge_interval_interval(
        left,
        right,
        right_interval,
        on,
        how=how,
        suffixes=suffixes,
    )

    df = df.drop(columns=list(right_interval))

    return df
