import numpy as np
from collections.abc import Sequence
from typing import TYPE_CHECKING
from pydantic import field_validator
from py_neuromodulation.utils.types import NMBaseModel, BoolSelector, NMFeature
from py_neuromodulation.utils.pydantic_extensions import (
    NMField,
    NMErrorList,
    create_validation_error,
)
from py_neuromodulation import logger

if TYPE_CHECKING:
    from py_neuromodulation.stream.settings import NMSettings
    from py_neuromodulation.filter import KalmanSettings


class BandpowerFeatures(BoolSelector):
    activity: bool = True
    mobility: bool = False
    complexity: bool = False


class BandPowerSettings(NMBaseModel):
    segment_lengths_ms: dict[str, int] = NMField(
        default={
            "theta": 1000,
            "alpha": 500,
            "low beta": 333,
            "high beta": 333,
            "low gamma": 100,
            "high gamma": 100,
            "HFA": 100,
        },
        custom_metadata={"field_type": "FrequencySegmentLength"},
    )
    bandpower_features: BandpowerFeatures = BandpowerFeatures()
    log_transform: bool = True
    kalman_filter: bool = False

    @field_validator("bandpower_features")
    @classmethod
    def bandpower_features_validator(cls, bandpower_features: BandpowerFeatures):
        if not len(bandpower_features.get_enabled()) > 0:
            raise create_validation_error(
                error_message="Set at least one bandpower_feature to True.",
                location=["bandpass_filter_settings", "bandpower_features"],
            )
        return bandpower_features

    def validate_fbands(self, settings: "NMSettings") -> NMErrorList:
        """_summary_

        :param settings: _description_
        :type settings: NMSettings
        :raises create_validation_error: _description_
        :raises create_validation_error: _description_
        :raises ValueError: _description_
        """
        errors: NMErrorList = NMErrorList()

        for fband_name, seg_length_fband in self.segment_lengths_ms.items():
            # Check that all frequency bands are defined in settings.frequency_ranges_hz
            if fband_name not in settings.frequency_ranges_hz:
                logger.warning(
                    f"Frequency band {fband_name} in bandpass_filter_settings.segment_lengths_ms"
                    " is not defined in settings.frequency_ranges_hz"
                )

            # Check that all segment lengths are smaller than settings.segment_length_features_ms
            if not seg_length_fband <= settings.segment_length_features_ms:
                errors.add_error(
                    f"segment length {seg_length_fband} needs to be smaller than "
                    f" settings['segment_length_features_ms'] = {settings.segment_length_features_ms}",
                    location=[
                        "bandpass_filter_settings",
                        "segment_lengths_ms",
                        fband_name,
                    ],
                )

        # Check that all frequency bands defined in settings.frequency_ranges_hz
        for fband_name in settings.frequency_ranges_hz.keys():
            if fband_name not in self.segment_lengths_ms:
                errors.add_error(
                    f"frequency range {fband_name} "
                    "needs to be defined in settings.bandpass_filter_settings.segment_lengths_ms",
                    location=[
                        "bandpass_filter_settings",
                        "segment_lengths_ms",
                        fband_name,
                    ],
                )

        return errors


class BandPower(NMFeature):
    def __init__(
        self,
        settings: "NMSettings",
        ch_names: Sequence[str],
        sfreq: float,
        use_kf: bool | None = None,
    ) -> None:
        settings.validate()

        self.bp_settings: BandPowerSettings = settings.bandpass_filter_settings
        self.kalman_filter_settings: KalmanSettings = settings.kalman_filter_settings
        self.sfreq = sfreq
        self.ch_names = ch_names
        self.KF_dict: dict = {}

        from py_neuromodulation.filter import MNEFilter

        self.bandpass_filter = MNEFilter(
            f_ranges=[
                tuple(frange) for frange in settings.frequency_ranges_hz.values()
            ],
            sfreq=self.sfreq,
            filter_length=self.sfreq - 1,
            verbose=False,
        )

        if use_kf or (use_kf is None and self.bp_settings.kalman_filter):
            self.init_KF("bandpass_activity")

        seglengths = self.bp_settings.segment_lengths_ms

        self.feature_params = []
        for ch_idx, ch_name in enumerate(self.ch_names):
            for f_band_idx, f_band in enumerate(settings.frequency_ranges_hz.keys()):
                seglength_ms = seglengths[f_band]
                seglen = int(np.floor(self.sfreq / 1000 * seglength_ms))
                for bp_feature in self.bp_settings.bandpower_features.get_enabled():
                    feature_name = "_".join([ch_name, "bandpass", bp_feature, f_band])
                    self.feature_params.append(
                        (
                            ch_idx,
                            f_band_idx,
                            seglen,
                            bp_feature,
                            feature_name,
                        )
                    )

    def init_KF(self, feature: str) -> None:
        from py_neuromodulation.filter import define_KF

        for f_band in self.kalman_filter_settings.frequency_bands:
            for channel in self.ch_names:
                self.KF_dict["_".join([channel, feature, f_band])] = define_KF(
                    self.kalman_filter_settings.Tp,
                    self.kalman_filter_settings.sigma_w,
                    self.kalman_filter_settings.sigma_v,
                )

    def update_KF(self, feature_calc: np.floating, KF_name: str) -> np.floating:
        if KF_name in self.KF_dict:
            self.KF_dict[KF_name].predict()
            self.KF_dict[KF_name].update(feature_calc)
            feature_calc = self.KF_dict[KF_name].x[0]
        return feature_calc

    def calc_feature(self, data: np.ndarray) -> dict:
        data = self.bandpass_filter.filter_data(data)

        feature_results = {}
        for (
            ch_idx,
            f_band_idx,
            seglen,
            bp_feature,
            feature_name,
        ) in self.feature_params:
            feature_results[feature_name] = self.calc_bp_feature(
                bp_feature, feature_name, data[ch_idx, f_band_idx, -seglen:]
            )

        return feature_results

    def calc_bp_feature(self, bp_feature, feature_name, data):
        match bp_feature:
            case "activity":
                feature_calc = np.var(data)
                if self.bp_settings.log_transform:
                    feature_calc = np.log10(feature_calc)
                if self.KF_dict:
                    feature_calc = self.update_KF(feature_calc, feature_name)
            case "mobility":
                feature_calc = np.sqrt(np.var(np.diff(data)) / np.var(data))
            case "complexity":
                feature_calc = self.calc_complexity(data)
            case _:
                raise ValueError(f"Unknown bandpower feature: {bp_feature}")

        return np.nan_to_num(feature_calc)

    @staticmethod
    def calc_complexity(data: np.ndarray) -> float:
        dat_deriv = np.diff(data)
        deriv_variance = np.var(dat_deriv)
        mobility = np.sqrt(deriv_variance / np.var(data))
        dat_deriv_2_var = np.var(np.diff(dat_deriv))
        deriv_mobility = np.sqrt(dat_deriv_2_var / deriv_variance)

        return deriv_mobility / mobility
