#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
A collection of functions that do not work directly with cube files, but with data produced by the other scripts

author: Oxana Andriuc

This code is under CC BY-NC-SA license: https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
"""

import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import sys


def OverlapKDEs(*fpaths,value_type='esp',xlim=None, ylim=None):
    """
    arguments: filepaths (str), value type (str, opt), xlim (None/list/tuple, opt), ylim (None/list/tuple, opt)

    calls: -

    returns: -
    
    this function plots overlapping KDEs as read from _kde.out files written by PlotHistogram from PlotCube (when save is True)
    
    fpaths can have however many filepaths of the type _kde.out as generated by PlotHistogram from PlotCube.py. 
    !!! The files have to contain "KDE data" in their first line, otherwise the execution of the script will be stopped.
    !!! If the units as they appear in the column name are not the same for all the files, the execution of the script will be stopped.
    
    value_type = can be ‘esp’/’potential’ or ‘dens’/’density’ (or any uppercase version of the aforementioned)    
    """
    
    val_units=""
    for fpath in fpaths:
        with open(fpath) as f:
            first_line=f.readline()
            if "KDE data" not in first_line:
                print("\nThe file "+fpath.split("/")[-1]+" does not seem to be a _kde.out file generated by PlotHistogram in PlotCube.py. Try running the program with files of the right format.\n")
                exit()
            second_line=f.readline()
            current_unit=second_line.split()[1][1:]
            if val_units!="" and current_unit!=val_units:
                print("\nYour data is not in the same units. Try running the program with files that have data in the same units.\n")
                exit()
            else:
                val_units=current_unit
            
    

    ##########################
    # Interpreting arguments #
    ##########################

    if value_type.lower()=='esp' or value_type.lower()=='potential':
        title = 'Electrostatic potential /'+val_units
    elif value_type.lower()=='dens' or value_type.lower()=='density':
        title = 'Density /'+val_units

    if xlim is None:
        xlim=[None,None]
    xmin=xlim[0]
    xmax=xlim[1]
    
    if ylim is None:
        ylim=[None,None]
    ymin=ylim[0]
    ymax=ylim[1]

    ###################
    # Creating figure #
    ###################

    fig = plt.figure()
    ax = fig.add_subplot(111)  # 1x1 grid, 1st subplot, 3D
    ax.set_facecolor((0.92,0.92,0.95))
    plt.grid(color=(1,1,1))
    ax.spines['bottom'].set_color([1, 1, 1])
    ax.spines['top'].set_color([1, 1, 1])
    ax.spines['right'].set_color([1, 1, 1])
    ax.spines['left'].set_color([1, 1, 1])

    #################
    # Plotting data #
    #################
    
    for fpath in fpaths:
        x=[]
        y=[]
        with open(fpath) as f:
            lines=f.readlines()[2:]
            for i in range(len(lines)):
                if len(lines[i].split())==2:
                    x.append(float(lines[i].split()[0]))
                    y.append(float(lines[i].split()[1]))
                else:
                    break
            ax.plot(x,y,label=fpath.split("/")[-1][:-8])

#    ax.legend(loc='best', fontsize='xx-small')
    ax.legend(loc='best')
    
    if title is None:
        ax.set_xlabel('Value')
        plt.title('Histogram')
    else:
        ax.set_xlabel(title)    
        plt.title(title.split("/")[0]+'histogram')
    ax.set_ylabel('Frequency')
    
    ax.set_axisbelow(True)
    
    if xmin is not None:
        ax.set_xlim(xmin=xmin)
    if xmax is not None:
        ax.set_xlim(xmax=xmax)
    if ymin is not None:
        ax.set_ylim(ymin=ymin)
    if ymax is not None:
        ax.set_ylim(ymax=ymax)

    plt.show()
            
def PlotHistogram(*fpaths, value_type='esp', colour='rainbow', cb_lim=None, xlim=None, ylim=None):
    """
    arguments: filepaths (str), value type (str, opt), colour (str, opt), colour bar limits (None/list/tuple, opt), xlim (None/list/tuple, opt), ylim (None/list/tuple, opt) 

    calls: -

    returns: -

    fpaths should include one _kde.out file and/or one _hist.out file as generated by PlotHistogram from PlotCube.py. 
    !!! If there is more than one of each file type (hist/kde), the execution of the script will be stopped
    !!! If the units as they appear in the column name are not the same for all the files, the execution of the script will be stopped.
    
    value_type = can be ‘esp’/’potential’ or ‘dens’/’density’ (or any uppercase version of the aforementioned)    
    """

    histpath=None
    kdepath=None
    val_units=""
    hist_no=0
    kde_no=0
    for fpath in fpaths:
        with open(fpath) as f:
            first_line=f.readline()
            if "Histogram" in first_line:
                histpath=fpath
                hist_no+=1
            elif "KDE" in first_line:
                kdepath=fpath
                kde_no+=1
            second_line=f.readline()    
            current_unit=second_line.split()[-2][1:]
            if val_units!="" and current_unit!=val_units:
                print("\nYour data is not in the same units. This function only works with files that have data in the same units.\n")
                exit()
            else:
                val_units=current_unit
    if hist_no>1:
        print("\nYou are trying to plot multiple histogram files. This function only works for one histogram file and/or one KDE file.\n")
        exit()
    elif kde_no>1:
        print("\nYou are trying to plot multiple KDE files. This function only works for one histogram file and/or one KDE file.\n")
        exit()
    
    filenames=""
    if histpath:
        filenames=histpath.split("/")[-1]
        with open(histpath) as h:
            hlines=h.readlines()[2:]
            
            b=[]
            n=[]
            
            for i in range(len(hlines)):
                b.append(float(hlines[i].split()[0]))
                n.append(float(hlines[i].split()[2]))
                
            b.append(float(hlines[-1].split()[1]))
            
            if cb_lim is None:
                cb_lim = [min(b), max(b)]
            elif cb_lim[0] > min(b) or cb_lim[1] < max(b):
                print('\nWarning: The colour bar limits you have selected do not cover the whole range of values\n')
                
            if xlim is None:
                xlim=[cb_lim[0]-(max(b)-min(b))/10,cb_lim[1]+(max(b)-min(b))/10]

                
    if kdepath:
        if filenames!="":
            filenames=filenames+"   "+kdepath.split("/")[-1]
        else:
            filenames=kdepath.split("/")[-1]
        with open(kdepath) as k:
            klines=k.readlines()[2:]
            x=[]
            y=[]
            for i in range(len(klines)):
                if len(klines[i].split())==2:
                    x.append(float(klines[i].split()[0]))
                    y.append(float(klines[i].split()[1]))
                else:
                    break

            
            
    ##########################
    # Interpreting arguments #
    ##########################

    if value_type.lower()=='esp' or value_type.lower()=='potential':
        title = 'Electrostatic potential /'+val_units
    elif value_type.lower()=='dens' or value_type.lower()=='density':
        title = 'Density /'+val_units
        
    if xlim is None:
        xlim=[None,None]
    xmin=xlim[0]
    xmax=xlim[1]
    
    if ylim is None:
        ylim=[0,None]
    ymin=ylim[0]
    ymax=ylim[1]
        
    ###################
    # Creating figure #
    ###################

    fig = plt.figure()
    plt.gcf().text(0.02, 0.97, filenames, fontsize=7)
    ax = fig.add_subplot(111)  # 1x1 grid, 1st subplot, 3D
    ax.set_facecolor((0.92,0.92,0.95))
    plt.grid(color=(1,1,1))
    ax.spines['bottom'].set_color([1, 1, 1])
    ax.spines['top'].set_color([1, 1, 1])
    ax.spines['right'].set_color([1, 1, 1])
    ax.spines['left'].set_color([1, 1, 1])

    # Plot histogram
    if histpath:
    # This is  the colormap I'd like to use.
    
        cmap = getattr(cm, colour)
        norm = matplotlib.colors.Normalize(vmin=cb_lim[0], vmax=cb_lim[1])

        n, b, patches = ax.hist(b[:-1], bins=b, weights=n)
            
        bin_centers = 0.5 * (b[:-1] + b[1:])
                
        # scale values to interval [0,1]
        col = bin_centers - min(bin_centers)
        col /= max(col)
    
        for c, p in zip(bin_centers, patches):
            plt.setp(p, 'facecolor', cmap(norm(c)))
    
    if kdepath:
        ax.plot(x,y,color=(0.1,0.1,0.1))
        
    if title is None:
        ax.set_xlabel('Value')
        plt.title('Histogram')
    else:
        ax.set_xlabel(title)    
        plt.title(title.split("/")[0]+'histogram')
    ax.set_ylabel('Frequency')
    
    ax.set_axisbelow(True)
    
    if xmin is not None:
        ax.set_xlim(xmin=xmin)
    if xmax is not None:
        ax.set_xlim(xmax=xmax)
    if ymin is not None:
        ax.set_ylim(ymin=ymin)
    if ymax is not None:
        ax.set_ylim(ymax=ymax)

    plt.show()


if __name__ == '__main__':
    # Map command line arguments to function arguments.
    fct = str(sys.argv[1])

    comp_args = []    # compulsory arguments
    arg_dict = {}     # optional arguments
    for argument in sys.argv[2:]:
        if "=" not in str(argument):
            comp_args.append(argument)
        else:
            try:
                arg_dict[str(argument).split('=')[0]] = eval(str(argument).split('=')[1])   # eval turns a string into the right format (list, float, boolean etc.)
            except NameError:
                arg_dict[str(argument).split('=')[0]] = str(argument).split('=')[1]

    globals()[fct](*comp_args, **arg_dict)
