# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/SequencePlot.ipynb (unless otherwise specified).

__all__ = ['format_uniprot_annotation', 'ptm_shape_dict', 'get_plot_data', 'plot_single_peptide_traces',
           'custom_color_palettes', 'uniprot_color_dict', 'aa_color_dict', 'get_quality_category',
           'get_exposure_category', 'get_alphafold_annotation', 'plot_peptide_traces', 'extract_annotation',
           'manipulate_cif', 'adjust_html', 'get_ms_concensus', 'plot_3d_structure', 'create_pdf_report']

# Cell
import os
import tempfile

# Cell
import pandas as pd

def format_uniprot_annotation(uniprot_ann: pd.DataFrame, uniprot_feature_dict: dict):
    """
    Function to format uniprot annotation for plotting.

    Args:
        uniprot_ann (pd.DataFrame): Formatted uniprot annotations from alphamap.
        uniprot_feature_dict (dict): Uniprot feature dictionary defined by alphamap.
    Returns:
        pd.DataFrame: Uniprot annotation with a combined structure entry for helix, strand and turn.

    """
    uniprot = uniprot_ann.copy(deep=True)
    uniprot.loc[uniprot.feature == "HELIX", "note"] = "Helix"
    uniprot.loc[uniprot.feature == "STRAND", "note"] = "Beta strand"
    uniprot.loc[uniprot.feature == "TURN", "note"] = "Turn"
    uniprot.loc[uniprot.feature.isin(["HELIX","STRAND","TURN"]), "feature"] = "STRUCTURE"

    uniprot_feature_dict_rev = {v: k for k, v in uniprot_feature_dict.items()}

    uniprot['annotation'] = uniprot['note']
    uniprot.loc[uniprot['annotation'].isnull(), 'annotation'] = uniprot['feature']
    uniprot = uniprot.replace({"annotation": uniprot_feature_dict_rev})
    return uniprot

# Cell

# The PTM shape dictionary defines the shapes for each modification, so they are consistent across datasets and analyses.

ptm_shape_dict = {
    '[Acetyl (K)]': 5,
    '[Acetyl (Protein N-term)]': 5,
    '[Carbamidomethyl (C)]': 3,
    '[Oxidation (M)]': 4,
    '[Phospho (STY)]': 0,
    '[GlyGly (K)]': 2,
    '[Methyl (KR)]': 6,
    '[Dimethyl (KR)]': 6,
    '[Trimethyl (K)]': 6,
    '[Pro5]': 9,
    '[Pro6]': 10,
    '[Glu->pyro-Glu]': 11,
    '[Gln->pyro-Glu]': 12,
    '[QQTGG (K)]': 13,
    '[Deamidation (N)]': 14,
    '[Deamidation 18O (N)]': 14,
    '[Deamidation (NQ)]': 14,
    '[Hydroxyproline]': 15,
    '[Carbamyl (N-term)]': 16,
    '[Delta:H(2)C(2) (N-term)]': 19,
    '[Dioxidation (MW)]': 4,
    '[Trioxidation (C)]': 4,
    '[Dethiomethyl (M)]': 20,
    '[Cation:Na (DE)]': 21,
    '[Methyl (E)]': 6,
    '[Dehydrated (ST)]': 23,
    '[Oxidation (P)]': 4,
    '[Dimethyl (K)]': 6,
    '[Amidated (Protein C-term)]': 22,
    '[Sulfo (STY)]': 24,
    '[Acetyl (N-term)]': 5,
    '[Amidated (C-term)]': 22,
    '[Sulfation (Y)]': 25,
    '[Phospho (ST)]': 0,
    '[Cys-Cys]': 26,
    '[Cysteinyl]': 27,
    '[Cysteinyl - carbamidomethyl]': 28,
    '[Oxidation (MP)]': 4
}

# Cell
import numpy as np
import pandas as pd
from pyteomics import fasta

def get_plot_data(protein,df,fasta):
    """
    Function to format experimental data for plotting.

    Args:
        protein (str): Uniprot protein accession.
        df (pd.DataFrame): Experimental data imported and formatted according to alphamap standards.
        fasta (fasta): Fasta file imported by pyteomics 'fasta.IndexedUniProt'.
    Returns:
        pd.DataFrame: Formatted dataframe for plotting.

    """
    protein_sequence = fasta[protein].sequence
    df_prot = df[df.unique_protein_id==protein]

    if df_prot.shape[0] == 0:
        df_plot = None
    else:
        df_peps = [np.arange(row['start'], row['end']+1) for _, row in df_prot.iterrows()]
        df_peps  = pd.DataFrame.from_records(data=df_peps)
        df_peps['modified_sequence'] = df_prot['modified_sequence'].values
        df_peps['all_protein_ids'] = df_prot['all_protein_ids'].values
        df_peps = df_peps.melt(id_vars=['modified_sequence','all_protein_ids'])
        df_peps = df_peps[['modified_sequence','all_protein_ids','value']].dropna()
        df_peps = df_peps.rename(columns={"value": "seq_position"})
        df_peps['marker_symbol'] = 1
        df_peps['marker_size'] = 8
        df_peps['PTM'] = np.NaN
        df_peps['PTMtype'] = np.NaN
        df_peps['PTMshape'] = np.NaN
        unique_pep = df_peps.modified_sequence.unique()
        for uid in unique_pep:
            df_peps_uid = df_peps[df_peps.modified_sequence==uid]
            start_uid = np.min(df_peps_uid.seq_position)
            end_uid = np.max(df_peps_uid.seq_position)
            df_peps['marker_symbol'] = np.where(df_peps.seq_position == start_uid, 7, df_peps.marker_symbol)
            df_peps['marker_symbol'] = np.where(df_peps.seq_position == end_uid, 8, df_peps.marker_symbol)
            df_peps['marker_size'] = np.where(df_peps.seq_position == start_uid, 6, df_peps.marker_size)
            df_peps['marker_size'] = np.where(df_peps.seq_position == end_uid, 6, df_peps.marker_size)

            df_PTMs_uid = df_prot[df_prot.modified_sequence==uid]
            #PTMsites = df_PTMs_uid.PTMsites.tolist()[0] + start_uid
            PTMsites = [x+start_uid for x in df_PTMs_uid.PTMsites.tolist()[0]]
            PTMtypes = df_PTMs_uid.PTMtypes.tolist()[0]

            #print(PTMsites)
            #print(PTMtypes)

            for i in range(0,len(PTMsites)):
                df_peps['PTM'] = np.where(df_peps["seq_position"]==PTMsites[i], 1, df_peps.PTM)
                df_peps['PTMtype'] = np.where(df_peps["seq_position"]==PTMsites[i], PTMtypes[i], df_peps.PTMtype)

            #print(df_peps)

            df_seq = pd.DataFrame({'seq_position':np.arange(0,len(protein_sequence))})

            df_plot = pd.merge(df_seq, df_peps, how='left', on='seq_position')
            df_plot['height']=0
            df_plot['color']="grey"

            unique_mods = df_plot['PTMtype'].dropna().unique()
            if len(unique_mods) > 0:
                for mod in df_plot['PTMtype'].dropna().unique():
                    if mod != 'nan':
                        #print(mod)
                        if mod not in ptm_shape_dict.keys():
                            ptm_shape_dict.update({mod : 17})

                        df_plot.loc[df_plot.PTMtype == mod, 'PTMshape'] = ptm_shape_dict[mod]
                        #df_plot.loc[df_plot.PTMtype == mod, 'PTMshape'] = 17

        #print(df_plot)

        df_plot['modified_sequence']= df_plot['modified_sequence'].astype(str)
        df_plot = df_plot.groupby('seq_position').agg({'modified_sequence':'; '.join,
                                                       'all_protein_ids': 'first',
                                                       'marker_symbol': 'first',
                                                       'marker_size': 'first',
                                                       'PTM': 'first',
                                                       'PTMtype': 'first',
                                                       'PTMshape': 'first',
                                                       'height': 'first',
                                                       'color': 'first'
                                                      }).reset_index()

        #print(df_plot)

    return(df_plot)

