#  Copyright (c) European Space Agency, 2017, 2018, 2019, 2020.
#
#  This file is subject to the terms and conditions defined in file 'LICENCE.txt', which
#  is part of this Pyxel package. No part of the package, including
#  this file, may be copied, modified, propagated, or distributed except according to
#  the terms contained in the file ‘LICENCE.txt’.

"""Pyxel wrapper class for NGHXRG - Teledyne HxRG Noise Generator model."""

import logging
import typing as t
from dataclasses import dataclass

import numpy as np
from typing_extensions import Literal

from pyxel.detectors import CMOS, CMOSGeometry
from pyxel.models.charge_measurement.nghxrg.nghxrg_beta import HXRGNoise
from pyxel.util import temporary_random_state


@dataclass
class KTCBiasNoise:
    """KTCBiasNoise parameters."""

    ktc_noise: float
    bias_offset: float
    bias_amp: float


@dataclass
class WhiteReadNoise:
    """WhiteReadNoise parameters."""

    rd_noise: float
    ref_pixel_noise_ratio: float


@dataclass
class CorrPinkNoise:
    """CorrPinkNoise parameters."""

    c_pink: float


@dataclass
class UncorrPinkNoise:
    """UncorrPinkNoise parameters."""

    u_pink: float


@dataclass
class AcnNoise:
    """AcnNoise parameters."""

    acn: float


@dataclass
class PCAZeroNoise:
    """PCAZeroNoise parameters."""

    pca0_amp: float


# Create an alias
NoiseType = t.Union[
    KTCBiasNoise,
    WhiteReadNoise,
    CorrPinkNoise,
    UncorrPinkNoise,
    AcnNoise,
    PCAZeroNoise,
]


# TODO: Use function `simcado.nghxrg.HXRGNoise` instead of
#       `pyxel.models.charge_measurement.nghxrg.nghxrg.HXRGNoise`
def compute_nghxrg(
    pixel_2d: np.ndarray,
    noise: t.Sequence[NoiseType],
    detector_shape: t.Tuple[int, int],
    window_pos: t.Tuple[int, int],
    window_size: t.Tuple[int, int],
    num_outputs: int,
    time_step: int,
    num_rows_overhead: int,
    num_frames_overhead: int,
    reverse_scan_direction: bool,
    reference_pixel_border_width: int,
) -> np.ndarray:
    """Compute NGHXRG.

    Parameters
    ----------
    pixel_2d : ndarray
    noise : sequence of NoiseType
    detector_shape : int, int
    window_pos : int, int
    window_size : int, int
    num_outputs : int
    time_step : int
    num_rows_overhead : int
    num_frames_overhead : int
    reverse_scan_direction : bool
    reference_pixel_border_width : int

    Returns
    -------
    ndarray
    """
    det_size_y, det_size_x = detector_shape
    window_y_start, window_x_start = window_pos
    window_y_size, window_x_size = window_size

    if window_pos == (0, 0) and window_size == detector_shape:
        window_mode = "FULL"  # type: Literal["FULL", "WINDOW"]
    else:
        window_mode = "WINDOW"

    data_hxrg_2d = np.asarray(pixel_2d, dtype=float)  # type: np.ndarray

    ng = HXRGNoise(
        n_out=num_outputs,
        time_step=time_step,
        nroh=num_rows_overhead,
        nfoh=num_frames_overhead,
        reverse_scan_direction=reverse_scan_direction,
        reference_pixel_border_width=reference_pixel_border_width,
        pca0=data_hxrg_2d,
        det_size_x=det_size_x,
        det_size_y=det_size_y,
        wind_mode=window_mode,
        wind_x_size=window_x_size,
        wind_y_size=window_y_size,
        wind_x0=window_x_start,
        wind_y0=window_y_start,
        verbose=True,
    )  # type: HXRGNoise

    final_data_2d = np.zeros(shape=window_size)

    for item in noise:  # type: NoiseType
        if isinstance(item, KTCBiasNoise):
            data = ng.add_ktc_bias_noise(
                ktc_noise=item.ktc_noise,
                bias_offset=item.bias_offset,
                bias_amp=item.bias_amp,
            )  # type: np.ndarray

        elif isinstance(item, WhiteReadNoise):
            data = ng.add_white_read_noise(
                rd_noise=item.rd_noise,
                reference_pixel_noise_ratio=item.ref_pixel_noise_ratio,
            )

        elif isinstance(item, CorrPinkNoise):
            data = ng.add_corr_pink_noise(c_pink=item.c_pink)

        elif isinstance(item, UncorrPinkNoise):
            data = ng.add_uncorr_pink_noise(u_pink=item.u_pink)

        elif isinstance(item, AcnNoise):
            data = ng.add_acn_noise(acn=item.acn)

        elif isinstance(item, PCAZeroNoise):
            data = ng.add_pca_zero_noise(pca0_amp=item.pca0_amp)

        else:
            raise TypeError(f"Unknown item: {item!r} !")

        data_2d = ng.format_result(data)  # type: np.ndarray
        if data_2d.any():
            if window_mode == "FULL":
                final_data_2d += data_2d
            elif window_mode == "WINDOW":
                window_y_end = window_y_start + window_y_size  # type: int
                window_x_end = window_x_start + window_x_size  # type: int

                final_data_2d[
                    window_y_start:window_y_end, window_x_start:window_x_end
                ] += data_2d

    return final_data_2d


