from typing import Literal, Self

from pydantic import BaseModel, model_validator

Host = tuple[str, int]


class Preprocessing(BaseModel):
    class Bandpass(BaseModel):
        lowcut: int
        highcut: int

        @property
        def bounds(self) -> tuple[int, int]:
            return (self.lowcut, self.highcut)

    bandpass: Bandpass | None = None
    subcarrier_step: int = 1
    denoising: Literal["none", "median", "mean"] = "median"
    denoising_period: float = 1

    steps: list[
        Literal[
            "fill_pilots",
            "skip_subcarriers",
            "remove_agc",
            "remove_sfo",
            "remove_sto",
            "bandpass",
        ]
    ] = ["fill_pilots", "skip_subcarriers", "remove_agc", "remove_sfo"]

    @model_validator(mode="after")
    def skip_in_steps(self) -> Self:
        if "skip_subcarriers" not in self.steps and self.subcarrier_step != 1:
            raise ValueError(
                "subcarrier_step must be 1 if skip_subcarriers is not in steps"
            )
        return self

    @model_validator(mode="after")
    def bandpass_in_steps(self) -> Self:
        if "bandpass" not in self.steps and self.bandpass is not None:
            raise ValueError(
                "bandpass must be in preprocessing steps if bandpass configuration is "
                "provided"
            )

        if "bandpass" in self.steps and self.bandpass is None:
            raise ValueError(
                "bandpass configuration must be provided if bandpass is in "
                "preprocessing steps"
            )
        return self


class MUSIC(BaseModel):
    eigval_threshold: float
    window_size: int

    class Heatmap(BaseModel):
        theta_resolution: int
        tof_resolution: int
        tof_max: float

    heatmap: Heatmap


class Antennas(BaseModel):
    spacing: float
    order: list[tuple[Host, int]]

    @property
    def count(self) -> int:
        return len(self.order)


class Config(BaseModel):
    receive_hosts: list[Host]
    transmit_host: Host

    antennas: Antennas

    collection_sample_rate: int
    processing_sample_rate: int
    central_freq: int
    channel_width: Literal[20, 40, 80, 160]
    frame_format: Literal["NOHT", "HT", "VHT", "HESU"]

    preprocessing: Preprocessing
    music: MUSIC

    @property
    def central_freq_hz(self) -> int:
        return self.central_freq * 1_000_000

    @property
    def band(self) -> Literal["2.4", "5", "6"]:
        if self.central_freq in range(2412, 2484):
            return "2.4"
        if self.central_freq in range(5180, 5320):
            return "5"
        if self.central_freq in range(5955, 7115):
            return "6"
        raise ValueError(f"{self.central_freq} is not a valid Wi-Fi channel")

    @property
    def num_guards(self) -> int:
        if self.channel_width == 20:
            return 7
        return 11

    @property
    def subcarriers(self) -> int:
        return (
            int((self.channel_width * 1e6) // self.delta_f) - self.num_guards
        ) // self.preprocessing.subcarrier_step

    @property
    def _delta_f_no_skipping(self) -> int:
        if self.frame_format == "HESU":
            return 78_125
        return 312_500

    @property
    def delta_f(self) -> int:
        return self._delta_f_no_skipping * self.preprocessing.subcarrier_step

    @property
    def subcarrier_frequencies(self) -> list[int]:
        num_subcarriers = self.channel_width * 1_000_000 // self._delta_f_no_skipping
        return [
            self.central_freq_hz + i * self._delta_f_no_skipping
            for i in range(-num_subcarriers // 2, num_subcarriers // 2)
        ]

    @model_validator(mode="after")
    def channels(self) -> Self:
        band_start = 2412 if self.band == "2.4" else 5180 if self.band == "5" else 5955
        if self.band == "2.4" and self.channel_width not in [20, 40]:
            raise ValueError(
                f"2.4GHz channel {self.central_freq} must "
                "have a channel width of 20 or 40 MHz"
            )

        if (self.band == "2.4" and (self.central_freq - band_start) % 5 != 0) or (
            self.band in ["5", "6"]
            and ((self.central_freq - band_start) % self.channel_width != 0)
        ):
            raise ValueError(
                f"Central frequency {self.central_freq} must be a channel as in https://en.wikipedia.org/wiki/List_of_WLAN_channels"
            )
        return self