"""Strict timeline, delayed-pair, and phase-conditioned alignment."""
from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable

import numpy as np


@dataclass(frozen=True)
class Pairing:
    left: np.ndarray
    right: np.ndarray
    sample_ids: np.ndarray
    dropped: int


def exact_time_pair(left: np.ndarray, left_steps: np.ndarray, right: np.ndarray, right_steps: np.ndarray) -> Pairing:
    a_steps = np.asarray(left_steps, dtype=np.int64)
    b_steps = np.asarray(right_steps, dtype=np.int64)
    if not np.array_equal(a_steps, b_steps):
        raise ValueError("time-paired artifacts require exactly identical lattice_steps")
    a, b = np.asarray(left), np.asarray(right)
    if a.shape != b.shape or a.shape[0] != len(a_steps):
        raise ValueError("time-paired fields must have identical shape and timeline length")
    return Pairing(a, b, a_steps.copy(), 0)


def valid_delayed_pairs(fields: np.ndarray, observables: np.ndarray, lattice_steps: np.ndarray, *, delay_steps: int, segment_ids: Iterable[int] | None = None) -> Pairing:
    q, y = np.asarray(fields), np.asarray(observables)
    steps = np.asarray(lattice_steps)
    if steps.ndim != 1 or not np.issubdtype(steps.dtype, np.integer):
        raise ValueError("lattice_steps must be a one-dimensional integer array")
    if len(steps) == 0 or np.any(np.diff(steps) <= 0):
        raise ValueError("lattice_steps must be unique and strictly increasing")
    if q.shape[0] != y.shape[0] or q.shape[0] != len(steps):
        raise ValueError("fields, observables, and lattice_steps must have equal length")
    if delay_steps < 0:
        raise ValueError("delay_steps must be non-negative")
    segments = np.zeros(len(steps), dtype=np.int64) if segment_ids is None else np.asarray(tuple(segment_ids))
    if segments.shape != steps.shape:
        raise ValueError("segment_ids must match timeline")
    index = {int(step): i for i, step in enumerate(steps)}
    left_indices, right_indices = [], []
    for i, step in enumerate(steps):
        j = index.get(int(step + delay_steps))
        if j is not None and segments[j] == segments[i]:
            left_indices.append(i)
            right_indices.append(j)
    li = np.asarray(left_indices, dtype=np.int64)
    ri = np.asarray(right_indices, dtype=np.int64)
    return Pairing(q[li], y[ri], steps[li], len(steps) - len(li))


def phase_bin_means(fields: np.ndarray, phase: np.ndarray, *, bins: int, min_count: int = 3) -> tuple[np.ndarray, np.ndarray]:
    q = np.asarray(fields)
    phi = np.mod(np.asarray(phase, dtype=float), 2 * np.pi)
    if q.shape[0] != len(phi) or bins < 2 or min_count < 1:
        raise ValueError("invalid phase-bin inputs")
    edges = np.linspace(0.0, 2 * np.pi, bins + 1)
    labels = np.minimum(np.searchsorted(edges, phi, side="right") - 1, bins - 1)
    means, counts = [], []
    for index in range(bins):
        selection = q[labels == index]
        counts.append(len(selection))
        if len(selection) < min_count:
            raise ValueError(f"phase bin {index} has {len(selection)} samples, requires {min_count}")
        means.append(selection.mean(axis=0))
    return np.asarray(means), np.asarray(counts, dtype=np.int64)


def phase_conditioned_correction(controlled: np.ndarray, controlled_phase: np.ndarray, baseline: np.ndarray, baseline_phase: np.ndarray, *, bins: int = 24, min_count: int = 3) -> tuple[np.ndarray, dict[str, np.ndarray]]:
    ctl_mean, ctl_count = phase_bin_means(controlled, controlled_phase, bins=bins, min_count=min_count)
    blk_mean, blk_count = phase_bin_means(baseline, baseline_phase, bins=bins, min_count=min_count)
    centers = (np.arange(bins) + 0.5) * (2 * np.pi / bins)
    return ctl_mean - blk_mean, {"phase_centers": centers, "controlled_counts": ctl_count, "baseline_counts": blk_count}
