#!/usr/bin/env python3

"""
Benchmark representation of unique side-chain contacts in the PDB for a fragment

"""

import os
import pickle
from pprint import pprint
import random
from multiprocessing import Pool

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
import pandas as pd
import seaborn as sns

from rdkit import Chem
from rdkit.Chem import Draw
from rdkit.Chem.Draw.MolDrawing import MolDrawing, DrawingOptions

from ..utils import *

def fragment_representation_bm(args):
    """
    Bootstrap AUC plot for unique side-chain fragment contact recovery for specific fragments

    Usage:
      bsff benchmark fragment_representation <user_defined_dir> <clusters_dir> [options]
      bsff benchmark fragment_representation figures <user_defined_dir> <clusters_dir> <bm_csv_dir> [options]
      bsff benchmark fragment_representation draw_fragments <fragment_inputs_csv>

    Arguments:
      <user_defined_dir>      Path to project root directory
      <clusters_dir>          Path to directory containing clusters generated by clusters command
      <bm_csv_dir>            Path to directory containing fragment recovery benchmark csvs
      <fragment_inputs_csv>   Path to directory containing defined fragment PDBs

    Options:
      --csv=<csv_path>        Previously generated benchmark data
      -e, --error_bands       Use error bands instead of bars for standard deviations
      -m, --maintext          Generate Figures for main text
      -p, --processes         Number of processes to run in parallel
    """

    """
    Before this function gets called:
      * Defined fragment (Use full amino acids to standardize atom names, need to figure out params for fragments)
      * Assemble fuzzball (SAVE REDUNDANT CONTACTS!!!)
    
    What happens here:
      * Assemble Cluster/PDB mapping dataframe from ProDy clusters
      * Remove singletons (keep track of this)
      * Find source of all structure codes in PDB
      * Bootstrap by taking 5% increments of PDB and intersect with dataframe
      * Save proportion of clusters recovered, plot
    
    Bam, done.
    
    All steps up to clustering to not rely on params files, so I might be able to get away with just generating fragments
    from ideal ligands from the PDB/LigandExpo.
    """

    user_defined_dir = args['<user_defined_dir>']
    clusters_dir = args['<clusters_dir>']
    processes = args['--processes'] if args['--processes'] else 12

    # Why.
    if args['draw_fragments']:
        draw_fragments(args['<fragment_inputs_csv>'])

    if args['figures']:
        png_dir = generate_fragment_pngs(user_defined_dir)
        # generate_cluster_counts_hist(clusters_dir, format_maintext=args['--maintext'])
        generate_auc_figures(args['<bm_csv_dir>'], png_dir, error_bands=args['--error_bands'], format_maintext=args['--maintext'])

    else:
        pool = Pool(processes=processes)
        cluster_dir_list = [cluster_dir for cluster_dir in directory_check(clusters_dir, base_only=False)]
        pool.map(perform_fragment_recovery_bm, cluster_dir_list)
        pool.close()
        pool.join()


def perform_fragment_recovery_bm(cluster_dir):
    """
    Perform fragment recovery benchmark
    :return:
    """

    # Import all PDBIDs from file
    pdb_entries_path = os.path.join(os.path.dirname(__file__), '..', '..', 'Additional_Files', 'entries.idx')
    pdb_df = pd.read_csv(pdb_entries_path,
                         delimiter='\t',
                         names=['IDCODE', 'HEADER', 'ACCESSION DATE', 'COMPOUND', 'SOURCE', 'AUTHOR LIST', 'RESOLUTION',
                                'EXPERIMENT TYPE'],
                         skiprows=2
                         )

    pdb_set = set([idcode.lower() for idcode in pdb_df['IDCODE']])
    all_pdbids_count = len(pdb_set)
    pdb_coverage_fractions = list(np.arange(0, 1.05, 0.05))

    recovery_list = list()

    # For loops all the way down...
    fragment_name = os.path.basename(cluster_dir)

    with open(os.path.join(cluster_dir, f'{fragment_name}-PDB_Sources.pickle'), 'rb') as pdb_source_pickle:
        cluster_source_dict = pickle.load(pdb_source_pickle)

    # todo: flatten cluster_source_dict!

    flattened_cluster_source_dict = dict()
    for cat_cluster_idx, contact_cluster_dict in cluster_source_dict.items():
        for contact_cluster_idx, contact_clusters in contact_cluster_dict.items():
            flattened_cluster_source_dict[(cat_cluster_idx, contact_cluster_idx)] = contact_clusters

    cluster_source_dict = flattened_cluster_source_dict
    pprint(cluster_source_dict)

    # Remove singletons
    singleton_cluster_ids = list()
    total_clusters = 0
    total_cluster_members = 0

    for cluster_id in cluster_source_dict:
        if len(cluster_source_dict[cluster_id]['list']) <= 1:
            singleton_cluster_ids.append(cluster_id)
        else:
            total_clusters += 1
            total_cluster_members += len(cluster_source_dict[cluster_id]['list'])

    for cluster_id in singleton_cluster_ids:
        del cluster_source_dict[cluster_id]

    print(total_clusters, total_cluster_members)

    for fraction in pdb_coverage_fractions:
        samples = int(all_pdbids_count * fraction)

        for i in range(1000):
            cluster_hits = 0
            selected_pdbids = set(random.sample(pdb_set, samples))

            for cluster_id, cluster_members in cluster_source_dict.items():
                cluster_pop = len(set([a.lower() for a in cluster_members['list']]) & selected_pdbids)
                if cluster_pop > 0:
                    cluster_hits += 1

            trial_dict = {'fragment': fragment_name,
                          'fraction': fraction,
                          'total_residues': total_cluster_members,
                          'total_clusters': total_clusters,
                          'cluster_hits': cluster_hits,
                          'recovered': cluster_hits / total_clusters
                          }
            recovery_list.append(trial_dict)

            print(fragment_name, fraction, total_clusters, cluster_hits)

    df = pd.DataFrame(recovery_list)
    df.to_csv(f'{os.path.basename(fragment_name)}-cluster_recovery.csv')

    return df


