import logging
from datetime import datetime

import numpy as np
import numpy.typing as npt

from ..config import config

logger = logging.getLogger(__name__)


class AoA:
    def __init__(self) -> None:
        self.historical_autocorr = np.array([])
        self.N_subcarriers = -1
        self.N_rx = -1
        self.timestamp = datetime.now()
        pass

    def smooth(self, data: npt.NDArray[np.complex128]) -> npt.NDArray[np.complex128]:
        assert len(data.shape) == 3

        M = data.shape[0]  # Number of subcarriers
        N = data.shape[1]  # Number of RX antennas
        T = data.shape[2]  # Number of TX antennas

        self.N_subcarriers = M
        self.N_rx = N

        logger.debug(f"Smoothing: Subcarriers: {M}, RX antennas: {N}, TX antennas: {T}")

        # This only works with 1 TX antenna (i.e. no MIMO) - see #4 for more details
        assert T == 1, "The current implementation only supports 1 TX antenna"

        H_n = np.zeros((N, M // 2, M // 2 + 1), dtype=np.complex128)

        for i in range(N):
            for j in range(M // 2):
                H_n[i, j] = data[j : j + M // 2 + 1, i, 0]

        H_sm_rows = [np.hstack(H_n[i : i + N // 2 + 1]) for i in range(N // 2)]
        H_sm = np.vstack(H_sm_rows)

        logger.debug(f"Smoothed: {H_sm.shape}")

        return H_sm

    def update(self, data: npt.NDArray[np.complex128]) -> None:
        self.timestamp = datetime.now()
        H_sm = self.smooth(data)

        auto_corr = np.matmul(H_sm, np.conj(H_sm).T)
        # This matrix is by definition Hermitian.
        # Therefore, all of its eigenvectors are orthogonal.

        if self.historical_autocorr.size == 0:
            self.historical_autocorr = np.expand_dims(auto_corr, 0)
        else:
            self.historical_autocorr = np.append(
                self.historical_autocorr, np.expand_dims(auto_corr, 0), axis=0
            )

        WINDOW_SIZE = config.music.window_size
        if self.historical_autocorr.shape[0] > WINDOW_SIZE:
            self.historical_autocorr = self.historical_autocorr[-WINDOW_SIZE:]

        # Is the moving average also Hermitian?
        R = np.mean(self.historical_autocorr, axis=0)

        # The smallest eigenvectors span the noise subspace,
        # and the largest span the signal subspace.
        eigvals, eigvecs = np.linalg.eigh(R)
        self.E_n = eigvecs[:, np.abs(eigvals) < config.music.eigval_threshold]

    def steering_vector(
        self, theta: float, tof: float
    ) -> npt.NDArray[np.complexfloating]:
        omega_t: npt.NDArray[np.complex128] = np.exp(-2j * np.pi * config.delta_f * tof)
        phi_theta: npt.NDArray[np.complex128] = np.exp(
            2j
            * np.pi
            * config.central_freq_hz
            * config.antennas.spacing
            * (1 - np.cos(theta))
            / 299_792_458
        )

        omega_t = np.expand_dims(omega_t, axis=-1)
        phi_theta = np.expand_dims(phi_theta, axis=-1)

        antenna_v = omega_t ** np.arange(self.N_subcarriers // 2)
        phis = phi_theta ** np.arange(self.N_rx // 2)
        antenna_v = np.expand_dims(antenna_v, axis=-1)
        steering = antenna_v * phis
        return steering.T.reshape(-1)

    def evaluate(self, theta: float, tof: float) -> float:
        try:
            steering = self.steering_vector(theta, tof)
            steering_h = np.conj(steering).T
        except Exception as e:
            logger.exception(e)
            return 0
        E_n = self.E_n
        E_n_H = np.conj(E_n).T
        c = 1 / (0.001 + (steering_h @ E_n @ E_n_H @ steering))
        return np.abs(c.real)


def test_smoothing() -> None:
    row, col = np.indices((4, 2))
    data = row + 1j * col
    aoa = AoA()
    smoothed = aoa.smooth(data)
    H_0 = np.array([[0 + 0j, 0 + 1j, 0 + 2j], [0 + 1j, 0 + 2j, 0 + 3j]])
    H_01 = np.vstack([H_0, H_0 + 1])
    H_12 = np.vstack([H_0 + 1, H_0 + 2])
    expected = np.hstack([H_01, H_12])
    print(expected)
    assert np.allclose(smoothed, expected)


def test_steering_vector() -> None:
    aoa = AoA()
    aoa.N_subcarriers = 10
    aoa.N_rx = 2
    tau = 1
    theta = 0
    print(aoa.steering_vector(theta, tau))
    assert False