import os

import numpy as np
import zmq

from common.runtime import load_root_env
from common.shared_stream_addrs import SHARED_868_ADDR, SHARED_915_ADDR
from core.multichannelswitcher import MultiChannel
from core.sig_n_medi_collect import Signal


load_root_env(__file__)


class Scheduler868:
    def __init__(self):
        self.signal_length = int(os.getenv('signal_length_868'))
        self.multi_channel = MultiChannel(
            [*map(float, os.getenv('f_step_868').split())],
            [*map(float, os.getenv('f_bases_868').split())],
            [*map(float, os.getenv('f_roofs_868').split())],
        )
        self.base_freq = float(self.multi_channel.init_f())
        self.signal = Signal()

    def get_current_freq(self):
        return float(self.multi_channel.get_cur_channel())

    def process(self, lvl):
        current_before = self.get_current_freq()
        metric = self.signal.fill_signal(lvl, self.signal_length)
        if metric == 0:
            return current_before, False

        next_freq = float(self.multi_channel.change_channel())
        self.signal.clear()
        lane_complete = next_freq == self.base_freq and current_before != self.base_freq
        return next_freq, lane_complete


class Scheduler915:
    def __init__(self):
        self.porog = float(os.getenv('POROG_915'))
        self.point_amount = 100_000
        self.split_size = 400_000
        self.show_amount = int(0.8 * self.point_amount)
        self.f_base = 0.91e9
        self.f_step = 20e6
        self.f_roof = 0.98e9
        self.f = self.f_base
        self.channel = 1
        self.flag = 0
        self.signal_arr = np.array([], dtype=np.complex64)

    def get_current_freq(self):
        return float(self.f)

    def _median(self, sig):
        samples = np.asarray(np.abs(np.array(sig, dtype=np.complex64)), dtype=np.float32)
        sorted_samples = sorted(samples)
        median = abs(float(np.median(sorted_samples[self.show_amount:])))
        self.flag = 0 if self.porog > median else 1

    def _advance(self):
        next_freq = self.f + self.f_step
        if next_freq >= self.f_roof:
            self.f = self.f_base
            self.channel = 1
            return float(self.f), True

        self.f = next_freq
        self.channel += 1
        return float(self.f), False

    def process(self, lvl):
        y = np.asarray(lvl, dtype=np.complex64).ravel()
        self.signal_arr = np.concatenate((self.signal_arr, y), axis=None)

        if self.flag == 0 and len(self.signal_arr) >= self.point_amount:
            self._median(self.signal_arr[:self.point_amount])
            self.signal_arr = np.array([], dtype=np.complex64)
            if self.flag == 0:
                return self._advance()

        if len(self.signal_arr) >= self.split_size:
            self.flag = 0
            self.signal_arr = np.array([], dtype=np.complex64)
            return self._advance()

        return float(self.f), False


class SharedRouter868915:
    def __init__(self):
        self.frontends = {
            '868': {
                'sample_rate': 20e6,
                'gain': 24,
                'if_gain': 24,
                'bb_gain': 100,
                'bandwidth': 0,
            },
            '915': {
                'sample_rate': 20e6,
                'gain': 16,
                'if_gain': 16,
                'bb_gain': 0,
                'bandwidth': 0,
            },
        }
        self.schedulers = {
            '868': Scheduler868(),
            '915': Scheduler915(),
        }
        self.context = zmq.Context.instance()
        self.sockets = {
            '868': self._make_socket(SHARED_868_ADDR),
            '915': self._make_socket(SHARED_915_ADDR),
        }
        self.active_name = '868'
        self.switch_counter = 0

    def _make_socket(self, addr):
        sock = self.context.socket(zmq.PUSH)
        sock.setsockopt(zmq.SNDHWM, 64)
        sock.setsockopt(zmq.LINGER, 0)
        sock.setsockopt(zmq.IMMEDIATE, 1)
        sock.bind(addr)
        return sock

    def close(self):
        for sock in self.sockets.values():
            try:
                sock.close(0)
            except Exception:
                pass

    def get_active_name(self):
        return self.active_name

    def get_active_freq(self):
        return self.schedulers[self.active_name].get_current_freq()

    def get_active_frontend(self):
        return dict(self.frontends[self.active_name])

    def get_start_freq(self):
        return self.get_active_freq()

    def route_vector(self, lvl):
        vector = np.asarray(lvl, dtype=np.complex64).ravel()
        active_name = self.active_name

        try:
            self.sockets[active_name].send(vector.tobytes(), zmq.NOBLOCK)
        except zmq.Again:
            pass

        next_freq, lane_complete = self.schedulers[active_name].process(vector)
        lane_switched = False

        if lane_complete:
            previous = self.active_name
            self.active_name = '915' if self.active_name == '868' else '868'
            next_freq = self.get_active_freq()
            self.switch_counter += 1
            lane_switched = True
            print(
                f'[shared-router-868-915] switch#{self.switch_counter}: '
                f'{previous} -> {self.active_name}, tune={next_freq}'
            )

        return float(next_freq), lane_switched