# from curses import set_tabsize
from pathlib import Path

import numpy as np
import xarray as xr

from gfatpy.lidar.file_manager import filename2path, search_dc

# from gfatpy.lidar.preprocessing.lidar_gluing_bravo_aranda import gluing
from gfatpy.lidar.preprocessing.gluing_de_la_rosa import gluing
from gfatpy.lidar.utils import LIDAR_INFO
from gfatpy.utils import calibration

# from .lidar_merge import apply_polarization_merge
# from .lidar_preprocessing_tools import *  # TODO: Remove wildcard importations within modules

# warnings.filterwarnings("ignore")

__author__ = "Bravo-Aranda, Juan Antonio"
__license__ = "GPL"
__version__ = "1.0"
__maintainer__ = "Bravo-Aranda, Juan Antonio"
__email__ = "jabravo@ugr.es"
__status__ = "Production"


""" DEFAULT AUXILIAR INFO
"""
# Root Directory (in NASGFAT)  according to operative system
""" LIDAR PREPROCESSING
"""


def preprocess(
    file_or_path: Path | str,
    channels: list[str] | None = None,
    crop_ranges: tuple[float, float] | None = (0, 20000),
    background_ranges: tuple[float, float] | None = None,
    # TODO: Use stage flags dc, dt, bz. Include add corrected in Dataset flag for False in each
    apply_dc: bool = True,
    apply_dt: bool = True,
    apply_bg: bool = True,
    apply_bz: bool = True,
    save_dc: bool = False,
    save_bg: bool = False,
    gluing_products: bool = False,
    force_dc_in_session=False,
    root_dir: Path | None = None,
) -> xr.Dataset:
    """Can preprocess a local NC file or search for it in NAS. Can apply optionally DC currection,
    DT correction, BG correction and calculate gluing products

    Args:
        file_or_path (Path | str): A relative (of where is executing) or absolute path to an nc file generated by the raw2l1. If not exists will try to search it in NAS
        channels (list[str] | None, optional): List of string channels, if not provided will return all available channels. Defaults to None.
        crop_ranges (tuple[float, float] | None, optional): outpur range. Defaults to (0, 20000).
        background_ranges (tuple[float, float] | None, optional): Range where background correction will be performed. Defaults to None.
        apply_dc (bool, optional): it allows the dark measurement correction. Defaults to True.
        apply_dt (bool, optional): it allows the dead time correction. Defaults to True.
        apply_bg (bool, optional): it allows the background correction. Defaults to True.
        apply_bz (bool, optional): it allows the zero bien correction. Defaults to True.
        save_dc (bool, optional): it saves the dark measurement in the xr.Dataset. Defaults to False.
        save_bg (bool, optional): it saves the background in the xr.Dataset. Defaults to False.
        gluing_products (bool, optional): it performs the detection-mode gluing. Defaults to False.
        force_dc_in_session (bool, optional): it forces to use the dark measurement linked to this measurement period. If not found, the correction cannot be performed. Defaults to False.
        root_dir (Path, None): Directory where lidar folder tree starts. Defaults to None. If None, DATA_DN is used —this points to the GFAT NAS—.
    Returns:
        xr.Dataset: xarray.Dataset with the lidar data and information.
    """

    _p = Path(file_or_path)
    if not _p.exists():
        print("File not found. Searching in DATA_DN.")
        filename = _p.name
        _p = filename2path(filename)
        if not _p.exists:
            raise ValueError("File not found.")

    dataset = xr.load_dataset(_p)
    dataset = drop_unwanted_channels(dataset, channels=channels)

    if apply_dc:
        dataset = apply_dark_current_correction(
            dataset,
            rs_path=_p,
            save_dc=save_dc,
            force_dc_in_session=force_dc_in_session,
        )

    if apply_dt:
        dataset = apply_dead_time_correction(
            dataset
        )  # TODO implement search_dt() in the same way of search_dc

    if apply_bg:
        dataset = apply_background_correction(
            dataset, background_ranges=background_ranges, save_bg=save_bg
        )

    if apply_bz:
        dataset = apply_bin_zero_correction(dataset, rs_path=_p)

    dataset = apply_crop_ranges_correction(dataset, crop_ranges=crop_ranges)

    if gluing_products:
        dataset = apply_detection_mode_merge(dataset)

    return dataset


