import copy
import os
import random
import shutil
import subprocess
import warnings

import natsort as ns
import numpy as np
import pandas as pd
import skimage.io as io
from alpineer import image_utils, io_utils, load_utils, misc_utils
from scipy.ndimage import gaussian_filter

from toffy.image_stitching import rescale_images
from toffy.json_utils import read_json_file
from toffy.streak_detection import streak_correction


def transform_compensation_json(json_path, comp_mat_path):
    """Converts the JSON file from ionpath into a compensation matrix.

    Args:
        json_path (str): path to json file
        comp_mat_path (str): path to comp matrix

    Returns:
        pd.DataTable: matrix with sources channels as rows and target channels as columns
    """
    data = read_json_file(json_path)["Data"]

    comp_mat = pd.read_csv(comp_mat_path, index_col=0)

    for i in range(len(data)):
        current_data = data[i]
        source = current_data["DonorMass"]
        target = current_data["RecipientMass"]
        val = current_data["Percent"] / 100
        comp_mat.loc[source, str(target)] = val

    comp_mat.to_csv(comp_mat_path)


def _compensate_matrix_simple(raw_inputs, comp_coeffs, out_indices):
    """Perform compensation on the raw data using the supplied compensation values.

    Args:
        raw_inputs (numpy.ndarray):
            array with shape [fovs, rows, cols, channels] containing the image data
        comp_coeffs (numpy.ndarray):
            2D array of coefficients with source channels as rows and targets as columns
        out_indices (numpy.ndarray):
            which indices to generate compensated outputs for

    Returns:
        numpy.ndarray: compensated copy of the raw inputs
    """
    outputs = np.copy(raw_inputs)

    # loop over each channel and construct compensation values
    for chan in out_indices:
        chan_coeffs = comp_coeffs[:, chan]

        # convert from 1D to 4D for broadcasting
        chan_coeffs = np.reshape(chan_coeffs, (1, 1, 1, len(chan_coeffs)))

        # broadcast across entire dataset and collapse into single set of values
        chan_vals = raw_inputs * chan_coeffs
        chan_vals = np.sum(chan_vals, axis=-1)

        # subtract compensated values from target channel
        outputs[..., chan] -= chan_vals

    # set negative values to zero
    outputs = np.where(outputs > 0, outputs, 0)

    # subset on specified indices
    outputs = outputs[..., out_indices]

    return outputs


def validate_inputs(
    raw_data_dir,
    comp_mat,
    acquired_masses,
    acquired_targets,
    input_masses,
    output_masses,
    all_masses,
    fovs,
    save_format,
    raw_data_sub_folder,
    batch_size,
    gaus_rad,
):
    """Helper function to validate inputs for compensate_image_data.

    Args:
        raw_data_dir (str): path to raw data
        comp_mat (pd.DataFrame): compensation matrix
        acquired_masses (list): masses in the supplied panel
        acquired_targets (list): targets in the supplied panel
        input_masses (list): masses to use for compensation
        output_masses (list): masses to compensate
        all_masses (list): masses in the compensation matrix
        fovs (list): fovs in the raw_data_dir
        save_format (str): format to save the data
        raw_data_sub_folder (string): sub-folder for raw images
        batch_size (int): number of images to process concurrently
        gaus_rad (int): radius for smoothing
    """
    # make sure panel is in increasing order
    if not np.all(acquired_masses == sorted(acquired_masses)):
        raise ValueError("Masses must be sorted numerically in the panel_info file")

    # make sure channels in comp matrix are same as those in panel csv
    misc_utils.verify_same_elements(acquired_masses=acquired_masses, compensation_masses=all_masses)

    # check first FOV to make sure all channels are present
    test_data = load_utils.load_imgs_from_tree(
        data_dir=raw_data_dir, fovs=fovs[0:1], img_sub_folder=raw_data_sub_folder
    )

    misc_utils.verify_in_list(
        listed_channels=acquired_targets, image_files=test_data.channels.values
    )

    # make sure supplied masses are present
    if input_masses is not None:
        misc_utils.verify_in_list(input_masses=input_masses, compensation_masses=all_masses)

    if output_masses is not None:
        misc_utils.verify_in_list(output_masses=output_masses, compensation_masses=all_masses)

    # make sure compensation matrix has valid values
    if comp_mat.isna().values.any():
        raise ValueError(
            "Compensation matrix must contain a value for every field; check to "
            "make sure there are no missing values"
        )

    # check for valid save_formats
    allowed_formats = ["raw", "rescaled", "both"]
    misc_utils.verify_in_list(save_format=save_format, allowed_formats=allowed_formats)

    if not isinstance(batch_size, int) or batch_size < 1:
        raise ValueError("batch_size parameter must be a positive integer")

    if not isinstance(gaus_rad, int) or gaus_rad < 0:
        raise ValueError("gaus_rad parameter must be a non-negative integer")


