# basic command line functions to use. Some of those functions can be implemented to
# scikit-beam later.
import os
import time
import json
import glob
import multiprocessing
import numpy as np
import h5py
import copy
from collections.abc import Iterable
import re

from ..core.quant_analysis import ParamQuantitativeAnalysis

from skbeam.core.fitting.xrf_model import linear_spectrum_fitting, define_range
from .fileio import output_data, read_hdf_APS, read_MAPS, sep_v
from .fit_spectrum import single_pixel_fitting_controller, save_fitdata_to_hdf

from ..core.map_processing import dask_client_create

import logging

logger = logging.getLogger(__name__)


def fit_pixel_data_and_save(
    working_directory,
    file_name,
    *,
    param_file_name,
    fit_channel_sum=True,
    fit_channel_each=False,
    param_channel_list=None,
    incident_energy=None,
    ignore_datafile_metadata=False,
    fln_quant_calib_data=None,
    quant_distance_to_sample=0,
    method="nnls",
    pixel_bin=0,
    raise_bg=0,
    comp_elastic_combine=False,
    linear_bg=False,
    use_snip=True,
    bin_energy=0,
    save_txt=False,
    save_tiff=True,
    scaler_name=None,
    use_average=False,
    interpolate_to_uniform_grid=False,
    data_from="NSLS-II",
    dask_client=None,
):
    """
    Do fitting for signle data set, and save data accordingly. Fitting can be performed on
    either summed data or each channel data, or both.

    Parameters
    ----------
    working_directory : str, required
        path folder
    file_names : str, required
        selected h5 file
    param_file_name : str, required
        param file name for summed data fitting
    fit_channel_sum : bool, optional
        fit summed data or not
    fit_channel_each : bool, optional
        fit each channel data or not
    param_channel_list : list, optional
        list of param file names for each channel
    incident_energy : float, optional
        use this energy as incident energy instead of the one in param file, i.e., XANES
        This value overrides the incident energy from metadata and from JSON parameter file
    ignore_datafile_metadata : bool
        tells whether to ignore metadata from the data file (if data file contains metadata).
        At the moment, only incident energy (monochromator energy) is used by the processing routine.
        If True, then the incident energy from the parameter JSON file is used; if False, then
        the energy from metadata is used. If the function parameter ``incident_energy`` is
        not None, then its value overrides the incident energy from metadata or JSON file.
    fln_quant_calib_data : str or list(str)
        file name or a list of file names that contain quantitative calibration data
    quant_distance_to_sample : float
        distance-to-sample used in quantitative calibration. If 0, then correction for
        distance is not applied (assumed that the standard and the sample were placed
        at the same distance from the detector).
    method : str, optional
        fitting method, default as nnls
    pixel_bin : int, optional
        bin pixel as 2by2, or 3by3
    raise_bg : int, optional
        add a constant value to each spectrum, better for fitting
    comp_elastic_combine : bool, optional
        combine elastic and compton as one component for fitting
    linear_bg : bool, optional
        use linear background instead of snip
    use_snip : bool, optional
        use snip method to remove background
    bin_energy : int, optional
        bin spectrum with given value
    save_txt : bool, optional
        save data to txt or not
    save_tiff : bool, optional
        save data to tiff or not
    scaler_name : str, optional
        name of the field representing the scaler (for example 'i0'), which must be present
        in the data files. If given, normalization will be performed before saving data
        as .txt and .tiff
    use_average : bool, optional
        if true, norm is performed as data/scaler*mean(scaler), otherwise just data/scaler
    interpolate_to_uniform_grid : bool
        interpolate the result to uniform grid before saving to tiff and txt files
        The grid dimensions match the dimensions of positional data for X and Y axes.
        The range of axes is chosen to fit the values of X and Y.
    data_from : str, optional
        where do data come from? Data format includes data from NSLS-II, or 2IDE-APS
    dask_client: dask.distributed.Client
        Dask client object. If None, then Dask client is created automatically.
        If a batch of files is processed, then creating Dask client and
        passing the reference to it to the processing functions will save
        execution time: `client = Client(processes=True, silence_logs=logging.ERROR)`
    """
    fpath = os.path.join(working_directory, file_name)

    # Load quantitative calibration files (if necessary)
    quant_norm = False  # Indicates if at least one calibration file is loaded
    param_quant_analysis = ParamQuantitativeAnalysis()
    if fln_quant_calib_data:
        if isinstance(fln_quant_calib_data, str):
            fln_quant_calib_data = [fln_quant_calib_data]
        for fln in fln_quant_calib_data:
            if os.path.isabs(fln):
                f = fln
            else:
                f = os.path.join(working_directory, fln)
            try:
                param_quant_analysis.load_entry(f)
                quant_norm = True
                logger.info(f"Quantitative calibration is loaded successfully from file '{f}'")
            except Exception as ex:
                logger.error(f"Error occurred while loading quantitative calibration from file '{f}': {ex}")

    t0 = time.time()
    prefix_fname = file_name.split(".")[0]
    if fit_channel_sum is True:
        if data_from == "NSLS-II":
            img_dict, data_sets, mdata = read_hdf_APS(working_directory, file_name, load_each_channel=False)
        elif data_from == "2IDE-APS":
            img_dict, data_sets, mdata = read_MAPS(working_directory, file_name, channel_num=1)
        else:
            print("Unknown data sets.")

        try:
            data_all_sum = data_sets[prefix_fname + "_sum"].raw_data
        except KeyError:
            data_all_sum = data_sets[prefix_fname].raw_data

        # load param file
        if not os.path.isabs(param_file_name):
            param_path = os.path.join(working_directory, param_file_name)
        else:
            param_path = param_file_name
        with open(param_path, "r") as json_data:
            param_sum = json.load(json_data)

        # update incident energy, required for XANES
        if incident_energy is not None:
            param_sum["coherent_sct_energy"]["value"] = incident_energy
            print("Using incident beam energy passed as the function parameter.")
        elif (
            mdata.is_metadata_available()
            and "instrument_mono_incident_energy" in mdata
            and not ignore_datafile_metadata
        ):
            param_sum["coherent_sct_energy"]["value"] = mdata["instrument_mono_incident_energy"]
            print(f"Using incident beam energy from the data file '{file_name}'.")
        else:
            print(f"Using incident beam energy from the parameter file '{param_path}'.")

        # The value of incident energy that is used for processing
        incident_energy_used = param_sum["coherent_sct_energy"]["value"]
        print(f"Incident beam energy: {incident_energy_used}.")

        result_map_sum, calculation_info = single_pixel_fitting_controller(
            data_all_sum,
            param_sum,
            incident_energy=incident_energy,
            method=method,
            pixel_bin=pixel_bin,
            raise_bg=raise_bg,
            comp_elastic_combine=comp_elastic_combine,
            linear_bg=linear_bg,
            use_snip=use_snip,
            bin_energy=bin_energy,
            dask_client=dask_client,
        )

        # output to .h5 file
        inner_path = "xrfmap/detsum"
        # fit_name = prefix_fname+'_fit'
        save_fitdata_to_hdf(fpath, result_map_sum, datapath=inner_path)

        def get_scaler_set(img_dict):
            sc_set_names = [_ for _ in img_dict if _.endswith("_scaler")]
            if sc_set_names:
                return img_dict[sc_set_names[0]]
            else:
                return {}

        def get_positions_set(img_dict):
            if "positions" in img_dict:
                return img_dict["positions"]
            else:
                return {}

        scaler_dict = get_scaler_set(img_dict)
        scaler_name_list = list(scaler_dict.keys())
        positions_dict = get_positions_set(img_dict)
        # Generate dataset
        dataset = copy.deepcopy(scaler_dict)
        dataset.update(result_map_sum)

        # Set parameters for quantitative normalization
        param_quant_analysis.experiment_incident_energy = incident_energy_used
        param_quant_analysis.experiment_distance_to_sample = quant_distance_to_sample
        param_quant_analysis.experiment_detector_channel = "sum"

        if save_txt is True:
            output_folder = "output_txt_" + prefix_fname
            output_path = os.path.join(working_directory, output_folder)
            output_data(
                output_dir=output_path,
                interpolate_to_uniform_grid=interpolate_to_uniform_grid,
                dataset_name="dataset_fit",  # Sum of all detectors: should end with '_fit'
                quant_norm=quant_norm,
                param_quant_analysis=param_quant_analysis,
                distance_to_sample=quant_distance_to_sample,
                dataset_dict=dataset,
                positions_dict=positions_dict,
                file_format="txt",
                scaler_name=scaler_name,
                scaler_name_list=scaler_name_list,
                use_average=use_average,
            )

        if save_tiff is True:
            output_folder = "output_tiff_" + prefix_fname
            output_path = os.path.join(working_directory, output_folder)
            output_data(
                output_dir=output_path,
                interpolate_to_uniform_grid=interpolate_to_uniform_grid,
                dataset_name="dataset_fit",  # Sum of all detectors: should end with '_fit'
                quant_norm=quant_norm,
                param_quant_analysis=param_quant_analysis,
                dataset_dict=dataset,
                positions_dict=positions_dict,
                file_format="tiff",
                scaler_name=scaler_name,
                scaler_name_list=scaler_name_list,
                use_average=use_average,
            )

    if fit_channel_each:
        img_dict, data_sets, mdata = read_hdf_APS(working_directory, file_name, load_each_channel=True)

        # Find the detector channels and the names of the channels
        det_channels = [_ for _ in data_sets.keys() if re.search(r"_det\d+$", _)]
        det_channel_names = [re.search(r"det\d+$", _)[0] for _ in det_channels]
        if param_channel_list is None:
            param_channel_list = [param_file_name] * 3
        else:
            if not isinstance(param_channel_list, list) and not isinstance(param_channel_list, tuple):
                raise RuntimeError("Parameter 'param_channel_list' must be a list or a tuple of strings")
            if len(param_channel_list) != len(det_channels):
                raise RuntimeError(
                    f"Parameter 'param_channel_list' must be 'None' or contain {len(det_channels)} file names."
                )

        channel_num = len(param_channel_list)
        for i in range(channel_num):
            inner_path = "xrfmap/" + det_channel_names[i]
            print(f"Processing data from detector channel {det_channel_names[i]} (#{i+1}) ...")

            # load param file
            param_file_name = param_channel_list[i]

            # load param file
            if not os.path.isabs(param_file_name):
                param_path = os.path.join(working_directory, param_file_name)
            else:
                param_path = param_file_name
            with open(param_path, "r") as json_data:
                param_det = json.load(json_data)

            # update incident energy, required for XANES
            if incident_energy is not None:
                param_det["coherent_sct_energy"]["value"] = incident_energy
                print("Using incident beam energy passed as the function parameter.")
            elif (
                mdata.is_metadata_available()
                and "instrument_mono_incident_energy" in mdata
                and not ignore_datafile_metadata
            ):
                param_det["coherent_sct_energy"]["value"] = mdata["instrument_mono_incident_energy"]
                print(f"Using incident beam energy from the data file '{file_name}'.")
            else:
                print(f"Using incident beam energy from the parameter file '{param_path}'.")

            # The value of incident energy that is used for processing
            incident_energy_used = param_det["coherent_sct_energy"]["value"]
            print(f"Incident beam energy: {incident_energy_used}.")

            data_all_det = data_sets[det_channels[i]].raw_data

            result_map_det, calculation_info = single_pixel_fitting_controller(
                data_all_det,
                param_det,
                incident_energy=incident_energy,
                method=method,
                pixel_bin=pixel_bin,
                raise_bg=raise_bg,
                comp_elastic_combine=comp_elastic_combine,
                linear_bg=linear_bg,
                use_snip=use_snip,
                bin_energy=bin_energy,
                dask_client=dask_client,
            )

            # output to .h5 file
            save_fitdata_to_hdf(fpath, result_map_det, datapath=inner_path)

            def get_scaler_set(img_dict):
                sc_set_names = [_ for _ in img_dict if _.endswith("_scaler")]
                if sc_set_names:
                    return img_dict[sc_set_names[0]]
                else:
                    return {}

            def get_positions_set(img_dict):
                if "positions" in img_dict:
                    return img_dict["positions"]
                else:
                    return {}

            scaler_dict = get_scaler_set(img_dict)
            scaler_name_list = list(scaler_dict.keys())
            positions_dict = get_positions_set(img_dict)
            # Generate dataset
            dataset = copy.deepcopy(scaler_dict)
            dataset.update(result_map_det)

            # Set parameters for quantitative normalization
            param_quant_analysis.experiment_incident_energy = incident_energy_used
            param_quant_analysis.experiment_distance_to_sample = quant_distance_to_sample
            param_quant_analysis.experiment_detector_channel = det_channel_names[i]

            if save_txt is True:
                output_folder = "output_txt_" + prefix_fname
                output_path = os.path.join(working_directory, output_folder)
                output_data(
                    output_dir=output_path,
                    interpolate_to_uniform_grid=interpolate_to_uniform_grid,
                    dataset_name=f"dataset_{det_channel_names[i]}_fit",  # ..._det1_fit, etc.
                    quant_norm=quant_norm,
                    param_quant_analysis=param_quant_analysis,
                    dataset_dict=dataset,
                    positions_dict=positions_dict,
                    file_format="txt",
                    scaler_name=scaler_name,
                    scaler_name_list=scaler_name_list,
                    use_average=use_average,
                )

            if save_tiff is True:
                output_folder = "output_tiff_" + prefix_fname
                output_path = os.path.join(working_directory, output_folder)
                output_data(