"""
Visualizes the landscape of a protein.
"""

import matplotlib.pyplot as plt
import matplotlib.ticker as tck
from matplotlib.ticker import MaxNLocator
import pandas as pd 
import numpy as np
import re
import ast
import logging
logger = logging.getLogger(__name__)


def plot_peptide_length_dist(first_df: pd.DataFrame, second_df: pd.DataFrame, first_explode: str, second_explode: str, first_column: str, second_column: str, first_label: str, second_label: str, mod_pattern: str) -> tuple[plt.figure,int,int]:
    """Visualize the distribution of lengths of sequences.
    
    Args:
        first_df: A pandas dataframe.
        second_df: A pandas dataframe.
        first_explode: The header of the column, that is used to explode the 
            dataframe. 
        second_column: The header of the column, that is used to explode the 
            dataframe. 
        first_column: The header of the column, that holds the sequences in 
            first_df, for which the length distribution will be plotted. 
        second_column: The header of the column, which holds the sequences in 
            second_df, for which the length distribution will be plotted.  
        first_label: Label for the values of first_column in the plot.
        second_label: Label for the values of second_column in the plot. 
        mod_pattern: A comma separated string with delimiters for peptide
            modifications


    Returns:
        A tuple including a matplotlib figure and two integers. The figure is a 
        histogram visualizing the length distribution of peptides and epitopes. 
        The two integers are the number of peptides and the number of epitopes. 
    """
    
    first_long = first_df.explode(first_explode)
    second_long = second_df.explode(second_explode)
    
    logger.info(f'{len(first_long.drop_duplicates("sequence"))} peptides were reduced to {len(second_long.drop_duplicates(["whole_epitopes", "consensus_epitopes"]))} epitopes.')

    # compute a histogram of the sequence lengths in first_column and second_column
    fig, ax = plt.subplots(layout='constrained')
    seq_first = first_long[first_column]
    seq_second = second_long[second_column]

    # remove modifications before accessing the sequence length
    pattern = r'\(.*?\)'
    peptide_seqs = seq_first.apply(lambda seq: re.sub(pattern,"",seq))
    pattern = r'\[.*?\]'
    peptide_seqs = peptide_seqs.apply(lambda seq: re.sub(pattern,"",seq))
    if mod_pattern:
        pattern = re.escape(mod_pattern.split(',')[0]) + r'.*?' + re.escape(mod_pattern.split(',')[1])
        peptide_seqs = peptide_seqs.apply(lambda seq: re.sub(pattern,"",seq))

    first_len = peptide_seqs.map(lambda pep: len(pep)).to_list()
    second_len = seq_second.map(lambda pep: len(pep)).to_list()
    
    ax.hist(first_len, bins=np.arange(5,50,1), color='grey', label=first_label, alpha=0.6)
    ax.hist(second_len, bins=np.arange(5,50,1), color='red', label=second_label, alpha=0.6)

    ax.legend()
    ax.set_xlabel('length')
    ax.set_ylabel('count')
    return fig, len(first_long), len(second_long)


def plot_protein_landscape(protein_df: pd.DataFrame, accession: str, proteome_dict: dict[str,str]) -> plt.figure:
    """Visualize the landscape of a protein.

    Args:
        protein_df: A pandas dataframe containing one protein per row.
        accession: The string of a protein accession.
        proteome_dict: A dictionary containing the reference proteome.

    Returns:
        A matplotlib bar plot that visualizes the peptide and core epitope 
        distribution across the sequence of the protein with the provided 
        accession. 
    """

    prot_row = protein_df[(protein_df['accession'] == accession)]
    if len(prot_row) == 0:
        raise Exception('The accession {} is not in your input data.'.format(accession))

    prot_seq = proteome_dict[accession]

    prot_landscape = [0 for _ in prot_seq]
    
    fig_width = max(15,round(len(prot_landscape)/50))
    fig_height = 3

    max_intens = 0

    fig, ax = plt.subplots(figsize=(fig_width, fig_height), layout='constrained')

    for group, landscape in enumerate(prot_row['landscape'].iloc[0]):

        if group % 3 == 0:
            color = 'red'
        elif group % 3 == 1:
            color = 'blue'
        else: 
            color = 'green'

        group_start = min(prot_row['grouped_peptides_start'].iloc[0][group])
        for idx, position in enumerate(landscape):
            ax.bar(group_start+int(idx),position,width=1, alpha=0.4, color=color)
        for pos in range(prot_row['core_epitopes_start'].iloc[0][group],prot_row['core_epitopes_end'].iloc[0][group]+1):
            ax.bar(pos,0.5,width=1,color=color)

        max_intens = max(max_intens, max(landscape))

    ybins=max_intens/10
    ax.yaxis.set_major_locator(MaxNLocator(nbins=max(ybins, 5)))
    ax.xaxis.set_major_locator(MaxNLocator(nbins=max(fig_width/10, 25)))

    ax.set_title('Landscape and consensus epitopes of the protein {}'.format(accession))
    ax.set_xlabel('Position in protein {}'.format(accession))
    ax.set_ylabel('Number of aligned peptides')
    return fig