def generate_fragment_pngs(user_defined_dir):
    """
    Create pngs of fragments for embedding into plots
    """
    fragment_inputs = os.path.join(user_defined_dir, 'Inputs', 'Fragment_Inputs')
    fragment_inputs_csv = os.path.join(fragment_inputs, 'Fragment_inputs.csv')
    df = pd.read_csv(fragment_inputs_csv)

    frag_smarts_dict = {f"Fragment_{row['Fragment']}": re.sub(r'[^\[\]]+(?=\])', lambda x: f"{x.group().split(';')[0]}",
                                              row['SMILES_fragment']).upper() for index, row in df.iterrows()}

    png_dir = os.path.join(user_defined_dir, 'Fragment_PNGs')
    os.makedirs(png_dir, exist_ok=True)

    for fragment_name, fragment_smarts in frag_smarts_dict.items():
        fragment_mol_H = Chem.MolFromSmarts(fragment_smarts)
        if fragment_mol_H is None:
            fragment_mol_H = Chem.MolFromPDBFile(os.path.join(fragment_inputs, f'{fragment_name}.pdb'), removeHs=True)

        fragment_mol = Chem.RemoveHs(fragment_mol_H)
        fragment_mol.UpdatePropertyCache(strict=False)  # https://github.com/rdkit/rdkit/issues/146
        Chem.rdDepictor.Compute2DCoords(fragment_mol)
        img = Draw.MolToMPL(fragment_mol, fitImage=True)
        plt.axis('off')
        img.savefig(os.path.join(png_dir, f'{fragment_name}.png'), bbox_inches='tight', transparent=True)

    return png_dir

def generate_cluster_counts_hist(clusters_dir, format_maintext=False):
    """Histogram to visualize cluster count/occupancy"""
    for cluster_dir in directory_check(clusters_dir, base_only=False):
        fragment_name = os.path.basename(cluster_dir)
        df = pd.read_csv(os.path.join(clusters_dir, fragment_name, f'{fragment_name}_report.csv'))
        df.sort_values(by=['cluster_members'], ascending=False, inplace=True)

        fig = plt.figure()
        ax = sns.barplot(x=list(range(1, len(df['cluster_members']) + 1)), y=df['cluster_members'], palette=['#007CBE']*len(df['cluster_members']))

        ax.set(ylim=(0, df['cluster_members'].max()))
        ax.set(xlim=(0, len(df['cluster_members'])))

        # 5% Vertical line
        plt.axvline(x=len(df['cluster_members'])/20, color='#EB093C')

        x_clusters = len(df['cluster_members'])
        x_cluster_ticks = len(df['cluster_members']) // 10

        plt.xticks(range(0, x_clusters, x_cluster_ticks), range(0, x_clusters, x_cluster_ticks))
        ax.set_xlabel('Clusters')
        ax.set_ylabel('Cluster Members')

        # Format axes for main text/suppmat
        if format_maintext:
            ax.tick_params(labelsize=16, labelrotation=45)
            plt.rcParams.update({'font.size': 16})

        fig.add_axes(ax)
        figure_name = f'{fragment_name}-cluster_occupancy.png'
        fig.savefig(figure_name, dpi=300, bbox_inches='tight', transparent=True)
        plt.close()

