#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Nov 16 2021

@author: rvangara
"""

import os
import sys
import string
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

plt.switch_backend("agg")
from matplotlib.backends.backend_pdf import PdfPages

import sigProfilerPlotting as plot
from SigProfilerAssignment.DecompositionPlots import PlotDecomposition as sp
from sigProfilerPlotting import plotActivity as plot_ac
from sigProfilerPlotting import tmbplot as tmb
import pypdf
import scipy
from pypdf import PdfWriter, PdfReader
import SigProfilerAssignment as spa
from SigProfilerAssignment import single_sample as ss
from scipy.spatial.distance import correlation as cor
from alive_progress import alive_bar
from joblib import Parallel, delayed
from multiprocessing import cpu_count


def getProcessAvg(
    samples,
    genome_build="GRCh37",
    cosmic_version=3.4,
    signature_database=None,
    connected_sigs=True,
    exome=False,
):
    paths = spa.__path__[0]

    if genome_build in ["GRCh37", "GRCh38", "mm9", "mm10", "mm39", "rn6", "rn7"]:
        genome_build = genome_build
    else:
        print(
            "The selected genome build is "
            + str(genome_build)
            + ". COSMIC signatures are available only for GRCh37/38, mm9/10/39 and rn6/7 genomes. So, the genome build is reset to GRCh37."
        )
        genome_build = "GRCh37"
    if samples.shape[0] == 96:
        if exome == False:
            sigDatabase = pd.read_csv(
                paths
                + "/data/Reference_Signatures/"
                + genome_build
                + "/COSMIC_v"
                + str(cosmic_version)
                + "_SBS_"
                + genome_build
                + ".txt",
                sep="\t",
                index_col=0,
            )
        else:
            sigDatabase = pd.read_csv(
                paths
                + "/data/Reference_Signatures/"
                + genome_build
                + "/COSMIC_v"
                + str(cosmic_version)
                + "_SBS_"
                + genome_build
                + "_exome.txt",
                sep="\t",
                index_col=0,
            )
        signames = sigDatabase.columns

    elif samples.shape[0] == 288:
        sigDatabase = pd.read_csv(
            paths
            + "/data/Reference_Signatures/GRCh37/COSMIC_v"
            + str(cosmic_version)
            + "_SBS"
            + str(samples.shape[0])
            + "_GRCh37.txt",
            sep="\t",
            index_col=0,
        )
        signames = sigDatabase.columns

    elif samples.shape[0] == 1536:
        sigDatabase = pd.read_csv(
            paths
            + "/data/Reference_Signatures/GRCh37/COSMIC_v"
            + str(cosmic_version)
            + "_SBS"
            + str(samples.shape[0])
            + "_GRCh37.txt",
            sep="\t",
            index_col=0,
        )
        signames = sigDatabase.columns

    elif samples.shape[0] == 78:
        if exome == False:
            sigDatabase = pd.read_csv(
                paths
                + "/data/Reference_Signatures/"
                + genome_build
                + "/COSMIC_v"
                + str(cosmic_version)
                + "_DBS_"
                + genome_build
                + ".txt",
                sep="\t",
                index_col=0,
            )
        else:
            sigDatabase = pd.read_csv(
                paths
                + "/data/Reference_Signatures/"
                + genome_build
                + "/COSMIC_v"
                + str(cosmic_version)
                + "_DBS_"
                + genome_build
                + "_exome.txt",
                sep="\t",
                index_col=0,
            )
        signames = sigDatabase.columns
        connected_sigs = False

    elif samples.shape[0] == 83:
        sigDatabase = pd.read_csv(
            paths
            + "/data/Reference_Signatures/GRCh37/COSMIC_v"
            + str(cosmic_version)
            + "_ID_GRCh37.txt",
            sep="\t",
            index_col=0,
        )
        signames = sigDatabase.columns
        connected_sigs = False
        # indel signatures have no re-normalization so "GRCh37" is used
        genome_build = "GRCh37"

    elif samples.shape[0] == 48:
        if cosmic_version < 3.3:
            print(
                f"The selected cosmic version is {cosmic_version}. However, CN signatures are "
                "available only for version 3.3 and newer. Therefore, the cosmic version has "
                "been reset to 3.4."
            )
            cosmic_version = 3.4
        sigDatabase = pd.read_csv(
            paths
            + "/data/Reference_Signatures/GRCh37/COSMIC_v"
            + str(cosmic_version)
            + "_CN_GRCh37.txt",
            sep="\t",
            index_col=0,
        )
        signames = sigDatabase.columns
        connected_sigs = False

    elif samples.shape[0] == 32:
        if cosmic_version < 3.4:
            print(
                f"The selected cosmic version is {cosmic_version}. However, SV signatures are "
                "available only for version 3.4 and newer. Therefore, the cosmic version has "
                "been reset to 3.4."
            )
            cosmic_version = 3.4
        sigDatabase = pd.read_csv(
            paths
            + "/data/Reference_Signatures/GRCh38/COSMIC_v"
            + str(cosmic_version)
            + "_SV_GRCh38.txt",
            sep="\t",
            index_col=0,
        )
        signames = sigDatabase.columns
        connected_sigs = False
    else:
        sigDatabase = pd.DataFrame(samples)
        sigDatabase.columns = sigDatabase.columns.astype(str)
        sigDatabase.index = sigDatabase.index.astype(str)
        signames = sigDatabase.columns
        connected_sigs = False

    return sigDatabase, signames, connected_sigs, genome_build

    if signature_database != None:  # pd.core.frame.DataFrame:
        print("################## USING CUSTOM SIGNATURE DATBASE ##################")
        signature_database = pd.read_csv(signature_database, sep="\t", index_col=0)
        if samples.shape[0] == signature_database.shape[0]:
            sigDatabase = signature_database
            signames = sigDatabase.columns
            # make_decomposition_plots=False
            del signature_database
        else:
            sys.exit(
                "The Signatures and the custom signature database have different context types."
            )
    sigDatabases = sigDatabase.reset_index()
    return sigDatabases


def signature_plotting_text(value, text, Type):
    name = text + ": "
    name_list = []
    total = np.sum(np.array(value))
    for i in value:
        if Type == "integer":
            i = int(i)
            p = round(i / total * 100, 1)
            i = format(i, ",d")
            tail = str(i) + "/" + str(p) + "%"
            name_list.append(name + tail)
        elif Type == "float":
            i = round(i, 2)
            name_list.append(name + str(i))
    return name_list


def make_letter_ids(idlenth=10, mtype="SBS96"):
    listOfSignatures = []
    letters = list(string.ascii_uppercase)
    letters.extend([i + b for i in letters for b in letters])
    letters = letters[0:idlenth]

    for j, l in zip(range(idlenth), letters):
        listOfSignatures.append(mtype + l)
    listOfSignatures = np.array(listOfSignatures)
    return listOfSignatures


def union(a, b):
    """return the union of two lists"""
    return list(set(a) | set(b))


def get_indeces(a, b):
    """
    Extracts the indices multiple items in a list.

    Parameters:
        a: list. where we want to get the index of the items.
        b: list. the items we want to get index of.
    #example:
    x = ['SBS1', 'SBS2', 'SBS3', 'SBS5', 'SBS8', 'SBS13', 'SBS40']
    y = ['SBS1',  'SBS5']
    get_indeces(x, y)
    #result
    >>> [1,3]
    """

    indeces = []
    for i in b:
        try:
            idx = a.index(i)
            indeces.append(idx)
        except:
            next

    return indeces


def get_items_from_index(x, y):
    """decipher the values of items in a list from their indices."""
    z = []
    for i in y:
        try:
            z.append(x[i])
        except:
            pass
    return z


def signature_decomposition(
    signatures,
    mtype,
    directory,
    genome_build="GRCh37",
    cosmic_version=3.4,
    signature_database=None,
    add_penalty=0.05,
    remove_penalty=0.01,
    mutation_context=None,
    connected_sigs=True,
    make_decomposition_plots=True,
    originalProcessAvg=None,
    new_signature_thresh_hold=0.8,
    sig_exclusion_list=[],
    exome=False,
    m_for_subgroups="SBS",
    volume=None,
):
    originalProcessAvg = originalProcessAvg.reset_index()
    if not os.path.exists(directory + "/Solution_Stats"):
        os.makedirs(directory + "/Solution_Stats")
    # open the log file for signature decomposition
    lognote = open(
        directory
        + "/Solution_Stats/Cosmic_"
        + mutation_context
        + "_Decomposition_Log.txt",
        "w",
    )
    lognote.write(
        "############################ Signature Decomposition Details ################################\n\n\n"
    )
    lognote.write("Context Type: {}\n".format(mtype))
    lognote.write("Genome Build: {}\n".format(genome_build))

    if signature_database == None:
        sigDatabase, signames, connected_sigs, genome_build = getProcessAvg(
            signatures,
            genome_build=genome_build,
            cosmic_version=cosmic_version,
            exome=exome,
        )
        # processAvg = processAvg.set_index('Type').rename_axis('MutationType')
    else:
        try:
            sigDatabase = pd.read_csv(signature_database, sep="\t", index_col=0)
            # indx = sigDatabase.index()
            if (
                sigDatabase.shape[0] == 1536
            ):  # collapse the 1596 context into 96 only for the deocmposition
                sigDatabase = sigDatabase.groupby(sigDatabase.index.str[1:8]).sum()

            elif (
                sigDatabase.shape[0] == 288
            ):  # collapse the 288 context into 96 only for the deocmposition
                # sigDatabase = pd.DataFrame(processAvg, index=index)
                sigDatabase = sigDatabase.groupby(sigDatabase.index.str[2:9]).sum()

            if (
                sigDatabase.shape[0] == 78
                or sigDatabase.shape[0] == 83
                or sigDatabase.shape[0] == 48
            ):
                connected_sigs = False
            lognote.write(
                "##### Using a custom signature database for decomposition #####"
            )
        except:
            sys.exit(
                "Wrong format of signature database for decompose_fit, Please pass a text file of signatures in the format of COSMIC sig database"
            )

    sig_exclusion_list = [m_for_subgroups + items for items in sig_exclusion_list]
    lognote.write(
        "The following signatures are excluded: "
        + " ".join(str(item) for item in sig_exclusion_list)
    )
    sigDatabase.drop(sig_exclusion_list, axis=1, inplace=True, errors="ignore")
    signames = sigDatabase.columns

    # if type(signature_database)==pd.core.frame.DataFrame:

    #     if signatures.shape[0]==signature_database.shape[0]:
    #         sigDatabase=signature_database
    #         signames = sigDatabase.columns
    #         #make_decomposition_plots=False
    #         del signature_database
    sigDatabases = sigDatabase.reset_index()
    letters = list(string.ascii_uppercase)
    letters.extend([i + b for i in letters for b in letters])
    letters = letters[0 : signatures.shape[1]]

    # replace the probability data of the process matrix with the number of mutation
    for i in range(signatures.shape[1]):
        signatures[:, i] = signatures[:, i] * 5000  # (np.sum(exposureAvg[i, :]))

    sigDatabase = np.array(sigDatabase)
    allsignatures = np.array([])
    newsig = list()  # will create the letter based id of newsignatures
    newsigmatrixidx = (
        list()
    )  # will create the original id of newsignature to help to record the original matrix
    fh = open(directory + "/De_Novo_map_to_COSMIC_" + mutation_context + ".csv", "w")
    fh.write(
        "De novo extracted, Global NMF Signatures, L1 Error %, L2 Error %, KL Divergence, Cosine Similarity, Correlation\n"
    )
    fh.close()
    dictionary = {}
    # bgsigs = (np.where(np.isin(signames.tolist(),['SBS1','SBS5']))[0]).tolist()

    # only for SBS96
    if mtype == "96" or mtype == "288" or mtype == "1536":
        bgsigs = get_indeces(list(signames), ["SBS1", "SBS5"])
    else:
        bgsigs = []

    # get the names of denovo signatures
    denovo_signature_names = make_letter_ids(
        signatures.shape[1], mtype=mutation_context
    )
    # lognote.write("\n********** Starting Signature Decomposition **********\n\n")
    activity_percentages = []
    merger = PdfWriter()

    for i, j in zip(range(signatures.shape[1]), denovo_signature_names):
        # Only for context SBS96
        if signatures.shape[0] == 96:
            lognote = open(
                directory
                + "/Solution_Stats/Cosmic_"
                + mutation_context
                + "_Decomposition_Log.txt",
                "a",
            )
            lognote.write(
                "\n\n\n\n######################## Decomposing "
                + j
                + " ########################\n"
            )
            lognote.close()
            if genome_build == "mm9" or genome_build == "mm10":
                check_rule_negatives = [1, 16]
                check_rule_penalty = 1.50
            else:
                check_rule_negatives = []
                check_rule_penalty = 1.0

            (
                _,
                exposures,
                L2dist,
                similarity,
                kldiv,
                correlation,
                cosine_similarity_with_four_signatures,
            ) = ss.add_remove_signatures(
                sigDatabase,
                signatures[:, i],
                metric="l2",
                solver="nnls",
                background_sigs=bgsigs,  # [0,4],
                permanent_sigs=bgsigs,  # [0,4],
                candidate_sigs="all",
                allsigids=signames,
                add_penalty=add_penalty,
                remove_penalty=remove_penalty,
                check_rule_negatives=check_rule_negatives,
                checkrule_penalty=check_rule_penalty,
                directory=directory
                + "/Solution_Stats/Cosmic_"
                + mutation_context
                + "_Decomposition_Log.txt",
                connected_sigs=connected_sigs,
                verbose=False,
            )

        else:
            lognote = open(
                directory
                + "/Solution_Stats/Cosmic_"
                + mutation_context
                + "_Decomposition_Log.txt",
                "a",
            )
            lognote.write(
                "\n\n\n\n######################## Decomposing "
                + j
                + " ########################\n"
            )
            lognote.close()

            (
                _,
                exposures,
                L2dist,
                similarity,
                kldiv,
                correlation,
                cosine_similarity_with_four_signatures,
            ) = ss.add_remove_signatures(
                sigDatabase,
                signatures[:, i],
                metric="l2",
                solver="nnls",
                background_sigs=[],
                candidate_sigs="all",
                add_penalty=add_penalty,
                remove_penalty=remove_penalty,
                check_rule_negatives=[],
                checkrule_penalty=[],
                directory=directory
                + "/Solution_Stats/Cosmic_"
                + mutation_context
                + "_Decomposition_Log.txt",
                connected_sigs=connected_sigs,
                verbose=False,
            )
        L1dist = np.linalg.norm(
            signatures[:, i] - np.dot(sigDatabase, exposures), ord=1
        ) / np.linalg.norm(signatures[:, i], ord=1)
        exposure_percentages = (
            exposures[np.nonzero(exposures)]
            / np.sum(exposures[np.nonzero(exposures)])
            * 100
        )
        listofinformation = list("0" * len(np.nonzero(exposures)[0]) * 3)

        count = 0
        decomposed_signatures = []
        contribution_percentages = []

        for j in np.nonzero(exposures)[0]:
            listofinformation[count * 3] = signames[j]
            listofinformation[count * 3 + 1] = round(exposure_percentages[count], 2)
            contribution_percentages.append(round(exposure_percentages[count], 2))
            listofinformation[count * 3 + 2] = "%"
            decomposed_signatures.append(signames[j])
            count += 1
        ListToTumple = tuple(
            [mtype, letters[i]]
            + listofinformation
            + [L1dist * 100]
            + [L2dist * 100]
            + [kldiv]
            + [similarity]
            + [correlation]
        )
        activity_percentages.append(contribution_percentages)

        weights = []
        basis_names = []
        nonzero_exposures = exposures[np.nonzero(exposures)]
        denovo_name = mutation_context + letters[i]
        for info in range(0, len(listofinformation), 3):
            # print(info)
            sigName = listofinformation[info]
            sigWeigt = str(listofinformation[info + 1]) + "%"
            weights.append(sigWeigt)
            basis_names.append(sigName)

        denovo_signames = []
        for letter in letters:
            denovo_signames.append(mutation_context + letter)

        sigDatabases_DF = sigDatabases

        if mtype == "1536":
            mtype_par = "1536"
        elif mtype == "288":
            mtype_par = "288"
        elif mtype == "96":
            mtype_par = "96"
        elif mtype == "DINUC" or mtype == "78":
            mtype_par = "78"
        elif mtype == "INDEL" or mtype == "83":
            mtype_par = "83"
        elif mtype == "CNV" or mtype == "48":
            mtype_par = "48"
        elif mtype == "SV" or mtype == "32":
            mtype_par = "32"
        else:
            mtype_par = "none"
        # only create decomposition plots for COSMIC signatures
        if mtype_par != "none" and make_decomposition_plots == True:
            # reformat the first column of cosmic signature dataframe
            cosmic_sigs_DF = sigDatabases_DF.copy(deep=True)
            cosmic_sigs_DF.columns = ["MutationType"] + cosmic_sigs_DF.columns[
                1:
            ].to_list()
            # Get the names of the columns for each dataframe
            denovo_col_names = originalProcessAvg.columns
            cosmic_col_names = cosmic_sigs_DF.columns
            # Get the name for the MutationType column
            cosmic_mut_types_col = cosmic_col_names[0]
            denovo_mut_types_col = denovo_col_names[0]
            # create lists of implemented columns
            basis_cols = basis_names.copy()
            basis_cols.insert(0, cosmic_mut_types_col)
            denovo_cols = [denovo_mut_types_col, denovo_name]
            byte_plot = sp.run_PlotDecomposition(
                originalProcessAvg[denovo_cols],
                denovo_name,
                cosmic_sigs_DF[basis_cols],
                basis_names,
                weights,
                nonzero_exposures / 5000,
                directory,
                "test",
                mtype_par,
                cosmic_version=cosmic_version,
                genome_build=genome_build,
                exome=exome,
                volume=volume,
                use_custom_basis=(signature_database is not None),
            )

            byte_plot.seek(0)
            reader = PdfReader(byte_plot)
            for page in reader.pages:
                merger.add_page(page)

            with alive_bar(
                1, ctrl_c=False, bar="blocks", title=f"Decompositon Plot:{denovo_name}"
            ) as bar:
                bar()

        strings = (
            "Signature %s-%s,"
            + " Signature %s (%0.2f%s) &" * (len(np.nonzero(exposures)[0]) - 1)
            + " Signature %s (%0.2f%s), %0.2f,  %0.2f, %0.3f, %0.2f, %0.2f\n"
        )
        # new_signature_thresh_hold = 0.8
        if (
            similarity > new_signature_thresh_hold
            and cosine_similarity_with_four_signatures > new_signature_thresh_hold
        ):  ########### minimum signtatures and cosine similarity needs to be fitted to become a unique signature
            allsignatures = np.append(allsignatures, np.nonzero(exposures))
            fh = open(
                directory + "/De_Novo_map_to_COSMIC_" + mutation_context + ".csv", "a"
            )
            fh.write(strings % (ListToTumple))
            fh.close()

            dictionary.update(
                {"{}".format(mutation_context + letters[i]): decomposed_signatures}
            )

        else:
            newsig.append(mutation_context + letters[i])
            newsigmatrixidx.append(i)
            fh = open(
                directory + "/De_Novo_map_to_COSMIC_" + mutation_context + ".csv", "a"
            )
            fh.write(
                "Signature {}-{}, Signature {}-{}, {}, {}, {}, {}, {}\n".format(
                    mtype, letters[i], mtype, letters[i], 0, 0, 0, 1, 1
                )
            )
            fh.close()
            dictionary.update(
                {
                    "{}".format(mutation_context + letters[i]): [
                        "{}".format(mutation_context + letters[i])
                    ]
                }
            )
            # dictionary.update({letters[i]:"Signature {}-{}, Signature {}-{}, {}\n".format(mtype, letters[i], mtype, letters[i], 1 )})

    try:
        if make_decomposition_plots and mtype_par != "none":
            # Write out the decomposition plots
            contexts = {
                "96": "SBS96",
                "288": "SBS288",
                "1536": "SBS1536",
                "78": "DBS78",
                "83": "ID83",
                "48": "CNV",
                "32": "SV",
            }
            merger.write(
                directory + "/" + contexts[mtype_par] + "_Decomposition_Plots.pdf"
            )
    except:
        print(
            "The context-"
            + str(mtype_par)
            + " decomposition pages were not able to be merged."
        )

    different_signatures = np.unique(allsignatures)
    different_signatures = different_signatures.astype(int)
    if mtype == "96" or mtype == "288" or mtype == "1536":
        different_signatures = list(set().union(different_signatures, bgsigs))
        different_signatures.sort()

    # get the name of the signatures
    try:
        detected_signatures = signames[different_signatures]
        globalsigmats = sigDatabases.loc[:, list(detected_signatures)]
    except:
        detected_signatures = [None]
        globalsigmats = None

    newsigsmats = signatures[:, newsigmatrixidx]

    # only for SBS96
    if mtype == "96" or mtype == "288" or mtype == "1536":
        background_sigs = get_indeces(list(detected_signatures), ["SBS1", "SBS5"])
        # add connected signatures
        different_signatures = ss.add_connected_sigs(
            different_signatures, list(signames)
        )
    # for other contexts
    else:
        background_sigs = []

    # close the lognote
    lognote.close()

    return {
        "globalsigids": list(detected_signatures),
        "newsigids": newsig,
        "globalsigs": globalsigmats,
        "newsigs": newsigsmats / 5000,
        "dictionary": dictionary,
        "background_sigs": background_sigs,
        "activity_percentages": activity_percentages,
    }


def process_sample(args):
    # (r, denovo_exposureAvg, attribution, allsigids, background_sigs, processAvg, allgenomes, layer_directory, solution_prefix, verbose, pcawg_rule, initial_remove_penalty, add_penalty, remove_penalty, check_rule_negatives, check_rule_penalty, connected_sigs) = args
    (
        r,
        denovo_exposureAvg,
        attribution,
        allsigids,
        background_sigs,
        processAvg,
        allgenomes,
        layer_directory,
        solution_prefix,
        verbose,
        pcawg_rule,
        initial_remove_penalty,
        add_penalty,
        remove_penalty,
        check_rule_negatives,
        check_rule_penalty,
        connected_sigs,
    ) = args
    if verbose:
        print(
            "\n\n\n\n\n                                        ################ Sample "
            + str(r + 1)
            + " #################"
        )

    lognote = open(
        layer_directory
        + "/Solution_Stats/"
        + solution_prefix
        + "_Signature_Assignment_log.txt",
        "a",
    )
    lognote.write(
        "\n\n\n\n\n                    ################ Sample "
        + str(r + 1)
        + " #################\n"
    )

    sample_exposure = np.array(denovo_exposureAvg.iloc[:, r])

    init_sig_idx = np.nonzero(sample_exposure)[0]
    init_sigs = denovo_exposureAvg.index[init_sig_idx]

    init_decomposed_sigs = []
    for de_novo_sig in init_sigs:
        init_decomposed_sigs = union(
            init_decomposed_sigs, list(attribution[de_novo_sig])
        )

    init_decomposed_sigs_idx = get_indeces(allsigids, init_decomposed_sigs)
    init_decomposed_sigs_idx.sort()
    init_decomposed_sigs_idx = list(
        set().union(init_decomposed_sigs_idx, background_sigs)
    )

    background_sig_idx = get_indeces(init_decomposed_sigs_idx, background_sigs)

    fit_signatures = processAvg[:, init_decomposed_sigs_idx]
    newExposure, newSimilarity = ss.fit_signatures(fit_signatures, allgenomes[:, r])

    exposureAvg = np.zeros([processAvg.shape[1], allgenomes.shape[1]])
    for nonzero_idx, nozero_exp in zip(init_decomposed_sigs_idx, newExposure):
        exposureAvg[nonzero_idx, r] = nozero_exp

    if pcawg_rule:
        maxmutation = np.sum(allgenomes[:, r])
        (
            exposureAvg[:, r],
            remove_distance,
            _,
        ) = ss.remove_all_single_signatures(
            processAvg,
            exposureAvg[:, r],
            allgenomes[:, r],
            metric="l2",
            verbose=False,
            cutoff=0.02,
        )
        maxcoef = max(list(exposureAvg[:, r]))
        idxmaxcoef = list(exposureAvg[:, r]).index(maxcoef)

        exposureAvg[:, r] = np.round(exposureAvg[:, r])

        if np.sum(exposureAvg[:, r]) != maxmutation:
            exposureAvg[:, r][idxmaxcoef] = (
                round(exposureAvg[:, r][idxmaxcoef])
                + maxmutation
                - sum(exposureAvg[:, r])
            )
    else:
        if verbose:
            print(
                "############################# Initial Composition #################################### "
            )
            print(pd.DataFrame(exposureAvg[:, r], index=allsigids).T)
            print("L2%: ", newSimilarity)

        lognote.write(
            "############################# Initial Composition ####################################\n"
        )
        exposures = pd.DataFrame(exposureAvg[:, r], index=allsigids).T
        lognote.write(
            "{}\n".format(exposures.iloc[:, exposures.to_numpy().nonzero()[1]])
        )
        lognote.write(
            "L2 Error %: {}\nCosine Similarity: {}\n".format(
                round(newSimilarity, 2),
                round(
                    cos_sim(
                        allgenomes[:, r],
                        np.dot(processAvg, exposureAvg[:, r]),
                    ),
                    2,
                ),
            )
        )
        (
            exposureAvg[:, r],
            L2dist,
            cosine_sim,
        ) = ss.remove_all_single_signatures(
            processAvg,
            exposureAvg[:, r],
            allgenomes[:, r],
            metric="l2",
            solver="nnls",
            cutoff=initial_remove_penalty,
            background_sigs=[],
            verbose=False,
        )
        if verbose:
            print(
                "############################## Composition After Initial Remove ############################### "
            )
            print(pd.DataFrame(exposureAvg[:, r], index=allsigids).T)
            print("L2%: ", L2dist)
        lognote.write(
            "############################## Composition After Initial Remove ###############################\n"
        )
        exposures = pd.DataFrame(exposureAvg[:, r], index=allsigids).T
        lognote.write(
            "{}\n".format(exposures.iloc[:, exposures.to_numpy().nonzero()[1]])
        )
        lognote.write(
            "L2 Error %: {}\nCosine Similarity: {}\n".format(
                round(L2dist, 2), round(cosine_sim, 2)
            )
        )
        lognote.write(
            "\n############################## Performing Add-Remove Step ##############################\n"
        )
        lognote.close()

        init_add_sig_idx = list(
            set().union(list(np.nonzero(exposureAvg[:, r])[0]), background_sigs)
        )

        if background_sigs != 0:
            background_sig_idx = get_indeces(allsigids, ["SBS1", "SBS5"])

        try:
            (
                _,
                exposureAvg[:, r],
                L2dist,
                similarity,
                kldiv,
                correlation,
                cosine_similarity_with_four_signatures,
            ) = ss.add_remove_signatures(
                processAvg,
                allgenomes[:, r],
                metric="l2",
                solver="nnls",
                background_sigs=init_add_sig_idx,
                permanent_sigs=background_sig_idx,
                candidate_sigs="all",
                allsigids=allsigids,
                add_penalty=add_penalty,
                remove_penalty=remove_penalty,
                check_rule_negatives=check_rule_negatives,
                checkrule_penalty=check_rule_penalty,
                connected_sigs=connected_sigs,
                directory=layer_directory
                + "/Solution_Stats/"
                + solution_prefix
                + "_Signature_Assignment_log.txt",
                verbose=False,
            )

            if verbose:
                print(
                    "####################################### Composition After Add-Remove #######################################\n"
                )
                print(exposureAvg[:, r])
                print("L2%: ", L2dist)
            lognote = open(
                layer_directory
                + "/Solution_Stats/"
                + solution_prefix
                + "_Signature_Assignment_log.txt",
                "a",
            )
            lognote.write(
                "####################################### Composition After Add-Remove #######################################\n"
            )
            exposures = pd.DataFrame(exposureAvg[:, r], index=allsigids).T
            lognote.write(
                "{}\n".format(exposures.iloc[:, exposures.to_numpy().nonzero()[1]])
            )
            lognote.write(
                "L2 Error %: {}\nCosine Similarity: {}\n".format(
                    round(L2dist, 2), round(similarity, 2)
                )
            )
            lognote.close()
        except:
            pass

    return exposureAvg[:, r]


def make_final_solution(
    processAvg,
    allgenomes,
    allsigids,
    layer_directory,
    m,
    index,
    allcolnames,
    process_std_error="none",
    signature_stabilities=" ",
    signature_total_mutations=" ",
    signature_stats="none",
    cosmic_sigs=False,
    attribution=0,
    denovo_exposureAvg="none",
    add_penalty=0.05,
    remove_penalty=0.01,
    initial_remove_penalty=0.05,
    de_novo_fit_penalty=0.02,
    background_sigs=0,
    genome_build="GRCh37",
    sequence="genome",
    export_probabilities=True,
    export_probabilities_per_mutation=False,
    refit_denovo_signatures=True,
    collapse_to_SBS96=True,
    connected_sigs=True,
    pcawg_rule=False,
    verbose=False,
    make_plots=True,
    samples="./",
    input_type="matrix",
    denovo_refit_option=True,
    exome=False,
    volume=None,
    cpu=-1,
):
    if processAvg.shape[0] == allgenomes.shape[0] and processAvg.shape[0] != 96:
        collapse_to_SBS96 = False

    solution_type = layer_directory.split("/")[-1]
    solution_prefix = solution_type.split("_")
    solution_prefix = "_".join(solution_prefix[0:2])
    if refit_denovo_signatures:
        solution_prefix_refit = solution_prefix + "_refit"

    if not os.path.exists(layer_directory + "/Signatures"):
        os.makedirs(layer_directory + "/Signatures")
    if not os.path.exists(layer_directory + "/Activities"):
        os.makedirs(layer_directory + "/Activities")
    if not os.path.exists(layer_directory + "/Solution_Stats"):
        os.makedirs(layer_directory + "/Solution_Stats")

    if refit_denovo_signatures:
        lognote = open(
            layer_directory
            + "/Solution_Stats/"
            + solution_prefix_refit
            + "_Signature_Assignment_log.txt",
            "w",
        )
    else:
        lognote = open(
            layer_directory
            + "/Solution_Stats/"
            + solution_prefix
            + "_Signature_Assignment_log.txt",
            "w",
        )
    lognote.write(
        "************************ Stepwise Description of Signature Assignment to Samples ************************"
    )
    lognote.close()

    if m == 83 or m == "83":
        signature_type = "INDEL83"
        connected_sigs = False
    elif m == 78 or m == "78":
        signature_type = "DINUC78"
        connected_sigs = False
    else:
        signature_type = "SBS" + str(m)

    allgenomes = np.array(allgenomes)
    if (
        (m == "96" or m == "1536" or m == "288")
        and (genome_build == "mm9" or genome_build == "mm10")
        and (collapse_to_SBS96 == True)
    ):
        check_rule_negatives = [1, 16]
        check_rule_penalty = 1.50
    else:
        check_rule_negatives = []
        check_rule_penalty = 1.0
    exposureAvg = np.zeros([processAvg.shape[1], allgenomes.shape[1]])
    if cosmic_sigs:
        denovo_exposureAvg = denovo_exposureAvg.T

    inputs = [
        (
            r,
            denovo_exposureAvg,
            attribution,
            allsigids,
            background_sigs,
            processAvg,
            allgenomes,
            layer_directory,
            solution_prefix,
            verbose,
            pcawg_rule,
            initial_remove_penalty,
            add_penalty,
            remove_penalty,
            check_rule_negatives,
            check_rule_penalty,
            connected_sigs,
        )
        for r in range(allgenomes.shape[1])
    ]

    # Passing inputs as a single argument and unpacking in the process_sample function
    # with tqdm(total=len(inputs)) as progress_bar:
    n_samples = allgenomes.shape[1]
    if cpu == -1:
        njobs = min(cpu_count(), allgenomes.shape[1])
    else:
        njobs = min(cpu, allgenomes.shape[1])
    # njobs=64
    # batchs= n_samples // njobs
    results = Parallel(n_jobs=njobs, mmap_mode="r", verbose=5, batch_size="auto")(
        delayed(process_sample)(input) for input in inputs
    )

    # Store the results in exposureAvg
    exposureAvg = np.zeros([processAvg.shape[1], allgenomes.shape[1]])
    for r, result in enumerate(results):
        exposureAvg[:, r] = result
    else:
        # Handle the non-cosmic_sigs case as needed
        pass

    processAvg = pd.DataFrame(processAvg.astype(float))
    processes = processAvg.set_index(index)
    processes.columns = allsigids
    processes = processes.rename_axis("MutationType", axis="columns")
    processes.to_csv(
        layer_directory
        + "/Signatures"
        + "/"
        + solution_prefix
        + "_"
        + "Signatures.txt",
        "\t",
        float_format="%.8f",
        index_label=[processes.columns.name],
    )
    exposureAvg = pd.DataFrame(exposureAvg.astype(int))
    allsigids = np.array(allsigids)
    exposures = exposureAvg.set_index(allsigids)
    exposures.columns = allcolnames
    exposures = exposures.T
    exposures = exposures.rename_axis("Samples", axis="columns")
    if refit_denovo_signatures:
        exposures.to_csv(
            layer_directory
            + "/Activities"
            + "/"
            + solution_prefix
            + "_"
            + "Activities_refit.txt",
            "\t",
            index_label=[exposures.columns.name],
        )
    else:
        exposures.to_csv(
            layer_directory
            + "/Activities"
            + "/"
            + solution_prefix
            + "_"
            + "Activities.txt",
            "\t",
            index_label=[exposures.columns.name],
        )

    tmb_exposures = pd.melt(exposures)
    if make_plots:
        if refit_denovo_signatures:
            tmb.plotTMB(
                tmb_exposures,
                scale=sequence,
                Yrange="adapt",
                output=layer_directory
                + "/Activities"
                + "/"
                + solution_prefix
                + "_"
                + "TMB_plot_refit.pdf",
            )
        else:
            tmb.plotTMB(
                tmb_exposures,
                scale=sequence,
                Yrange="adapt",
                output=layer_directory
                + "/Activities"
                + "/"
                + solution_prefix
                + "_"
                + "TMB_plot.pdf",
            )
        del tmb_exposures

    if make_plots:
        if refit_denovo_signatures:
            plot_ac.plotActivity(
                layer_directory
                + "/Activities"
                + "/"
                + solution_prefix
                + "_"
                + "Activities_refit.txt",
                output_file=layer_directory
                + "/Activities/"
                + solution_prefix
                + "_"
                + "Activity_Plots_refit.pdf",
                bin_size=50,
                log=False,
            )
        else:
            plot_ac.plotActivity(
                layer_directory
                + "/Activities"
                + "/"
                + solution_prefix
                + "_"
                + "Activities.txt",
                output_file=layer_directory
                + "/Activities/"
                + solution_prefix
                + "_"
                + "Activity_Plots.pdf",
                bin_size=50,
                log=False,
            )

    est_genomes = np.dot(processAvg, exposureAvg)
    all_similarities, cosine_similarities = calculate_similarities(
        allgenomes, est_genomes, allcolnames
    )
    all_similarities.iloc[:, [3, 5]] = (
        all_similarities.iloc[:, [3, 5]].astype(str) + "%"
    )

    if refit_denovo_signatures:
        all_similarities.to_csv(
            layer_directory
            + "/Solution_Stats/"
            + solution_prefix
            + "_Samples_Stats_refit.txt",
            sep="\t",
        )
    else:
        all_similarities.to_csv(
            layer_directory
            + "/Solution_Stats/"
            + solution_prefix
            + "_Samples_Stats.txt",
            sep="\t",
        )

    if refit_denovo_signatures:
        try:
            process_std_error = pd.DataFrame(process_std_error)
            processSTE = process_std_error.set_index(index)
            processSTE.columns = allsigids
            processSTE = processSTE.rename_axis("MutationType", axis="columns")
            processSTE.to_csv(
                layer_directory
                + "/Signatures"
                + "/"
                + solution_prefix
                + "_"
                + "Signatures_SEM_Error.txt",
                "\t",
                float_format="%.2E",
                index_label=[processes.columns.name],
            )
        except:
            pass
    if refit_denovo_signatures:
        try:
            signature_stats = signature_stats.set_index(allsigids)
            signature_stats = signature_stats.rename_axis("Signatures", axis="columns")
            signature_stats.to_csv(
                layer_directory
                + "/Solution_Stats"
                + "/"
                + solution_prefix
                + "_"
                + "Signatures_Stats.txt",
                "\t",
                index_label=[exposures.columns.name],
            )
            signature_total_mutations = np.sum(exposureAvg, axis=1).astype(int)
            signature_total_mutations = signature_plotting_text(
                signature_total_mutations, "Sig. Mutations", "integer"
            )
        except:
            pass
    else:
        signature_total_mutations = np.sum(exposureAvg, axis=1).astype(int)
        signature_total_mutations = signature_plotting_text(
            signature_total_mutations, "Sig. Mutations", "integer"
        )
        if (m == "1536" or m == "288") and collapse_to_SBS96 == True:
            m = "96"
    if make_plots:
        if m == "DINUC" or m == "78":
            plot.plotDBS(
                layer_directory
                + "/Signatures/"
                + solution_prefix
                + "_"
                + "Signatures.txt",
                layer_directory + "/Signatures/",
                solution_prefix,
                "78",
                True,
                custom_text_upper=signature_stabilities,
                custom_text_middle=signature_total_mutations,
                volume=volume,
            )
        elif m == "INDEL" or m == "83":
            plot.plotID(
                layer_directory
                + "/Signatures/"
                + solution_prefix
                + "_"
                + "Signatures.txt",
                layer_directory + "/Signatures/",
                solution_prefix,
                "94",
                True,
                custom_text_upper=signature_stabilities,
                custom_text_middle=signature_total_mutations,
                volume=volume,
            )
        elif m == "CNV" or m == "48":
            plot.plotCNV(
                layer_directory
                + "/Signatures/"
                + solution_prefix
                + "_"
                + "Signatures.txt",
                layer_directory + "/Signatures/",
                solution_prefix,
                percentage=True,
                aggregate=False,
            )
        elif m == "SV" or m == "32":
            plot.plotSV(
                layer_directory
                + "/Signatures/"
                + solution_prefix
                + "_"
                + "Signatures.txt",
                layer_directory + "/Signatures/",
                solution_prefix,
                percentage=True,
                aggregate=False,
            )
        elif (
            m == "96" or m == "288" or m == "384" or m == "1536" or m == "4608"
        ) and collapse_to_SBS96 == True:
            plot.plotSBS(
                layer_directory
                + "/Signatures/"
                + solution_prefix
                + "_"
                + "Signatures.txt",
                layer_directory + "/Signatures/",
                solution_prefix,
                m,
                True,
                custom_text_upper=signature_stabilities,
                custom_text_middle=signature_total_mutations,
                volume=volume,
            )
        elif m == "96":
            plot.plotSBS(
                layer_directory
                + "/Signatures/"
                + solution_prefix
                + "_"
                + "Signatures.txt",
                layer_directory + "/Signatures/",
                solution_prefix,
                m,
                True,
                custom_text_upper=signature_stabilities,
                custom_text_middle=signature_total_mutations,
                volume=volume,
            )
        elif m == "288":
            plot.plotSBS(
                layer_directory
                + "/Signatures/"
                + solution_prefix
                + "_"
                + "Signatures.txt",
                layer_directory + "/Signatures/",
                solution_prefix,
                m,
                True,
                custom_text_upper=signature_stabilities,
                custom_text_middle=signature_total_mutations,
                volume=volume,
            )
        elif m == "384":
            plot.plotSBS(
                layer_directory
                + "/Signatures/"
                + solution_prefix
                + "_"
                + "Signatures.txt",
                layer_directory + "/Signatures/",
                solution_prefix,
                m,
                True,
                custom_text_upper=signature_stabilities,
                custom_text_middle=signature_total_mutations,
                volume=volume,
            )
        elif m == "1536":
            plot.plotSBS(
                layer_directory
                + "/Signatures/"
                + solution_prefix
                + "_"
                + "Signatures.txt",
                layer_directory + "/Signatures/",
                solution_prefix,
                m,
                True,
                custom_text_upper=signature_stabilities,
                custom_text_middle=signature_total_mutations,
                volume=volume,
            )
        elif m == "4608":
            plot.plotSBS(
                layer_directory
                + "/Signatures/"
                + solution_prefix
                + "_"
                + "Signatures.txt",
                layer_directory + "/Signatures/",
                solution_prefix,
                m,
                True,
                custom_text_upper=signature_stabilities,
                custom_text_middle=signature_total_mutations,
                volume=volume,
            )
        else:
            custom_signatures_plot(processes, layer_directory + "/Signatures")

    if export_probabilities:
        probability = probabilities(
            processAvg, exposureAvg, index, allsigids, allcolnames
        )
        probability = probability.set_index("Sample Names")

        if denovo_refit_option:
            if refit_denovo_signatures:
                probability.to_csv(
                    layer_directory
                    + "/Activities"
                    + "/"
                    + "De_Novo_MutationType_Probabilities_refit.txt",
                    "\t",
                )
            else:
                probability.to_csv(
                    layer_directory
                    + "/Activities"
                    + "/"
                    + "De_Novo_MutationType_Probabilities.txt",
                    "\t",
                )
        if not denovo_refit_option:
            probability.to_csv(
                layer_directory
                + "/Activities"
                + "/"
                + "Decomposed_MutationType_Probabilities.txt",
                "\t",
            )

    if export_probabilities_per_mutation:
        if export_probabilities:
            if input_type == "vcf":
                if m == "96" or m == "78" or m == "83":
                    (
                        probability_per_mutation,
                        samples_prob_per_mut,
                    ) = probabilities_per_mutation(probability, samples, m, exome)

                    if denovo_refit_option:
                        if refit_denovo_signatures:
                            ppm_file_name = "De_Novo_Mutation_Probabilities_refit"
                            output_path_prob_per_mut = (
                                layer_directory + "/Activities" + "/" + ppm_file_name
                            )
                        else:
                            ppm_file_name = "De_Novo_Mutation_Probabilities"
                            output_path_prob_per_mut = (
                                layer_directory + "/Activities" + "/" + ppm_file_name
                            )
                    else:
                        ppm_file_name = "Decomposed_Mutation_Probabilities"
                        output_path_prob_per_mut = (
                            layer_directory + "/Activities" + "/" + ppm_file_name
                        )

                    if not os.path.exists(output_path_prob_per_mut):
                        os.makedirs(output_path_prob_per_mut)
                    for matrix, sample in zip(
                        probability_per_mutation, samples_prob_per_mut
                    ):
                        matrix = matrix.set_index("Sample Names")
                        matrix = matrix.sort_values(by=["Chr", "Pos"])
                        matrix.to_csv(
                            layer_directory
                            + "/Activities"
                            + "/"
                            + ppm_file_name
                            + "/"
                            + ppm_file_name
                            + "_"
                            + sample
                            + ".txt",
                            "\t",
                        )
                else:
                    print(
                        "Probabilities per mutation are only calculated for SBS96, DBS78 and ID83 mutational contexts."
                    )
            else:
                print(
                    'Probabilities per mutation are only calculated if input_type is "vcf".'
                )
        else:
            print(
                "Probabilities per mutation require to calculate probabilities per context type. Please re-run your analysis setting export_probabilites=True."
            )

    return exposures


################################################################### FUNCTION ONE ###################################################################
# function to calculate multiple similarities/distances
def calculate_similarities(genomes, est_genomes, sample_names=False):
    from numpy import inf

    if sample_names is False:
        sample_names = ["None"] * genomes.shape[1]

    cosine_similarity_list = []
    kl_divergence_list = []
    correlation_list = []
    l1_norm_list = []
    l2_norm_list = []
    total_mutations_list = []
    relative_l1_list = []
    relative_l2_list = []

    for i in range(genomes.shape[1]):
        p_i = genomes[:, i]
        q_i = est_genomes[:, i]
        cosine_similarity_list.append(round(cos_sim(p_i, q_i), 3))
        kl_divergence_list.append(round(scipy.stats.entropy(p_i, q_i), 5))
        correlation_list.append(round(scipy.stats.pearsonr(p_i, q_i)[0], 3))
        l1_norm_list.append(round(np.linalg.norm(p_i - q_i, ord=1), 3))
        relative_l1_list.append(
            round((l1_norm_list[-1] / np.linalg.norm(p_i, ord=1)) * 100, 3)
        )
        l2_norm_list.append(round(np.linalg.norm(p_i - q_i, ord=2), 3))
        relative_l2_list.append(
            round((l2_norm_list[-1] / np.linalg.norm(p_i, ord=2)) * 100, 3)
        )
        total_mutations_list.append(np.sum(p_i))

    kl_divergence_list = np.array(kl_divergence_list)
    kl_divergence_list[kl_divergence_list == inf] = 1000
    similarities_dataframe = pd.DataFrame(
        {
            "Sample Names": sample_names,
            "Total Mutations": total_mutations_list,
            "Cosine Similarity": cosine_similarity_list,
            "L1 Norm": l1_norm_list,
            "L1_Norm_%": relative_l1_list,
            "L2 Norm": l2_norm_list,
            "L2_Norm_%": relative_l2_list,
            "KL Divergence": kl_divergence_list,
            "Correlation": correlation_list,
        }
    )
    similarities_dataframe = similarities_dataframe.set_index("Sample Names")
    return [similarities_dataframe, cosine_similarity_list]


############################################################### FUNCTIONS TO CALCULATE DISTANCES BETWEEN VECTORS ##################################################
################################################################### 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 cor_sim(a, b):
    """Takes 2 vectors a, b and returns the corrilation 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
    corr = 1 - cor(a, b)
    return corr