# Cell
import plotly.graph_objects as go

def plot_single_peptide_traces(df_plot,protein,fasta):
    """
    Function to plot single peptide trace.

    Args:
        df_plot (pd.DataFrame): Formatted dataframe for plotting, generated by get_plot_data.
        protein (str): Uniprot protein accession.
        fasta (fasta): Fasta file imported by pyteomics 'fasta.IndexedUniProt'.
    Returns:
        go.Figure: Figure data for a single dataset.

    """
    protein_sequence = fasta[protein].sequence
    entry_name = fasta[protein].description['GN']
    protein_name = fasta[protein].description['name']

    plot0 = go.Scatter(y=[None],
                       name='',
                       xaxis='x1',
                       showlegend=False)

    ## Peptide backbone
    df_plot_pep = df_plot.dropna(subset=['modified_sequence'])
    df_plot_pep = df_plot_pep[~df_plot_pep.modified_sequence.str.contains('nan')]
    plot1 = go.Scatter(x=df_plot_pep.seq_position+1,
                       y=df_plot.height,
                       xaxis='x2',
                       mode='markers',
                       marker_size=df_plot_pep.marker_size,
                       marker_symbol=df_plot_pep.marker_symbol,
                       marker_line_color=df_plot_pep.color,
                       marker_color=df_plot_pep.color,
                       marker_opacity=1,
                       meta=df_plot_pep.modified_sequence,
                       text=df_plot_pep.all_protein_ids,
                       hovertemplate ='Peptide: %{meta}<br>' +
                       'Protein IDs: %{text}',
                       name='',
                       showlegend=False)

    covered_AA = len(df_plot_pep.seq_position.unique())
    percent_AA_coverage = int(np.round(100/len(protein_sequence)*covered_AA))
    #print(percent_AA_coverage)

    ## PTM dots
    df_plot_ptm = df_plot.dropna(subset=['PTM'])
    #print(df_plot_ptm)
    plot2 = go.Scatter(x=df_plot_ptm.seq_position+1,
                       y=df_plot_ptm.height+0.3,
                       xaxis='x2',
                       mode='markers',
                       marker_size=8,
                       marker_symbol=df_plot_ptm.PTMshape,
                       marker_line_color=df_plot_ptm.color,
                       marker_color=df_plot_ptm.color,
                       marker_opacity=1,
                       text=df_plot_ptm.PTMtype,
                       hovertemplate = 'PTM: %{text}',
                       #hoverinfo='text',
                       name='',
                       showlegend=False)

    layout = go.Layout(
            yaxis=dict(
                title = "",
                ticks = None,
                showticklabels=False,
                range=[-1, 2],
                showgrid=False,
                zeroline=False
                ),
            xaxis1=dict(
                title= 'protein sequence',
                tickmode = 'array',
                range=[-10, len(protein_sequence)+10],
                tickvals = np.arange(1,len(protein_sequence)+1),
                ticktext = list(protein_sequence),
                tickangle=0,
                matches="x2",
                type="linear",
                anchor="y",
                showgrid=False,
                zeroline=False
                ),
            xaxis2=dict(
                title= 'AA position',
                tickmode = 'auto',
                range=[-10, len(protein_sequence)+10],
                tickangle=0,
                matches="x1",
                side="top",
                type="linear",
                anchor="y",
                showgrid=False,
                zeroline=False,
                tickformat = '.d'
                ),
        #showlegend=False,
        #height=400,
        #width=1000,
        plot_bgcolor='rgba(0,0,0,0)',
        title=f"Sequence plot for: {protein_name}<br>{entry_name} - {protein}",
        meta=percent_AA_coverage,
        margin = dict(l=20, r=20, t=150, b=20)
        )

    fig = go.Figure(data=[plot1,plot2,plot0], layout=layout)

    #print(fig.layout.meta)

    for i in range(0, df_plot_ptm.shape[0]):
            fig.add_shape(
                    dict(
                        type="line",
                        xref="x2",
                        x0=df_plot_ptm.seq_position.values[i] +1,
                        y0=df_plot_ptm.height.values[i],
                        x1=df_plot_ptm.seq_position.values[i] +1,
                        y1=df_plot_ptm.height.values[i]+0.3,
                        line=dict(
                            color=df_plot_ptm.color.values[i],
                            width=1
                        )
                    )
            )

    return fig