def generate_auc_figures(bm_csv_dir, png_dir, error_bands=False, format_maintext=False):
    """
    Generate figures for fragment recovery benchmark
    :param df:
    :return:
    """
    global fraction
    fraction_recovery_list = []

    annotation_x_positions = (0.3, 0.3, 0.3, 0.3, 0.3) if format_maintext else (0.3, 0.3, 0.3, 0.2, 0.2)
    annotation_y_positions = (0.33, 0.26, 0.19, 0.12, 0.05) if format_maintext else (0.3, 0.25, 0.2, 0.15, 0.1, 0.05)
    annotation_fontsize = 16 if format_maintext else 10

    for csv in os.listdir(bm_csv_dir):
        if csv.endswith('.csv'):
            df = pd.read_csv(os.path.join(bm_csv_dir, csv))

            for fragment, fragment_df in df.groupby('fragment'):
                print(fragment)

                fig = plt.figure()
                plt.rcParams.update({'font.size': annotation_fontsize})

                if error_bands:
                    muh_plot = sns.lineplot(x='fraction', y='recovered', data=fragment_df, ci='sd', err_style='band', estimator=np.mean)
                    sns.regplot(x='fraction', y='recovered', data=fragment_df, x_estimator=np.mean, x_ci=None, fit_reg=False, ax=muh_plot, scatter_kws={'color': list(sns.color_palette("Blues_r")[1])})
                else:
                    muh_plot = sns.regplot(x='fraction', y='recovered', data=fragment_df, x_ci='sd', x_estimator=np.mean, fit_reg=False)

                # Format axes for main text/suppmat
                if format_maintext:
                    muh_plot.tick_params(labelsize=16, labelrotation=45)

                muh_plot.set(ylim=(0, 1.1))
                muh_plot.set(xlim=(0, 1.1))

                muh_plot.set_xlabel('PDB Fraction')
                muh_plot.set_ylabel('Clusters Recovered')

                # Fraction of PDB for 80% recovery
                for fraction, fraction_df in fragment_df.groupby('fraction'):
                    avg_recovered = fraction_df['recovered'].mean()
                    if avg_recovered > 0.8: break
                fraction_recovery_list.append({'fraction': fraction, 'fragment': fragment})

                # AUC calculation
                # auc_info = [(fraction, recovered['recovered'].mean()) for fraction, recovered in fragment_df.groupby('fraction')]
                # x_coords = [a[0] for a in auc_info]
                # y_coords = [a[1] for a in auc_info]
                # muh_plot.text(0.7, 0.05, f'AUC: {np.trapz(y_coords, x_coords):.4f}')

                # embed fragment RDKit image
                if not format_maintext:
                    fragment_png = mpimg.imread(os.path.join(png_dir, f'{fragment}.png'))
                    fragment_inlay = OffsetImage(fragment_png, zoom=0.15)
                    ab = AnnotationBbox(fragment_inlay, (0.915, 0.25), pad=0)
                    muh_plot.add_artist(ab)

                # Add annotations
                # if not format_maintext:
                #     muh_plot.text(annotation_x_positions[0], annotation_y_positions[0], f'{fragment}')

                muh_plot.text(annotation_x_positions[1], annotation_y_positions[1], f'Total Residues: {fragment_df["total_residues"].median():.0f}')
                muh_plot.text(annotation_x_positions[2], annotation_y_positions[2], f'Total Clusters: {fragment_df["total_clusters"].median():.0f}')
                muh_plot.text(annotation_x_positions[3], annotation_y_positions[3], f'Fraction PDB sampled for')
                muh_plot.text(annotation_x_positions[4], annotation_y_positions[4], f'>80% contact recovery: {fraction:.2f}')

                if format_maintext:
                    muh_plot.tick_params(labelsize=16)

                fig.add_axes(muh_plot)
                figure_name = f'{fragment}-cluster_recovery-{"band" if error_bands else "bars"}.png'
                fig.savefig(figure_name, dpi=600, bbox_inches='tight', transparent=True)
                plt.close()

    # Fraction recovered Histogram
    df = pd.DataFrame(fraction_recovery_list)

    sns.set(style="whitegrid", palette="hls")
    sns.set_context("paper", font_scale=1.5)

    fig, ax = plt.subplots(figsize=(6,6))
    sns.countplot(df['fraction'].round(2), color='#007CBE', ax=ax)
    fig.suptitle('Fraction of PDB Required for >80%\nContact Recovery Across All Tested Fragments')
    ax.set_xlabel('Fraction', fontsize=12)
    ax.set_ylabel('Count', fontsize=12)

    fig.savefig('RecoveryHist.tiff', dpi=300)
    df.to_csv('FractionRecovered.csv')


# todo: how do I embed these into the plots?
def draw_fragments(fragment_inputs_csv):
    """
    Use RDKit to draw fragments and output to .png
    :return:
    """

    df = pd.read_csv(fragment_inputs_csv)
    frag_smarts_dict = {f"Fragment_{row['Fragment']}": re.sub(r'[^\[\]]+(?=\])', lambda x: f"{x.group().split(';')[0]}", row['SMILES_fragment']).upper() for index, row in df.iterrows()}

    name_list = list()
    mol_list = list()

    for fragment_name, fragment_smarts in frag_smarts_dict.items():
        print(fragment_smarts)
        fragment_mol = Chem.RemoveHs(Chem.MolFromSmarts(fragment_smarts))
        fragment_mol.UpdatePropertyCache(strict=False)  # https://github.com/rdkit/rdkit/issues/146
        Chem.rdDepictor.Compute2DCoords(fragment_mol)
        Draw.MolToFile(fragment_mol, f'{fragment_name}.png')

        name_list.append(fragment_name)
        mol_list.append(fragment_mol)

    img = Draw.MolsToGridImage(mol_list, molsPerRow=4, subImgSize=(200, 200), legends=name_list)
    img.save('AUC_Fragments.png')