def clean_rosetta_test_dir(folder_path):
    """Remove the unnecessary intermediate folders created by rosetta test data computation.

    Args:
        folder_path (str): base dir for testing, image subdirs will be stored here
    """
    # remove any files beginning with ._, needed to ensure external drive hidden files clear
    _ = subprocess.call(["find", folder_path, "-type", "f", "-name", "._*", "-delete"])

    # remove the compensated data folders
    comp_folders = io_utils.list_folders(folder_path, substrs="compensated_data_")
    for cf in comp_folders:
        shutil.rmtree(os.path.join(folder_path, cf))

    # remove the stitched image folder
    shutil.rmtree(os.path.join(folder_path, "stitched_images"))


def combine_compensation_files(comp_matrix_path, compensation_matrix_names, final_matrix_name):
    """Combine a list of round two compensation matrix files in a given cohort folder.
    This is done additively since round two compensation files are mutually exclusive w.r.t.
    output channels.

    Args:
        comp_matrix_path (str):
            Path to the compensation matrix files to combine
        compensation_matrix_names (list):
            List of files inside `cohort_folder_path` to combine
        final_matrix_name (str):
            Where to write the combined compensation matrix to
    """
    # load in the first matrix inside compensation_matrix_names
    final_compensation_matrix = pd.read_csv(
        os.path.join(comp_matrix_path, compensation_matrix_names[0])
    )

    # loop over the rest and add them in
    # NOTE: skip the first column, since that just delineates the channels
    for matrix in compensation_matrix_names[1:]:
        final_compensation_matrix.iloc[:, 1:] = final_compensation_matrix.iloc[:, 1:].add(
            pd.read_csv(os.path.join(comp_matrix_path, matrix)).iloc[:, 1:]
        )

    # save the final compensation matrix to final_matrix_name
    final_compensation_matrix.to_csv(os.path.join(comp_matrix_path, final_matrix_name), index=False)


def flat_field_correction(img, gaus_rad=100):
    """Apply flat field correction to an image.

    Args:
        img (np.ndarray): image to be corrected
        gaus_rad (int): radius for smoothing

    Returns:
        np.ndarray: corrected image
    """
    # if image is empty, return empty image
    if not np.any(img):
        warnings.warn("Image for flatfield correction is empty")
        return img

    # smooth image
    img_smooth = gaussian_filter(img, sigma=gaus_rad)

    # calculate mean to preserve overall intensity
    img_mean = np.mean(img)

    # apply correction
    img_corr = (img / img_smooth) * img_mean

    return img_corr


def get_masses_from_channel_names(names, panel_df):
    """Get the weights for the given names from the input dataframe.

    Args:
        names (list): the channels whose masses will be returned
        panel_df (pd.DataFrame): the panel containing the masses and channel names

    Returns:
        list: the masses for the given channels
    """
    misc_utils.verify_in_list(
        supplied_channel_names=names, panel_channel_names=panel_df["Target"].values
    )

    weights = panel_df.loc[np.isin(panel_df["Target"], names)]["Mass"].values

    return weights