# TODO: copyright
@temporary_random_state
def nghxrg(
    detector: CMOS,
    noise: t.Sequence[
        t.Mapping[
            Literal[
                "ktc_bias_noise",
                "white_read_noise",
                "corr_pink_noise",
                "uncorr_pink_noise",
                "acn_noise",
                "pca_zero_noise",
            ],
            t.Mapping[str, float],
        ]
    ],
    window_position: t.Optional[t.Tuple[int, int]] = None,
    window_size: t.Optional[t.Tuple[int, int]] = None,
    seed: t.Optional[int] = None,  # This parameter is used by '@temporary_random_state'
) -> None:
    """Generate fourier noise power spectrum on HXRG detector.

    For more information see :cite:p:`2015:rauscher`.

    Parameters
    ----------
    detector: Detector
    noise: list
    window_position: t.Sequence, optional
        [x0 (columns), y0 (rows)].
    window_size: t.Sequence, optional
        [x (columns), y (rows)].
    seed: int, optional
    """
    # Converter
    params = []  # type: t.List[NoiseType]
    for item in noise:
        if "ktc_bias_noise" in item:
            sub_item = item["ktc_bias_noise"]  # type: t.Mapping[str, float]
            param = KTCBiasNoise(
                ktc_noise=sub_item["ktc_noise"],
                bias_offset=sub_item["bias_offset"],
                bias_amp=sub_item["bias_amp"],
            )  # type: NoiseType

        elif "white_read_noise" in item:
            sub_item = item["white_read_noise"]
            param = WhiteReadNoise(
                rd_noise=sub_item["rd_noise"],
                ref_pixel_noise_ratio=sub_item["ref_pixel_noise_ratio"],
            )

        elif "corr_pink_noise" in item:
            sub_item = item["corr_pink_noise"]
            param = CorrPinkNoise(c_pink=sub_item["c_pink"])

        elif "uncorr_pink_noise" in item:
            sub_item = item["uncorr_pink_noise"]
            param = UncorrPinkNoise(u_pink=sub_item["u_pink"])

        elif "acn_noise" in item:
            sub_item = item["acn_noise"]
            param = AcnNoise(acn=sub_item["acn"])

        elif "pca_zero_noise" in item:
            sub_item = item["pca_zero_noise"]
            param = PCAZeroNoise(pca0_amp=sub_item["pca0_amp"])

        else:
            raise KeyError(f"Unknown key: item = {item!r}")

        params.append(param)

    logging.getLogger("nghxrg").setLevel(logging.WARNING)

    # Prepare the parameters
    geo = detector.geometry  # type: CMOSGeometry

    if window_position is None:
        window_position = (0, 0)
    if window_size is None:
        window_size = (geo.row, geo.col)

    if detector.is_dynamic:
        time_step = int(detector.time / detector.time_step)  # type: int
    else:
        time_step = 1

    # Compute new pixels
    result_2d = compute_nghxrg(
        pixel_2d=detector.pixel.array,
        noise=params,
        detector_shape=(geo.row, geo.col),
        window_pos=window_position,
        window_size=window_size,
        num_outputs=geo.n_output,
        time_step=time_step,
        num_rows_overhead=geo.n_row_overhead,
        num_frames_overhead=geo.n_frame_overhead,
        reverse_scan_direction=geo.reverse_scan_direction,
        reference_pixel_border_width=geo.reference_pixel_border_width,
    )  # type: np.ndarray

    # Add the pixels
    detector.pixel.array = detector.pixel.array + result_2d