# Cell

# Color palettes used for the different uniprot annotation features
custom_color_palettes = {
    'col_greens':["#5C965D","#6AA16B","#77AC78","#84B786","#91C193","#9FCCA1","#B3DCB5","#C6EBC9"],
    'col_ornages':["#ff4800","#ff5400","#ff6000","#ff6d00","#ff7900","#ff8500","#ff9100","#ff9e00","#ffaa00","#ffb600"],
    'col_purples':["#ffa69e","#febaae","#fcb088","#d9f3e2","#b8f2e6","#aed9e0","#9baed9","#9199d5","#8783d1"],
    'col_turquises':["#00a9a5","#4e8098","#90c2e7"],
    'col_darkpinks':["#42033d","#6f0c59","#901468","#7c238c","#924ea6","#9c5eae"],
    'col_browns':["#5a2a27","#5c4742","#8d5b4c","#a5978b","#c4bbaf"]
}

# Cell

# The uniprot_color_dict maps each uniprot annotation feature to a color in the custom_color_palettes.
uniprot_color_dict = {'CHAIN': custom_color_palettes['col_greens'][0],
                      'INIT_MET': custom_color_palettes['col_greens'][1],
                      'PEPTIDE': custom_color_palettes['col_greens'][2],
                      'PROPEP': custom_color_palettes['col_greens'][3],
                      'SIGNAL': custom_color_palettes['col_greens'][4],
                      'TRANSIT': custom_color_palettes['col_greens'][5],

                      'COILED': custom_color_palettes['col_purples'][0],
                      'COMPBIAS': custom_color_palettes['col_purples'][1],
                      'DOMAIN': custom_color_palettes['col_purples'][2],
                      'MOTIF': custom_color_palettes['col_purples'][3],
                      'REGION': custom_color_palettes['col_purples'][4],
                      'REPEAT': custom_color_palettes['col_purples'][5],
                      'ZN_FING': custom_color_palettes['col_purples'][6],

                      'INTRAMEM': custom_color_palettes['col_turquises'][0],
                      'TOPO_DOM': custom_color_palettes['col_turquises'][1],
                      'TRANSMEM': custom_color_palettes['col_turquises'][2],

                      'STRUCTURE': 'black',
                      # extra structures
                      'Helix': '#5dabe8',
                      'Turn': '#e094bc',
                      'Beta strand': '#8cdbad',

                      'CROSSLNK': custom_color_palettes['col_ornages'][2],
                      'DISULFID': custom_color_palettes['col_ornages'][3],
                      'CARBOHYD': custom_color_palettes['col_ornages'][4],
                      'LIPID': custom_color_palettes['col_ornages'][5],
                      'MOD_RES': custom_color_palettes['col_ornages'][6],

                      'BINDING': custom_color_palettes['col_darkpinks'][0],
                      'CA_BIND': custom_color_palettes['col_darkpinks'][1],
                      'DNA_BIND': custom_color_palettes['col_darkpinks'][2],
                      'METAL': custom_color_palettes['col_darkpinks'][3],
                      'NP_BIND': custom_color_palettes['col_darkpinks'][4],
                      'SITE': custom_color_palettes['col_darkpinks'][5],

                      'NON_STD': custom_color_palettes['col_browns'][0],
                      'NON_CONS': custom_color_palettes['col_browns'][1],
                      'NON_TER': custom_color_palettes['col_browns'][2],
                      'VARIANT': custom_color_palettes['col_browns'][3],
                      'CONFLICT': custom_color_palettes['col_browns'][4],

                      'VAR_SEQ': '#fae7b1',
                      'UNSURE': 'grey',
                      'MUTAGEN': 'darkgrey',
                     }

# Cell

# Dictionary that maps one-letter amino acid abbreviations to their full name.
aa_color_dict = {'A':'Alanine',
                 'R':'Arginine',
                 'N':'Asparagine',
                 'D':'Aspartic acid',
                 'C':'Cysteine',
                 'E':'Glutamic acid',
                 'Q':'Glutamine',
                 'G':'Glycine',
                 'H':'Histidine',
                 'I':'Isoleucine',
                 'L':'Leucine',
                 'K':'Lysine',
                 'M':'Methionine',
                 'F':'Phenylalanine',
                 'P':'Proline',
                 'S':'Serine',
                 'T':'Threonine',
                 'W':'Tryptophan',
                 'Y':'Tyrosine',
                 'V':'Valine',
                 'X':'nan',
                 'U':'nan'}

# Cell
from structuremap.processing import download_alphafold_cif, download_alphafold_pae, format_alphafold_data, annotate_accessibility, get_smooth_score

# Cell
def get_quality_category(s):
    if s >= 90:
        n='Very high'
    elif s >= 70:
        n='Confident'
    elif s >= 50:
        n='Low'
    else:
        n='Very low'
    return(n)

# Cell
def get_exposure_category(s):
    if s > 10:
        n='Surrounded'
    elif s > 5:
        n='Many neighbors'
    elif s > 0:
        n='Few neighbors'
    else:
        n='No neighbors'
    return(n)