def compensate_image_data(
    raw_data_dir,
    comp_data_dir,
    comp_mat_path,
    panel_info,
    input_masses=None,
    output_masses=None,
    save_format="rescaled",
    raw_data_sub_folder="",
    batch_size=1,
    gaus_rad=1,
    norm_const=200,
    ffc_masses=[39],
    correct_streaks=False,
    streak_chan="Noodle",
):
    """Function to compensate MIBI data with a flow-cytometry style compensation matrix.

    Args:
        raw_data_dir: path to directory containing raw images
        comp_data_dir: path to directory where compensated images will be saved
        comp_mat_path: path to compensation matrix, nxn with channel labels
        panel_info: array with information about the panel
        input_masses (list): masses from the compensation matrix to use for compensation. If None,
            all masses will be used
        output_masses (list): masses from the compensation matrix that will have tifs generated. If
            None, all masses will be used
        save_format: flag to control how the processed tifs are saved. Must be one of:
            'raw': Direct output from the compensation matrix corresponding to number of ion events
                detected per pixel. These will not be viewable in many image processing programs
            'rescaled': all images are divided by 200 to enable visualization. This transform
                has no effect on downstream analysis as it preserves relative expression values
            'both': saves both 'raw' and 'rescaled' images
        raw_data_sub_folder (string): sub-folder for raw images
        batch_size: number of images to process at a time
        gaus_rad: radius for blurring image data. Passing 0 will result in no blurring
        norm_const: constant used for rescaling if save_format == 'rescaled'
        ffc_masses (list): masses that need to be flat field corrected.
        correct_streaks (bool): whether to correct streaks in the image
        streak_chan (str): the channel to use for streak correction
    """
    io_utils.validate_paths([raw_data_dir, comp_data_dir, comp_mat_path])

    # get list of all fovs
    # TODO: list_folders does not handle cases such as "fov0" correctly
    # need to add a fix in to alpineer to deal with this
    fovs = [f for f in os.listdir(raw_data_dir) if "fov" in f]

    # load csv files
    comp_mat = pd.read_csv(comp_mat_path, index_col=0)
    acquired_masses = panel_info["Mass"].values
    acquired_targets = panel_info["Target"].values
    all_masses = comp_mat.columns.values.astype("int")

    # convert ffc mass into ffc channel names
    if ffc_masses is not None:
        ffc_channels = [
            panel_info.loc[panel_info.Mass == mass].Target.values[0] for mass in ffc_masses
        ]
    else:
        ffc_channels = None

    validate_inputs(
        raw_data_dir,
        comp_mat,
        acquired_masses,
        acquired_targets,
        input_masses,
        output_masses,
        all_masses,
        fovs,
        save_format,
        raw_data_sub_folder,
        batch_size,
        gaus_rad,
    )

    # set unused masses to zero
    if input_masses is not None:
        zero_idx = ~np.isin(all_masses, input_masses)
        comp_mat.loc[zero_idx, :] = 0

    if output_masses is None:
        out_indices = np.arange(len(all_masses))
    else:
        out_indices = np.where(np.isin(all_masses, output_masses))[0]

    # loop over each set of FOVs in the batch
    for i in range(0, len(fovs), batch_size):
        print(f"Processing {fovs[i]}")

        # load batch of fovs
        batch_fovs = fovs[i : i + batch_size]
        batch_data = load_utils.load_imgs_from_tree(
            data_dir=raw_data_dir,
            fovs=batch_fovs,
            channels=acquired_targets,
            img_sub_folder=raw_data_sub_folder,
        )

        # convert to float32 for gaussian_filter and rosetta compatibility
        batch_data = batch_data.astype(np.float32)

        # blur data
        if gaus_rad > 0:
            for j in range(batch_data.shape[0]):
                for k in range(batch_data.shape[-1]):
                    batch_data[j, :, :, k] = gaussian_filter(batch_data[j, :, :, k], sigma=gaus_rad)

        if correct_streaks:
            corrected_channels, _ = streak_correction(
                fov_data=batch_data, streak_channel=streak_chan
            )
            batch_data[0] = corrected_channels.values

        # apply flat field correction if specified
        if ffc_channels is not None:
            misc_utils.verify_in_list(
                flat_field_correction_masses=ffc_channels, all_masses=acquired_targets
            )
            for fov in batch_fovs:
                for chan in ffc_channels:
                    raw_img = batch_data.loc[fov, :, :, chan].values
                    batch_data.loc[fov, :, :, chan] = flat_field_correction(raw_img)

        comp_data = _compensate_matrix_simple(
            raw_inputs=batch_data, comp_coeffs=comp_mat.values, out_indices=out_indices
        )

        # save data
        for j in range(batch_data.shape[0]):
            fov_folder = os.path.join(comp_data_dir, batch_data.fovs.values[j])
            os.makedirs(fov_folder)

            # create directories for saving tifs
            if save_format in ["rescaled", "both"]:
                rescale_folder = os.path.join(fov_folder, "rescaled")
                os.makedirs(rescale_folder)

            if save_format in ["raw", "both"]:
                raw_folder = os.path.join(fov_folder, "raw")
                os.makedirs(raw_folder)

            # this may be only a subset of masses, based on output_masses
            for idx, val in enumerate(out_indices):
                channel_name = batch_data.channels.values[val] + ".tiff"

                # save tifs to appropriate directories
                if save_format in ["rescaled", "both"]:
                    save_path = os.path.join(rescale_folder, channel_name)
                    image_utils.save_image(save_path, comp_data[j, :, :, idx] / norm_const)

                if save_format in ["raw", "both"]:
                    save_path = os.path.join(raw_folder, channel_name)
                    image_utils.save_image(save_path, comp_data[j, :, :, idx])


