import numpy as np
from creme import shuffle
from tqdm import tqdm
import operator


############################################################################################
# TSS Context Dependence Test
############################################################################################

def context_dependence_test(model, x, tile_pos, num_shuffle, mean=True, drop_wt=False):
    """
    This test embeds a sequence pattern bounded by start and end in shuffled
    background contexts -- in line with a global importance analysis.
    The background contexts are generated by dinucleotide shuffling the original sequence.

    Parameters
    ----------
        model : keras.Model 
            A keras model.
        x : np.array
            Single one-hot sequence shape (L, A).
        tile_pos : list
            List with start index and end index of pattern-of-interest along L (i.e. [start, end]).
        num_shuffle : int
            Number of shuffles to apply and average over.
        mean : bool
            If True, return the mean predictions across shuffles, otherwise return full predictions.
        drop_wt : bool
            If true, do not run predictions on the WT sequence. Use this to avoid the computational
            cost if predictions are already available.

    Returns
    -------
        np.array : prediction of wild type sequence.
        np.array : prediction of mutant sequences.
    """

    # get wild-type prediction
    if drop_wt:
        pred_wt = [None]
    else:
        pred_wt = model.predict(x[np.newaxis])

    # crop pattern of interest
    start, end = tile_pos
    x_pattern = x[start:end, :]

    # loop over shuffles
    pred_mut = []
    for n in range(num_shuffle):
        x_mut = shuffle.dinuc_shuffle(x)
        x_mut[start:end, :] = x_pattern
        pred_mut.append(model.predict(x_mut)[np.newaxis][0])
    pred_mut = np.concatenate(pred_mut, axis=0)

    if mean:
        return pred_wt[0], np.mean(pred_mut, axis=0), np.std(pred_mut, axis=0)
    else:
        return pred_wt, pred_mut


############################################################################################
# TSS Context Swap Test
############################################################################################

def context_swap_test(model, x_source, x_target, tile_pos):
    """
    This test places a source sequence pattern bounded by start and end in a 
    target sequence context at the same position -- in line with a global importance analysis.

    inputs:
        model: keras model 
            A keras model.
        x_source(s) : np.array
            Source sequence (one-hot with shape (L, A) or (N, L, A) from which a pattern will be taken.
        x_target(s) : np.array
            Target sequence with shape (L, A) or (N, L, A) that will inherit a source pattern.
        tile_pos : list
            List of start and end index of pattern along L.
        mean : bool
            If True, return the mean predictions across shuffles, otherwise return full predictions.

    Returns
    -------
        np.array : prediction of wild type sequence.
        np.array : prediction of mutant sequences.
    """
    # Fix the shapes of sequences for input into the model if one 1 sequence is given (of shape (L, A)
    if len(x_target.shape) == 2:
        x_target = np.expand_dims(x_target, axis=0)
    if len(x_source.shape) == 2:
        x_source = np.expand_dims(x_source, axis=0)
    # get start and end coordinates
    start, end = tile_pos

    # place source pattern in target sequence
    x_mut = np.copy(x_target)
    x_mut[:, start:end, :] = x_source[:, start:end, :]

    # predict mutant sequence
    pred_mut = model.predict(x_mut)

    return pred_mut


############################################################################################
# CRE Necessity Test
############################################################################################

def generate_tile_shuffles(x, tile_set, num_shuffle):
    """
    inputs:
        x: np.array
            Source sequenc (one-hot with shape (L, A) from which a pattern will be taken.
        tile_set: list
            List of start and end positions.
        num_shuffle: int
            Number of shuffles to apply and average over.

    Returns
    -------
        Mutant sequences with shuffled tile(s).
    """
    seq_mut = np.empty(([num_shuffle] + list(x.shape)))
    for n in range(num_shuffle):
        x_mut = np.copy(x)  # start a new mutant
        for pos in tile_set:  # tile set can include more than one start
            start, end = pos

            # shuffle tile
            x_mut[start:end, :] = shuffle.dinuc_shuffle(x_mut[start:end, :])

            # save mutant
            seq_mut[n] = x_mut
    return seq_mut


