import numpy as np
import pandas as pd
from scipy import spatial
import time
import json
import matplotlib.pyplot as plt
import matplotlib
import matplotlib.gridspec as gridspec
import warnings
warnings.filterwarnings("ignore")


def get_gt(query_file):
    """Load ground truth file

    Args:
        query_file (str): json file with the ground truth data

    Returns:
        dict: ground truth dictionary
    """    
    with open(query_file, "r") as gt_file:
        gt = json.load(gt_file)
    return gt


def calculate_distances(search_results, dataset):
    """Calculate Euclidean distances and cosine similarity in saerch results

    Args:
        search_results (DataFrame): search results from evaluator
        dataset (DataFrane): protein descriptors

    Returns:
        DataFrame: search results with distances
    """    
    print("Calculating distances.")
    lmi_data_dict = dict(zip(dataset.index,range(len(dataset))))
    np_lmi = dataset.to_numpy()
    print("Calculating euclid.")
    start = time.time()
    search_results['euclid'] = search_results[['query', 'result']].apply(lambda row: np.linalg.norm(np_lmi[lmi_data_dict[row['query']]] - np_lmi[lmi_data_dict[row['result']]]), axis=1)
    end = time.time()
    print(f"{search_results.shape[0]} euclidian distances calculated in {end - start:.4f}s")
    print("Calculating cosine.")
    start = time.time()
    search_results['cosine'] = search_results[['query', 'result']].apply(lambda row: spatial.distance.cosine(np_lmi[lmi_data_dict[row['query']]], np_lmi[lmi_data_dict[row['result']]]), axis=1)
    end = time.time()
    print(f"{search_results.shape[0]} cosine similarities calculated in {end - start:.4f}s")
    return search_results


def calculate_gt_counts(query_file, ranges, search_results):
    """Get ground truth counts for specified ranges"""    
    queries = search_results[['query']].drop_duplicates().reset_index(drop=True)
    
    gt = get_gt(query_file)
        
    all_distances = []
    lmi_found = []
    for rng in ranges:
        lmi_found.append(search_results[search_results['distance'] < rng].groupby("query")[['distance']].agg(lambda x: len(x)).rename(columns={'distance': f'LMI_{rng}'}))
    for _, query in queries.iterrows():
        neighbours = gt[query['query']]
        distances = []
        for _, distance in neighbours.items():
            distances.append(distance)
        new_row = {"query": query['query']}
        for rng in ranges:
            new_row[f"GT_{rng}"] = sum(np.array(distances) > (1 - rng))
            
        all_distances.append(new_row)
    queries_gt = pd.DataFrame(all_distances)
    for i in range(len(ranges)):
        queries_gt = queries_gt.merge(right=lmi_found[i], how='left', on=['query'])
    return queries_gt


def filter_with_cutoff(search_results, queries_gt, ranges, metric, cutoffs):
    """Evaluate the filtering step with different ranges and cutoff values"""    
    result = {}
    for q_distance in ranges:
        cutoff_list = []
        recall_list = []
        precision_list = []
        for cutoff in cutoffs:
            queries_results = queries_gt
            dst = search_results[(search_results['distance'] < q_distance) & (search_results[metric] < cutoff)].groupby("query")[['distance']].agg(lambda x: len(x))
            dst = dst.rename(columns={'distance': metric})
            total = search_results[(search_results[metric] < cutoff)].groupby("query")[['distance']].agg(lambda x: len(x))
            total = total.rename(columns={'distance': f'{metric}_total'})
            queries_results = queries_results.merge(right=dst, how='left', on=['query'])
            queries_results = queries_results.merge(right=total, how='left', on=['query'])
            queries_results[f'Recall_{metric}'] = queries_results[metric] / queries_results[f'GT_{q_distance}']
            queries_results[f'Precision_{metric}'] = queries_results[metric] / queries_results[f'{metric}_total']
            cutoff_list.append(cutoff)
            recall = queries_results[f'Recall_{metric}'].mean()
            recall_list.append(recall)
            precision = queries_results[f'Precision_{metric}'].mean()
            precision_list.append(precision)
        result[q_distance] = {}
        result[q_distance]['cutoff'] = cutoff_list
        result[q_distance]['recall'] = recall_list
        result[q_distance]['precision'] = precision_list
    return result