def copy_round_one_compensated_images(
    runs, round_one_comp_folder, round_two_comp_folder, channels_to_copy
):
    """Copies channels that don't need round two compensation to the round two comp folder.

    Args:
        runs (list):
            runs in `round_two_comp_folder` that need copying from `round_one_comp_folder`
        round_one_comp_folder (str):
            path to the round one Rosetta compensated images
        round_two_comp_folder (str):
            path to the round two Rosetta compensated images
        channels_to_copy (list):
            channels to copy from `round_one_comp_folder` to `round_two_comp_folder`
    """
    io_utils.validate_paths([round_one_comp_folder, round_two_comp_folder])

    # verify runs found in both round one and round two Rosetta folders
    r1_runs = io_utils.list_folders(round_one_comp_folder)
    r2_runs = io_utils.list_folders(round_two_comp_folder)
    misc_utils.verify_in_list(provided_runs=runs, round_one_comp_runs=r1_runs)
    misc_utils.verify_in_list(provided_runs=runs, round_two_comp_runs=r2_runs)

    # for each FOV, copy the channel from their r1_runs folder to r2_runs folder
    for run in runs:
        fovs = io_utils.list_folders(os.path.join(round_one_comp_folder, run), substrs="fov")

        for fov in fovs:
            channel_files = io_utils.list_files(
                os.path.join(round_one_comp_folder, run, fov, "rescaled"), substrs=channels_to_copy
            )

            for cf in channel_files:
                shutil.copy(
                    os.path.join(round_one_comp_folder, run, fov, "rescaled", cf),
                    os.path.join(round_two_comp_folder, run, fov, "rescaled", cf),
                )