def necessity_test(model, x, tiles, num_shuffle, mean=True, return_seqs=False):
    """
    This test systematically measures how tile shuffles affects model predictions. 

    Parameters
    ----------
        model : keras.Model 
            A keras model.
        x : np.array
            Single one-hot sequence shape (L, A).
        tiles : list
            List of tile positions (start, end) to shuffle (i.e. [[start1, end1], [start2, end2],...]).
        num_shuffle : int
            Number of shuffles to apply and average over.
        mean : bool
            If True, return the mean predictions across shuffles, otherwise return full predictions.
        return_seqs : bool
            If True, return generated sequences for future use.

    Returns
    -------
        list : WT sequence prediction, mean and standard deviation of mutant predictions (with shuffled tile) or
        all mutant predictions without averaging and (optionally) return generated sequences.
    """

    # get wild-type prediction
    pred_wt = model.predict(x[np.newaxis])

    # loop over shuffle positions list
    pred_mut = []
    all_muts = np.empty((len(tiles), num_shuffle, x.shape[-2], x.shape[-1]))
    for tile_i, pos in tqdm(enumerate(tiles), total=len(tiles)):
        start, end = pos

        # loop over number of shuffles
        pred_shuffle = []

        for n in range(num_shuffle):
            x_mut = np.copy(x)

            # shuffle tile
            x_mut[start:end, :] = shuffle.dinuc_shuffle(x_mut[start:end, :])
            all_muts[tile_i, n, :, :] = x_mut
            # predict mutated sequence
            pred_shuffle.append(model.predict(x_mut[np.newaxis])[0])
        pred_mut.append(pred_shuffle)
    pred_mut = np.array(pred_mut)

    if mean:
        test_res = [pred_wt, np.mean(pred_mut, axis=1), np.std(pred_mut, axis=1)]
    else:
        test_res = [pred_wt, pred_mut]
    if return_seqs:
        test_res.append(all_muts)
    return test_res


############################################################################################
# CRE Sufficiency Test
############################################################################################

def sufficiency_test(model, x, tss_tile, tiles, num_shuffle, tile_seq=None, mean=True, return_seqs=False):
    """
    This test measures if a region of the sequence together with the TSS tile is sufficient to get model
    predictions same as in the WT case.

    Parameters
    ----------
        model : keras.Model 
            A keras model.
        x : np.array
            Single one-hot sequence shape (L, A).
        tss_tile : list
            List of the tss_tile position to embed in shuffled sequences, i.e. [start, end].
        tiles : list
            List of tile positions (start, end) to embed in shuffled sequences.
            (i.e. [[start1, end1], [start2, end2],...]).
        num_shuffle : int
            Number of dinuc shuffles to apply to sequence context and average over.
        tile_seq : numpy
            Sequence of the TSS to be embedded. If provided this overrules the TSS tile coordinates.
        mean : bool
            If True, return the mean predictions across shuffles, otherwise return full predictions.
        return_seqs : bool
            If True, return the generated mutant sequences.

    Returns
    -------
        list of numpy arrays. Depending on arguments returns either the WT prediction, mean and standard deviation of
        mutant sequence (dinuc shuffled sequence with TSS and tile) predictions, mean and standard deviation of
        control sequences (dinuc shuffled sequence with TSS only) or all the predictions without averaging and
        (optionally) the constructed mutant sequences.
    """

    # get wild-type prediction
    pred_wt = model.predict(x[np.newaxis])

    # loop over number of shuffles
    pred_mut = []
    pred_control = []
    for pos in tiles:
        start, end = pos

        pred_mut_shuffle = []
        pred_control_shuffle = []
        sequences = np.empty((num_shuffle, model.seq_length, 4))
        for n in range(num_shuffle):
            x_mut = shuffle.dinuc_shuffle(x)

            # embed tss tile
            x_mut[tss_tile[0]:tss_tile[1], :] = x[tss_tile[0]:tss_tile[1], :]
            sequences[n] = x_mut.copy()
            # predict shuffled context with just TSS
            pred_control_shuffle.append(model.predict(x_mut[np.newaxis])[0])

            # embed tile of interest in
            if tile_seq:
                x_mut[start:end, :] = tile_seq
            else:
                x_mut[start:end, :] = x[start:end, :]

            # predict mutated sequence
            pred_mut_shuffle.append(model.predict(x_mut[np.newaxis])[0])

        # store results
        pred_mut.append(pred_mut_shuffle)
        pred_control.append(pred_control_shuffle)

    pred_mut = np.array(pred_mut)
    pred_control = np.array(pred_control)

    if mean:
        test_res = [pred_wt[0], np.mean(pred_mut, axis=1), np.std(pred_mut, axis=1), np.mean(pred_control, axis=1),
                    np.std(pred_control, axis=1)]
    else:
        test_res = [pred_wt, pred_mut, pred_control]
    if return_seqs:
        test_res.append(sequences)
    return test_res