def drop_unwanted_channels(
    dataset: xr.Dataset, channels: list[str] | None
) -> xr.Dataset:
    if channels is None:
        return dataset

    channel_names = list(filter(lambda c: c not in channels, dataset.channel.values))

    dataset = dataset.drop_vars(
        map(lambda c: f"signal_{c}", channel_names), errors="raise"
    )

    dataset = dataset.drop_sel(channel=channel_names, errors="raise")

    return dataset


def apply_dark_current_correction(
    dataset: xr.Dataset, rs_path: Path, save_dc: bool = False, force_dc_in_session=False
) -> xr.Dataset:
    groups = calibration.split_continous_measurements(dataset.time.values)
    analog_channels = list(filter(lambda c: c.endswith("a"), dataset.channel.values))

    for group in groups:
        dc_path = search_dc(
            rs_path,
            session_period=group[[0, -1]],
            force_dc_in_session=force_dc_in_session,
        )
        dc = xr.open_dataset(dc_path)

        lower_idx = np.where(dataset.time == group[0])[0][0]
        upper_idx = np.where(dataset.time == group[-1])[0][0] + 1

        for channel in analog_channels:
            signal_str = f"signal_{channel}"
            dc_mean = dc[signal_str].values.mean(0)
            dataset[signal_str][lower_idx:upper_idx] -= dc_mean

            if save_dc:
                if (
                    f"dc_{channel}" not in list(dataset.variables.keys())
                    and lower_idx == 0
                ):
                    dataset[f"dc_{channel}"] = dataset[signal_str] * np.nan

                dataset[f"dc_{channel}"][lower_idx:upper_idx] = dc_mean

    dataset.attrs["dc_corrected"] = True

    return dataset


def apply_dead_time_correction(dataset: xr.Dataset) -> xr.Dataset:
    # dt_path = search_dt(rs_path, session_period=dataset.time.values[[0,-1]]) #
    # dt_dict = open_dataset(dt_path) #TODO

    dt_dict = None

    if dt_dict is None:
        lidar_name = dataset.attrs["system"].upper()
        try:
            dt_dict = {
                key: value["dead_time_ns"]
                for (key, value) in LIDAR_INFO["lidars"][lidar_name]["channels"].items()
                if value.get("dead_time_ns", False)
            }
        except Exception:
            raise ValueError(
                f"No dead time value defined in LIDAR_INFO->{lidar_name}]."
            )

    photocounting_channels = list(
        filter(lambda c: c.endswith("p"), dataset.channel.values)
    )

    for channel in photocounting_channels:
        # tau from ns to us
        tau_us = dt_dict[channel] * 1e-3

        signal_str = f"signal_{channel}"

        # Eq 4 [D'Amico et al., 2016]
        dataset[signal_str] = c = dataset[signal_str] / (
            1 - dataset[signal_str] * tau_us
        )

        # No infinites nor negative values
        c = np.where(np.logical_or(np.isinf(c), c < 0), np.nan, c)

        dataset.attrs["dt_corrected"] = True

    return dataset


def apply_background_correction(
    dataset: xr.Dataset,
    background_ranges: tuple[float, float] | None = None,
    save_bg: bool = False,
) -> xr.Dataset:

    if background_ranges is None:
        background_ranges = (
            dataset.attrs["BCK_MIN_ALT"],
            dataset.attrs["BCK_MAX_ALT"],
        )

    if background_ranges[1] <= background_ranges[0]:
        raise ValueError("background_ranges should be in order (min, max)")

    ranges_between = (background_ranges[0] < dataset.range) & (
        dataset.range < background_ranges[1]
    )

    for channel in dataset.channel:
        signal_str = f"signal_{channel}"
        background = dataset[signal_str].loc[:, ranges_between].mean(axis=1)
        dataset[signal_str] -= background

        if save_bg:
            dataset[f"bg_{channel}"] = background

    dataset.attrs["bg_corrected"] = True

    return dataset