def create_tiled_comparison(
    input_dir_list,
    output_dir,
    max_img_size,
    img_sub_folder="rescaled",
    channels=None,
    img_size_scale=0.5,
):
    """Creates a tiled image comparing FOVs from all supplied runs for each channel.

    Args:
        input_dir_list: list of directories to compare
        output_dir: directory where tifs will be saved
        max_img_size (int): largest fov image size
        img_sub_folder: subfolder within each input directory to load images from
        channels: list of channels to compare.
        img_size_scale (int/float): amount to scale down image, set to None for no scaling
    """
    test_dir = input_dir_list[0]
    test_fov = io_utils.list_folders(test_dir)[0]
    test_data = load_utils.load_imgs_from_tree(
        data_dir=test_dir,
        fovs=[test_fov],
        img_sub_folder=img_sub_folder,
        channels=channels,
    )

    if not channels:
        channels = test_data.channels.values

    misc_utils.verify_in_list(
        provided_channels=channels, test_data_channels=test_data.channels.values
    )

    channel_num = len(channels)

    # check that all dirs have the same number of fovs and correct subset of channels
    fov_names = io_utils.list_folders(input_dir_list[0])
    for dir_name in input_dir_list[1:]:
        current_folders = io_utils.list_folders(dir_name)
        misc_utils.verify_same_elements(fov_names1=fov_names, fov_names2=current_folders)
        current_channels = load_utils.load_imgs_from_tree(
            data_dir=dir_name, img_sub_folder=img_sub_folder, fovs=current_folders[:1]
        ).channels.values
        misc_utils.verify_in_list(specified_channels=channels, current_channels=current_channels)

    fov_num = len(fov_names)

    # loop over each channel
    for j in range(channel_num):
        # create tiled array of dirs x fovs
        tiled_image = np.zeros(
            (max_img_size * len(input_dir_list), max_img_size * fov_num),
            dtype=test_data.dtype,
        )

        # loop over each fov, and place into columns of tiled array
        for i in range(fov_num):
            start = i * max_img_size
            end = (i + 1) * max_img_size

            # go through each of the directories, read in the images, and place in the right spot
            for idx, key in enumerate(input_dir_list):
                dir_data = load_utils.load_imgs_from_tree(
                    key,
                    channels=channels[j : j + 1],
                    img_sub_folder=img_sub_folder,
                    max_image_size=max_img_size,
                )
                tiled_image[(max_img_size * idx) : (max_img_size * (idx + 1)), start:end] = (
                    dir_data.values[i, :, :, 0]
                )
        # scale image size
        if img_size_scale:
            tiled_image = rescale_images(tiled_image, img_size_scale)
        fname = os.path.join(output_dir, channels[j] + "_comparison.tiff")

        # save, compress to 3 decimal points for space optimization
        image_utils.save_image(fname, np.round(tiled_image, 3))


def add_source_channel_to_tiled_image(
    raw_img_dir,
    tiled_img_dir,
    output_dir,
    source_channel,
    max_img_size,
    img_sub_folder="",
    percent_norm=99,
    img_size_scale=0.5,
):
    """Adds the specified source_channel to the first row of previously generated tiled images.

    Args:
        raw_img_dir (str): path to directory containing the raw images
        tiled_img_dir (str): path to directory contained the tiled images
        output_dir (str): path to directory where outputs will be saved
        img_sub_folder (str): subfolder within raw_img_dir to load images from
        max_img_size (int): largest fov image size
        source_channel (str): the channel which will be prepended to the tiled images
        percent_norm (int): percentile normalization param to enable easy visualization, set to
            None to skip this step
        img_size_scale (int/float): amount to scale down image, set to None for no scaling
    """
    # load source images
    source_imgs = load_utils.load_imgs_from_tree(
        raw_img_dir,
        channels=[source_channel],
        img_sub_folder=img_sub_folder,
        max_image_size=max_img_size,
    )

    # convert stacked images to concatenated row
    source_list = [source_imgs.values[fov, :, :, 0] for fov in range(source_imgs.shape[0])]
    source_row = np.concatenate(source_list, axis=1)
    source_scaled = rescale_images(source_row, img_size_scale)

    # get percentile of source row if percent_norm set, otherwise leave unset
    perc_source = np.percentile(source_scaled, percent_norm) if percent_norm else None

    # normalize by max if percent_norm is 0 for source in case of sparse signal
    if percent_norm and perc_source == 0:
        perc_source = np.percentile(source_scaled, 100)

    # confirm tiled images have expected shape
    tiled_images = io_utils.list_files(tiled_img_dir)
    test_file = io.imread(os.path.join(tiled_img_dir, tiled_images[0]))
    if test_file.shape[1] != source_scaled.shape[1]:
        raise ValueError(
            "Tiled image {} has shape {}, but source image {} hasshape {}".format(
                tiled_images[0], test_file.shape, source_channel, source_scaled.shape
            )
        )

    # loop through each tiled image, prepend source row scaled, and save
    for tile_name in tiled_images:
        current_tile = io.imread(os.path.join(tiled_img_dir, tile_name))

        # if percent_norm set, normalize the source row to be in the same range as the current tile
        # otherwise, just leave as is (divide by 1)
        perc_ratio = 1
        if percent_norm:
            perc_tile = np.percentile(current_tile, percent_norm)

            # normalize by max if percent_norm is 0 for tile in case of sparse signal
            if perc_tile == 0:
                perc_tile = np.percentile(current_tile, 100)

            perc_ratio = perc_source / perc_tile

        rescaled_source = source_scaled / perc_ratio

        # combine together and save, compress to 3 decimal points for space optimization
        combined_tile = np.concatenate([rescaled_source, current_tile])
        save_name = tile_name.split(".tiff")[0] + "_source_" + source_channel + ".tiff"
        image_utils.save_image(os.path.join(output_dir, save_name), np.round(combined_tile, 3))