# Cell
def get_alphafold_annotation(protein: str,
                             selected_features: list,
                             download_folder: str = tempfile.gettempdir()) -> pd.DataFrame:

    alphafold_feature_dict = dict({'AlphaFold confidence':'quality',
        'AlphaFold exposure':'nAA_12_70_pae',
        'AlphaFold IDR':'AlphaFold IDR',
        'AlphaFold secondary structures':'structure_group'})

    alphafold_features = [i for i in [
        'AlphaFold confidence',
        'AlphaFold exposure',
        'AlphaFold IDR',
        'AlphaFold secondary structures'
    ] if i in selected_features]

    if len(alphafold_features) > 0:

        # download cif file
        valid_proteins, invalid_proteins, existing_proteins = download_alphafold_cif(
            proteins=[protein],
            out_folder=download_folder)
        is_available = np.sum([len(valid_proteins),len(existing_proteins)])

        if is_available > 0:

            af_data = format_alphafold_data(directory = download_folder, protein_ids = [protein])

            if len([i for i in ['AlphaFold exposure',
                                'AlphaFold IDR'] if i in selected_features]) > 0:
                # download pae file
                valid_pae, invalid_pae, existing_pae = download_alphafold_pae(
                    proteins=[protein],
                    out_folder=download_folder)
                is_available_pae = np.sum([len(valid_pae),len(existing_pae)])
                #@ToDo
                #if is_available == 0:
                #    print('error')

                # annotate exposure
                exposure_df = annotate_accessibility(df = af_data,
                                                     max_dist = 12,
                                                     max_angle = 70,
                                                     error_dir = download_folder)

                af_data = af_data.merge(exposure_df, how='left', on=['protein_id','AA','position'])

                # annotate exposure
                full_exposure_df = annotate_accessibility(df = af_data,
                                                          max_dist = 24,
                                                          max_angle = 180,
                                                          error_dir = download_folder)

                af_data = af_data.merge(full_exposure_df, how='left', on=['protein_id','AA','position'])

            # smooth scores
            smoothing_scores = []
            if len([i for i in ['AlphaFold IDR'] if i in alphafold_features]) > 0:
                smoothing_scores.append('nAA_24_180_pae')
            if len(smoothing_scores) > 0:
                af_data = get_smooth_score(af_data, np.array(smoothing_scores),[10])

            # annotate predicted IDRs
            if 'AlphaFold IDR' in alphafold_features:
                af_data['AlphaFold IDR'] = np.where(af_data['nAA_24_180_pae_smooth10']<=34.27, "IDR", "folded")

            af_columns = ['position','AA']
            af_columns = af_columns + [alphafold_feature_dict[i] for i in alphafold_features]
            af_data = af_data[af_columns]

            if 'quality' in af_data.columns:
                af_data["AlphaFold confidence"] = [get_quality_category(x) for x in af_data.quality]

            if 'nAA_12_70_pae' in af_data.columns:
                af_data["AlphaFold exposure"] = [get_exposure_category(x) for x in af_data.nAA_12_70_pae]

            if 'structure_group' in af_data.columns:
                af_data = af_data.rename(columns={"structure_group": "AlphaFold secondary structures"})

        else:
            af_data = pd.DataFrame({'AA' : []})
    else:
        af_data = pd.DataFrame({'AA' : []})

    return(af_data)

# Cell

import plotly.graph_objects as go
from .proteolytic_cleavage import get_cleavage_sites

