import logging
from datetime import datetime

import numpy as np
import numpy.typing as npt
import torch

from ..config import config

logger = logging.getLogger(__name__)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.set_default_device(device)


class AoA:
    def __init__(self) -> None:
        self.historical_autocorr = torch.tensor([], dtype=torch.complex64)
        self.N_subcarriers = -1
        self.N_rx = -1
        self.timestamp = datetime.now()
        pass

    def smooth(self, data: torch.Tensor) -> torch.Tensor:
        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 = torch.zeros((N, M // 2, M // 2 + 1), dtype=torch.complex64)

        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 = [torch.hstack(list(H_n[i : i + N // 2 + 1])) for i in range(N // 2)]
        H_sm = torch.vstack(H_sm_rows)

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

        return H_sm

    def update(self, data: torch.Tensor) -> None:
        self.timestamp = datetime.now()
        H_sm = self.smooth(data)
        logger.debug(f"Calculated smoothed CSI matrix: {H_sm.shape}")

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

        if len(self.historical_autocorr.shape) <= 1:
            self.historical_autocorr = torch.unsqueeze(auto_corr, 0)
        else:
            self.historical_autocorr = torch.cat(
                (self.historical_autocorr, torch.unsqueeze(auto_corr, 0))
            )

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

        logger.debug("Finished updating autocorrelation matrix")

    def steering_vector(self, theta: torch.Tensor, tof: torch.Tensor) -> torch.Tensor:
        assert theta.shape == tof.shape
        assert len(theta.shape) == 1
        N = theta.shape[0]

        omega_t: torch.Tensor = torch.exp(-2j * np.pi * config.delta_f * tof)
        phi_theta: torch.Tensor = torch.exp(
            2j
            * np.pi
            * config.central_freq_hz
            * config.antennas.spacing
            * (1 - torch.cos(theta))
            / 299_792_458
        )
        assert omega_t.shape == phi_theta.shape == (N,)

        omega_t = torch.unsqueeze(omega_t, dim=-1)
        phi_theta = torch.unsqueeze(phi_theta, dim=-1)

        assert omega_t.shape == phi_theta.shape == (N, 1)

        antenna_v = omega_t ** torch.arange(
            self.N_subcarriers // 2, dtype=torch.float32
        )
        phis = phi_theta ** torch.arange(self.N_rx // 2, dtype=torch.float32)

        assert antenna_v.shape == (N, self.N_subcarriers // 2)
        assert phis.shape == (N, self.N_rx // 2)

        antenna_v = torch.unsqueeze(antenna_v, dim=1)
        phis = torch.unsqueeze(phis, dim=-1)

        assert antenna_v.shape == (N, 1, self.N_subcarriers // 2)
        assert phis.shape == (N, self.N_rx // 2, 1)

        steering = torch.bmm(phis, antenna_v)

        assert steering.shape == (N, self.N_rx // 2, self.N_subcarriers // 2)
        return steering.reshape(N, -1)

    def evaluate(self, theta: torch.Tensor, tof: torch.Tensor) -> torch.Tensor:
        R = torch.mean(self.historical_autocorr, dim=0)
        assert (R == torch.conj(R).T).all()

        # The smallest eigenvectors span the noise subspace,
        # and the largest span the signal subspace.
        logger.debug(f"Calculating eigenvectors of R: {R.shape}")
        eigvals, eigvecs = torch.linalg.eigh(R)
        assert isinstance(eigvals, torch.Tensor)
        assert isinstance(eigvecs, torch.Tensor)
        logger.info(f"Eigenvalues: {eigvals}")
        E_n = eigvecs[:, torch.abs(eigvals) < config.music.eigval_threshold]

        logger.debug(f"Signal subspace: {E_n.shape}")
        steering = torch.unsqueeze(self.steering_vector(theta, tof), dim=-1)
        steering_h = torch.conj(steering).permute(0, 2, 1)

        E_n = E_n.unsqueeze(0)
        E_n_H = torch.conj(E_n).permute(0, 2, 1)
        logger.debug(
            f"Heatmap multiplication: {steering_h.shape}, {E_n.shape}, "
            f"{E_n_H.shape}, {steering.shape}"
        )
        c: torch.Tensor = 1 / (0.001 + (steering_h @ E_n @ E_n_H @ steering))
        return torch.abs(c.real)

    def heatmap(self) -> npt.NDArray[np.float32]:
        thetas = np.linspace(
            0, np.pi, config.music.heatmap.theta_resolution, dtype=np.float32
        )
        tofs = np.linspace(
            0,
            config.music.heatmap.tof_max,
            config.music.heatmap.tof_resolution,
            dtype=np.float32,
        )
        thetas_mesh, tofs_mesh = np.meshgrid(thetas, tofs)
        logger.debug(
            f"Calculating heatmap with {thetas_mesh.shape} and {tofs_mesh.shape}"
        )
        evaluated = self.evaluate(
            torch.tensor(thetas_mesh.reshape(-1)),
            torch.tensor(tofs_mesh.reshape(-1)),
        )
        logger.debug(f"Evaluated heatmap: {evaluated.shape}")
        heatmap: npt.NDArray[np.float32] = evaluated.reshape(
            config.music.heatmap.theta_resolution,
            config.music.heatmap.tof_resolution,
        ).numpy(force=True)
        return heatmap


def test_smoothing() -> None:
    row, col = np.indices((6, 4))
    data = row + 1j * col
    data = np.expand_dims(data, axis=2)
    np.set_printoptions(linewidth=200)
    print(data.shape)
    aoa = AoA()
    aoa.N_subcarriers = 6
    aoa.N_rx = 4
    smoothed = aoa.smooth(data)
    print(smoothed)


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