"""Merge operations for collapsing overlapping regions in one table."""

__all__ = [
    'merge_depth',
]

from typing import Optional, Iterable

import polars as pl

from .col import CoordCol, get_coord_cols, make_unique_col


def merge_depth(
        df: pl.LazyFrame | pl.DataFrame,
        distance: int = 0,
        col_names: Optional[CoordCol | Iterable[str] | str] = None,
        col_max_depth: Optional[str] = 'max_depth'
):
    """Merge overlapping regions in a table.

    :param df: Table of coordinates.
    :param distance: Merge records within this distance.
    :param col_names: Names of columns (defaults to chromosome columns).
    :param col_max_depth: Column name for maximum depth. Must not be in `col_names`.

    :return: A table of depths.
    """
    col_names = get_coord_cols(col_names)

    if col_max_depth in col_names:
        raise ValueError(f'Column name {col_max_depth} is in column names: {col_names!r}')

    col_loc = make_unique_col('_loc', col_names)
    col_depth = make_unique_col('_depth', col_names)
    col_rle_id = make_unique_col('_rle_id', col_names)

    if isinstance(df, pl.DataFrame):
        df = df.lazy()

    return (
        melt_depth(
            df=df.with_columns(pl.col(col_names.end) + distance),
            col_names=col_names,
            col_loc=col_loc,
            col_depth=col_depth,
        )
        .with_columns(
            (pl.col(col_depth) > 0).rle_id().over(pl.col(col_names.chrom)).alias(col_rle_id),
        )
        .filter(pl.col(col_depth) > 0)
        .group_by([pl.col(col_names.chrom), col_rle_id])
        .agg(
            pl.col(col_loc).min().alias(col_names.pos),
            pl.col(col_loc).max().alias(col_names.end) - distance,
            pl.col(col_depth).max().alias(col_max_depth) if col_max_depth else None,
        )
        .select(list(col_names) + ([col_max_depth] if col_max_depth else []))
        .sort(list(col_names))
    )


def melt_depth(
        df: pl.LazyFrame,
        col_names: CoordCol,
        col_loc: str = '_loc',
        col_depth: str = '_depth',
):
    """Get a table of depths by locations of changes.

    The table will have one row for each location of change with end positions having two entries (one for the
    previous depth, and one for the new depth).

    This is typically used as an intermediate step for depth or merging operations. The default names for the location
    and depth columns have a leading underscore suggesting this is an intermediate table. The caller should ensure
    these columns do not collide with names in `col_names`.

    For example:

    ===== ==== ======
    chrom _loc _depth
    ===== ==== ======
    chr1  100  1
    chr1  120  2
    chr1  160  2
    chr1  160  1
    chr1  200  1
    chr1  200  0
    ===== ==== ======

    This would be generated by one of two input tables:

    ===== ==== ======
    chrom pos  end
    ===== ==== ======
    chr1  100  160
    chr1  120  200
    ===== ==== ======

    or

    ===== ==== ======
    chrom pos  end
    ===== ==== ======
    chr1  100  200
    chr1  120  160
    ===== ==== ======

    Note that adjoining records will be represented with rows that go up and down:

    ===== ==== ======
    chrom _loc _depth
    ===== ==== ======
    chr1  100  1
    chr1  200  2
    chr1  200  2
    chr1  200  1
    chr1  300  1
    chr1  300  0
    ===== ==== ======

    Would result from:

    ===== ==== ======
    chrom pos  end
    ===== ==== ======
    chr1  100  200
    chr1  200  300
    ===== ==== ======

    This could be used to decide to collapse neighboring records or not.

    :param df: Table to melt to depth change locations.
    :param col_names: Names of chrom, pos, end columns to melt.
    :param col_loc: Name of the location column.
    :param col_depth: Name of the depth column to return.

    :return: A LazyFrame of depth change locations.
    """
    col_names = get_coord_cols(col_names)

    if isinstance(df, pl.DataFrame):
        df = df.lazy()

    assert not (col_loc in col_names or col_depth in col_names), (
        'Column names "%s" or "%s" are in column names "%s"' % (col_loc, col_depth, col_names)
    )

    assert col_loc != col_depth, 'Column names "%s" and "%s" are the same' % (col_loc, col_depth)

    return (
        pl.concat(
            [
                (  # Start
                    df
                    .select(
                        pl.col(col_names.chrom),
                        pl.col(col_names.pos).alias(col_loc),
                        pl.lit(1).alias(col_depth)
                    )
                ),
                (  # End position
                    df
                    .select(
                        pl.col(col_names.chrom),
                        pl.col(col_names.end).alias(col_loc),
                        pl.lit(0).alias(col_depth)
                    )
                ),
                (  # Adjust depth after end position
                    df
                    .select(
                        pl.col(col_names.chrom),
                        pl.col(col_names.end).alias(col_loc),
                        pl.lit(-1).alias(col_depth)
                    )
                ),
            ]
        )
        .sort([col_names.chrom, col_loc, col_depth], descending=[False, False, True])
        .with_columns(
            pl.col(col_depth).cum_sum().over(col_names.chrom).alias(col_depth)
        )
    )