def draw_graphs(data_euclid, data_cosine, rng, filename=None):
    """Filtering step graphs"""    
    fig, axs = plt.subplots(figsize=(8.5, 12), ncols=2, nrows=3)
    
    for index, rng in enumerate(rng):
    
        axs[index, 0].plot(data_euclid[rng]["cutoff"], data_euclid[rng]["recall"], label="Recall")
        axs[index, 0].plot(data_euclid[rng]["cutoff"], data_euclid[rng]["precision"], label="Precision")
        axs[index, 0].set(ylim=(0.0, 1.0))
        axs[index, 0].set_xlabel('Euclidean distance cutoff')
        axs[index, 0].legend()

        axs[index, 1].plot(data_cosine[rng]["cutoff"], data_cosine[rng]["recall"], label="Recall")
        axs[index, 1].plot(data_cosine[rng]["cutoff"], data_cosine[rng]["precision"], label="Precision")
        axs[index, 1].set(ylim=(0.0, 1.0))
        axs[index, 1].set_xlabel('Cosine similarity cutoff')
        axs[index, 1].legend()
    
    if filename is not None:
        fig.savefig(filename, bbox_inches='tight')

    
def best_f1(recall_precission_euclid, recall_precission_cosine, distances):
    """Get the best F1-score"""    
    def print_best_f1(recall_precision_array, range_, metric=""):
        cutoff = np.array(recall_precision_array[range_]['cutoff'][1:])
        recall = np.array(recall_precision_array[range_]['recall'][1:])
        precision = np.array(recall_precision_array[range_]['precision'][1:])
        f1 = 2 * (recall * precision) / (recall + precision)
        print(f"{metric} Best f1: {max(f1):.03f} (Recall={recall[np.argmax(f1)]:.03f}) using cutoff {cutoff[np.argmax(f1)]:.02f}")

    for dst in distances:
        print_best_f1(recall_precission_euclid, dst, "EUCLID")

    for dst in distances:
        print_best_f1(recall_precission_cosine, dst, "COSINE")


def draw_heatmap(queries_results, ranges, filename=None):
    """Draw the heatmap of the result query counts compared to the ground truth query counts"""    
    gs = gridspec.GridSpec(4, 4, wspace=0.8, hspace=0.5)
    figure = plt.figure(figsize=(8,8))
    
    for index, rng in enumerate(ranges):
        total_query_results_lmi = []
        total_query_results_gt = []
        for _, row in queries_results.iterrows():
            try:
                lmi_total = int(row[f'LMI_{rng}'])
            except:
                lmi_total = 0
            gt_total = int(row[f'GT_{rng}'])
            if lmi_total < 2 or gt_total < 2:
                continue
            total_query_results_lmi.append(lmi_total)
            total_query_results_gt.append(gt_total)
        
        points=list(set(zip(total_query_results_gt,total_query_results_lmi))) 
        count=[len([x for x,y in zip(total_query_results_gt,total_query_results_lmi) if x==p[0] and y==p[1]]) for p in points]
        plot_x=[i[0] for i in points]
        plot_y=[i[1] for i in points]
        count=np.array(count)

        if index == 0:
            ax1 = figure.add_subplot(gs[:2, :2])
        elif index == 1:
            ax1 = figure.add_subplot(gs[:2, 2:])
        else:
            ax1 = figure.add_subplot(gs[2:4, 1:3])
            
        scatter1 = ax1.scatter(plot_x,plot_y,c=count,s=10*count**0.75,cmap='Spectral_r')
        ax1.set_yscale('symlog')
        ax1.set_xscale('symlog')
        ax1.set_xlabel('GT results')
        ax1.set_ylabel('LMI results')
        ax1.set_title(f'Range={rng}')

        ax1.grid(which='major')
        ax1.xaxis.set_major_formatter(matplotlib.ticker.ScalarFormatter())
        ax1.yaxis.set_major_formatter(matplotlib.ticker.ScalarFormatter())
        ax1.xaxis.set_minor_locator(matplotlib.ticker.SymmetricalLogLocator(base=10, linthresh=2, subs=list(range(1, 10))))
        ax1.yaxis.set_minor_locator(matplotlib.ticker.SymmetricalLogLocator(base=10, linthresh=2, subs=list(range(1, 10))))
        plt.colorbar(scatter1, ax=ax1, ticks=list(range(1, max(count) + 1, (max(count) // 5) if max(count) == 5 else (max(count) // 9))))

    if filename is not None:  
        plt.savefig(filename, bbox_inches='tight')