################################################### Generation of probabilities for each processes given to A mutation type ############################################
def probabilities(W, H, index, allsigids, allcolnames):
    # setting up the indices
    rows = index
    cols = allcolnames
    sigs = allsigids

    W = np.array(W)
    H = np.array(H)
    # rebuild the original matrix from the estimated W and H
    genomes = np.dot(W, H)

    result = 0
    for i in range(H.shape[1]):  # here H.shape is the number of sample
        M = genomes[:, i][np.newaxis]
        probs = W * H[:, i] / M.T
        probs = pd.DataFrame(probs)
        probs.columns = sigs
        col1 = [cols[i]] * len(rows)
        probs.insert(loc=0, column="Sample Names", value=col1)
        probs.insert(loc=1, column="MutationType", value=rows)
        if i != 0:
            result = pd.concat([result, probs], axis=0)
        else:
            result = probs

    return result


################################################### Generation of probabilities for each processes given to A mutation ############################################
def probabilities_per_mutation(probability_matrix, samples_path, m, exome=False):

    probability_matrix = probability_matrix.reset_index()

    if m == "96":
        seqinfo_path = samples_path + "/output/vcf_files/SNV/"
        interval_low = 3
        interval_high = -1
    if m == "78":
        seqinfo_path = samples_path + "/output/vcf_files/DBS/"
        interval_low = 4
        interval_high = -2
    if m == "83":
        seqinfo_path = samples_path + "/output/vcf_files/ID/"
        interval_low = 2
        interval_high = 100

    seqinfo_files = os.listdir(seqinfo_path)
    seqinfo_files.sort()

    all_mutations = pd.DataFrame()
    for file in seqinfo_files:
        if "exome" in file:
            exome_df = pd.read_csv(seqinfo_path + file, sep="\t", header=None)
        else:
            try:
                new = pd.read_csv(seqinfo_path + file, sep="\t", header=None)
                all_mutations = pd.concat([all_mutations, new])
            except pd.errors.EmptyDataError:
                pass

    all_mutations[3] = all_mutations[3].str[interval_low:interval_high]
    if m == "96" or m == "78":
        del all_mutations[4]
    else:
        del all_mutations[6]
        del all_mutations[5]
        del all_mutations[4]

    all_mutations.columns = ["Sample Names", "Chr", "Pos", "MutationType"]
    if exome == True:
        del exome_df[2]
        del exome_df[3]
        del exome_df[4]
        exome_df.columns = ["Chr", "Pos"]
        exome_df["Chr"] = [str(x) for x in (exome_df["Chr"]).to_list()]
        all_mutations["Chr"] = [str(x) for x in (all_mutations["Chr"]).to_list()]
        all_mutations = pd.merge(all_mutations, exome_df)

    all_samples_mutations = [y for x, y in all_mutations.groupby("Sample Names")]

    prob_per_mut = []
    sample_names = []
    for sample_mutations in all_samples_mutations:
        new = sample_mutations.merge(probability_matrix)
        prob_per_mut.append(new)
        sample_names.append(new["Sample Names"][0])

    result = [prob_per_mut, sample_names]

    return result


