import os

# ==============================
# FILE Patterns for NuMuCC e NuE
# ==============================

# -- Reconstruction
NUMUCC_FILE_RECO = "km3_v4_numuCC_{}.evt.JEnergy.rootreco_trees.npz"
NUE_FILE_RECO = "km3_v4_nueCC_{}.evt.JEnergy.rootreco_trees.npz"

# -- Doms Map
NUMUCC_FILE_DOMS_MAP = "numu_{}_sel5_doms_map.npz"
NUE_FILE_DOMS_MAP = "nue_{}_sel_5_doms_map.npz"


# -- Multi-data
NUMUCC_FILE_MULTIDATA = "Xy_numu_{}_multi_data.npz"
NUE_FILE_MULTIDATA = "Xy_nue_{}_multi_data.npz"

# -- Training/Validation/Test/Indices
TRAINING_FILE = "Xy_train{}_sel5_doms.npz"
VALIDATION_FILE = "Xy_val{}_sel5_doms.npz"
TEST_FILE = "Xy_test{}_sel5_doms.npz"
INDEX_SPLIT_FILE = "Xy_indx{}_sel5_doms.npz"


def _get_numu_nue_filepaths(dir_path, numu_filep, nue_filep, n_files=100):
    numu_filelist = [os.path.join(dir_path, numu_filep.format(i + 1))
                     for i in range(n_files)]
    nue_filelist = [os.path.join(dir_path, nue_filep.format(i + 1))
                    for i in range(n_files)]
    return numu_filelist, nue_filelist


def get_reco_files(directory_reco, merge=False, n_files=100):
    """
    Utility function to get the list of paths of reconstructed
    files for NuMucc e NuE
    
    Parameters
    ----------
    directory_reco: str
        Path to the folder containing reconstruction files
        
    merge: bool (default False)
        whether joining file-paths in one unique lists of
        tuples (couples)

    n_files: int (default 100)
        Number of Files to extract

    Returns
    -------
        reco_files_numu: list
            List of reconstruction file paths for NuMuCC of length `n_files`
        reco_files_nue: list
            List of reconstruction file paths for NuE of length `n_files`

    One single list is returned if `merge == True`.
    """
    reco_files_numu, reco_files_nue = _get_numu_nue_filepaths(directory_reco,
                                                              NUMUCC_FILE_RECO, NUE_FILE_RECO, n_files)
    if not merge:
        return reco_files_numu, reco_files_nue
    return [(reco_numu, reco_nue) for reco_numu, reco_nue in zip(reco_files_numu, reco_files_nue)]


def get_doms_maps(directory_doms_map, merge=False, n_files=100):
    """
    Utility function to get the list of paths of Doms Maps files
    for NuMucc e NuE

    Parameters
    ----------
    directory_doms_map: str
        Path to the folder containing doms maps

    merge: bool (default False)
        whether joining file-paths in one unique lists of
        tuples (couples)

    n_files: int (default 100)
        Number of Files to extract

    Returns
    -------
        doms_maps_numu: list
            List of reconstruction file paths for NuMuCC of length `n_files`
        doms_maps_nue: list
            List of reconstruction file paths for NuE of length `n_files`

    One single list is returned if `merge == True`.
    """
    doms_maps_numu, doms_maps_nue = _get_numu_nue_filepaths(directory_doms_map, NUMUCC_FILE_DOMS_MAP,
                                                            NUE_FILE_DOMS_MAP, n_files)
    if not merge:
        return doms_maps_numu, doms_maps_nue
    return [(map_numu, maps_nue) for map_numu, maps_nue in zip(doms_maps_numu, doms_maps_nue)]


def get_multi_data_files(multi_data_folder, merge=True, n_files=100):
    """
    Utility function to get the list of paths of Multi-data files
    (i.e. files containing both data and metadata) for NuMucc e NuE

    Parameters
    ----------
    multi_data_folder: str
        Path to the folder containing multidata files

    merge: bool (default True)
        whether joining file-paths in one unique lists of
        tuples (couples). True by default

    n_files: int (default 100)
        Number of Files to extract

    Returns
    -------
        multi_data_numu: list
            List of reconstruction file paths for NuMuCC of length `n_files`
        multi_data_nue: list
            List of reconstruction file paths for NuE of length `n_files`

    One single list is returned if `merge == True` (default).
    """
    multi_data_numu, multi_data_nue = _get_numu_nue_filepaths(multi_data_folder, NUMUCC_FILE_MULTIDATA,
                                                              NUE_FILE_MULTIDATA, n_files)
    if not merge:
        return multi_data_numu, multi_data_nue
    return [(multi_numu, multi_nue) for multi_numu, multi_nue in zip(multi_data_numu, multi_data_nue)]


def get_train_validation_test_files(train_test_dir, n_files=100):
    """
    Collect paths to files in training, validation, and tets sets, respectively,
    along with the corresponding index_split_files containing
    indices of events belonging in the different sets.

    Parameters
    ----------
    train_test_dir: str
        path to the folder containing files as generated by the train/test split process

    n_files: int (default 100)
        Number of Files to extract

    Returns
    -------
        fnames_train: list
            List of the data files in the Training set
        fnames_val: list
            List of the data files in the Validation set
        fnames_test: list
            List of the data files in the Test set
        index_split_filelist: list
            List of index files containing mapping indices for the events in each of the training/val/test sets.

    See Also
    --------
    export_train_test.export_train_validation_test

    """
    fnames_train = [os.path.join(train_test_dir, TRAINING_FILE.format(i + 1)) for i in range(n_files)]
    fnames_test = [os.path.join(train_test_dir, TEST_FILE.format(i + 1)) for i in range(n_files)]
    fnames_val = [os.path.join(train_test_dir, VALIDATION_FILE.format(i + 1)) for i in range(n_files)]
    index_split_filelist = [os.path.join(train_test_dir, INDEX_SPLIT_FILE.format(i + 1)) for i in range(n_files)]
    return fnames_train, fnames_val, fnames_test, index_split_filelist
