import numpy as np
from scipy import signal

from .myrandom import GaussianNoise


class ChannelEqualization(object):
    def __init__(self, data_gen, channel, N=0, input_delay=0, noise=GaussianNoise(), SNR=None):
        if input_delay > N:
            raise ValueError('input delay cannot be higher than the N')

        self._gen = data_gen
        self._channel = channel
        self._N = N
        self._input_delay = int(input_delay)
        self._iter = None
        self._noise = noise or 0
        self._SNR = SNR

    def init(self, *d):
        self._tx = self._gen(*d)
        self._rx = self._channel(self._tx)

        noise = self._noise
        if self._SNR is not None:
            noise = self._noise(std=np.sqrt(self._rx.var() / (2 * self._SNR)),
                                complex=np.iscomplex(self._rx).any())

        self._rx += noise(*d)
        self._iter = self._N

    def next(self):
        rx = self._rx[(self._iter - self._N):(self._iter + 1)][::-1]
        tx = self._tx[self._iter - self._input_delay]

        self._iter += 1

        return rx, tx

    def var(self, axis=-1):
        return self._tx.var(axis=axis)
