#! /usr/bin/env python
"""Module for the post-processing FMMF algorithm."""

__author__ = "Thomas Bédrine"
__all__ = ["FMMFBuilder", "PPFMMF"]

from typing import Optional
from dataclasses import dataclass

import numpy as np
from dataclass_builder import dataclass_builder

from .dataset import Dataset
from .postproc import PostProc
from ..invprob import fmmf, FMMF_Params
from ..config.utils_conf import algo_calculates_decorator as calculates


@dataclass
class PPFMMF(PostProc, FMMF_Params):
    """Post-processing forward model matching filter algorithm."""

    _algo_name: str = "fmmf"
    snr_map: np.ndarray = None

    # TODO: write test
    @calculates("frame_final", "snr_map")
    def run(
        self,
        dataset: Optional[Dataset] = None,
        nproc: Optional[int] = None,
    ):
        """
        Run the post-processing FMMF algorithm for model PSF subtraction.

        Parameters
        ----------
        dataset : Dataset object
            An Dataset object to be processed.
        model: {'KLIP', 'LOCI'}, optional
            If you want to change the default model. See documentation above for more
            information.
        nproc : None or int, optional
            Number of processes for parallel computing. If None the number of
            processes will be set to cpu_count()/2. By default the algorithm works
            in single-process mode.
        verbose : bool, optional
            If True prints to stdout intermediate info.

        """
        self.snr_map = None
        self._update_dataset(dataset)

        if self.dataset.fwhm is None:
            raise ValueError("`fwhm` has not been set")

        self._explicit_dataset()
        if nproc is not None:
            self.nproc = nproc

        params_dict = self._create_parameters_dict(FMMF_Params)
        all_params = {"algo_params": self}
        res = fmmf(**all_params)

        self.frame_final, self.snr_map = res

        if self.results is not None:
            self.results.register_session(
                params=params_dict,
                frame=self.frame_final,
                snr_map=self.snr_map,
                algo_name=self._algo_name,
            )

    def make_snrmap(self):
        """
        Do nothing, snr_map is already generated by ``fmmf``.

        This function must overload the inherited one to avoid misusage.

        """


FMMFBuilder = dataclass_builder(PPFMMF)
