"""Строгий программный приём QPSK из непрерывного комплексного IQ."""

from __future__ import annotations

from dataclasses import dataclass
import math
import time

import numpy as np
from scipy.signal import oaconvolve

from protocol import punctured_fec
from protocol.video003_modem import (
    MAX_FULL_PACKET_BYTES,
    PHY_HEADER_SYMBOLS,
    PREAMBLE_SYMBOLS,
    SAMPLE_RATE_HZ,
    _decode_phy_header,
    receive_frame,
    samples_per_symbol,
)
from protocol.video005_carrier import (
    correct_pilot_estimate,
    estimate_from_preamble,
    track_qpsk_body,
)
from protocol.video005_iq_waveform import NOMINAL_CARRIER_HZ, SYMBOL_RATE, rrc_taps
from protocol.video005_timing import refine_timing, sample_symbols, search_preamble_grid


@dataclass(frozen=True)
class ReceivedIqPacket:
    packet: bytes
    payload: bytes
    start_sample: float
    timing_phase_symbols: float
    carrier_offset_hz: float
    initial_phase_rad: float
    quadrant: int
    preamble_score: float
    tracking_error_rms_rad: float


@dataclass(frozen=True)
class IqStreamResult:
    packets: tuple[ReceivedIqPacket, ...]
    candidate_count: int
    rejected_candidates: int
    stage_time_s: dict[str, float]


def receive_iq_stream(
    iq: np.ndarray,
    *,
    symbol_rate: int = SYMBOL_RATE,
    carrier_hz: float = NOMINAL_CARRIER_HZ,
    decoder_backend: str = "numba",
) -> IqStreamResult:
    """Получает только IQ и параметры радиорежима, без подсказок передатчика.

    Внешние метаданные сигнала и ожидаемое число пакетов не являются аргументами.
    Возвращаются только пакеты, прошедшие PHY/FEC/CRC.
    """
    samples_per_symbol_count = samples_per_symbol(symbol_rate)
    if not math.isfinite(carrier_hz) or abs(carrier_hz) >= SAMPLE_RATE_HZ / 2:
        raise ValueError("неверная номинальная несущая")
    if decoder_backend == "numba":
        from protocol.video004_viterbi import decode_soft as fec_decoder
    elif decoder_backend == "legacy":
        fec_decoder = punctured_fec.decode_soft
    else:
        raise ValueError("неизвестный выбранный декодер")
    samples = np.asarray(iq)
    if samples.ndim != 1 or not np.iscomplexobj(samples):
        raise ValueError("приёмнику нужен одномерный комплексный IQ")
    if not np.all(np.isfinite(samples)):
        raise ValueError("IQ содержит неконечные значения")
    stage: dict[str, float] = {}
    if samples.size < PREAMBLE_SYMBOLS * samples_per_symbol_count:
        return IqStreamResult((), 0, 0, stage)

    started = time.perf_counter()
    period = round(SAMPLE_RATE_HZ / abs(carrier_hz)) if carrier_hz else 0
    if carrier_hz and abs(abs(carrier_hz) * period - SAMPLE_RATE_HZ) < 1e-8:
        cycle = np.exp(-1j * np.sign(carrier_hz) * 2 * np.pi
                       * np.arange(period) / period).astype(np.complex64)
        oscillator = np.resize(cycle, samples.size)
        baseband = samples.astype(np.complex64, copy=False) * oscillator
    else:
        indexes = np.arange(samples.size, dtype=np.float64)
        oscillator = np.exp(-1j * (2 * np.pi * carrier_hz / SAMPLE_RATE_HZ) * indexes)
        baseband = samples * oscillator
    stage["digital_downconversion"] = time.perf_counter() - started

    started = time.perf_counter()
    matched = oaconvolve(baseband, rrc_taps(symbol_rate).astype(np.float32), mode="full")
    stage["rrc"] = time.perf_counter() - started

    started = time.perf_counter()
    candidates = search_preamble_grid(matched, samples_per_symbol_count)
    stage["preamble_search"] = time.perf_counter() - started
    maximum_symbols = (
        PREAMBLE_SYMBOLS + PHY_HEADER_SYMBOLS
        + math.ceil(punctured_fec.coded_bit_count(MAX_FULL_PACKET_BYTES * 8, "2/3") / 2)
    )
    stage.update(timing_recovery=0.0, carrier_estimation=0.0, phy_header=0.0,
                 carrier_tracking=0.0, strict_phy_fec_crc=0.0)
    accepted: list[ReceivedIqPacket] = []
    rejected = 0
    for candidate in candidates:
        started = time.perf_counter()
        refined = refine_timing(matched, candidate, samples_per_symbol_count)
        symbols = sample_symbols(matched, refined.start_sample,
                                 samples_per_symbol_count, maximum_symbols)
        stage["timing_recovery"] += time.perf_counter() - started
        if symbols.size < PREAMBLE_SYMBOLS + PHY_HEADER_SYMBOLS:
            rejected += 1
            continue
        started = time.perf_counter()
        estimate = estimate_from_preamble(symbols, symbol_rate)
        corrected = correct_pilot_estimate(symbols, estimate)
        stage["carrier_estimation"] += time.perf_counter() - started
        started = time.perf_counter()
        packet_bytes = _decode_phy_header(
            corrected[PREAMBLE_SYMBOLS:PREAMBLE_SYMBOLS + PHY_HEADER_SYMBOLS].real
        )
        stage["phy_header"] += time.perf_counter() - started
        if packet_bytes is None:
            rejected += 1
            continue
        frame_symbols = (
            PREAMBLE_SYMBOLS + PHY_HEADER_SYMBOLS
            + math.ceil(punctured_fec.coded_bit_count(packet_bytes * 8, "2/3") / 2)
        )
        corrected = corrected[:frame_symbols]
        started = time.perf_counter()
        tracked, error_rms = track_qpsk_body(corrected)
        stage["carrier_tracking"] += time.perf_counter() - started
        started = time.perf_counter()
        reception = receive_frame(
            tracked, "QPSK", "2/3", symbol_rate,
            maximum_start=0, fec_decoder=fec_decoder,
        )
        stage["strict_phy_fec_crc"] += time.perf_counter() - started
        if not reception.crc_pass or reception.packet is None or reception.payload is None:
            rejected += 1
            continue
        accepted.append(ReceivedIqPacket(
            reception.packet, reception.payload, refined.start_sample,
            (refined.start_sample % samples_per_symbol_count) / samples_per_symbol_count,
            estimate.frequency_hz, estimate.phase_rad, estimate.quadrant,
            refined.score, error_rms,
        ))
    return IqStreamResult(tuple(accepted), len(candidates), rejected, stage)