import pandas as pd
from pathlib import Path
import numpy as np
import os
import matplotlib.pyplot as plt


"""
    Forked from Demo_loadandanalyzeDLCdata.ipynb by Federico Claudi
    https://github.com/FedeClaudi
"""


def get_cmap(n, name='hsv'):
    return plt.cm.get_cmap(name, n)


def Histogram(vector, color, bins):
    dvector = np.diff(vector)
    dvector = dvector[np.isfinite(dvector)]
    plt.hist(dvector, color=color, histtype='step', bins=bins)


def PlottingResults(Dataframe,
                    bodyparts2plot,
                    alphavalue=.2,
                    pcutoff=.5,
                    colormap='jet',
                    fs=(4, 3)):
    ''' Plots poses vs time; pose x vs pose y; histogram of differences and likelihoods.'''
    plt.figure(figsize=fs)
    colors = get_cmap(len(bodyparts2plot), name=colormap)
    scorer = Dataframe.columns.get_level_values(0)[
        0]  #you can read out the header to get the scorer name!

    for bpindex, bp in enumerate(bodyparts2plot):
        Index = Dataframe[scorer][bp]['likelihood'].values > pcutoff
        plt.plot(Dataframe[scorer][bp]['x'].values[Index],
                 Dataframe[scorer][bp]['y'].values[Index],
                 '.',
                 color=colors(bpindex),
                 alpha=alphavalue)

    plt.gca().invert_yaxis()

    sm = plt.cm.ScalarMappable(cmap=plt.get_cmap(colormap),
                               norm=plt.Normalize(vmin=0,
                                                  vmax=len(bodyparts2plot) -
                                                  1))
    sm._A = []
    cbar = plt.colorbar(sm, ticks=range(len(bodyparts2plot)))
    cbar.set_ticklabels(bodyparts2plot)
    #plt.savefig(os.path.join(tmpfolder,"trajectory"+suffix))
    plt.figure(figsize=fs)
    Time = np.arange(np.size(Dataframe[scorer][bodyparts2plot[0]]['x'].values))

    for bpindex, bp in enumerate(bodyparts2plot):
        Index = Dataframe[scorer][bp]['likelihood'].values > pcutoff
        plt.plot(Time[Index],
                 Dataframe[scorer][bp]['x'].values[Index],
                 '--',
                 color=colors(bpindex),
                 alpha=alphavalue)
        plt.plot(Time[Index],
                 Dataframe[scorer][bp]['y'].values[Index],
                 '-',
                 color=colors(bpindex),
                 alpha=alphavalue)

    sm = plt.cm.ScalarMappable(cmap=plt.get_cmap(colormap),
                               norm=plt.Normalize(vmin=0,
                                                  vmax=len(bodyparts2plot) -
                                                  1))
    sm._A = []
    cbar = plt.colorbar(sm, ticks=range(len(bodyparts2plot)))
    cbar.set_ticklabels(bodyparts2plot)
    plt.xlabel('Frame index')
    plt.ylabel('X and y-position in pixels')
    #plt.savefig(os.path.join(tmpfolder,"plot"+suffix))

    plt.figure(figsize=fs)
    for bpindex, bp in enumerate(bodyparts2plot):
        Index = Dataframe[scorer][bp]['likelihood'].values > pcutoff
        plt.plot(Time,
                 Dataframe[scorer][bp]['likelihood'].values,
                 '-',
                 color=colors(bpindex),
                 alpha=alphavalue)

    sm = plt.cm.ScalarMappable(cmap=plt.get_cmap(colormap),
                               norm=plt.Normalize(vmin=0,
                                                  vmax=len(bodyparts2plot) -
                                                  1))
    sm._A = []
    cbar = plt.colorbar(sm, ticks=range(len(bodyparts2plot)))
    cbar.set_ticklabels(bodyparts2plot)
    plt.xlabel('Frame index')
    plt.ylabel('likelihood')

    #plt.savefig(os.path.join(tmpfolder,"plot-likelihood"+suffix))

    plt.figure(figsize=fs)
    bins = np.linspace(0, np.amax(Dataframe.max()), 100)

    for bpindex, bp in enumerate(bodyparts2plot):
        Index = Dataframe[scorer][bp]['likelihood'].values < pcutoff
        X = Dataframe[scorer][bp]['x'].values
        X[Index] = np.nan
        Histogram(X, colors(bpindex), bins)
        Y = Dataframe[scorer][bp]['x'].values
        Y[Index] = np.nan
        Histogram(Y, colors(bpindex), bins)

    sm = plt.cm.ScalarMappable(cmap=plt.get_cmap(colormap),
                               norm=plt.Normalize(vmin=0,
                                                  vmax=len(bodyparts2plot) -
                                                  1))
    sm._A = []
    cbar = plt.colorbar(sm, ticks=range(len(bodyparts2plot)))
    cbar.set_ticklabels(bodyparts2plot)
    plt.ylabel('Count')
    plt.xlabel('DeltaX and DeltaY')

    #plt.savefig(os.path.join(tmpfolder,"hist"+suffix))
