import logging
from datetime import datetime
from typing import Any, Callable

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

from ..config import config
from ..visualise import server as visualise

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 = config.subcarriers
        self.N_rx = config.antennas.count
        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
            * (torch.sin(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, dtype=torch.float32)
        phis = phi_theta ** torch.arange(self.N_rx, dtype=torch.float32)

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

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

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

        steering = torch.bmm(phis, antenna_v)

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

    def evaluate(
        self,
        theta: torch.Tensor,
        tof: torch.Tensor,
        visualiser: None
        | Callable[[npt.NDArray[Any], visualise.figures.Figure], None] = None,
    ) -> torch.Tensor:
        R = torch.mean(self.historical_autocorr, dim=0)

        # 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.eig(R)
        if visualiser:
            visualiser(eigvals.numpy(), visualise.figures.Figure.MUSIC_EIGENVALUES)
        assert isinstance(eigvals, torch.Tensor)
        assert isinstance(eigvecs, torch.Tensor)
        logger.debug(f"Eigenvalues: {eigvals}")
        E_n = eigvecs[:, torch.abs(eigvals) < config.music.eigval_threshold]

        logger.info(f"Signal subspace: {E_n.shape}")
        steering = torch.unsqueeze(self.steering_vector(theta, tof), dim=-1)
        print(steering.shape)
        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 / (steering_h @ E_n @ E_n_H @ steering)
        return torch.abs(c)[:, 0, 0]

    def heatmap(
        self,
        visualiser: None | Callable[[npt.NDArray[Any], visualise.figures.Figure], None],
    ) -> 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)),
            visualiser=visualiser,
        )
        logger.debug(f"Evaluated heatmap: {evaluated.shape}")
        heatmap: npt.NDArray[np.float32] = evaluated.reshape(
            config.music.heatmap.tof_resolution,
            config.music.heatmap.theta_resolution,
        ).numpy(force=True)

        if visualiser:
            visualiser(heatmap, visualise.figures.Figure.AOA_HEATMAP)
        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