def plot_peptide_traces(df: pd.DataFrame or list,
                        name: str or list,
                        protein: str,
                        fasta: fasta,
                        uniprot: pd.DataFrame,
                        selected_features: list,
                        uniprot_feature_dict: dict,
                        uniprot_color_dict: dict,
                        selected_proteases: list = [],
                        selected_alphafold_features: list = [],
                        dashboard: bool = False,
                        trace_colors: list = [],
                        download_folder: str = tempfile.gettempdir()):

    """
    Function to generate the sequence plot.

    Args:
        df (pd.DataFrame/list): Single dataframe or list of dataframes containing the datasets to plot.
        name (str/list): Single string or list of strings containing the names for each dataset in df.
        protein (str): Uniprot protein accession.
        fasta (fasta): Fasta file imported by pyteomics 'fasta.IndexedUniProt'.
        uniprot (pd.DataFrame): Uniprot annotations formatted by alphamap.
        selected_features (list): List of uniprot features to plot.
        uniprot_feature_dict (dict): Uniprot feature dictionary.
        uniprot_color_dict (dict): Uniprot color dictionary.
        selected_proteases (list, optional): List of proteases to plot. Default is an empty list.
        dashboard (bool, optional): Flag if the function is called from the dashboard. Default is 'False'.
        trace_colors (list, optional): List of manualy selected colors for each dataset in df. Default is an empty list.

    Returns:
        go.Figure: Sequence plot.

    """

    figure_height = 200

    protein_sequence = fasta[protein].sequence

    # colors for experimental data traces
    if len(trace_colors) == 0:
        colors = ["#023e8a","#0096c7","#90e0ef","#7fd14d","#26a96c",
                  "#0D7F8B", "#24B4AB", "#9DE7BE", "#1B9CC6", "#016699"]
    else:
        colors = trace_colors

    # generation of a reverse uniprot_feature_dict
    uniprot_feature_dict_rev = {v: k for k, v in uniprot_feature_dict.items()}
    #uniprot_feature_dict_rev["STRUCTURE"] = "Secondary structure"

    # subsetting of the uniprot annotation to the selected features
    uniprot_annotation_p = uniprot[uniprot.protein_id==protein]
    # formatting of uniprot annotations
    uniprot_annotation_p_f = format_uniprot_annotation(uniprot_annotation_p, uniprot_feature_dict)
    # subset for selected features
    uniprot_annotation_p_f_f = uniprot_annotation_p_f[uniprot_annotation_p_f.feature.isin(selected_features)]

    #@ToDo if

    if isinstance(df, pd.DataFrame):
        df_plot = get_plot_data(protein=protein,
                              df = df,
                              fasta = fasta)

        df_plot.color = colors[0]

        observed_mods = list(set(df_plot.PTMtype))
        ptm_shape_dict_sub = {key: ptm_shape_dict[key] for key in observed_mods if key in ptm_shape_dict}

        fig = plot_single_peptide_traces(df_plot,protein=protein,fasta = fasta)

        AA_coverage = fig.layout.meta
        trace_name = [name + "<br> (" + str(AA_coverage) + "% coverage)"]

        fig.update_layout(yaxis=dict(showticklabels=True,
                                     tickmode = 'array',
                                     tickvals = [0],
                                     ticktext = [name + "(" + str(AA_coverage) + "%)"],
                                     showgrid=False))

        y_max = 1

    elif isinstance(df, list):

        df_plot = [get_plot_data(protein=protein,
                               df = d,
                               fasta = fasta) for d in df]

        # Subset data and annotations for the samples where the selected protein was detected
        valid_idx = []
        for i in range(len(df_plot)):
            if df_plot[i] is not None:
                valid_idx.append(i)
        df_plot = [df_plot[i] for i in valid_idx]
        name = [name[i] for i in valid_idx]
        colors = [colors[i] for i in valid_idx]
        #observed_mods = set([df_plot[i].PTMtype for i in valid_idx])
        observed_mods = []
        for i in range(len(df_plot)):
            observed_mods.extend(list(set(df_plot[i].PTMtype)))
        observed_mods = list(set(observed_mods))
        ptm_shape_dict_sub = {key: ptm_shape_dict[key] for key in observed_mods if key in ptm_shape_dict}

        for i in range(len(df_plot)):
            df_plot[i].color = colors[i]
            df_plot[i].height = 1+i

        plot_list = [plot_single_peptide_traces(df,protein=protein,fasta = fasta) for df in df_plot]
        new_data = [p.data for p in plot_list]
        new_data = sum(new_data, ())
        new_layout = plot_list[0].layout
        shapes = [p.layout.shapes for p in plot_list]
        shapes = sum(shapes, ())
        new_layout.shapes = new_layout.shapes + tuple(shapes)
        AA_coverage = [p.layout.meta for p in plot_list]
        trace_name = [n + "<br> (" + str(c) + "% coverage)" for n,c in zip(name,AA_coverage)]

        fig = go.Figure(data=new_data, layout=new_layout)
        fig.update_layout(yaxis=dict(range=[0,len(df_plot)+1],
                                     showticklabels=True,
                                     tickmode = 'array',
                                     tickvals = np.arange(0, len(df_plot))+1,
                                     ticktext = np.array(trace_name),
                                     showgrid=False))

        y_max = len(df_plot)+1

        figure_height = figure_height + (len(df_plot)*50)


    ptm_shape_dict_sub = dict(sorted(ptm_shape_dict_sub.items()))
    for i in range(len(ptm_shape_dict_sub)):
        fig.add_trace(go.Scatter(y=[None],
                                 mode='markers',
                                 xaxis='x2',
                                 marker=dict(symbol=list(ptm_shape_dict_sub.values())[i],
                                             color='black'),
                                 name=list(ptm_shape_dict_sub.keys())[i],
                                 showlegend=True))

    all_uniprot_features = list(uniprot_color_dict.keys())
    available_features = list(set(uniprot_annotation_p_f_f.feature))
    unique_features = [x for x in all_uniprot_features if x in available_features]
    if len(unique_features) > 0:

        y_max = y_max+1

        for j in range(0,len(unique_features)):

            figure_height = figure_height + 50

            domain = unique_features[j]
            domain_info_sub = uniprot_annotation_p_f_f[uniprot_annotation_p_f_f.feature==domain].reset_index(drop=True)
            for i in range(0, domain_info_sub.shape[0]):
                start = int(domain_info_sub.start[i])
                end = domain_info_sub.end[i]
                if np.isnan(domain_info_sub.end[i]):
                    end=start #+1
                else:
                    end=int(end)

                if domain_info_sub.feature[i] == "STRUCTURE":
                    marker_col = uniprot_color_dict[domain_info_sub.annotation[i]]
                else:
                    marker_col = uniprot_color_dict[domain_info_sub.feature[i]]

                fig.add_trace(go.Bar(x=list(range(start,end+1)),
                                     y=list(np.repeat(0.2,end-start+1)),
                                     base=list(np.repeat(y_max+(j/2),end-start+1)-0.1),
                                     marker_color=marker_col,
                                     marker_line_width=0,
                                     opacity=0.8,
                                     showlegend=False,
                                     xaxis='x2',
                                     name='',
                                     text=np.repeat(domain_info_sub.annotation[i],len(range(start,end+1))),
                                     hovertemplate ='%{text}'
                                     #hovertext=domain_info_sub.annotation[i],
                                     #hoverinfo='text'
                                    ))
        fig.update_layout(barmode='stack', bargap=0, hovermode='x unified',hoverdistance=1)

        y_max = y_max + (len(unique_features)/2)

    selected_proteases = sorted(selected_proteases)
    if len(selected_proteases) > 0:

        y_max = y_max+1

        for u in range(0,len(selected_proteases)):

            figure_height = figure_height + 50

            protease = selected_proteases[u]
            sites = get_cleavage_sites(protein_sequence,protease)
            for s in sites:
                fig.add_trace(go.Bar(x=list(range(s+1,s+2)),
                                     y=[0.2],
                                     base=y_max+(u/2)-0.1,
                                     marker_color="grey",
                                     opacity=0.8,
                                     showlegend=False,
                                     xaxis='x2',
                                     name='',
                                     text=np.repeat(protease,len(range(s+1,s+2))),
                                     hovertemplate ='%{text}'
                                     #hovertext=protease,
                                     #hoverinfo='text'
                                    ))

        y_max = y_max + (len(selected_proteases)/2)


    alphafold_annotation = get_alphafold_annotation(protein = protein,
                                                    selected_features = selected_alphafold_features,
                                                    download_folder = download_folder)

    if alphafold_annotation.empty:
        #print('no alphafold')
        alphafold_features = []
    else:
        #print('alphafold')

        alphafold_features = [i for i in alphafold_annotation.columns if i in [
            "AlphaFold confidence",
            "AlphaFold exposure",
            "AlphaFold IDR",
            "AlphaFold secondary structures"]]

        alphafold_color_dict = {'HELX': '#5dabe8',
                                'STRN': '#8cdbad',
                                'TURN': '#e094bc',
                                'BEND': '#fac682',
                                'unstructured': '#aaaaaa',
                                'Very high':'#4035de',
                                'Confident':'#6aaff7',
                                'Low':'#f5d44e',
                                'Very low':'#ff8f4a',
                                'No neighbors': "#8FF5A2",
                                'Few neighbors': "#38BFA7",
                                'Many neighbors': "#367BC3",
                                'Surrounded': "#524582",
                                'IDR': "#B5B5B5",
                                'folded': "#367BC3"}

        y_max = y_max+1

        for j in range(0,len(alphafold_features)):

            figure_height = figure_height + 50

            af_column = alphafold_features[j]
            if af_column == "AlphaFold confidence":
                extra_annotation = alphafold_annotation["quality"].values
            elif af_column == "AlphaFold exposure":
                extra_annotation = alphafold_annotation["nAA_12_70_pae"].values
            else:
                extra_annotation = np.repeat("",alphafold_annotation.shape[0])

            for i in range(0, alphafold_annotation.shape[0]):
                start = alphafold_annotation.position[i]
                end = alphafold_annotation.position[i]

                alpha_col = alphafold_color_dict[alphafold_annotation[af_column][i]]

                fig.add_trace(go.Bar(x=list(range(start,end+1)),
                                     y=list(np.repeat(0.2,end-start+1)),
                                     base=y_max+(j/2)-0.1,
                                     marker_color=alpha_col,
                                     marker_line_width=0,
                                     opacity=0.8,
                                     showlegend=False,
                                     xaxis='x2',
                                     name='',
                                     text=np.repeat(str(alphafold_annotation[af_column][i]) + " " + str(extra_annotation[i]),len(range(start,end+1))),
                                     hovertemplate ='%{text}'
                                     #hovertext=domain_info_sub.annotation[i],
                                     #hoverinfo='text'
                                    ))

        y_max = y_max + (len(alphafold_features)/2)

    fig.add_trace(go.Scatter(x=np.arange(1,len(protein_sequence)+1,1),
                        y=np.repeat(0,len(protein_sequence)),
                        marker=dict(color='rgba(135, 206, 250, 0)'),
                        name='',
                        mode='markers',
                        xaxis='x2',
                        text=[aa_color_dict[x] for x in list(protein_sequence)],
                        #text=np.arange(1,len(protein_sequence)+1,1),
                        meta=list(protein_sequence),
                        hovertemplate ='<b>%{meta}: %{text}<b>',
                        showlegend=False))

    if figure_height < 500:
        figure_height = 500

    if dashboard:
        plot_width = None
    else:
        plot_width = 1500

    fig.update_layout(barmode='stack', bargap=0, hovermode='x unified',hoverdistance=1,
                      width=plot_width, height=figure_height)

    mapped_feature_names = [uniprot_feature_dict_rev.get(key) for key in unique_features]
    if isinstance(df, pd.DataFrame):
        trace_ax = list([0])
        trace_count = 1
    elif isinstance(df, list):
        trace_ax = list(np.arange(1,len(df_plot)+1,1))
        trace_count = len(df_plot) + 1
    #print(trace_ax)

    if len(unique_features) > 0:
        trace_count += 1
        uniprot_ax = list(np.arange(0,len(unique_features)/2, 0.5) + trace_count)
        trace_count = max(uniprot_ax)+0.5
        #print(uniprot_ax)
    else:
        uniprot_ax = []
    if len(selected_proteases) > 0:
        trace_count += 1
        protease_ax = list(np.arange(0, len(selected_proteases)/2,0.5) + trace_count)
        trace_count = max(protease_ax)+0.5
        #print(protease_ax)
    else:
        protease_ax = []
    if len(alphafold_features) > 0:
        trace_count += 1
        alphafold_ax = list(np.arange(0, len(alphafold_features)/2,0.5) + trace_count)
        trace_count = max(alphafold_ax)+0.5
        #print(alphafold_ax)
    else:
        alphafold_ax = []

    fig.update_yaxes(showticklabels=True,
                     #tickvals= np.arange(0, 1+len(unique_features)+len(selected_proteases)),
                     tickvals= trace_ax + uniprot_ax + protease_ax + alphafold_ax,
                     ticktext=np.hstack((np.array(trace_name),np.array(mapped_feature_names),np.array(selected_proteases),np.array(alphafold_features))),
                     automargin=True,
                     range=[-1, y_max + 0.2],
                     showgrid=False)

    #config = {'toImageButtonOptions': {'format': 'svg', # one of png, svg, jpeg, webp
    #                                   'filename': 'custom_image',
    #                                   'height': 500,
    #                                   'width': 1500,
    #                                   'scale': 1 # Multiply title/legend/axis/canvas sizes by this factor
    #                                  }
    #         }

    return fig #.show(config=config)

