import logging
import time
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Iterator

import numpy as np
import numpy.typing as npt

from where_fi.collection.protocols import CSIProducer, MergedCSI
from where_fi.config import config

C = 299_792_458
MAX_DEPTH = 2
LOSS_EXPONENT = 0.8

logger = logging.getLogger(__name__)


@dataclass
class PathComponent:
    delay: float
    phase_offset: float
    attenuation: float


class ChannelImpulseResponse:
    def __init__(self, path_components: list[PathComponent]) -> None:
        self.path_components = path_components

    @staticmethod
    def delayed(
        other: "ChannelImpulseResponse", delay: float, reflect: bool = False
    ) -> "ChannelImpulseResponse":
        new_path_components: list[PathComponent] = []
        distance = delay * C

        for component in other.path_components:
            new_path_components.append(
                PathComponent(
                    delay=component.delay + delay,
                    phase_offset=component.phase_offset + (np.pi if reflect else 0),
                    attenuation=component.attenuation
                    * (1 / (1 + distance) ** LOSS_EXPONENT),
                    # attenuation=component.attenuation
                    # * (np.exp(-distance * LOSS_EXPONENT)),
                )
            )
        return ChannelImpulseResponse(new_path_components)

    def __add__(self, other: "ChannelImpulseResponse") -> "ChannelImpulseResponse":
        return ChannelImpulseResponse(self.path_components + other.path_components)


class PathObject:
    def __init__(self, x: float, y: float, z: float) -> None:
        self.x = x
        self.y = y
        self.z = z
        self.cir = ChannelImpulseResponse([])

    def distance(self, other: "PathObject") -> float:
        return (
            (self.x - other.x) ** 2 + (self.y - other.y) ** 2 + (self.z - other.z) ** 2
        ) ** 0.5

    def reflect(
        self,
        incoming_cir: ChannelImpulseResponse,
        target: list["PathObject"],
        depth: int = 0,
    ) -> None:
        self.cir = self.cir + incoming_cir

        if depth > MAX_DEPTH:
            return

        for obj in target:
            if id(obj) == id(self):
                continue
            distance = self.distance(obj)
            assert distance > 0
            delay = distance / C
            obj.reflect(
                ChannelImpulseResponse.delayed(incoming_cir, delay, reflect=True),
                target,
                depth + 1,
            )


class Transmitter(PathObject):
    def __init__(self, x: float, y: float, z: float) -> None:
        super().__init__(x, y, z)


DELTA_T = 10
GAMMA = np.pi / 4


class Receiver(PathObject):
    def __init__(self, x: float, y: float, z: float, ideal: bool = True) -> None:
        super().__init__(x, y, z)
        self.delta_t = 0 if ideal else DELTA_T
        self.gamma = 0 if ideal else GAMMA

    def get_cfr(self) -> npt.NDArray[np.complex64]:
        """
        Calculate the CSI matrix for this rx-tx pair. This is computed by the Fourier
        Transform of the Channel Impulse Response (CIR). The CIR is calculated as in
        [1], [2].

        To get the FT of the CIR, we use the sifting property of the Dirac delta

        [1] - https://tns.thss.tsinghua.edu.cn/wst/docs/pre
        [2] - https://dl.acm.org/doi/10.1145/2543581.2543592, Equation 4
        """
        cfr = np.zeros(len(config.subcarrier_frequencies), dtype=np.complex64)
        for i_sub, f_sub in enumerate(config.subcarrier_frequencies):
            cfr[i_sub] = sum(
                [
                    component.attenuation
                    * np.exp(1j * component.phase_offset)
                    * np.exp(-1j * 2 * np.pi * f_sub * component.delay)
                    for component in self.cir.path_components
                ]
            ) * np.exp(
                1j
                * (
                    2
                    * np.pi
                    * (i_sub / len(config.subcarrier_frequencies))
                    * self.delta_t
                    + self.gamma
                )
            )
        return cfr


class Environment:
    def __init__(self, objects: list[PathObject]) -> None:
        self.transmitters = [obj for obj in objects if isinstance(obj, Transmitter)]
        self.receivers = [obj for obj in objects if isinstance(obj, Receiver)]
        self.objects = [
            obj for obj in objects if obj not in self.transmitters + self.receivers
        ]
        self.count = 0

    def move(self) -> None:
        self.count += 1
        if self.count > 4000:
            print("Moving objects")
            for obj in self.objects:
                obj.x += np.sin(self.count / 1000 * 2 * np.pi) * 2

    @staticmethod
    def from_config(filename: Path) -> "Environment":
        """
        Read a config file that includes a scene description and create an envionment
        based on that

        The config file should be a text file where each line corresponds to an object.
        The first word of each line should be the type of object (TX/RX/OBJ), followed
        by the x and y coordinates of the object.
        """
        objects: list[PathObject] = []
        with open(filename, "r") as f:
            for line in f.readlines():
                if line.startswith("#"):
                    continue
                parts = line.split(" ")
                x, y, z = map(float, parts[1:])
                if parts[0] == "TX":
                    objects.append(Transmitter(x, y, z))
                elif parts[0] == "RX":
                    objects.append(Receiver(x, y, z))
                else:
                    objects.append(PathObject(x, y, z))

        return Environment(objects)

    def add_awgn(
        self, signal: npt.NDArray[np.complex64], snr_dB: float
    ) -> npt.NDArray[np.complex64]:
        signal_power = np.mean(np.abs(signal) ** 2)
        snr_linear = 10 ** (snr_dB / 10)
        noise_power = signal_power / snr_linear
        noise = np.sqrt(noise_power / 2) * (
            np.random.randn(*signal.shape) + 1j * np.random.randn(*signal.shape)
        )
        return signal + noise.astype(np.complex64)

    def get_csi(self) -> npt.NDArray[np.complex64]:
        """
        Calculate the Channel State Information (CSI) matrix for the simulated
        environment.

        The returned matrix is of shape (num_subcarriers, num_receivers,
                                         num_transmitters)

        This is calculated by finding all paths leading to each receiver, and
        calculating the CFR evaluated at each subcarrier.
        """
        csi = np.zeros(
            (
                len(config.subcarrier_frequencies),
                len(self.receivers),
                len(self.transmitters),
            ),
            dtype=np.complex64,
        )
        for i_tx, transmitter in enumerate(self.transmitters):
            for obj in self.objects + self.receivers + self.transmitters:
                obj.cir = ChannelImpulseResponse([])
            transmitter.reflect(
                ChannelImpulseResponse(
                    [
                        PathComponent(
                            delay=0,
                            phase_offset=0,
                            # attenuation=100000000,
                            attenuation=100,
                        )
                    ]
                ),
                self.objects + self.receivers,
            )
            for i_rx, receiver in enumerate(self.receivers):
                logger.debug(f"RX {i_rx} paths: {len(receiver.cir.path_components)}")
                csi[:, i_rx, i_tx] = self.add_awgn(receiver.get_cfr(), snr_dB=30)
        return csi


class SimulatedCSIProducer(CSIProducer):
    pass

    def __init__(self, environment: Environment) -> None:
        self.environment = environment
        self.active = True

    def __call__(self) -> Iterator[MergedCSI]:
        while self.active:
            self.environment.move()
            start = datetime.now()
            csi = self.environment.get_csi()
            yield MergedCSI(
                frames={},
                matrix=csi,
            )
            time.sleep(
                max(
                    0,
                    1 / config.collection_sample_rate
                    - (datetime.now() - start).total_seconds(),
                )
            )

    def stop(self) -> None:
        self.active = False
