"""Оценка несущей по преамбуле и слежение за фазой QPSK-пакета."""

from __future__ import annotations

from dataclasses import dataclass
import math

import numpy as np

from protocol.video003_modem import (
    PHY_HEADER_SYMBOLS,
    PREAMBLE_SYMBOLS,
    _estimate_phase_frequency,
    bpsk_symbols,
    known_preamble_bits,
)


BODY_START = PREAMBLE_SYMBOLS + PHY_HEADER_SYMBOLS
TRACK_BLOCK_SYMBOLS = 64


@dataclass(frozen=True)
class PilotEstimate:
    phase_rad: float
    frequency_hz: float
    phase_step_rad: float
    quadrant: int


def estimate_from_preamble(symbols: np.ndarray, symbol_rate: int) -> PilotEstimate:
    """Оценивает частоту и абсолютную фазу, включая поворот на 90°/180°."""
    if symbols.ndim != 1 or len(symbols) < PREAMBLE_SYMBOLS:
        raise ValueError("недостаточно символов известной преамбулы")
    known = bpsk_symbols(known_preamble_bits())
    phase, phase_step = _estimate_phase_frequency(symbols[:PREAMBLE_SYMBOLS], known)
    return PilotEstimate(
        phase,
        phase_step * symbol_rate / (2 * math.pi),
        phase_step,
        round(phase / (math.pi / 2)) % 4,
    )


def correct_pilot_estimate(symbols: np.ndarray, estimate: PilotEstimate) -> np.ndarray:
    """Убирает оценённые фазу и CFO из всего найденного пакета."""
    indexes = np.arange(len(symbols), dtype=np.float64)
    correction = np.exp(-1j * (estimate.phase_rad + estimate.phase_step_rad * indexes))
    return symbols * correction.astype(np.complex64)


def track_qpsk_body(symbols: np.ndarray) -> tuple[np.ndarray, float]:
    """Причинное слежение: решение строится из текущего QPSK, без эталонных битов."""
    corrected = symbols.copy()
    phase_state = 0.0
    errors: list[float] = []
    for start in range(BODY_START, len(symbols), TRACK_BLOCK_SYMBOLS):
        stop = min(start + TRACK_BLOCK_SYMBOLS, len(symbols))
        block = symbols[start:stop] * np.exp(-1j * phase_state)
        decisions = (np.where(block.real >= 0, 1.0, -1.0)
                     + 1j * np.where(block.imag >= 0, 1.0, -1.0)).astype(np.complex64)
        error = float(np.angle(np.sum(block * np.conj(decisions))))
        error = float(np.clip(error, -0.2, 0.2))
        phase_state += 0.35 * error
        corrected[start:stop] = symbols[start:stop] * np.exp(-1j * phase_state)
        errors.append(error)
    root_mean_square = math.sqrt(sum(error * error for error in errors) / len(errors)) if errors else 0.0
    return corrected, root_mean_square