def replace_with_intensity_image(run_dir, channel="Au", replace=True, fovs=None):
    """Replaces the specified channel with the intensity image of that channel.

    Args:
        run_dir (str): directory containing extracted run data
        channel (str): the channel whose intensity image will be copied over
        fovs (list or None): the subset of fovs within run_dir which will have their
            intensity image copied over. If None, applies to all fovs
        replace (bool): controls whether intensity image is copied over with _intensity appended
            or if it will overwrite existing channel
    """
    all_fovs = io_utils.list_folders(run_dir)

    # ensure supplied folders are valid
    if fovs is not None:
        misc_utils.verify_in_list(specified_folders=fovs, all_folders=all_fovs)
        all_fovs = fovs

    # ensure channel is valid
    test_file = os.path.join(run_dir, all_fovs[0], "intensities", channel + "_intensity.tiff")
    if not os.path.exists(test_file):
        raise FileNotFoundError("Could not find specified file {}".format(test_file))

    # loop through each fov
    for fov in all_fovs:
        # control whether intensity image overwrites previous image or is copied with new name
        if replace:
            suffix = ".tiff"
        else:
            suffix = "_intensity.tiff"
        shutil.copy(
            os.path.join(run_dir, fov, "intensities", channel + "_intensity.tiff"),
            os.path.join(run_dir, fov, channel + suffix),
        )


def remove_sub_dirs(run_dir, sub_dirs, fovs=None):
    """Removes specified sub-folders from fovs in a run.

    Args:
        run_dir (str): path to directory containing fovs
        sub_dirs (list): directories to remove from each fov
        fovs (list): list of fovs to remove dirs from, otherwise removes from all fovs
    """
    all_fovs = io_utils.list_folders(run_dir)

    # ensure supplied folders are valid
    if fovs is not None:
        misc_utils.verify_in_list(specified_folders=fovs, all_folders=all_fovs)
        all_fovs = fovs

    # ensure all sub_dirs exist
    for sub_dir in sub_dirs:
        if not os.path.isdir(os.path.join(run_dir, all_fovs[0], sub_dir)):
            raise ValueError("Did not find {} in {}".format(sub_dir, all_fovs[0]))

    for fov in all_fovs:
        for sub_dir in sub_dirs:
            shutil.rmtree(os.path.join(run_dir, fov, sub_dir))


