from typing import Literal, Self

from pydantic import BaseModel, model_validator

Host = tuple[str, int]


class Preprocessing(BaseModel):
    moving_average_alpha: float

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

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

    bandpass: Bandpass
    subcarrier_step: int


class MUSIC(BaseModel):
    eigval_threshold: int
    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]]


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 delta_f(self) -> int:
        if self.frame_format == "HESU":
            return 78_125 * self.preprocessing.subcarrier_step
        return 312_500 * self.preprocessing.subcarrier_step

    @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