"""Lab050: независимый приём IQ-блоков с ограниченной памятью, без RF.

Использует корреляцию и строгий PHY/FEC/CRC Lab048. Принятые пакеты сразу
возвращаются вызывающему коду; история IQ/пакетов внутри не накапливается.
Адаптер констант старых лабораторий рассчитан на последовательный вызов.
"""

from __future__ import annotations

from collections import Counter
from contextlib import contextmanager
from dataclasses import dataclass
import math
import time

import numpy as np

from experiments import lab046c_offline_recovery as offline
from experiments import lab049_bpsk_symbol_rate as phy
from experiments.lab049_independent_iq import receiver_mode


@contextmanager
def streaming_mode(mode: phy.PhyMode):
    """Оставить не менее двух выборок/символ в грубом поиске."""
    module = phy.lab045.lab044.lab043
    original = module.DETECTION_DECIMATION
    with receiver_mode(mode):
        module.DETECTION_DECIMATION = min(20, mode.samples_per_symbol // 2)
        try:
            yield
        finally:
            module.DETECTION_DECIMATION = original


@dataclass(frozen=True)
class ReceivedPacket:
    start_sample: int
    end_sample: int
    emitted_after_sample: int
    payload: bytes


class StreamingReceiver:
    """push_iq_block выдаёт пакеты до конца потока; finish только опустошает хвост.

    Размер кадра для удержания IQ определяется максимальной допустимой PHY
    length (1024 байта), а не TX-границами/числом ожидаемых пакетов.
    Порог Lab047 остаётся фиксированным. Источник и размеры объекта неизвестны.
    """

    def __init__(self, symbol_rate: int, *, block_samples: int = 65_536):
        self.mode = phy.PhyMode(symbol_rate)
        if self.mode.samples_per_symbol < 4:
            raise ValueError("Для этого адаптера нужно не менее четырёх выборок/символ")
        if type(block_samples) is not int or block_samples < 1024:
            raise ValueError("Размер блока должен быть целым и не меньше 1024")
        self.block_samples = block_samples
        self.decimation = min(20, self.mode.samples_per_symbol // 2)
        self.core_samples = math.ceil(block_samples / self.decimation) * self.decimation
        margin = offline.REFINEMENT_MARGIN_SAMPLES
        self.margin = math.ceil(margin / self.decimation) * self.decimation
        maximum_symbols = (phy.lab045.PREAMBLE_BIT_COUNT + phy.lab045.PHY_HEADER_BIT_COUNT
                           + phy.lab045.conv.encoded_bit_count(phy.lab045.MAX_PACKET_BYTES * 8)
                           + 2 * phy.lab045.GUARD_SYMBOL_COUNT + phy.lab045.RRC_SPAN_SYMBOLS)
        self.lookahead = maximum_symbols * self.mode.samples_per_symbol + self.margin
        self.buffer_limit_samples = self.core_samples + self.lookahead + self.margin + block_samples
        self.buffer = np.empty(0, np.complex64)
        self.base_sample = self.scan_sample = self.total_samples = self.accepted_until = 0
        self.peak_buffer_samples = self.peak_candidate_count = 0
        self.packet_count = self.candidate_count = self.blocks_processed = 0
        self.processing_seconds = self.coarse_seconds = self.strict_seconds = 0.0
        self.score_count = 0
        self.score_sum = 0.0
        self.rejected_reasons: Counter = Counter()
        self.closed = False
        with streaming_mode(self.mode):
            self.reference = offline.lab046b._decimate_then_downconvert(
                phy.lab045.lab044._preamble_reference_samples()).copy()

    @property
    def memory_estimate_bytes(self) -> int:
        """Консервативная оценка массивов, не измеренный RSS Python/SciPy."""
        window_bytes = self.buffer_limit_samples * np.dtype(np.complex64).itemsize
        # Два буфера при копировании + прежний коэффициент временных массивов
        # Lab048 + запас для строгого matched filter и индексов кандидатов.
        return 2 * window_bytes + 12 * window_bytes + 8 * window_bytes + self.reference.nbytes

    def _scan(self, owned_end: int) -> list[ReceivedPacket]:
        started = time.perf_counter()
        scores = offline._correlation_scores(self.buffer, self.reference)
        self.coarse_seconds += time.perf_counter() - started
        indexes = offline._all_local_peak_indexes(scores)
        offsets = self.base_sample + indexes * self.decimation
        indexes = indexes[(offsets >= self.scan_sample) & (offsets < owned_end)]
        self.score_count += len(indexes)
        self.score_sum += float(np.sum(scores[indexes], dtype=np.float64))
        indexes = indexes[scores[indexes] >= np.float32(phy.lab047.FIXED_HARDWARE_CANDIDATE_THRESHOLD)]
        self.peak_candidate_count = max(self.peak_candidate_count, len(indexes))
        emitted = []
        for index in indexes:
            absolute = self.base_sample + int(index) * self.decimation
            if absolute < self.accepted_until:
                continue
            self.candidate_count += 1
            started = time.perf_counter()
            outcome = offline.decode_candidate_from_memmap(
                self.buffer.view(np.float32), self.buffer.size,
                offline.CoarseCandidate(int(index) * self.decimation, float(scores[index])),
                self.candidate_count,
                read_iq_window=lambda raw, start, end: raw[2*start:2*end].view(np.complex64),
            )
            self.strict_seconds += time.perf_counter() - started
            if outcome.accepted and outcome.crc == phy.lab045.CRC_PASS and outcome.frame is not None:
                end = self.base_sample + outcome.next_sample
                self.accepted_until = end
                self.packet_count += 1
                emitted.append(ReceivedPacket(
                    self.base_sample + outcome.refined_start_sample, end, self.total_samples,
                    bytes(outcome.frame.reception.payload_bytes),
                ))
            else:
                self.rejected_reasons[outcome.reason] += 1
        self.blocks_processed += 1
        self.scan_sample = owned_end
        retain_start = max(0, self.scan_sample - self.margin)
        trim = retain_start - self.base_sample
        if trim > 0:
            self.buffer = self.buffer[trim:].copy()
            self.base_sample = retain_start
        return emitted

    def push_iq_block(self, samples: np.ndarray) -> list[ReceivedPacket]:
        if self.closed:
            raise RuntimeError("Поток уже закрыт")
        samples = np.asarray(samples, dtype=np.complex64)
        if samples.ndim != 1 or samples.size > self.block_samples or not np.all(np.isfinite(samples)):
            raise ValueError("Нужен конечный одномерный IQ-блок не больше block_samples")
        started = time.perf_counter()
        self.total_samples += samples.size
        self.buffer = np.concatenate((self.buffer, samples))
        self.peak_buffer_samples = max(self.peak_buffer_samples, self.buffer.size)
        if self.buffer.size > self.buffer_limit_samples:
            raise RuntimeError("Превышен предел буфера IQ")
        emitted = []
        with streaming_mode(self.mode):
            while self.total_samples >= self.scan_sample + self.core_samples + self.lookahead:
                emitted.extend(self._scan(self.scan_sample + self.core_samples))
        self.processing_seconds += time.perf_counter() - started
        return emitted

    def finish(self) -> list[ReceivedPacket]:
        if self.closed:
            raise RuntimeError("Поток уже закрыт")
        started = time.perf_counter()
        emitted = []
        with streaming_mode(self.mode):
            if self.scan_sample < self.total_samples:
                emitted = self._scan(self.total_samples)
        self.buffer = np.empty(0, np.complex64)
        self.closed = True
        self.processing_seconds += time.perf_counter() - started
        return emitted