"""
This module contains a Helper class to make config files.

The processing config is still evolving and this class and its methods may change.

Development Notes:
 - Initially, we only supported EMTF style band specification, but this should be updated
     to allow logarithmically spaced bands of constant Q.

"""

import pathlib
from typing import Optional, Union

from loguru import logger
from mt_metadata.processing.window import Window
from mth5.processing.kernel_dataset import KernelDataset

from aurora.config import BANDS_DEFAULT_FILE
from aurora.config.metadata import Processing
from aurora.sandbox.io_helpers.emtf_band_setup import EMTFBandSetupFile


SUPPORTED_BAND_SPECIFICATION_STYLES = ["EMTF", "band_edges"]


class ConfigCreator:
    def __init__(
        self,
        emtf_band_file: Optional[Union[str, pathlib.Path, None]] = None,
        band_edges: Optional[Union[dict, None]] = None,
    ):
        """
        Constructor

        Parameters
        ----------
        emtf_band_file: Optional[Union[str, pathlib.Path, None]]
            Allows the specification of an EMTF "band setup file" for defining frequency bands.
        band_edges: dict
            Keys are integers corresponding to decimation level.  Values are numpy arrays.
            Numpy arrays are one-row-per-band.  Array shape n_bands x 2. array[i_band,0] is the
            lower edge of band, and array[i_band,0] is the upper.

        """
        self._emtf_band_file = emtf_band_file
        self._band_edges = band_edges
        self._band_specification_style = None

    # TODO: Delete or update with config info ... currently kernel_dataset.processing_id is used.
    # def processing_id(self, kernel_dataset):
    #     """
    #     Generates a string id label for the processing config: WIP.

    #     In the past, we used f"{local}-{remote}" or  f"{local}-{run_id}"
    #     Neither of these is sufficiently unique.  In fact, they only describe the
    #     dataset, and not the processing config.  It is difficult to see how to make a
    #     comprehensive, unique id without it being very long or involving hash
    #     functions.

    #     For now, will try to use {local}-{local_runs}-{remote}-{remote_runs},
    #     which at least describes the dataset, then a string can be generated by the
    #     config and appended if needed.

    #     Parameters
    #     ----------
    #     kernel_dataset: aurora.transfer_function.kernel_dataset.KernelDataset
    #         An object that defines the data to be processed.

    #     Returns
    #     -------
    #     id: str
    #         A label for the processing config.
    #     """
    #     id = f"{kernel_dataset.local_station_id}-{kernel_dataset.remote_station_id}"
    #     return id

    @property
    def band_specification_style(self):
        """return a description of the scheme used to define the bands."""
        return self._band_specification_style

    @band_specification_style.setter
    def band_specification_style(self, value: str) -> None:
        """
        Sets the band_specification_style

        Parameters
        ----------
        value: str
            The label for the scheme used to define the bands.

        Returns
        -------

        """
        if value not in SUPPORTED_BAND_SPECIFICATION_STYLES:
            msg = f"Won't set band specification style to unrecognized value {value}"
            logger.warning(msg)
            raise NotImplementedError(msg)
        else:
            self._band_specification_style = value

    def determine_band_specification_style(self) -> None:
        """
            Tries to identify which scheme was used to define the bands

        TODO: Should emtf_band_file path be stored in config to support reproducibility?

        """

        if (self._emtf_band_file is None) & (self._band_edges is None):
            msg = "Bands not defined; setting to EMTF BANDS_DEFAULT_FILE"
            logger.info(msg)
            self._emtf_band_file = BANDS_DEFAULT_FILE
            self._band_specification_style = "EMTF"
        elif (self._emtf_band_file is not None) & (self._band_edges is not None):
            msg = "Bands defined twice, and possibly inconsistently"
            logger.error(msg)
            raise ValueError(msg)
        elif self._band_edges is not None:
            self._band_specification_style = "band_edges"
        elif self._emtf_band_file is not None:
            self._band_specification_style = "EMTF"

    def create_from_kernel_dataset(
        self,
        kernel_dataset: KernelDataset,
        input_channels: Optional[list] = None,
        output_channels: Optional[list] = None,
        remote_channels: Optional[list] = None,
        estimator: Optional[str] = None,
        emtf_band_file: Optional[Union[str, pathlib.Path]] = None,
        band_edges: Optional[dict] = None,
        decimation_factors: Optional[list] = None,
        num_samples_window: Optional[int] = None,
        **kwargs,
    ) -> Processing:
        """
        Creates a processing config from a kernel dataset.

        TODO: Make this a method of kernel_dataset.

        **Development Notes:**

        1. The number of decimation levels must be defined either by:

        - decimation_factors argument (normally accompanied by a bands_dict)
        - number of decimations implied by EMTF band setup file.

        Theoretically, you could also use the number of decimations implied by bands_dict but this is sloppy, because it would assume the decimation factor.

        2. 2024-12-29 Added setting of decimation_obj.stft.per_window_detrend_type = "linear"

        This makes tests pass following a refactoring of mt_metadata.  Could use more testing.

        Parameters
        ----------
        kernel_dataset: aurora.transfer_function.kernel_dataset.KernelDataset
            An object that defines the data to be processed.
        input_channels: list
            List of the input channels that will be used in TF estimation (usually "hx", "hy")
        output_channels: list
            List of the output channels that will be estimated by TF (usually "ex", "ey", "hz")
        remote_channels: list
            List of the remote reference channels (usually "hx", "hy" at remote site)
        estimator:  Optional[Union[str, None]]
            The name of the regression estimator to use for TF estimation.
        emtf_band_file: Optional[Union[str, pathlib.Path, None]]
            The emtf and setup file if used.
        band_edges: Optional[Union[dict, None]]
            The band edges if emtf_band_file not used
        decimation_factors: Optional[Union[list, None]]
            List of decimation factors, normally [1, 4, 4, 4, ... 4]
        num_samples_window: Optional[Union[int, None]]
            The size of the window (usually for FFT)

        **kwargs:
            Additional keyword arguments passed to Processing constructor. Could contain:

            - save_fcs: bool
                If True, save Fourier coefficients during processing.
            - save_fcs_type: str
                File type for saving Fourier coefficients. Options are "h5" or "csv".

        Returns
        -------
        processing_obj: aurora.config.metadata.processing.Processing
            Object storing the processing parameters.
        """

        processing_obj = Processing(id=kernel_dataset.processing_id)  # , **kwargs)

        # pack station and run info into processing object
        processing_obj.stations.from_dataset_dataframe(kernel_dataset.df)

        # Unpack optional arguments
        self._emtf_band_file = emtf_band_file
        self._band_edges = band_edges
        decimation_factors = decimation_factors
        num_samples_window = num_samples_window

        # Determine if band_setup or edges dict is to be used for bands
        self.determine_band_specification_style()

        # Set Frequency Bands
        if self.band_specification_style == "EMTF":
            emtf_band_setup_file = EMTFBandSetupFile(
                filepath=self._emtf_band_file,
                sample_rate=kernel_dataset.sample_rate,
            )
            num_decimations = emtf_band_setup_file.num_decimation_levels
            # Assign optional arguments if they have not been passed
            if decimation_factors is None:
                # set default values to EMTF default values [1, 4, 4, 4, ..., 4]
                decimation_factors = num_decimations * [4]
                decimation_factors[0] = 1
            if num_samples_window is None:
                default_window = Window()
                num_samples_window = num_decimations * [default_window.num_samples]
            elif isinstance(num_samples_window, int):
                num_samples_window = num_decimations * [num_samples_window]
            # now you can define the frequency bands
            band_edges = emtf_band_setup_file.compute_band_edges(
                decimation_factors, num_samples_window
            )
            self._band_edges = band_edges

        processing_obj.assign_bands(
            self._band_edges,
            kernel_dataset.sample_rate,
            decimation_factors,
            num_samples_window,
        )
        processing_obj.band_specification_style = self.band_specification_style
        if self.band_specification_style == "EMTF":
            processing_obj.band_setup_file = str(self._emtf_band_file)
        for key, decimation_obj in processing_obj.decimations_dict.items():
            if input_channels is None:
                decimation_obj.input_channels = kernel_dataset.input_channels
            else:
                decimation_obj.input_channels = input_channels

            if output_channels is None:
                decimation_obj.output_channels = kernel_dataset.output_channels
            else:
                decimation_obj.output_channels = output_channels

            if remote_channels is None:
                if kernel_dataset.remote_channels is not None:
                    decimation_obj.reference_channels = kernel_dataset.remote_channels

            if num_samples_window is not None:
                decimation_obj.stft.window.num_samples = num_samples_window[key]

            if kwargs.get("save_fcs", False):
                decimation_obj.save_fcs = True
                decimation_obj.save_fcs_type = kwargs.get("save_fcs_type", "h5")

            # set estimator if provided as kwarg
            if estimator:
                try:
                    decimation_obj.estimator.engine = estimator["engine"]
                except KeyError:
                    pass
            decimation_obj.stft.per_window_detrend_type = "linear"

        return processing_obj