# Cell
def extract_annotation(df: pd.DataFrame):
    df = df[df.marker_symbol>0]
    df = df.astype({"PTMtype": str})
    df['MS data'] = ['detected' if m=='nan' else m for m in df.PTMtype]
    df = df[['seq_position','MS data']]
    df['MS data'] = [x.replace(" ", "") for x in df['MS data']]
    df = df.rename(columns={"seq_position": "position"})
    df['position'] = df['position']+1
    return df

# Cell
def manipulate_cif(protein: str,
                   MS_data: pd.DataFrame,
                   download_folder: str = tempfile.gettempdir(),
                   ):


    print(MS_data)

    alphafold_annotation = get_alphafold_annotation(protein = protein,
                                                    selected_features = ["AlphaFold confidence",
                                                                         "AlphaFold exposure",
                                                                         "AlphaFold IDR",
                                                                         "AlphaFold secondary structures"],
                                                    download_folder = download_folder)

    if alphafold_annotation.empty:
        return False

    else:
        mod_table = False
        AA_position = 0
        infile = os.path.join(download_folder, protein+'.cif')
        outfile = os.path.join(download_folder, protein+'_mod.cif')

        with open(infile) as f:
            lines = []
            for line in f:
                line = line.rstrip()
                if mod_table:
                    if line == '#':
                        mod_table = False
                        lines.append(line)
                    else:
                        AA_position += 1

                        # append exposure
                        exposure = str(alphafold_annotation[alphafold_annotation.position==AA_position].nAA_12_70_pae.values[0])
                        n_extra_blank = 4-len(exposure)
                        appended_line = line+' '+exposure+(n_extra_blank*' ')
                        # append MS data
                        if AA_position in list(MS_data.position.values):
                            appended_line = appended_line+' '+'detected'+' '
                            ms_evidence = MS_data[MS_data.position==AA_position]['MS data'].values[0]
                            n_extra_blank = 40-len(ms_evidence)
                            appended_line = appended_line+' '+ms_evidence+(n_extra_blank*' ')
                        else:
                            appended_line = appended_line+' '+'absent  '+' '
                            n_extra_blank = 40-len('absent')
                            appended_line = appended_line+' '+'absent'+(n_extra_blank*' ')
                        # append PTM data
                        ptm_sub = MS_data[MS_data['MS data'] != 'detected']
                        if AA_position in list(ptm_sub.position.values):
                            ptm_evidence = ptm_sub[ptm_sub.position==AA_position]['MS data'].values[0]
                            n_extra_blank = 40-len(ptm_evidence)
                            appended_line = appended_line+' '+ptm_evidence+(n_extra_blank*' ')
                        else:
                            n_extra_blank = 40-len('none')
                            appended_line = appended_line+' '+'none'+(n_extra_blank*' ')
                        # append IDR
                        IDR_info = alphafold_annotation[alphafold_annotation.position==AA_position]['AlphaFold IDR'].values[0]
                        n_extra_blank = 10-len(IDR_info)
                        appended_line = appended_line+' '+IDR_info+(n_extra_blank*' ')
                        # append IDR + PTM data
                        if AA_position in list(ptm_sub.position.values):
                            n_extra_blank = 40-len(ptm_evidence)
                            appended_line = appended_line+' '+ptm_evidence+(n_extra_blank*' ')
                        else:
                            n_extra_blank = 40-len(IDR_info)
                            appended_line = appended_line+' '+IDR_info+(n_extra_blank*' ')
                        # append structure group
                        structure_info = alphafold_annotation[alphafold_annotation.position==AA_position]['AlphaFold secondary structures'].values[0]
                        n_extra_blank = 15-len(structure_info)
                        appended_line = appended_line+' '+structure_info+(n_extra_blank*' ')

                        lines.append(appended_line)
                else:
                    if line == '_ma_qa_metric_local.ordinal_id':
                        lines.append(line)
                        lines.append('_ma_qa_metric_local.exposure')
                        lines.append('_ma_qa_metric_local.aa')
                        lines.append('_ma_qa_metric_local.aa_ptm')
                        lines.append('_ma_qa_metric_local.ptm')
                        lines.append('_ma_qa_metric_local.idr')
                        lines.append('_ma_qa_metric_local.idr_ptm')
                        lines.append('_ma_qa_metric_local.structure_group')
                        mod_table = True
                    else:
                        lines.append(line)

        try:
            os.remove(outfile)
        except:
            pass

        with open(outfile, 'a') as outfile:
            for l in lines:
                outfile.write(l+'\n')

        return True