############################################################################################
# TSS-CRE Distance Test
############################################################################################

def distance_test(model, x, tile_fixed_coord, tile_var_coord, test_positions, num_shuffle, mean=True, seed=False):
    """
    This test maps out the distance dependence of tile1 (anchored) and tile 2 (variable position).
    Tiles are placed in dinuc shuffled background contexts, in line with global importance analysis. 

    Parameters
    ----------
        model : keras.Model 
            A keras model.
        x : np.array
            Single one-hot sequence shape (L, A).
        tile_fixed_coord : list
            List with start index and end index of tile that is anchored (i.e. [start, end]).
        tile_var_coord : list
            List with start index and end index of tile that is to be tested.
        test_positions : list
            List with start index of positions to test tile_var.
        num_shuffle : int
            Number of shuffles to apply and average over.
        mean : bool
            If True, return the mean predictions across shuffles, otherwise return full predictions.
        seed: bool
            If Ture, set a seed for the random dinuc shuffle of sequence and use the same background sequences
            for all position tests (per sequence).

    Returns
    -------
        dict: results organized as dictionary containing control (i.e. sequence with TSS and tile in original position)
        and mutant (variable tile location) predictions (either summarized as mean and standard deviation or all the
        predictions).

    """

    # crop pattern of interest
    x_tile_fixed = x[tile_fixed_coord[0]:tile_fixed_coord[1], :]  # fixed tile sequence
    x_tile_var = x[tile_var_coord[0]:tile_var_coord[1], :]  # variable position tile sequence

    # get sufficiency of tiles in original positions
    pred_control = []
    for n in range(num_shuffle):
        # shuffle sequence and place tiles in respective positions
        if seed:
            x_mut = shuffle.dinuc_shuffle(x, seed=n)
        else:
            x_mut = shuffle.dinuc_shuffle(x)
        x_mut[tile_fixed_coord[0]:tile_fixed_coord[1], :] = x_tile_fixed
        x_mut[tile_var_coord[0]:tile_var_coord[1], :] = x_tile_var

        # predict mutant sequence
        pred_control.append(model.predict(x_mut[np.newaxis])[0])
    pred_control = np.array(pred_control)

    # loop over embedding tile_var in available position list
    pred_mut = []
    tile_len = tile_var_coord[1] - tile_var_coord[0]
    for start in tqdm(test_positions):

        # loop over number of shuffles
        pred_shuffle = []
        for n in range(num_shuffle):

            # shuffle sequence
            if seed:
                x_mut = shuffle.dinuc_shuffle(x, seed=n)
            else:
                x_mut = shuffle.dinuc_shuffle(x)

            # place tile 1 in original location
            x_mut[tile_fixed_coord[0]:tile_fixed_coord[1], :] = x_tile_fixed

            # place tile 2 in new position
            x_mut[start:start + tile_len, :] = x_tile_var

            # predict mutant sequence
            pred_shuffle.append(model.predict(x_mut[np.newaxis])[0])
        pred_mut.append(pred_shuffle)
    pred_mut = np.array(pred_mut)

    if mean:
        res = {"mean_control": np.mean(pred_control, axis=0), "std_control": np.std(pred_control, axis=0),
               "mean_mut": np.mean(pred_mut, axis=1), "std_mut": np.std(pred_mut, axis=1)}
    else:
        res = {'control': pred_control, 'mut': pred_mut}
    return res