# TODO: This generates plot. It should be in class `Output`
# def display_noisepsd(
#     array: np.ndarray,
#     nb_output: int,
#     dimensions: t.Tuple[int, int],
#     noise_name: str,
#     path: Path,
#     step: int,
#     # mode: Literal["plot"] = "plot",
# ) -> t.Tuple[t.Any, np.ndarray]:
#     """Display noise PSD from the generated FITS file.
#
#     For power spectra, need to be careful of readout directions.
#
#     For the periodogram, using Welch's method to smooth the PSD
#     > Divide data in N segment of length nperseg which overlaps at nperseg/2
#     (noverlap argument)
#     nperseg high means less averaging for the PSD but more points.
#     """
#     # Conversion gain
#     conversion_gain = 1
#     data_corr = array * conversion_gain
#
#     if nb_output <= 1:
#         raise ValueError("Parameter 'nb_output' must be >= 1.")
#
#     # Should be in the simulated data header
#     dimension = dimensions[0]  # type: int
#     pix_p_output = dimension ** 2 / nb_output  # Number of pixels per output
#     nbcols_p_channel = dimension / nb_output  # Number of columns per channel
#     nperseg = int(pix_p_output / 10.0)  # Length of segments for Welch's method
#     read_freq = 100000  # Frame rate [Hz]
#
#     # Initializing table of nb_outputs periodogram +1 for dimensions
#     pxx_outputs = np.zeros((nb_output, int(nperseg / 2) + 1))
#
#     # For each output
#     for i in np.arange(nb_output):
#         # If i odd, flatten data since its the good reading direction
#         if i % 2 == 0:
#             start_x_idx = int(i * nbcols_p_channel)
#             end_x_idx = int((i + 1) * nbcols_p_channel)
#
#             output_data = data_corr[:, start_x_idx:end_x_idx].flatten()
#         # Else, flip it left/right and then flatten it
#         else:
#             start_x_idx = int(i * nbcols_p_channel)
#             end_x_idx = int((i + 1) * nbcols_p_channel)
#
#             output_data = np.fliplr(data_corr[:, start_x_idx:end_x_idx])
#             output_data.flatten()
#         # output data without flipping
#         # output_data = data_corr[:,int(i*(nbcols_p_channel)):
#         #                         int((i+1)*(nbcols_p_channel))].flatten()
#
#         # print(output_data, read_freq, nperseg)
#         # Add periodogram to the previously initialized array
#         f_vect, pxx_outputs[i] = signal.welch(output_data, read_freq, nperseg=nperseg)
#
#     # For the detector
#     # detector_data = data_corr.flatten()
#     # Add periodogram to the previously initialized array
#     # test, Pxx_detector = signal.welch(detector_data,
#     #                                   read_freq,
#     #                                   nperseg=nperseg)
#
#     # if mode == "plot":
#     #     fig, ax1 = plt.subplots(1, 1, figsize=(8, 8))
#     #     fig.canvas.set_window_title("Power Spectral Density")
#     #     fig.suptitle(
#     #         noise_name
#     #         + " Power Spectral Density\n"
#     #         + "Welch seg. length / Nb pixel output: "
#     #         + str("{:1.2f}".format(nperseg / pix_p_output))
#     #     )
#     #     ax1.plot(f_vect, np.mean(pxx_outputs, axis=0), ".-", ms=3, alpha=0.5, zorder=32)
#     #     for _, ax in enumerate([ax1]):
#     #         ax.set_xlim([1, 1e5])
#     #         ax.set_xscale("log")
#     #         ax.set_yscale("log")
#     #         ax.set_xlabel("Frequency [Hz]")
#     #         ax.set_ylabel("PSD [e-${}^2$/Hz]")
#     #         ax.grid(True, alpha=0.4)
#     #
#     #     filename = path.joinpath(
#     #         "nghxrg_" + noise_name + "_" + str(step) + ".png"
#     #     )  # type: Path
#     #     plt.savefig(filename, dpi=300)
#     #     plt.close()
#
#     result = np.asarray(np.mean(pxx_outputs, axis=0))  # type: np.ndarray
#     return f_vect, result