# Cell
def adjust_html(protein: str,
                coloring: str,
                data_name: str):

    html_template = """
                    <!DOCTYPE html>
                    <html lang="en">
                      <head>
                        <meta charset="utf-8" />
                        <meta name="viewport" content="width=device-width, user-scalable=no, minimum-scale=1.0, maximum-scale=1.0">
                        <title>PDBe Molstar</title>

                        <!-- Molstar CSS & JS -->
                        <link rel="stylesheet" type="text/css" href="js/alphafold-viz.css">
                        <script type="text/javascript" src="js/alphafold-viz.js"></script>
                        <style>
                          * {
                              margin: 0;
                              padding: 0;
                              box-sizing: border-box;
                          }
                          .msp-plugin ::-webkit-scrollbar-thumb {
                              background-color: #5b5b68 !important;
                          }
                          .viewerSection {
                            padding-top: 100px;
                          }
                          .controlsSection {
                            width: 300px;
                            float: left;
                            padding: 40px 0 0 40px;
                            margin-right: 30px;
                          }
                          .controlBox {
                            border: 1px solid lightgray;
                            padding: 10px;
                            margin-bottom: 20px;
                          }
                          #myViewer{
                            float:left;
                            width: 800px;
                            height: 800px;
                            position:relative;
                          }
                        </style>
                      </head>

                      <body>



                        <div class="viewerSection">
                          <!-- <h4>3D viewer</h4> -->
                          <h4>   </h4>

                          <!-- Molstar container -->
                          <div id="myViewer"></div>

                        </div>
                        <script>

                          globalThis.color_choice = your_choice //'af', 'access', 'ptm'

                          //Create plugin instance
                          var viewerInstance = new PDBeMolstarPlugin();

                          //Set options (Checkout available options list in the documentation)
                          var options = {
                            customData: {
                              url: 'cif/your_file.cif',
                              format: 'cif'
                            },
                            hideControls: false,
                            // expanded: true,
                            isAfView: true,
                            hideCanvasControls: [] //['selection', 'animation', 'controlToggle', 'controlInfo']
                          }

                          //Get element from HTML/Template to place the viewer
                          var viewerContainer = document.getElementById('myViewer');

                          //Call render method to display the 3D view
                          viewerInstance.render(viewerContainer, options);
                          viewerInstance.canvas.setBgColor({r:255, g:255, b:255});

                        </script>
                      </body>

                    </html>
                    """

    html_template = html_template.replace('your_file.cif',protein+'_'+data_name+'_mod.cif')

    if coloring == 'AlphaFold confidence':
        color = 'af'
    elif coloring == 'AlphaFold exposure':
        color = 'exposure'
    elif coloring == 'MS peptides':
        color = 'aa'
    elif coloring == 'MS modified peptides':
        color = 'aa_ptm'
    elif coloring == 'MS PTMs':
        color = 'ptm'
    elif coloring == 'AlphaFold IDR':
        color = 'idr'
    elif coloring == 'MS PTMs on AlphaFold IDR':
        color = 'idr_ptm'
    elif coloring == 'AlphaFold secondary structures':
        color = 'structure_group'
    else:
        color = 'af'

    html = html_template.replace('globalThis.color_choice = your_choice',f"globalThis.color_choice = '{color}'")

    return html