############################################################################################
# CRE Higher-order Interaction Test
############################################################################################


def higher_order_interaction_test(model, x, cre_tiles_to_test, optimization, num_shuffle=10, num_rounds=None):
    """
    This test performs a greedy search to identify which tile sets lead to optimal changes
    in model predictions. In each round, a new tile is identified, given the previous sets 
    of tiles by shuffling each tile and selecting the tile with biggest effect (similar to
    necessity test)

    Parameters
    ----------
        model : keras.Model 
            A keras model.
        x : np.array
            Single one-hot sequence shape (L, A).
        cre_tiles_to_test : list
            List with tile coordinates to be tested, each with a list that consists of start index and end index.
        optimization : np.argmax or np.argmin
            Function that identifies/selects tile index for each round of greedy search.
        num_shuffle : int
            Number of shuffles to apply and average over.
        num_rounds : int
            Number of rounds to perform greedy search.

    Returns
    -------

        dictionary with keys as iteration number, values as another dictionary with results for that iteration.
        These include: initial predictions for that iteration (in iteration 0 this is WT, in iteration 2 this
        is for a sequence with 2 tiles shuffled already); predictions for newly generated mutants; selected tile
        based on predictions of this iteration; per tile mean of shuffles for the selected best tile


    """

    result_summary = {}
    if not num_rounds:
        num_rounds = len(cre_tiles_to_test)

    for iteration_i in tqdm(range(num_rounds)):
        result_summary[iteration_i] = {}
        # run one sweep of tile shuffles and keep shuffled seqs
        pred_wt, pred_mut, all_muts = necessity_test(model, x, cre_tiles_to_test, num_shuffle, False, True)

        # get per tile predictions (average across bins)
        per_tile_preds = pred_mut[..., 0].mean(-1)  # [tile number, shuffle n]

        result_summary[iteration_i]['initial_pred'] = pred_wt.mean()  # keep track of initial seq prediction
        result_summary[iteration_i]['preds'] = per_tile_preds  # save all tile preds for comparing to hypothetical model
        per_tile_mean = per_tile_preds.mean(axis=-1)  # average across shuffles

        # find optimal tile
        selected_tile_i = optimization(per_tile_mean)  # find best tile index
        best_tile_preds = per_tile_preds[selected_tile_i, :]  # all shuffle outputs for best tile
        keep_shuffled = cre_tiles_to_test[selected_tile_i]  # tile coords of the best tile
        result_summary[iteration_i]['selected_tile'] = keep_shuffled  # keep record
        cre_tiles_to_test.remove(keep_shuffled)  # remove this for next iteration

        selected_mean_pred = per_tile_mean[selected_tile_i]  # select the best tile prediction for trace
        result_summary[iteration_i]['selected_mean_pred'] = selected_mean_pred
        # update seq for next iteration selecting sequence yielding the closest prediction to mean
        x = all_muts[selected_tile_i, np.argmin(np.abs(best_tile_preds - selected_mean_pred)), :, :]
    return result_summary


