"""Pairwise overlap constants."""

__all__ = [
    'RESERVED_COLS',
    'AUTOGEN_COLS',
    'DEFAULT_CHUNK_SIZE',
    'EXPR_OVERLAP_RO',
    'EXPR_SZRO',
    'EXPR_OFFSET_DIST',
    'EXPR_OFFSET_PROP',
    'JOIN_COL_EXPR',
    'INVARIANT_JOIN_COLS',
    'DEFAULT_JOIN_COLS',
]

from collections.abc import Callable
from typing import Any

import polars as pl


RESERVED_COLS: frozenset[str] = frozenset({
    '_index',
    '_end_ro',
})
"""Reserved columns are added automatically to input tables."""

AUTOGEN_COLS: frozenset[str] = frozenset({
    'varlen',
    'id',
})
"""Columns automatically generated by the join process for input tables if they are not present."""

DEFAULT_CHUNK_SIZE: int = 1_500
"""
Default size to chunk tables before joining. A value of 10,000 works well to balance combinatorial
explosion without exploding the number of chunks to merge when overlapping large variant tables.
"""

EXPR_OVERLAP_RO = (
    (
        (
            pl.min_horizontal('_end_ro_a', '_end_ro_b')
            - pl.max_horizontal('pos_a', 'pos_b')
        ) / (
            pl.max_horizontal('varlen_a', 'varlen_b')
        )
    )
    .clip(0.0, 1.0)
    .cast(pl.Float32)
)
"""Expression for overlap reciprocal overlap (RO)."""

EXPR_SZRO = (
    (
        (
            pl.min_horizontal('varlen_a', 'varlen_b')
        ) / (
            pl.max_horizontal('varlen_a', 'varlen_b')
        )
    )
    .cast(pl.Float32)
)
"""Expression for computing reciprocal size overlap (SZRO)."""

EXPR_OFFSET_DIST = (
    pl.max_horizontal(
        (pl.col('pos_a') - pl.col('pos_b')).abs(),
        (pl.col('end_a') - pl.col('end_b')).abs()
    )
    .cast(pl.Int32)
)
"""Expression for computing offset distance."""

EXPR_OFFSET_PROP = (
    (
        EXPR_OFFSET_DIST / pl.min_horizontal('varlen_a', 'varlen_b')
    )
    .cast(pl.Float32)
)
"""Expression for computing the offset proportion."""

JOIN_COL_EXPR: dict[str, pl.Expr | Callable[[Any], pl.Expr]] = {
    'index_a': pl.col('_index_a'),
    'index_b': pl.col('_index_b'),
    'id_a': pl.col('id_a'),
    'id_b': pl.col('id_b'),
    'ro': EXPR_OVERLAP_RO,
    'offset_dist': EXPR_OFFSET_DIST,
    'offset_prop': EXPR_OFFSET_PROP,
    'size_ro': EXPR_SZRO,
}
"""Known join columns."""

INVARIANT_JOIN_COLS: list[str] = ['index_a', 'index_b']
"""Columns that appear in all joins regardless of parameters."""

DEFAULT_JOIN_COLS: list[str] = [
    col for col in list(JOIN_COL_EXPR.keys()) + ['match_prop', 'weight']
    if col not in INVARIANT_JOIN_COLS
]
"""List of the default join columns."""