def plot_core_mapping_peptides_hist(epitope_df: pd.DataFrame) ->  plt.figure:
    """A histogram of the number of peptides mapped to each epitope. 
    
    Args:
        epitope_df: A dataframe containing one epitope and 
            information of that epitope per row. 

    Returns:
        A histogram visualizing the number of peptides mapped to 
        each epitope. 
    """
    fig, ax = plt.subplots(layout='constrained')
    n_peps = epitope_df['grouped_peptides_sequence'].apply(lambda sequences: len(sequences.split(',')))
    ax.hist(n_peps,bins=np.arange(1,max(n_peps)+2,1))
    ax.set_yscale('log')
    ax.set_xlabel('number of peptides contributing to each consensus epitope')
    ax.set_ylabel('count')
    return fig


def compute_coverage(pep_start, pep_end, consensus_start, consensus_end):
    '''Compute the coverage of the consensus sequence by the peptide contributing to it.

    Args:
        peptide: Peptide sequence.
        consensus: Consensus sequence.
        entire: Sequence of entire consensus sequence.

    Returns:
        The coverage of the consensus sequence by the peptide sequence
    '''

    if pep_start < consensus_start:
        overlap = max(min(consensus_end-consensus_start+1, pep_end-consensus_start+1),0)
    else:
        overlap = max(min(pep_end-pep_start+1, consensus_end-pep_start+1),0)
    return overlap/(consensus_end-consensus_start+1)


def compute_epitopes_coverage(epitope_df: pd.DataFrame, out):
    '''Compute the coverage of consensus sequence epitopes.

    Args:
        cores_df: A pandas DataFrame containing consensus sequence epitopes..

    Returns:
        A tuple containing the list of coverages for all consensus epitope sequences and the list of coverages for consensus sequence epitope sequences that are defined by at least two peptides.
    '''
    epitope_df['grouped_peptides_start'] = epitope_df['grouped_peptides_start'].apply(lambda starts: [re.sub(r'[^a-zA-Z0-9]', '',start) for start in starts.split(',')])
    epitope_df['grouped_peptides_end'] = epitope_df['grouped_peptides_end'].apply(lambda starts: [re.sub(r'[^a-zA-Z0-9]', '',start) for start in starts.split(',')])
    epitope_df['grouped_peptides_sequence'] = epitope_df['grouped_peptides_sequence'].apply(lambda starts: [re.sub(r'[^a-zA-Z0-9]', '',start) for start in starts.split(',')])
    epitope_df = epitope_df.explode(['grouped_peptides_start', 'grouped_peptides_end', 'grouped_peptides_sequence'])
    epitope_df['coverage'] = epitope_df.apply(lambda row: compute_coverage(int(row['grouped_peptides_start']), int(row['grouped_peptides_end']), int(row['core_epitopes_start']), int(row['core_epitopes_end'])), axis=1)
    coverage_all = epitope_df['coverage'].to_list()
    epitope_df = epitope_df.drop_duplicates(['grouped_peptides_sequence', 'whole_epitopes', 'consensus_epitopes'])
    epitope_df.to_csv(f'{out}/coverage.csv')
    return coverage_all


def plot_coverage(coverage_hist_all: list[float], out:str):
    '''Plot the coverages of the consensus sequences by the peptides.

    Args:
        coverage_hist_all: List containing all coverages.
        out: Path were the plot gets saved.
    '''
    figure, axis = plt.subplots(1, 1, dpi=150)
    n_all, bins_all, _ = axis.hist(coverage_hist_all, 100, color='red', label='all consensus sequences')
    axis.set_yscale('log')
    axis.set_ylabel('Number of peptides')
    axis.set_xlabel(f'Consensus sequence coverage')
    axis.grid()
    plt.tight_layout()
    figure.savefig(f'{out}/consensus_sequence_coverage.png')


def plot_consensus_sequence_coverage(epitopes_df: pd.DataFrame, out: str):
    '''Plot the consensus sequence coverage of the computed consensus epitope sequences.
    
    Args:
        epitopes_df: A pandas dataframe containing the calculated consensus epitope sequences.
        out: Path were the plot gets saved.
    '''
    epitopes_df_copy = epitopes_df.copy()
    all = compute_epitopes_coverage(epitopes_df_copy, out)
    plot_coverage(all, out)
    logger.info(f'The calculated consensus epitope sequences have an average peptide coverage of {np.mean(all)}')


def create_html(out_name):
    ''' Create html version of plot.

    Args: 
        out_name: Location where the figure should be stored.    
    '''
    with open(f'{out_name[:-4]}svg', 'r') as svg_file:
        svg_content = svg_file.read()
        svg_content = re.sub(r'<\?xml[^>]+\?>', '', svg_content)
        svg_content = re.sub(r'<!DOCTYPE[^>]+>', '', svg_content)
        html = f'<!DOCTYPE html> <html> <body>{svg_content}</body></html>'
        with open(f'{out_name}','w') as f:
            f.write(html)