############################################################################################
# CRE Multiplicity Test
############################################################################################
def multiplicity_test(model, x, tss_tile_coord, cre_tile_coord, cre_tile_seq, test_coords, num_shuffle, num_copies,
                      optimization):
    """
    Parameters
    ----------
        model : keras.Model
            A keras model.
        x : np.array
            Single one-hot sequence shape (L, A).
        tss_tile_coord : list
            Start and end coordinates of the TSS tile (which is fixed from the beginning)
        cre_tile_coord : list
            Start and end coordinates for where to insert the CRE as a control.
        cre_tile_seq : np.array
            Single one-hot sequence of the CRE shape (L, A) where L equals the length of the CRE.
        test_coords : np.array
            Tile start positions to test. In iteration 0 all the positions in the array will be tested and the one
            with the most optimal prediction will be selected (and removed from the set of positions for subsequent
            iterations).
        num_shuffle : int
            Number of shuffles to apply and average over.
        num_copies : int
            Number of copies to insert, i.e. iterations to run.
        optimization : np.argmax or np.argmin
            Function that identifies tile index for each round of greedy search.

    Returns
    ----------
    dict: results organized as dictionary containing: (i) TSS activity on its own (in dinuc shuffled backgrounds),
    (ii) TSS and CRE activity when CRE is positioned at the specified position,
    (iii) list of the most optimal prediction in each iteration showing the steps of the optimization process
    (iv) list of tile positions that were selected in each iteration, (v) list all mutant predictions -
    for each iteration the predictions for each of the tested positions.

    """
    shuffled_seqs = shuffle.dinuc_shuffle(x, num_shuffle)  # create backgroun seqs
    # re-insert TSS sequence
    shuffled_seqs[:, tss_tile_coord[0]:tss_tile_coord[1], :] = x[tss_tile_coord[0]:tss_tile_coord[1], :].copy()
    # get TSS only predictions
    only_tss_pred = model.predict(shuffled_seqs).mean()
    # get predictions for when CRE is inserted in specified position
    tss_and_cre = shuffled_seqs.copy()
    tss_and_cre[:, cre_tile_coord[0]: cre_tile_coord[1]] = cre_tile_seq
    tss_and_cre_pred = model.predict(tss_and_cre).mean()

    tile_positions_to_test = test_coords.copy()
    current_seq_version = shuffled_seqs.copy()  # start with the TSS only sequence in the first round
    # lists for saving predictions for mutants, optimal signal and selected tiles
    all_mutants = []
    best_tss_signal = []
    selected_tile_order = []
    for _ in tqdm(range(num_copies)):  # per iteration
        # placeholder array for sequences of shape [N, T, L, A] where N = shuffles to average over, T = tiles to test
        test_seqs = np.empty((num_shuffle, len(tile_positions_to_test), model.seq_length, 4))
        # placeholder for saving mutant predictions
        mutant_preds = np.empty((num_shuffle, len(tile_positions_to_test)))
        for s, shuffled_seq in enumerate(current_seq_version): # per shuffled background sequence
            for t, (tile_start, tile_end) in enumerate(tile_positions_to_test): # per tile position remaining to test
                test_seq = shuffled_seq.copy()
                test_seq[tile_start: tile_end] = cre_tile_seq.copy() # embed CRE in test position
                mutant_preds[s, t] = model.predict(test_seq).mean() # prediction for current mutant
                test_seqs[s, t, ...] = test_seq.copy()
        # pick the optimal tile position index based on prediction means across shuffles
        best_index = optimization(mutant_preds.mean(axis=0))
        selected_tile = tile_positions_to_test[best_index]  # pick the tile coordinate
        # save the mean TSS signal at optimal tile position
        best_tss_signal.append(mutant_preds.mean(axis=0)[best_index])
        tile_positions_to_test.remove(selected_tile)  # remove selected tile from set to test
        selected_tile_order.append(selected_tile)  # save the tile coordinate selected
        all_mutants.append(mutant_preds)  # save the mutant predictions
        # re-set the starting sequences to the selected mutant which has the CRE sequence embedded at the newly
        # selected best position (and all the previous ones)
        current_seq_version = test_seqs[:, best_index, ...].copy()
    return {'only_tss_pred': only_tss_pred, 'tss_and_cre_pred': tss_and_cre_pred, 'best_tss_signal': best_tss_signal,
            'selected_tile_order': selected_tile_order, 'all_mutants': all_mutants}


########################################################################################
# Pruning function
########################################################################################