def create_rosetta_matrices(
    default_matrix,
    save_dir,
    multipliers,
    current_channel_name,
    output_channel_names,
    masses=None,
):
    """Creates a series of compensation matrices for evaluating coefficients.

    Args:
        default_matrix (str): path to the rosetta matrix to use as the default
        save_dir (str): output directory
        multipliers (list): the range of values to multiply the default matrix by
            to get new coefficients
        current_channel_name (str): channel being adjusted
        output_channel_names (list): subset of the channels to compensate for
        masses (list | None): an optional list of masses to include in the multiplication. If
            only a subset of masses are specified, other masses will retain their values
            in all iterations. If None, all masses are included
    """
    # Read input matrix
    comp_matrix = pd.read_csv(default_matrix, index_col=0)
    comp_masses = comp_matrix.index

    # Check that all entries of comp_matrix are numeric
    if not np.issubdtype(comp_matrix.values.dtype, np.number):
        raise ValueError("Compensation matrix must include only numeric entries")

    # Check channel input
    if masses is None:
        masses = comp_masses
    else:
        try:
            masses = [int(x) for x in masses]
        except ValueError:
            raise ValueError("Masses must be provided as integers")
        misc_utils.verify_in_list(specified_masses=masses, rosetta_masses=comp_masses)

    # define the file prefix used for each compensation matrix file
    default_mat_prefix = os.path.basename(default_matrix).split(".csv")[0]
    output_chan_str = "_".join(output_channel_names) if output_channel_names is not None else "all"
    comp_file_prefix = f"{current_channel_name}_{output_chan_str}_{default_mat_prefix}_mult"

    # loop over each specified multiplier and create separate compensation matrix
    for i in multipliers:
        mult_matrix = copy.deepcopy(comp_matrix)

        for j in range(len(comp_matrix)):
            # multiply specified channel by multiplier
            if comp_masses[j] in masses:
                mult_matrix.iloc[j, :] = comp_matrix.iloc[j, :] * i
        comp_name = f"{comp_file_prefix}_{i}.csv"
        mult_matrix.to_csv(os.path.join(save_dir, comp_name))


# TODO: anything with [f for f in os.listdir(...) ...] needs to be changed
# after list_folders with substrs specified is fixed
def copy_image_files(
    cohort_name, run_names, rosetta_testing_dir, extracted_imgs_dir, fovs_per_run=5
):
    """Creates a new directory for rosetta testing and copies over a random subset of
        previously extracted images.

    Args:
        cohort_name (str): name for all combined runs
        run_names (list): gives names of run folders to retrieve extracted images from
        rosetta_testing_dir (str): directory where to create cohort rosetta testing folder
        extracted_imgs_dir (str): directory containing images from each run,
        fovs_per_run: number of fovs from each run to use for testing, default 5

    """
    # path validation
    io_utils.validate_paths([rosetta_testing_dir, extracted_imgs_dir])

    # validate provided run names
    small_runs = []
    for run in run_names:
        if not os.path.exists(os.path.join(extracted_imgs_dir, run)):
            raise ValueError(f"{run} is not a valid run name found in {extracted_imgs_dir}")
        fovs_in_run = [f for f in os.listdir(os.path.join(extracted_imgs_dir, run)) if "fov" in f]
        # check number of fovs in each run
        if len(fovs_in_run) < fovs_per_run:
            small_runs.append(run)

    run_names_process = run_names[:]
    if len(small_runs) > 0:
        # edge case if none of the runs in run_names have length equal to fovs_per_run
        if len(small_runs) == len(run_names_process):
            raise ValueError(
                f"None of the runs specified contain the minimum amount of FOVs ({fovs_per_run}) "
                "defined by the fovs_per_run given"
            )

        # warn user that runs below fovs_per_run threshold will be skipped
        warnings.warn(
            "The following runs will be skipped because they do not contain the minimum amount "
            f"of FOVs ({fovs_per_run}) defined by the fovs_per_run given: {small_runs}."
        )

        run_names_process = list(set(run_names_process).difference(small_runs))

    # make rosetta testing dir and extracted images subdir
    cohort_rosetta_dir = os.path.join(rosetta_testing_dir, cohort_name)
    os.makedirs(os.path.join(cohort_rosetta_dir, "extracted_images"))

    # randomly choose fovs from a run and copy them to the img subdir in rosetta testing dir
    for i, run in enumerate(ns.natsorted(run_names_process)):
        run_path = os.path.join(extracted_imgs_dir, run)

        fovs_in_run = [f for f in os.listdir(run_path) if "fov" in f]
        fovs_in_run = ns.natsorted(fovs_in_run)
        rosetta_fovs = random.sample(fovs_in_run, k=fovs_per_run)

        for fov in rosetta_fovs:
            fov_path = os.path.join(os.path.join(extracted_imgs_dir, run, fov))
            # prepend the run name to each fov
            new_path = os.path.join(
                os.path.join(cohort_rosetta_dir, "extracted_images", run + "_" + fov)
            )
            shutil.copytree(fov_path, new_path)