def apply_bin_zero_correction(dataset: xr.Dataset, rs_path: Path) -> xr.Dataset:
    # bz_path = search_bz(rs_path, session_period=dataset.time.values[[0,-1]]) #
    # bz_dict = open_dataset(bz_path) #TODO
    bz_dict = None

    if bz_dict is None:
        lidar_name = dataset.attrs["system"].upper()
        try:
            bz_dict = {
                key: value["bin_zero"]
                for (key, value) in LIDAR_INFO["lidars"][lidar_name]["channels"].items()
            }
        except Exception:
            raise ValueError(f"No bin zero value defined in LIDAR_INFO->{lidar_name}].")

    for channel in dataset.channel.values:
        signal_str = f"signal_{channel}"
        # set_trace()
        dataset[signal_str] = dataset[signal_str].shift(
            range=-bz_dict[channel], fill_value=0.0
        )

    dataset.attrs["bz_corrected"] = True

    return dataset


def apply_crop_ranges_correction(
    dataset: xr.Dataset, crop_ranges: tuple[float, float] | None = (0, 20000)
) -> xr.Dataset:
    # TODO: Apply crop ranges. With dataset.sel(range=slice(*crop_ranges))

    if crop_ranges is None:
        return dataset

    if crop_ranges[0] > crop_ranges[-1]:
        raise ValueError("crop_ranges should be in order (min, max)")

    dataset = dataset.sel(range=slice(*crop_ranges))

    return dataset


def apply_detection_mode_merge(dataset: xr.Dataset):
    LIDAR_INFO["metadata"]["code_telescope_str2number"]
    LIDAR_INFO["metadata"]["code_mode_str2number"]
    code_polarization_str2number = LIDAR_INFO["metadata"][
        "code_polarization_str2number"
    ]

    channels_pc: list[str] = list(
        filter(lambda c: c.endswith("p"), dataset.channel.values)
    )

    # wavelength, telescope, polarization, mode, channel_name = [], [], [], [], []
    range_m = dataset["range"].values

    glued_list: list[dict] = []

    glued_list: list[dict] = []

    for channel_pc in channels_pc:
        channel_an = f"{channel_pc[0:-1]}a"

        if f"signal_{channel_an}" not in list(dataset.variables.keys()):
            continue

        # rcs_an = signal_to_rcs(dataset[f"signal_{channel_an}"], range_m)  # type: ignore
        # rcs_pc = signal_to_rcs(dataset[f"signal_{channel_pc}"], range_m)  # type: ignore

        # if type(rcs_an) != xr.DataArray or type(rcs_pc) != xr.DataArray:
        #     raise TypeError("RCS must be xarray.DataArray")

        signal_gl = gluing(
            dataset[f"signal_{channel_an}"], dataset[f"signal_{channel_pc}"]
        )

        channel_sel = dataset.sel(channel=channel_an)

        glued_list.append(
            {
                "name": f"{channel_pc[0:-1]}g",
                "signal": signal_gl,
                "polarization": channel_sel.polarization,
                "telescope": channel_sel.telescope,
                "wavelength": channel_sel.telescope,
            }
        )

    glued_signal = {
        f"signal_{glued['name']}": (["time", "range"], glued["signal"])
        for glued in glued_list
    }

    other_var = {
        "wavelength": (["channel"], [g["wavelength"] for g in glued_list]),
        "polarization": (["channel"], [g["polarization"] for g in glued_list]),
        "telescope": (["channel"], [g["telescope"] for g in glued_list]),
        "bin_shift": (["channel"], [0 for _ in glued_list]),
    }

    glued_dataset = xr.Dataset(
        glued_signal | other_var,
        coords={
            "range": range_m,
            "time": dataset["time"].values,
            "channel": list(map(lambda i: i["name"], glued_list)),
        },
    )

    return xr.merge([dataset, glued_dataset])