# Cell
def get_ms_concensus(ms_list):
    unique_ms_list = set(ms_list)
    if len(unique_ms_list)==1:
        out = unique_ms_list
    else:
        out = [i for i in unique_ms_list if i != 'detected']

    out = ';'.join(out)

    return out

# Cell
import alphamap
def plot_3d_structure(df: pd.DataFrame or list,
                        name: str,
                        protein: str,
                        fasta: fasta,
                        selected_coloring: str,
                        dashboard: bool = False,
                        download_folder: str = tempfile.gettempdir()):

    """
    Function to generate the sequence plot.

    Args:
        df (pd.DataFrame/list): Single dataframe or list of dataframes containing the datasets to plot.
        name (str): Single string containing the name of the MS dataset in df.
        protein (str): Uniprot protein accession.
        fasta (fasta): Fasta file imported by pyteomics 'fasta.IndexedUniProt'.
        selected_coloring (str): Coloring to show.
        dashboard (bool, optional): Flag if the function is called from the dashboard. Default is 'False'.

    Returns:
        go.Figure: 3D plot.

    """

    if isinstance(df, pd.DataFrame):
        df_plot = get_plot_data(protein=protein,
                                df = df,
                                fasta = fasta)
    elif isinstance(df, list):
        df_plot = [get_plot_data(protein=protein,
                                 df = d,
                                 fasta = fasta) for d in df]
        df_plot = pd.concat(df_plot)

    if df_plot is None:
        MS_annotation =  pd.DataFrame({'position' : [], 'MS data': []})
    else:
        MS_annotation = extract_annotation(df_plot)
        MS_annotation = MS_annotation.groupby('position')['MS data'].apply(list).reset_index(name='MS data')
        MS_annotation['MS data'] = MS_annotation['MS data'].apply(lambda x: get_ms_concensus(x))

    cif_available = manipulate_cif(protein = protein,
                                   MS_data = MS_annotation,
                                   download_folder = download_folder)

    BASE_PATH = os.path.dirname(alphamap.__file__)
    js_path = os.path.join(BASE_PATH, 'js')
    cif_path = download_folder


    import shutil
    new_file = os.path.join(cif_path, protein+'_'+name+'_mod.cif')
    try:
        os.remove(new_file)
    except:
        pass
    shutil.copyfile(src = os.path.join(cif_path, protein+'_mod.cif'),
                    dst = new_file)

    if cif_available:
        mod_html = adjust_html(protein=protein,
                               coloring=selected_coloring,
                               data_name=name)
    else:
        mod_html = f"""
            <!DOCTYPE html>
            <center> No AlphaFold structure available for {protein}. <center>
            """



    return mod_html, js_path, cif_path

# Cell
from .pdflib import *

def create_pdf_report(proteins: list,
                      df: pd.DataFrame or list,
                      name: str or list,
                      fasta: fasta,
                      uniprot: pd.DataFrame,
                      selected_features: list,
                      uniprot_feature_dict: dict,
                      uniprot_color_dict: dict,
                      selected_proteases: list = [],
                      trace_colors: list = []):
    """
    Function to write pdf reports for selected proteins

    Args:
        proteins (list): List of uniprot protein accessions.
        df (pd.DataFrame/list): Single dataframe or list of dataframes containing the datasets to plot.
        name (str/list): Single string or list of strings containing the names for each dataset in df.
        fasta (fasta): Fasta file imported by pyteomics 'fasta.IndexedUniProt'.
        uniprot (pd.DataFrame): Uniprot annotations formatted by alphamap.
        selected_features (list): List of uniprot features to plot.
        uniprot_feature_dict (dict): Uniprot feature dictionary.
        uniprot_color_dict (dict): Uniprot color dictionary.
        selected_proteases (list, optional): List of proteases to plot. Default is an empty list.
        trace_colors (list, optional): List of manualy selected colors for each dataset in df. Default is an empty list.

    Returns:
        BytesIO: BytesIO object for writing a pdf report.
    """

    if isinstance(df, pd.DataFrame):
        max_height = 200 + 50 + (len(selected_features)*50) + (len(selected_proteases)*50)
    else:
        max_height = 200 + (len(df)*50) + (len(selected_features)*50) + (len(selected_proteases)*50)


    if max_height < 700:
        max_height = 700

    footer_text = '<font size="20">This report was generated by <a href="https://github.com/MannLabs/alphamap" color="darkblue"><b>AlphaMap</b></a>.</font>'

    pdf_buf = BytesIO()
    pdf_report = canvas.Canvas(pdf_buf, pagesize=(1600,max_height))

    if isinstance(proteins, list):
        for p in range(0,len(proteins)):
            plot = plot_peptide_traces(df=df, name=name, protein=proteins[p], fasta=fasta,
                                       uniprot=uniprot, selected_features=selected_features,
                                       uniprot_feature_dict=uniprot_feature_dict,
                                       uniprot_color_dict=uniprot_color_dict,
                                       selected_proteases=selected_proteases,
                                       trace_colors=trace_colors)
            draw_content(pdf_report, plot, width=1600, height=max_height,
                         spacing=5, border=20)
            draw_content(pdf_report, footer_text, width=1600, height=100,
                         spacing=20, border=30)
            pdf_report.showPage()
    else:
        plot = plot_peptide_traces(df=df, name=name, protein=proteins, fasta=fasta,
                                   uniprot=uniprot, selected_features=selected_features,
                                   uniprot_feature_dict=uniprot_feature_dict,
                                   uniprot_color_dict=uniprot_color_dict,
                                   selected_proteases=selected_proteases,
                                   trace_colors=trace_colors)
        draw_content(pdf_report, plot, width=1600, height=max_height,
                         spacing=20, border=30)
        draw_content(pdf_report, footer_text, width=1600, height=100,
                         spacing=20, border=30)
        pdf_report.showPage()

    pdf_report.save()
    pdf_buf.seek(0)

    #with open("file.pdf", "wb") as file:
    #    file.write(pdf_buf.getvalue())

    return pdf_buf