def rescale_raw_imgs(img_out_dir, scale=200):
    """Rescale image data to be between 0 and 1.

    Args:
        img_out_dir (str): the directory containing extracted images
        scale (int): how much to rescale the image data by

    Returns:
        saves the rescaled images in a subdir
    """
    io_utils.validate_paths(img_out_dir)

    fovs = io_utils.list_folders(img_out_dir)
    for fov in fovs:
        fov_dir = os.path.join(img_out_dir, fov)
        # create subdirectory for the new images
        sub_dir = os.path.join(fov_dir, "rescaled")
        if not os.path.exists(sub_dir):
            os.makedirs(sub_dir)
        chans = io_utils.list_files(fov_dir)
        # rescale each channel image
        for chan in chans:
            fname = os.path.join(sub_dir, chan)
            if not os.path.exists(fname):
                img = io.imread(os.path.join(fov_dir, chan))
                img = (img / scale).astype("float32")
                image_utils.save_image(fname, img)


def generate_rosetta_test_imgs(
    rosetta_mat_path,
    img_out_dir,
    multipliers,
    folder_path,
    panel,
    current_channel_name="Noodle",
    output_channel_names=None,
    gaus_rad=1,
    norm_const=1,
    ffc_masses=[39],
):
    """Compensate example FOV images based on given multipliers.

    Args:
        rosetta_mat_path (str): path to rosetta compensation matrix
        img_out_dir (str): directory where extracted images are stored
        multipliers (list): list of coeffient multipliers to create different matrices for
        folder_path (str): base dir for testing, image subdirs will be stored here
        panel (pd.DataFrame): the panel containing the masses and channel names
        current_channel_name (str): channel being adjusted, default Noodle
        output_channel_names (list): subset of the channels to compensate for, default None is all
        gaus_rad: radius for blurring image data. Passing 0 will result in no blurring
        norm_const: constant used for rescaling
        ffc_masses (list): masses that need to be flat field corrected.

    Returns:
        Create subdirs containing rosetta compensated images for each multiplier and stitched imgs
    """
    io_utils.validate_paths([rosetta_mat_path, img_out_dir, folder_path])

    # get mass information
    current_channel_mass = get_masses_from_channel_names([current_channel_name], panel)

    if output_channel_names is not None:
        output_masses = get_masses_from_channel_names(output_channel_names, panel)
    else:
        output_masses = None

    # generate rosetta matrices for each multiplier
    create_rosetta_matrices(
        default_matrix=rosetta_mat_path,
        save_dir=folder_path,
        multipliers=multipliers,
        current_channel_name=current_channel_name,
        output_channel_names=output_channel_names,
        masses=current_channel_mass,
    )

    # define the file prefix used for each compensation matrix file
    matrix_name = io_utils.remove_file_extensions([os.path.basename(rosetta_mat_path)])[0]
    output_chan_str = "_".join(output_channel_names) if output_channel_names is not None else "all"
    comp_file_prefix = f"{current_channel_name}_{output_chan_str}_{matrix_name}_mult"

    # loop over each multiplier and compensate the data
    rosetta_dirs = [img_out_dir]
    for multiplier in multipliers:
        print(f"Processing images with multiplier {multiplier}")
        rosetta_mat_path = os.path.join(folder_path, f"{comp_file_prefix}_{multiplier}.csv")
        rosetta_out_dir = os.path.join(folder_path, "compensated_data_{}".format(multiplier))
        rosetta_dirs.append(rosetta_out_dir)
        os.makedirs(rosetta_out_dir)
        compensate_image_data(
            raw_data_dir=img_out_dir,
            comp_data_dir=rosetta_out_dir,
            comp_mat_path=rosetta_mat_path,
            raw_data_sub_folder="rescaled",
            panel_info=panel,
            batch_size=1,
            gaus_rad=gaus_rad,
            norm_const=norm_const,
            output_masses=output_masses,
            ffc_masses=ffc_masses,
        )