def custom_signatures_plot(signatures, output):
    with PdfPages(output + "/Custom_Signature_Plots.pdf") as pdf:
        plt.figure(figsize=(10, 3))
        plt.bar(list(range(1, 1 + len(signatures.iloc[:, 0]))), signatures.iloc[:, 0])
        plt.title("Custom Signature {}".format(0 + 1))
        plt.xticks([])
        plt.xlabel("Mutation Types")
        plt.ylabel("Probabilities")
        pdf.savefig()  # saves the current figure into a pdf page
        plt.close()
        for i in range(1, signatures.shape[1]):
            # if LaTeX is not installed or error caught, change to `usetex=False`
            plt.rc("text", usetex=False)
            plt.figure(figsize=(10, 3))
            plt.bar(
                list(range(1, 1 + len(signatures.iloc[:, i]))), signatures.iloc[:, i]
            )
            plt.title("Custom Signature {}".format(i + 1))
            plt.xticks([])
            plt.xlabel("Mutation Types")
            plt.ylabel("Probabilities")
            pdf.attach_note("signature plots")
            pdf.savefig()
            plt.close()


def merge_pdf(input_folder, output_file):
    pdf2merge = []
    for filename in os.listdir(input_folder):
        # print(filename)
        if filename.endswith(".pdf"):
            pdf2merge.append(filename)

    pdf2merge.sort()
    pdfWriter = pypdf.PdfFileWriter()
    for filename in pdf2merge:
        pdfFileObj = open(input_folder + "/" + filename, "rb")
        pdfReader = pypdf.PdfFileReader(pdfFileObj)
        for pageNum in range(pdfReader.numPages):
            pageObj = pdfReader.getPage(pageNum)
            pdfWriter.addPage(pageObj)

    pdfOutput = open(output_file + ".pdf", "wb")
    pdfWriter.write(pdfOutput)
    # Outputting the PDF
    pdfOutput.close()
