# from curses import set_tabsize
from pathlib import Path
from pdb import set_trace

import numpy as np
import xarray as xr

from gfatpy.lidar.file_manager import search_dc

# from gfatpy.lidar.preprocessing.lidar_gluing_bravo_aranda import gluing
from gfatpy.lidar.preprocessing.gluing_proportional import gluing
from gfatpy.lidar.utils import LIDAR_INFO
from gfatpy.utils import calibration
from gfatpy.utils.io import unzip_file_to_temp

# 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,
    apply_ov: bool = False,
    save_dc: bool = False,
    save_bg: bool = False,
    gluing_products: bool = False,
    force_dc_in_session=False,
    dir: Path | None = None,
    overlap_path: Path | str | 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 linc. 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.
        apply_ov (bool, optional): it allows the overlap correction. Defaults to False.
        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.
        dir (Path, None): Directory where lidar folder tree starts.
        overlap_path (Path, str, None): Path to the overlap file. Defaults to None. If None and apply_ov is TRUE, the overlap file is searched in the dir/<lidar_name>/QA/overlap/.
    Returns:
        xr.Dataset: xarray.Dataset with the lidar data and information.
    """

    _p = Path(file_or_path)
    if not _p.exists():
        raise ValueError(f"{file_or_path} not found")

    with xr.open_dataset(_p, chunks={}) as _nc:
        dataset = _nc
    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,
        )
    else:
        dataset.attrs["dc_corrected"] = str(False)

    if apply_dt:
        dataset = apply_dead_time_correction(
            dataset
        )  # TODO implement search_dt() in the same way of search_dc
    else:
        dataset.attrs["dt_corrected"] = str(False)

    if apply_bg:
        dataset = apply_background_correction(
            dataset, background_ranges=background_ranges, save_bg=save_bg
        )
    else:
        dataset.attrs["bg_corrected"] = str(False)

    if apply_bz:
        dataset = apply_bin_zero_correction(dataset, rs_path=_p)
    else:
        dataset.attrs["bz_corrected"] = str(False)

    if apply_ov:
        if overlap_path is None:
            raise
        elif isinstance(overlap_path, str):
            overlap_path = Path(overlap_path)
        dataset = apply_overlap_correction(dataset, overlap_path)
        dataset.attrs["ov_corrected"] = "True"
    else:
        dataset.attrs["ov_corrected"] = "False"

    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)
    channels = dataset.channel.values

    analog_channels: list[str] = list(filter(lambda c: c.endswith("a"), channels))

    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}"

            if dc[signal_str].sel(range=slice(0, 5000)).std("range").std("time") > 0.07:
                raise Warning(
                    f"Dark current std is too high for channel {channel} in {dc_path}. Check if DC was measured correctly."
                )

            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"] = str(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

    lidar_name = dataset.attrs["system"].lower()
    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[str] = list(
        filter(lambda c: c.endswith("p"), dataset.channel.values)
    )

    for channel in photocounting_channels:
        # tau from ns to us
        try:
            tau_us = dt_dict[channel] * 1e-3
        except KeyError:
            raise ValueError(  # TODO
                f"No dead time value defined in LIDAR_INFO->{lidar_name}->{channel}."
            )

        signal_str = f"signal_{channel}"

        # Eq 4 [D'Amico et al., 2016]
        # No infinites nor negative values
        # condition = np.logical_and(~np.isinf(dataset[signal_str]), dataset[signal_str] > 0)
        # dataset[signal_str] = dataset[signal_str].where(condition)/ ( 1 - dataset[signal_str].where(condition) * tau_us )

        dataset[signal_str] = dataset[signal_str] / (1 - dataset[signal_str] * tau_us)
        dataset.attrs["dt_corrected"] = str(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]
    )
    channels: list[str] = dataset.channel.values
    for channel in channels:
        signal_str = f"signal_{channel}"
        try:
            background = dataset[signal_str].loc[:, ranges_between].mean(axis=1)
        except:
            background = np.ones(1)
        dataset[signal_str] -= background

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

    dataset.attrs["bg_corrected"] = str(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"].lower()
        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}.")

    channels: list[str] = dataset.channel.values
    for channel in channels:
        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"] = str(True)

    return dataset


def apply_overlap_correction(
    dataset: xr.Dataset,
    ff_overlap_path: Path | str | None = None,
    nf_overlap_path: Path | str | None = None,
) -> xr.Dataset:
    """_summary_

    Args:
        dataset (xr.Dataset): _description_
        overlap_path (Path): _description_

    Returns:
        xr.Dataset: _description_
    """

    def apply_correction(
        dataset: xr.Dataset,
        overlap: xr.DataArray,
        channels2correct: list[str],
        overlap_channel: str,
    ) -> xr.Dataset:
        try:
            for channel_ in channels2correct:
                if channel_ in dataset.channel.values:
                    signal_str = f"signal_{channel_}"

                    # Perform overlap correction
                    dataset[signal_str] = dataset[signal_str] / overlap
                    dataset[signal_str].attrs["overlap_applied"] = overlap_channel
                    dataset["overlap_corrected"][dataset.channel.values == channel_] = 1
        except:
            raise Warning("Could not apply overlap correction")
        return dataset

    # New array to store if overlap correction was applied
    overlap_corrected = xr.DataArray(
        np.zeros(dataset.channel.size),
        dims=("channel"),
        coords={"channel": dataset.channel},
    )
    dataset["overlap_corrected"] = overlap_corrected

    # Apply overlap correction for far-field channels
    if ff_overlap_path is not None:
        linked_channels = {
            "355fpa": [
                "355fpa",
                "355fpp",
                "355fsa",
                "355fsp",
                "354fta",
                "354ftp",
                "408fta",
                "408ftp",
            ],
            "532fta": ["532fta", "532ftp", "531fta", "531ftp"],
            "1064fta": ["1064fta", "1061fta", "1064ftp"],
        }

        # Load overlap dataset
        overlap = xr.open_dataarray(ff_overlap_path)

        for channel_ in overlap.channel.values:
            # Select overlap array for channel
            overlap_ = overlap.sel(channel=channel_)

            # Check size of arrays
            if dataset.range.size != overlap_.range.size:
                # Adapt size of overlap array to dataset
                overlap_ = overlap_.interp(range=dataset.range, method="linear")

            # Apply overlap correction
            dataset = apply_correction(
                dataset, overlap_, linked_channels[channel_], overlap_channel=channel_
            )

            # Save overlap array
            if dataset["overlap_corrected"].values.any():
                dataset[f"overlap_{channel_}"] = overlap_

    # Apply overlap correction for far-field channels
    if nf_overlap_path is not None:
        linked_channels = {
            "355npa": [
                "355npa",
                "355npp",
                "355nsa",
                "355nsp",
                "387nta",
                "387ntp",
                "408nta",
                "408ntp",
            ],
            "532npa": ["532npa", "532npp", "607nta", "607ntp"],
        }
        # Load overlap dataset
        overlap = xr.open_dataarray(nf_overlap_path)

        for channel_ in overlap.channel.values:
            # Select overlap array for channel
            overlap_ = overlap.sel(channel=channel_)

            # Check size of arrays
            if dataset.range.size != overlap_.range.size:
                # Adapt size of overlap array to dataset
                overlap_ = overlap_.interp(range=dataset.range, method="linear")

            # Apply overlap correction
            dataset = apply_correction(
                dataset, overlap_, linked_channels[channel_], overlap_channel=channel_
            )

            # Save overlap array
            if dataset["overlap_corrected"].values.any():
                dataset[f"overlap_{channel_}"] = overlap_

    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])
