"""The sole interval overlap projection kernel used by both coordinate modes."""
from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from scipy import sparse


@dataclass
class Projection:
    x: np.ndarray
    observed: np.ndarray
    coverage: np.ndarray
    source_count: np.ndarray
    first_source: np.ndarray
    last_source: np.ndarray
    source_weights: sparse.csr_matrix
    observed_dimensions: np.ndarray
    coverage_dimensions: np.ndarray
    quality_mean: np.ndarray
    quality_available_fraction: np.ndarray


def project_intervals(
    values: np.ndarray,
    source_intervals: np.ndarray,
    target_intervals: np.ndarray,
    observed: np.ndarray,
    quality: np.ndarray | None = None,
    quality_available: np.ndarray | None = None,
) -> Projection:
    """Project source cells using overlap * quality * observed validity.

    Row provenance uses any valid dimension. Feature values use validity per
    dimension, so partially observed physical features remain partially missing.
    """
    source = np.asarray(values, dtype=np.float32)
    src = np.asarray(source_intervals, dtype=np.float64)
    dst = np.asarray(target_intervals, dtype=np.float64)
    if source.ndim != 2 or src.shape != (len(source), 2) or dst.ndim != 2 or dst.shape[1] != 2:
        raise ValueError("inconsistent source features or interval dimensions")
    if not np.isfinite(source).all() or not np.isfinite(src).all() or not np.isfinite(dst).all():
        raise ValueError("non-finite source features or intervals")
    if np.any(src[:, 1] < src[:, 0]) or np.any(dst[:, 1] <= dst[:, 0]):
        raise ValueError("source widths must be nonnegative and target widths positive")
    obs = np.asarray(observed, bool)
    if obs.shape == (len(source),):
        obs_dim = np.broadcast_to(obs[:, None], source.shape)
    elif obs.shape == source.shape:
        obs_dim = obs
        obs = obs.any(axis=1)
    else:
        raise ValueError("observed must have source-row or source-feature shape")
    q = np.ones(len(source), dtype=np.float64) if quality is None else np.asarray(quality, dtype=np.float64)
    if q.shape != (len(source),) or not np.isfinite(q).all() or np.any(q < 0):
        raise ValueError("quality must be finite and nonnegative per source row")
    overlap = np.maximum(0.0, np.minimum(dst[:, None, 1], src[None, :, 1])
                         - np.maximum(dst[:, None, 0], src[None, :, 0]))
    available = np.zeros(len(source), bool) if quality_available is None else np.asarray(quality_available, bool)
    if available.shape != (len(source),):
        raise ValueError("quality availability must be per source row")
    # Keep the same multiplication and accumulation order as the original
    # official-unaligned projection when quality is uniformly one.
    physical = overlap.copy()
    physical *= obs[None, :]
    row_weight = physical.copy()
    row_weight *= q[None, :]
    mass = row_weight.sum(axis=1)
    row_valid = mass > 0
    normalized = np.zeros_like(row_weight, dtype=np.float32)
    normalized[row_valid] = (row_weight[row_valid] / mass[row_valid, None]).astype(np.float32)
    support = row_weight > 0
    count = support.sum(axis=1).astype(np.uint16)
    first = np.full(len(dst), -1, dtype=np.int32)
    last = np.full(len(dst), -1, dtype=np.int32)
    if row_valid.any():
        first[row_valid] = support[row_valid].argmax(axis=1)
        last[row_valid] = len(source) - 1 - support[row_valid, ::-1].argmax(axis=1)
    width = dst[:, 1] - dst[:, 0]
    physical_mass = physical.sum(axis=1)
    coverage = np.clip(physical_mass / width, 0.0, 1.0).astype(np.float32)
    qmean = np.ones(len(dst), np.float32)
    qavailable = np.zeros(len(dst), np.float32)
    physical_valid = physical_mass > 0
    qmean[physical_valid] = (mass[physical_valid] / physical_mass[physical_valid]).astype(np.float32)
    qavailable[physical_valid] = ((physical[physical_valid] @ available.astype(np.float64))
                                  / physical_mass[physical_valid]).astype(np.float32)

    # The common full-dimension case follows the original matrix product
    # exactly; this is also much faster for 500 x 768 input.
    if np.array_equal(obs_dim, np.broadcast_to(obs[:, None], source.shape)):
        x = np.zeros((len(dst), source.shape[1]), np.float32)
        x[row_valid] = ((row_weight[row_valid] @ source) / mass[row_valid, None]).astype(np.float32)
        observed_dimensions = np.broadcast_to(row_valid[:, None], x.shape).copy()
        coverage_dimensions = np.broadcast_to(coverage[:, None], x.shape).copy()
    else:
        dim_physical = overlap[:, :, None] * obs_dim[None, :, :]
        dim_weight = dim_physical * q[None, :, None]
        dim_mass = dim_weight.sum(axis=1)
        observed_dimensions = dim_mass > 0
        x = np.zeros((len(dst), source.shape[1]), np.float32)
        numerator = np.einsum("ksd,sd->kd", dim_weight, source, optimize=True)
        x[observed_dimensions] = (numerator[observed_dimensions] / dim_mass[observed_dimensions]).astype(np.float32)
        coverage_dimensions = np.clip(dim_physical.sum(axis=1) / width[:, None], 0, 1).astype(np.float32)
    return Projection(x, row_valid, coverage, count, first, last,
                      sparse.csr_matrix(normalized), observed_dimensions, coverage_dimensions,
                      qmean, qavailable)