#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri May 10 10:55:51 2019

@author: mishugeb
"""
import numpy as np
from numpy import linalg as LA
from scipy.optimize import nnls
from scipy.optimize import minimize
from SigProfilerAssignment import decompose_subroutines as sub
import scipy.stats
import copy
import os

#############################################################################################################
#################################### Functions For Single Sample Algorithms #############################


################################################################### FUNCTION ONE ###################################################################
# function to calculate the cosine similarity
def cos_sim(a, b):
    """Takes 2 vectors a, b and returns the cosine similarity according
    to the definition of the dot product

    Dependencies:
    *Requires numpy library.
    *Does not require any custom function (constructed by me)

    Required by:
    * pairwise_cluster_raw
    """
    if np.sum(a) == 0 or np.sum(b) == 0:
        return 0.0
    dot_product = np.dot(a, b)
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    return dot_product / (norm_a * norm_b)


def add_connected_sigs(background_sigs, allsigids):
    connected_sigs = [
        ["SBS2", "SBS13"],
        ["SBS7a", "SBS7b", "SBS7c", "SBS7d"],
        ["SBS10a", "SBS10b"],
        ["SBS17a", "SBS17b"],
    ]

    backround_sig_names = sub.get_items_from_index(allsigids, background_sigs)
    connect_the_sigs = []
    for i in range(len(connected_sigs)):
        if len(set(connected_sigs[i]).intersection(set(backround_sig_names))) > 0:
            connect_the_sigs = connect_the_sigs + connected_sigs[i]

    backround_sig_names = list(set(backround_sig_names).union(set(connect_the_sigs)))
    background_sigs = sub.get_indeces(allsigids, backround_sig_names)
    background_sigs.sort()
    return background_sigs


def parameterized_objective2_custom(x, signatures, samples):
    rec = np.dot(signatures, x)
    try:
        y = LA.norm(samples - rec[:, np.newaxis])
    except:
        y = LA.norm(samples - rec)
    return y


def constraints1(x, samples):
    sumOfSamples = np.sum(samples, axis=0)
    # print(sumOfSamples)
    # print(x)
    result = sumOfSamples - (np.sum(x))
    # print (result)
    return result


def create_bounds(idxOfZeros, samples, numOfSignatures):
    total = np.sum(samples)
    b = (0.0, float(total))
    lst = [b] * numOfSignatures

    for i in idxOfZeros:
        lst[i] = (0.0, 0.0)

    return lst


# required for removing signatures with background
def get_changed_background_sig_idx(exposures, background_sigs):
    background_sigs_values = sub.get_items_from_index(exposures, background_sigs)
    temp_exposures = exposures[:]
    temp_exposures[:] = (value for value in temp_exposures if value != 0)

    # remove the background signatures with zero values
    background_sigs[:] = (value for value in background_sigs_values if value != 0)

    # get the new indices of the background signatures
    background_sigs = sub.get_indeces(temp_exposures, background_sigs_values)

    return background_sigs


def roundConserveSum(x):

    # Total to be conserved
    total = np.round(np.sum(x))

    # Round to nearest integer above
    x_out = np.ceil(x)

    # Order integer residuals
    order = np.argsort(x - x_out)

    # Take one off first n residuals to correct total
    x_out[order[: int(np.sum(x_out) - total + 1e-10)]] -= 1

    return x_out


# Fit signatures
def fit_signatures(W, genome, metric="l2"):
    genome = np.array(genome)
    maxmutation = round(np.sum(genome))

    # initialize the guess
    x0 = np.random.rand(W.shape[1], 1) * maxmutation
    x0 = x0 / np.sum(x0) * maxmutation

    # the optimization step
    # set the bounds and constraints
    # bnds = create_bounds([], genome, W.shape[1])
    # cons1 ={'type': 'eq', 'fun': constraints1, 'args':[genome]}

    # sol = scipy.optimize.minimize(parameterized_objective2_custom, x0, args=(W, genome),  bounds=bnds, constraints =cons1, tol=1e-30)
    # solution = sol.x

    ### using NNLS algorithm
    reg = nnls(W, genome)
    weights = reg[0]
    est_genome = np.dot(W, weights)
    normalised_weights = weights / sum(weights)
    solution = normalised_weights * sum(genome)

    # print(W1)
    # convert the newExposure vector into list type structure
    newExposure = list(solution)

    # get the maximum value of the new Exposure
    # maxcoef = max(newExposure)
    # idxmaxcoef = newExposure.index(maxcoef)

    # newExposure = np.round(newExposure)

    # We may need to tweak the maximum value of the new exposure to keep the total number of mutation equal to the original mutations in a genome
    # import pdb; pdb.set_trace()
    newExposure = roundConserveSum(newExposure)
    # # if np.sum(newExposure) <= maxmutation:
    # newExposure[idxmaxcoef] = (
    #     round(newExposure[idxmaxcoef]) + maxmutation - sum(newExposure)
    #     )
    # # else:
    # #     newExposure = np.round()
    if np.sum(newExposure) != maxmutation:
        raise ValueError("Mutation count not conserved")

    if metric == "cosine":
        newSimilarity = cos_sim(genome, est_genome)
    else:
        newSimilarity = np.linalg.norm(genome - est_genome, ord=2) / np.linalg.norm(
            genome, ord=2
        )

    return (newExposure, newSimilarity)


# to do the calculation, we need: W, samples(one column), H for the specific sample, tolerence
# Fit signatures with multiprocessing
def fit_signatures_pool(total_genome, W, index, metric="l2"):
    genome = total_genome[:, index]
    genome = np.array(genome)
    maxmutation = round(np.sum(genome))

    # initialize the guess
    x0 = np.random.rand(W.shape[1], 1) * maxmutation
    x0 = x0 / np.sum(x0) * maxmutation

    # the optimization step
    # set the bounds and constraints
    # bnds = create_bounds([], genome, W.shape[1])
    # cons1 ={'type': 'eq', 'fun': constraints1, 'args':[genome]}

    # sol = scipy.optimize.minimize(parameterized_objective2_custom, x0, args=(W, genome),  bounds=bnds, constraints =cons1, tol=1e-30)
    # solution = sol.x

    ### using NNLS algorithm
    reg = nnls(W, genome)
    weights = reg[0]
    est_genome = np.dot(W, weights)
    normalised_weights = weights / sum(weights)
    solution = normalised_weights * sum(genome)

    # print(W1)
    # convert the newExposure vector into list type structure
    newExposure = list(solution)

    # get the maximum value of the new Exposure
    maxcoef = max(newExposure)
    idxmaxcoef = newExposure.index(maxcoef)

    newExposure = np.round(newExposure)

    # We may need to tweak the maximum value of the new exposure to keep the total number of mutation equal to the original mutations in a genome
    if np.sum(newExposure) != maxmutation:
        newExposure[idxmaxcoef] = (
            round(newExposure[idxmaxcoef]) + maxmutation - sum(newExposure)
        )

    if metric == "cosine":
        newSimilarity = cos_sim(genome, est_genome)
    else:
        newSimilarity = np.linalg.norm(genome - est_genome, ord=2) / np.linalg.norm(
            genome, ord=2
        )

    return (newExposure, newSimilarity)


def add_signatures(
    W,
    genome,
    cutoff=0.05,
    presentSignatures=[],
    toBeAdded="all",
    metric="l2",
    solver="nnls",
    check_rule_negatives=[],
    check_rule_penalty=1.0,
    verbose=False,
):
    # This function takes an array of signature and a single genome as input, returns a dictionray of cosine similarity, exposures and presence
    # of signatures according to the indices of the original signature array

    if toBeAdded == "all":
        notToBeAdded = []
    else:
        # print(len(list(range(W.shape[1]))))
        # print(len(toBeAdded))
        notToBeAdded = list(
            set(list(range(W.shape[1]))) - set(toBeAdded)
        )  # get the indices of the signatures to be excluded

        # print(len(notToBeAdded))
    originalSimilarity = (
        np.inf
    )  # it can be also written as oldsimilarity. wanted to put a big number
    maxmutation = round(np.sum(genome))
    init_listed_idx = presentSignatures

    init_nonlisted_idx = list(
        set(list(range(W.shape[1]))) - set(init_listed_idx) - set(notToBeAdded)
    )

    finalRecord = [
        ["similarity place-holder"],
        ["newExposure place-holder"],
        ["signatures place-holder"],
    ]  # for recording the cosine difference, similarity, the new exposure and the index of the best signauture

    # get the initial original similarity if some signatures are already given to be present
    if len(init_listed_idx) != 0:
        loop_liststed_idx = init_listed_idx
        loop_liststed_idx.sort()

        W1 = W[:, loop_liststed_idx]

        # print (W1.shape)
        # initialize the guess
        x0 = np.random.rand(W1.shape[1], 1) * maxmutation
        x0 = x0 / np.sum(x0) * maxmutation

        # set the bounds and constraints
        bnds = create_bounds([], genome, W1.shape[1])
        cons1 = {"type": "eq", "fun": constraints1, "args": [genome]}

        # the optimization step
        if solver == "slsqp":
            sol = minimize(
                parameterized_objective2_custom,
                x0,
                args=(W1, genome),
                bounds=bnds,
                constraints=cons1,
                tol=1e-30,
            )
            newExposure = list(sol.x)
            est_genome = np.dot(W1, np.array(newExposure))
        if solver == "nnls":
            ### using NNLS algorithm
            reg = nnls(W1, genome[:, 0])
            weights = reg[0]
            est_genome = np.dot(W1, weights)
            normalised_weights = weights / sum(weights)
            solution = normalised_weights * sum(np.sum(genome, axis=0))
            newExposure = list(solution)
        # print(W1)
        # convert the newExposure vector into list type structure

        # newExposure = list(solution) #for nnls

        # get the maximum value of the new Exposure
        maxcoef = max(newExposure)
        idxmaxcoef = newExposure.index(maxcoef)

        newExposure = np.round(newExposure)

        # We may need to tweak the maximum value of the new exposure to keep the total number of mutation equal to the original mutations in a genome
        if np.sum(newExposure) != maxmutation:
            newExposure[idxmaxcoef] = (
                round(newExposure[idxmaxcoef]) + maxmutation - sum(newExposure)
            )

        if metric == "cosine":
            originalSimilarity = 1 - cos_sim(genome[:, 0], est_genome)
        elif metric == "l2":
            originalSimilarity = np.linalg.norm(
                genome[:, 0] - est_genome, ord=2
            ) / np.linalg.norm(genome[:, 0], ord=2)

        finalRecord = [originalSimilarity, newExposure, init_listed_idx]

    while True:
        bestDifference = -np.inf  #### wanted to put a big initila number
        bestSimilarity = np.inf  #### actually bestdistance
        loopRecord = [
            ["newExposure place-holder"],
            ["signatures place-holder"],
            ["best loop signature place-holder"],
        ]

        for sig in init_nonlisted_idx:
            if len(init_listed_idx) != 0:
                loop_liststed_idx = init_listed_idx + [sig]
                loop_liststed_idx.sort()
                # print(loop_liststed_idx)
                W1 = W[:, loop_liststed_idx]
                # print (W1.shape)
                # initialize the guess
                x0 = np.random.rand(W1.shape[1], 1) * maxmutation
                x0 = x0 / np.sum(x0) * maxmutation

                # set the bounds and constraints
                bnds = create_bounds([], genome, W1.shape[1])
                cons1 = {"type": "eq", "fun": constraints1, "args": [genome]}
            # for the first time addintion
            else:
                W1 = W[:, sig][:, np.newaxis]
                # print (W1.shape)
                # initialize the guess
                x0 = np.ones((1, 1)) * maxmutation

                # set the bounds and constraints
                bnds = create_bounds([], genome, 1)
                cons1 = {"type": "eq", "fun": constraints1, "args": [genome]}

            if solver == "slsqp":
                # the optimization step
                sol = minimize(
                    parameterized_objective2_custom,
                    x0,
                    args=(W1, genome),
                    bounds=bnds,
                    constraints=cons1,
                    tol=1e-30,
                )
                newExposure = list(sol.x)
                # compute the estimated genome
                est_genome = np.dot(W1, np.array(newExposure))
            if solver == "nnls":
                reg = nnls(W1, genome[:, 0])
                weights = reg[0]
                est_genome = np.dot(W1, weights)
                normalised_weights = weights / sum(weights)
                solution = normalised_weights * sum(genome)
                newExposure = list(solution)  # for nnls
            # print(W1)
            # convert the newExposure vector into list type structure

            # get the maximum value of the new Exposure
            maxcoef = max(newExposure)
            idxmaxcoef = newExposure.index(maxcoef)

            newExposure = np.round(newExposure)

            # We may need to tweak the maximum value of the new exposure to keep the total number of mutation equal to the original mutations in a genome
            if np.sum(newExposure) != maxmutation:
                newExposure[idxmaxcoef] = (
                    round(newExposure[idxmaxcoef]) + maxmutation - sum(newExposure)
                )

            if metric == "cosine":
                newSimilarity = 1 - cos_sim(genome[:, 0], est_genome)
            elif metric == "l2":
                newSimilarity = np.linalg.norm(
                    genome[:, 0] - est_genome, ord=2
                ) / np.linalg.norm(genome[:, 0], ord=2)

            if sig in check_rule_negatives:
                newSimilarity = newSimilarity * check_rule_penalty

            difference = originalSimilarity - newSimilarity

            # record the best values so far that creates the best difference
            if difference > bestDifference:
                bestDifference = difference
                bestSimilarity = newSimilarity
                loopRecord = [
                    newExposure,
                    W1,
                    sig,
                ]  # recording the cosine difference, the new exposure and the index of the best signauture

        if originalSimilarity - bestSimilarity > cutoff:
            originalSimilarity = bestSimilarity
            init_listed_idx.append(
                loopRecord[2]
            )  # add the qualified signature in the signature list
            init_nonlisted_idx.remove(
                loopRecord[2]
            )  # remover the qualified signature from the to be added list
            init_listed_idx.sort()
            # print(originalSimilarity)
            finalRecord = [
                originalSimilarity,
                loopRecord[0],
                init_listed_idx,
                loopRecord[1],
                genome,
            ]
            if verbose == True:
                print("Current signaure indices: ", init_listed_idx)
                if metric == "cosine":
                    print("Cosine similarity: ", 1 - originalSimilarity)
                elif metric == "l2":
                    print("L2 norm: ", originalSimilarity)
                print("\n")
            if len(init_nonlisted_idx) != 0:
                continue
            else:
                break
        else:
            break

    # print("Good so far")
    # print(finalRecord)

    dictExposure = {
        "similarity": finalRecord[0],
        "exposures": finalRecord[1],
        "signatures": finalRecord[2],
    }
    addExposure = np.zeros([W.shape[1]])
    addExposure[dictExposure["signatures"]] = dictExposure["exposures"]
    # print(cos_sim(genome[:,0], np.dot(W,addExposure)))
    if metric == "cosine":
        return (
            addExposure,
            cos_sim(genome[:, 0], np.dot(W, addExposure)),
            cos_sim(genome[:, 0], np.dot(W, addExposure)),
        )  # first value is the exprosure, second value is the similarity (based on the similarity matic), third value is the cosine similarity
    elif metric == "l2":
        return (
            addExposure,
            originalSimilarity,
            cos_sim(genome[:, 0], np.dot(W, addExposure)),
        )  # first value is the exprosure, second value is the similarity (based on the similarity matic), third value is the cosine similarity


def remove_all_single_signatures(
    W,
    H,
    genomes,
    metric="l2",
    solver="nnls",
    cutoff=0.05,
    background_sigs=[],
    verbose=False,
):
    # make the empty list of the successfull combinations
    signature_ids = sub.make_letter_ids(idlenth=W.shape[1], mtype="Signature ")

    current_signatures = signature_ids
    # print(current_signatures)
    successList = [0, [], 0]
    background_sig = copy.deepcopy(background_sigs)

    # setting the base_similarity
    base_H_index = list(np.nonzero(H)[0])
    reg = nnls(W[:, base_H_index], genomes)
    base_weights = reg[0]

    if metric == "cosine":
        # get the cos_similarity with sample for the oringinal W and H[:,i]
        originalSimilarity = 1 - cos_sim(
            genomes, np.dot(W[:, base_H_index], base_weights)
        )
        if verbose == True:
            print("originalSimilarity", 1 - originalSimilarity)
    elif metric == "l2":
        originalSimilarity = np.linalg.norm(
            genomes - np.dot(W[:, base_H_index], base_weights), ord=2
        ) / np.linalg.norm(genomes, ord=2)
        if verbose == True:
            print("originalSimilarity", originalSimilarity)
    # make the original exposures of specific sample round
    # oldExposures = np.round(H)
    oldExposures = roundConserveSum(H)
    # import pdb; pdb.set_trace()
    # set the flag for the while loop
    if len(oldExposures[np.nonzero(oldExposures)]) > 1:
        Flag = True
    else:
        Flag = False
        if metric == "cosine":
            return (
                oldExposures,
                1 - originalSimilarity,
                cos_sim(genomes, np.dot(W, H)),
            )  # first value is the exprosure, second value is the similarity (based on the similarity matic), third value is the cosine similarity
        else:
            return (
                oldExposures,
                originalSimilarity,
                cos_sim(genomes, np.dot(W, H)),
            )  # first value is the exprosure, second value is the similarity (based on the similarity matic), third value is the cosine similarity
    # The while loop starts here
    current_signatures = sub.get_items_from_index(
        signature_ids, np.nonzero(oldExposures)[0]
    )
    while Flag:
        # get the list of the indices those already have zero values
        if len(successList[1]) == 0:
            initialZerosIdx = list(np.where(oldExposures == 0)[0])
            # get the indices to be selected
            selectableIdx = list(np.where(oldExposures > 0)[0])
        elif len(successList[1]) > 1:
            initialZerosIdx = list(np.where(successList[1] == 0)[0])
            # get the indices to be selected
            selectableIdx = list(np.where(successList[1] > 0)[0])
        else:
            print("iteration is completed")
            # break

        # get the total mutations for the given sample
        maxmutation = round(np.sum(genomes))

        # new signature matrix omiting the column for zero
        # Winit = np.delete(W, initialZerosIdx, 1)
        Winit = W[:, selectableIdx]

        # set the initial cos_similarity or other similarity distance
        record = [np.inf, [], 1]  #
        # get the number of current nonzeros
        l = Winit.shape[1]

        background_sig = get_changed_background_sig_idx(
            list(oldExposures), background_sig
        )

        for i in range(l):
            if i in background_sig:
                continue

            loopSelection = list(range(l))
            del loopSelection[i]

            W1 = Winit[:, loopSelection]

            # initialize the guess
            x0 = np.random.rand(l - 1, 1) * maxmutation
            x0 = x0 / np.sum(x0) * maxmutation

            # set the bounds and constraints
            bnds = create_bounds([], genomes, W1.shape[1])
            cons1 = {"type": "eq", "fun": constraints1, "args": [genomes]}

            # the optimization step
            if solver == "slsqp":
                sol = minimize(
                    parameterized_objective2_custom,
                    x0,
                    args=(W1, genomes),
                    bounds=bnds,
                    constraints=cons1,
                    tol=1e-15,
                )

                # print (sol.success)
                # print (sol.x)

                # convert the newExposure vector into list type structure
                newExposure = list(sol.x)
            if solver == "nnls":
                ### using NNLS algorithm
                reg = nnls(W1, genomes)
                weights = reg[0]
                newSample = np.dot(W1, weights)
                normalised_weights = weights / sum(weights)
                solution = normalised_weights * sum(genomes)
                newExposure = list(solution)

            # insert the loopZeros in its actual position
            newExposure.insert(i, 0)

            # insert zeros in the required position the newExposure matrix
            initialZerosIdx.sort()

            for zeros in initialZerosIdx:
                newExposure.insert(zeros, 0)

            # get the maximum value the new Exposure
            # maxcoef = max(newExposure)
            # idxmaxcoef = newExposure.index(maxcoef)

            # newExposure = np.round(newExposure)

            # if np.sum(newExposure)!=maxmutation:
            # newExposure[idxmaxcoef] = round(newExposure[idxmaxcoef])+maxmutation-sum(newExposure)

            newExposure = np.array(newExposure)

            if verbose == True:
                # print(newExposure)
                print("\nRemoving {}".format(current_signatures[i]))
            if metric == "cosine":
                newSimilarity = 1 - cos_sim(genomes, newSample)
                if verbose == True:
                    print("newSimilarity", 1 - newSimilarity)
            elif metric == "l2":
                newSimilarity = np.linalg.norm(
                    genomes - newSample, ord=2
                ) / np.linalg.norm(genomes, ord=2)
                if verbose == True:
                    print("newSimilarity", newSimilarity)
            difference = newSimilarity - originalSimilarity
            if difference < 0:
                difference = cutoff + 1e-100
            if verbose == True:
                print("difference", difference)
            if difference < record[0]:
                record = [difference, newExposure, newSimilarity]

                # print("difference", difference)
            if verbose == True:
                print("--------------------------------------")

        if verbose == True:
            print("\n############################\n############################")
            # print("Selected Exposure")
            # print(record[1])
            selected_sigs = sub.get_items_from_index(
                signature_ids, np.nonzero(record[1])[0]
            )
            dropped_sig = list(set(current_signatures) - set(selected_sigs))
            print("Dropped Signature: {}".format(dropped_sig))
            current_signatures = selected_sigs
            print("New Similarity", record[2])
            print("Similarity Difference: {}".format(record[2] - originalSimilarity))
            print("Current Signatures: {}".format(selected_sigs))
            print("****************************\n****************************\n\n")
            # print ("This loop's selection is {}".format(record))

        if record[0] > cutoff:
            Flag = False
        elif len(record[1][np.nonzero(record[1])]) == 1:
            successList = record
            Flag = False
        else:
            successList = record
            background_sig = get_changed_background_sig_idx(
                list(record[1]), background_sig
            )
            originalSimilarity = record[2]

        # print("The loop selection is {}".format(successList))

        # print (Flag)
        # print ("\n\n")

    # print ("The final selection is {}".format(successList))

    if len(successList[1]) == 0:
        successList = [0.0, oldExposures, originalSimilarity]

    if verbose == True:
        print("\n")
        print("Final Exposure")
        print(successList[1])
        if metric == "cosine":
            print("Final Similarity: ", 1 - successList[2])
        if metric == "l2":
            print("Final Similarity: ", successList[2])
    H = successList[1]
    return (
        H,
        successList[2],
        cos_sim(genomes, np.dot(W, H)),
    )  # first value is the exprosure, second value is the similarity (based on the similarity matic), third value is the cosine similarity


def remove_all_single_signatures_pool(indices, W, exposures, totoalgenomes):
    i = indices
    H = exposures[:, i]
    genomes = totoalgenomes[:, i]
    # make the empty list of the successfull combinations
    successList = [0, [], 0]
    # get the cos_similarity with sample for the oringinal W and H[:,i]
    originalSimilarity = cos_sim(genomes, np.dot(W, H))
    # make the original exposures of specific sample round
    oldExposures = np.round(H)
    oldExposures = roundConserveSum(H)
    # set the flag for the while loop
    if len(oldExposures[np.nonzero(oldExposures)]) > 1:
        Flag = True
    else:
        Flag = False
        return oldExposures
    # The while loop starts here
    while Flag:
        # get the list of the indices those already have zero values
        if len(successList[1]) == 0:
            initialZerosIdx = list(np.where(oldExposures == 0)[0])
            # get the indices to be selected
            selectableIdx = list(np.where(oldExposures > 0)[0])
        elif len(successList[1]) > 1:
            initialZerosIdx = list(np.where(successList[1] == 0)[0])
            # get the indices to be selected
            selectableIdx = list(np.where(successList[1] > 0)[0])
        else:
            print("iteration is completed")
            # break

        # get the total mutations for the given sample
        maxmutation = round(np.sum(genomes))

        # new signature matrix omiting the column for zero
        # Winit = np.delete(W, initialZerosIdx, 1)
        Winit = W[:, selectableIdx]

        # set the initial cos_similarity
        record = [0.11, []]
        # get the number of current nonzeros
        l = Winit.shape[1]

        for i in range(l):
            # print(i)
            loopSelection = list(range(l))
            del loopSelection[i]
            # print (loopSelection)
            W1 = Winit[:, loopSelection]

            # initialize the guess
            x0 = np.random.rand(l - 1, 1) * maxmutation
            x0 = x0 / np.sum(x0) * maxmutation

            # set the bounds and constraints
            bnds = create_bounds([], genomes, W1.shape[1])
            cons1 = {"type": "eq", "fun": constraints1, "args": [genomes]}

            # the optimization step
            sol = minimize(
                parameterized_objective2_custom,
                x0,
                args=(W1, genomes),
                bounds=bnds,
                constraints=cons1,
                tol=1e-15,
            )

            # print (sol.success)
            # print (sol.x)

            # convert the newExposure vector into list type structure
            newExposure = list(sol.x)

            # insert the loopZeros in its actual position
            newExposure.insert(i, 0)

            # insert zeros in the required position the newExposure matrix
            initialZerosIdx.sort()
            for zeros in initialZerosIdx:
                newExposure.insert(zeros, 0)

            # get the maximum value the new Exposure
            maxcoef = max(newExposure)
            idxmaxcoef = newExposure.index(maxcoef)

            newExposure = np.round(newExposure)
            newExposure = roundConserveSum(newExposure)
            if np.sum(newExposure) != maxmutation:
                raise ValueError("Mutation count not conserved")

            # if np.sum(newExposure) != maxmutation:
            #     newExposure[idxmaxcoef] = (
            #         round(newExposure[idxmaxcoef]) + maxmutation - sum(newExposure)
            #     )

            newSample = np.dot(W, newExposure)
            newSimilarity = cos_sim(genomes, newSample)

            difference = originalSimilarity - newSimilarity
            # print(originalSimilarity)
            # print(newSample)
            # print(newExposure)
            # print(newSimilarity)

            # print(difference)
            # print (newExposure)
            # print (np.round(H))
            # print ("\n\n")

            if difference < record[0]:
                record = [difference, newExposure, newSimilarity]

        # print ("This loop's selection is {}".format(record))

        if record[0] > 0.01:
            Flag = False
        elif len(record[1][np.nonzero(record[1])]) == 1:
            successList = record
            Flag = False
        else:
            successList = record
        # print("The loop selection is {}".format(successList))

        # print (Flag)
        # print ("\n\n")

    # print ("The final selection is {}".format(successList))

    if len(successList[1]) == 0:
        successList = [0.0, oldExposures, originalSimilarity]
    # print ("one sample completed")
    return successList[1]


def add_remove_signatures(
    W,
    sample,
    metric="l2",
    solver="nnls",
    background_sigs=[],
    permanent_sigs=[],
    candidate_sigs="all",
    add_penalty=0.05,
    remove_penalty=0.01,
    check_rule_negatives=[],
    checkrule_penalty=1.00,
    allsigids=False,
    directory=os.getcwd() + "/Signature_assaignment_logfile.txt",
    connected_sigs=True,
    verbose=False,
):
    lognote = open(directory, "a")

    maxmutation = np.sum(np.array(sample))

    always_background = copy.deepcopy(permanent_sigs)
    M = sample
    if candidate_sigs == "all":
        candidate_sigs = list(range(W.shape[1]))
    # first add each signatures
    original_distance = np.inf  # a big number
    layer = 0

    # check the cosine_similarity with 4 signatures (highest)
    cosine_similarity_with_four_signatures = (
        1.0  # a value that will allow the signature not be reported as novel signature
    )

    # set the signature's name
    if type(allsigids) != bool:
        allsigids = sub.get_items_from_index(allsigids, candidate_sigs)
    else:
        allsigids = candidate_sigs

    while True:
        if verbose:
            print("\n\n\n\n!!!!!!!!!!!!!!!!!!STARTING LAYER: ", layer)
        lognote.write(
            "\n\n!!!!!!!!!!!!!!!!!!!!!!!!! LAYER: {} !!!!!!!!!!!!!!!!!!!!!!!!!\n".format(
                layer
            )
        )
        layer = layer + 1
        layer_original_distance = np.inf  # a big number
        sigsToBeAdded = list(set(candidate_sigs) - set(background_sigs))
        # set the signature's name
        if type(allsigids) != bool:
            allsigidsToBeAdded = sub.get_items_from_index(allsigids, sigsToBeAdded)
        else:
            allsigidsToBeAdded = sigsToBeAdded

        for i, j in zip(sigsToBeAdded, allsigidsToBeAdded):
            loop_sig = [i]

            background_sigs = list(set(background_sigs).union(set(always_background)))
            if connected_sigs == True:
                background_sigs = add_connected_sigs(background_sigs, list(allsigids))

            add_exposures, add_distance, _ = add_signatures(
                W,
                M[:, np.newaxis],
                presentSignatures=copy.deepcopy(background_sigs),
                toBeAdded=loop_sig,
                metric="l2",
                verbose=False,
                check_rule_negatives=check_rule_negatives,
                check_rule_penalty=checkrule_penalty,
                cutoff=add_penalty,
            )

            if verbose:
                print(
                    "\n\n\n################## Add Index {} ########################".format(
                        j
                    )
                )
                print(np.nonzero(add_exposures)[0])
                print(sigsToBeAdded)
                print(add_distance)

            remove_exposures, remove_distance, _ = remove_all_single_signatures(
                W,
                add_exposures,
                M,
                background_sigs=always_background,
                metric="l2",
                verbose=False,
                cutoff=remove_penalty,
            )

            if verbose:
                print("\n################## Remove ########################")
                print(np.nonzero(remove_exposures)[0])
                print(remove_distance)
            # check if there is any change between the add and remove signatures and assign the distance and exposure accoridingly
            if (
                np.nonzero(add_exposures)[0].all()
                == np.nonzero(remove_exposures)[0].all()
            ) and np.nonzero(add_exposures)[0].shape == np.nonzero(remove_exposures)[
                0
            ].shape:
                distance = add_distance
                exposures = add_exposures
            else:
                distance = remove_distance
                exposures = remove_exposures

            if distance < layer_original_distance:
                selected_signatures = np.nonzero(exposures)[0]
                layer_original_distance = distance
                activities = exposures
        if verbose:
            print(
                "\n\n#################### After Add-Remove #################################"
            )
            print(selected_signatures)
            print(layer_original_distance)

        lognote.write(
            "Best Signature Composition {}\n".format(
                sub.get_items_from_index(allsigids, selected_signatures)
            )
        )
        lognote.write("L2 Error % {}\n".format(round(layer_original_distance, 2)))
        lognote.write(
            "Cosine Similarity {}\n".format(round(cos_sim(M, np.dot(W, activities)), 2))
        )
        if layer_original_distance < original_distance:
            original_distance = layer_original_distance
            background_sigs = list(selected_signatures)
            finalactivities = activities
            if len(background_sigs) == 4:
                cosine_similarity_with_four_signatures = cos_sim(
                    M, np.dot(W, finalactivities)
                )
                # print(cosine_similarity_with_four_signatures)
        else:
            N = np.dot(W, finalactivities)
            cosine_similarity = cos_sim(M, N)
            kldiv = round(scipy.stats.entropy(M, N), 3)
            correlation, _ = scipy.stats.pearsonr(M, N)
            correlation = round(correlation, 2)
            break

    # get the maximum value of the new Exposure
    # maxcoef = np.max(finalactivities)
    # idxmaxcoef = list(finalactivities).index(maxcoef)

    # finalactivities = np.round(finalactivities)
    finalactivities = roundConserveSum(finalactivities)

    # We may need to tweak the maximum value of the new exposure to keep the total number of mutation equal to the original mutations in a genome
    # if np.sum(finalactivities) != maxmutation:
    #     finalactivities[idxmaxcoef] = (
    #         round(finalactivities[idxmaxcoef]) + maxmutation - np.sum(finalactivities)
    #     )
    if np.sum(finalactivities) != round(maxmutation):
        raise ValueError("Mutation count not conserved")

    if verbose:
        print("\n########################## Final ###########################")
        print(background_sigs)
        print(original_distance)
        print(finalactivities)
    lognote.write(
        "\n#################### Final Composition #################################\n"
    )
    lognote.write(
        "{}\n".format(sub.get_items_from_index(allsigids, selected_signatures))
    )
    lognote.write("L2 Error % {}\n".format(round(original_distance, 2)))
    lognote.write(
        "Cosine Similarity {}\n".format(
            round(cos_sim(M, np.dot(W, finalactivities)), 2)
        )
    )
    # close lognote
    lognote.close()
    # newExposure, newSimilarity = fit_signatures(W[:,list(background_sigs)], M, metric="l2")
    # print(newExposure, newSimilarity)
    return (
        background_sigs,
        finalactivities,
        original_distance,
        cosine_similarity,
        kldiv,
        correlation,
        cosine_similarity_with_four_signatures,
    )
