"""
    DataUtils.XrLoader.py
"""
from glob import glob
import numpy as np

def load_xr(folder_path):
    """
    Load X-ray scattering data from a folder containing .dat files.

    Parameters
    ----------
    folder_path : str
        Path to the folder containing .dat files.

    Returns
    -------
    xr_array : np.ndarray
        3D array containing the X-ray scattering data.

    datafiles : list of str
        List of data file paths corresponding to the loaded data.

    Notes
    -----
    The function assumes that each .dat file contains data in a format compatible with np.loadtxt.
    The first dimension corresponds to the number of files, the second to the number of points, and the third to the data columns.
    """
    input_list = []
    datafiles = []
    for path in sorted(glob(folder_path + "/*.dat")):
        try:
            input_list.append(np.loadtxt(path))
            datafiles.append(path)
        except Exception as e:
            print(f"Error loading {path}: {e}")
    try:
        xr_array = np.array(input_list)
    except ValueError as e:
        print(f"Error converting input list to array: {e}")
        from molass_legacy.SerialAnalyzer.SerialDataUtils import convert_to_the_least_shape
        xr_array = np.array(convert_to_the_least_shape(input_list)[0])
        print(f"Converted to least shape array with shape {xr_array.shape}")
    except Exception:
        raise
    return xr_array, datafiles

def xr_remove_bubbles(xr_array, logger=None, debug=False):
    """Remove bubbles from the XR data array.

    Parameters
    ----------
    xr_array : np.ndarray
        3D array containing the X-ray scattering data.
    logger : logging.Logger, optional
        Logger for logging messages. If None, print to console.
    debug : bool, optional
        If True, enable debug mode with additional output.
    """
    from molass.DataObjects.Curve import create_icurve
    from molass.DataUtils.AnomalyHandlers import bubble_check_impl, remove_bubbles_impl
    qv = xr_array[0,:,0]
    xrM = xr_array[:,:,1].T
    x = np.arange(xrM.shape[1])
    icurve = create_icurve(x, xrM, qv, pickvalue=0.02)
    bubbles = bubble_check_impl(icurve.y)

    if debug:
        from copy import deepcopy
        print(f"bubbles: {bubbles}")
        icurve_orig = deepcopy(icurve)

    if len(bubbles) > 0:
        excluded_set = set()
        remove_bubbles_impl(xr_array, bubbles, excluded_set)
        if logger is None:
            print("bubbles have been removed at %s" % bubbles)
        else:
            logger.warning("bubbles have been removed at %s", bubbles)

    if debug:
        import matplotlib.pyplot as plt
        xrM_ = xr_array[:,:,1].T
        icurve_ = create_icurve(x, xrM_, qv, pickvalue=0.02)
        fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(12, 5))
        ax1.plot(*icurve_orig.get_xy())
        ax2.plot(*icurve_.get_xy())
        plt.show()

def load_xr_with_options(folder_path, remove_bubbles=False, logger=None, debug=False):
    """Load X-ray scattering data from a folder with options to preprocess.

    Parameters
    ----------
    folder_path : str
        Path to the folder containing .dat files.
    remove_bubbles : bool, optional
        If True, remove bubbles from the data, by default False.
    logger : logging.Logger, optional
        Logger for logging messages. If None, print to console.
    debug : bool, optional
        If True, enable debug mode with additional output, by default False.
        
    Returns
    -------
    xr_array : np.ndarray
        3D array containing the X-ray scattering data.
    datafiles : list of str
        List of data file paths corresponding to the loaded data.
    """
    xr_array, datafiles = load_xr(folder_path)
    if remove_bubbles:
        xr_remove_bubbles(xr_array, logger=logger, debug=debug)
    return xr_array, datafiles