"""Поиск преамбулы и восстановление дробного момента выборки из IQ."""

from __future__ import annotations

from dataclasses import dataclass
import math

import numpy as np
from scipy.signal import fftconvolve, find_peaks

from protocol.video003_modem import PREAMBLE_SYMBOLS, bpsk_symbols, known_preamble_bits


@dataclass(frozen=True)
class PreambleCandidate:
    start_sample: float
    score: float


def search_preamble_grid(
    matched: np.ndarray,
    samples_per_symbol: int,
    minimum_score: float = 0.5,
) -> tuple[PreambleCandidate, ...]:
    """Ищет пилот во всём сигнале, не получая границ кадров от передатчика."""
    if matched.ndim != 1 or samples_per_symbol < 2 or not 0 < minimum_score < 1:
        raise ValueError("неверные параметры поиска преамбулы")
    known = bpsk_symbols(known_preamble_bits())
    all_peaks: list[PreambleCandidate] = []
    # Две тактовые фазы дают грубый поиск; дробный пик уточняется отдельно.
    phase_step = max(1, samples_per_symbol // 2)
    for phase in range(0, samples_per_symbol, phase_step):
        series = matched[phase::samples_per_symbol]
        if series.size < PREAMBLE_SYMBOLS:
            continue
        correlation = fftconvolve(series, known[::-1].conj(), mode="valid")
        power = np.abs(series) ** 2
        running = np.concatenate(([0.0], np.cumsum(power, dtype=np.float64)))
        window_power = running[PREAMBLE_SYMBOLS:] - running[:-PREAMBLE_SYMBOLS]
        score = np.abs(correlation) / np.sqrt(
            np.maximum(window_power * PREAMBLE_SYMBOLS, 1e-18)
        )
        peaks, _ = find_peaks(
            score, height=minimum_score, distance=PREAMBLE_SYMBOLS // 2
        )
        for index in peaks:
            all_peaks.append(PreambleCandidate(float(phase + index * samples_per_symbol),
                                               float(score[index])))
    all_peaks.sort(key=lambda item: item.start_sample)
    merged: list[PreambleCandidate] = []
    for candidate in all_peaks:
        if merged and candidate.start_sample - merged[-1].start_sample < 2 * samples_per_symbol:
            if candidate.score > merged[-1].score:
                merged[-1] = candidate
        else:
            merged.append(candidate)
    return tuple(merged)


def refine_timing(
    matched: np.ndarray,
    candidate: PreambleCandidate,
    samples_per_symbol: int,
) -> PreambleCandidate:
    """Уточняет пик корреляции до доли отсчёта по известной преамбуле."""
    known = bpsk_symbols(known_preamble_bits())
    center = int(candidate.start_sample)
    half_phase_step = max(2, samples_per_symbol // 4)
    if (center < half_phase_step + 1
            or center + (PREAMBLE_SYMBOLS - 1) * samples_per_symbol
            + half_phase_step + 1 >= matched.size):
        return candidate
    indexes = center + np.arange(PREAMBLE_SYMBOLS) * samples_per_symbol
    shifts = range(-half_phase_step - 1, half_phase_step + 2)
    scores: list[float] = [float(abs(np.vdot(known, matched[indexes + shift])))
                           for shift in shifts]
    best = int(np.argmax(scores))
    if best in (0, len(scores) - 1):
        fractional = 0.0
    else:
        curvature = scores[best - 1] - 2 * scores[best] + scores[best + 1]
        fractional = (0.0 if curvature >= -1e-12 else
                      0.5 * (scores[best - 1] - scores[best + 1]) / curvature)
    fractional = float(np.clip(fractional, -0.5, 0.5))
    return PreambleCandidate(candidate.start_sample + shifts[best] + fractional,
                             candidate.score)


def sample_symbols(
    matched: np.ndarray,
    start_sample: float,
    samples_per_symbol: int,
    maximum_symbols: int,
) -> np.ndarray:
    """Интерполирует найденные такты без входного oracle timing."""
    if not math.isfinite(start_sample) or start_sample < 0 or maximum_symbols < 1:
        raise ValueError("неверные границы выборки символов")
    available = int((matched.size - 2 - start_sample) // samples_per_symbol) + 1
    count = min(maximum_symbols, max(0, available))
    if count == 0:
        return np.empty(0, dtype=np.complex64)
    positions = start_sample + np.arange(count, dtype=np.float64) * samples_per_symbol
    lower = np.floor(positions).astype(np.int64)
    fraction = positions - lower
    return (matched[lower] * (1 - fraction)
            + matched[lower + 1] * fraction).astype(np.complex64)