def prune_sequence(model, wt_seq, control_sequences, mut, whole_tile_start, whole_tile_end, scales, thresholds, frac,
                   N_batches, cre_type='enhancer'):
    """
    This function prunes a tile through greedy search to find the most enhancing subset of sub-tiles, explaining a
    set fraction of the original enhancement. It's done in stages where sub-tiles of a specified scale are shuffled,
    keeping the least enhancing N sub-tiles (N defined by N_batches). The TSS activity with only the remaining
    sub-tiles is computed and compared to the case when the entire tile is inserted, using a ratio, i.e. the 'score'.
    If the score is above a set threshold, the optimization continues with more iterations. When the threshold is
    reached, the last step is reverted, and a new stage can begin if defined.

    Parameters
    ----------
        model : keras.Model
            A keras model.
        wt_seq : np.array
            Single one-hot sequence shape (L, A).
        control_sequences : np.array
            One-hot background sequences of shape (N, L, A).
        mut : float
            Prediction when only TSS and the entire CRE are embedded in background sequences. This is used
            to compute fraction restored by a subset of tile sequences embedded.
        whole_tile_start : int
            Start coordinate of the CRE to prune.
        whole_tile_end : int
            End coordinate of the CRE to prune.
        scales : list
            Window sizes to use for sub-tiles. Each stage of pruning can have a different window size, e.g.
            500bp in the first and 50bp in the second, to speed up the optimization.
        thresholds : list
            Score thresholds to use to determine when to stop a given optimization.
        frac : float
            Fraction of scale or window size to use to compute step size.
        N_batches : list
            List of integer batch sizes to use in each stage of the optimization. Batch size determines the
            number of tiles that are pruned out in each iteration. For example, batch size of 1 means that only 1 tile
            (if searching for enhancers then the most silencing tile) will be pruned. This parameter allows to prune
            discontinuous patches of sequences.
        cre_type : string
            'enhancer' or 'silencer' - defines the optimization type, ie. to either prune the least enhancing or
            least silencing elements.

    Returns
    ----------
    dict: returns a dictionary with a summary of results for each iteration. The information of each stage is saved as
    window size (key) and corresponding dictionary of 'scores' - the fraction tile activity recovered, bps - number of
    bps embedded, 'all_removed_tiles' - np.array of all the removed sub-tiles, 'insert_coords' - set of
    remaining/surviving sub-tiles.

    """
    remove_tiles = [] #  set of sub-tiles that are pruned out

    # save what to put back from wt sequence in form of coordinates
    insert_coords = [[whole_tile_start, whole_tile_end]]
    pruned_seqs = control_sequences.copy()  # pruned sequences with optimized set of sub-tiles preserved
    bps = np.zeros((whole_tile_end - whole_tile_start))  # number of bps embedded
    result_summary = {}
    # for each stage of optimization using a set window size for sub-tiles, threshold for stopping the current stage
    # and N_batch number of sub-tiles to remove
    for (window, threshold, N_batch) in zip(scales, thresholds, N_batches):
        result_summary[window] = {'scores': [], 'bps': []} # create a new entry in the output dictionary
        print(f"Tile size = {window}, threshold = {threshold}")

        step = int(window * frac)  # determine step size as fraction of window size
        test_coords = []

        # for each coordinate to be re-inserted
        for insert_coord in insert_coords:
            # recover WT sequence
            pruned_seqs[:, insert_coord[0]: insert_coord[1], :] = wt_seq[insert_coord[0]: insert_coord[1]].copy()
            bps[insert_coord[0] - whole_tile_start: insert_coord[1] - whole_tile_start] = 1  # count added bps
            # add sub-tile chunks for the current interval
            test_coords += [[s, s + window] for s in list(range(insert_coord[0], insert_coord[1] - step + 1, step))]

        # fraction recovered with all the sub-tiles re-inserted. Note, in the first stage this is the entire CRE.
        score = model.predict(pruned_seqs).mean() / mut
        print(f"Starting score: {score}")

        test_coords = np.array(test_coords)
        print(len(test_coords))

        final_check_seq = pruned_seqs.copy()
        all_removed_tiles = np.array([[], []]).T

        print("Starting optimization...")
        # select optimization type
        if cre_type == 'enhancer':
            comp = operator.gt
        elif cre_type == 'silencer':
            comp = operator.lt
        # continue pruning while threshold of score is not crossed
        while comp(score, threshold) and len(test_coords):

            print(f'score = {score}')

            pruned_seqs = final_check_seq.copy()  # save removed seq tiles
            # remove test coordinates if already pruned
            new_test_coords = []
            for test_coord in test_coords:
                if test_coord not in all_removed_tiles:
                    new_test_coords.append(test_coord)
            test_coords = new_test_coords
            print(f"Number of tiles to test: {len(test_coords)}")
            results = []
            for test_coord in test_coords:
                # shuffle subtile and get TSS activity
                test_seqs = pruned_seqs.copy()  # don't change pruned_seqs yet
                test_seqs[:, test_coord[0]: test_coord[1], :] = control_sequences[:, test_coord[0]: test_coord[1],
                                                                :].copy()
                results.append(model.predict(test_seqs).mean())

            # select the batch of the least enhancing or least silencing sub-tiles
            if cre_type == 'enhancer':  # prune out silencers, ie. tiles that when shuffled lead to higher pred
                remove_tiles = np.array(test_coords)[np.argsort(results)[-N_batch:]]  # choose N useless
            elif cre_type == 'silencer':  # remove enhancers = tiles the shuffling of which leads to drop in TSS
                remove_tiles = np.array(test_coords)[np.argsort(results)[:N_batch]]  # choose N useless

            all_removed_tiles = np.concatenate([all_removed_tiles, remove_tiles])  # add to the list of pruned sub-tiles

            # final check needed if batch size > 1
            final_check_seq = pruned_seqs.copy()
            for tile in remove_tiles:
                final_check_seq[:, tile[0]: tile[1], :] = control_sequences[:, tile[0]: tile[1],
                                                          :].copy()  # prune out selected tiles
                bps[tile[0] - whole_tile_start: tile[1] - whole_tile_start] = 0
            # TSS activitiy with pruned sub-tiles / TSS activity with entire CRE
            score = model.predict(final_check_seq).mean() / mut
            print(f"Number of tiles at the end of iteration: {len(test_coords)}, score = {score}, bps = {bps.sum()}")
            result_summary[window]['scores'].append(score)  # save score
            result_summary[window]['bps'].append(bps.sum())  # save bps

        result_summary[window]['all_removed_tiles'] = all_removed_tiles # save the pruned tile list for this stage
        insert_coords = test_coords.copy()  # the remaining sub-tile coordinates
        result_summary[window]['insert_coords'] = insert_coords
    return result_summary


########################################################################################
# Normalization functions
########################################################################################


def context_effect_on_tss(pred_wt, pred_mut):
    """Normalization based on difference between the effect size of the mutation and wt divided by wt"""

    return (pred_wt - pred_mut) / pred_wt


def fold_change_over_control(pred_wt, pred_mut, bin_index):
    """Normalization based on difference between the effect size of the mutation and wt divided by wt"""
    if len(pred_mut.shape) == 1:
        return pred_mut[bin_index] / pred_wt[bin_index]
    else:
        return pred_mut[:, bin_index] / pred_wt[bin_index]


def normalized_tile_effect(pred_wt, pred_mut, pred_control, bin_index):
    """Normalization used for sufficiency test"""
    return (pred_mut[:, bin_index] - pred_control[:, bin_index]) / pred_wt[bin_index]


def reduce_pred_index(pred, bin_index):
    """Reduce multivariate prediction by selecting an index"""
    return pred[:, bin_index]


def remove_tss_tile(tiles, tile_index):
    """Remove a tile form a list of tile coordinates."""
    del tiles[tile_index]
