#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
A collection of functions that plot data from cube files

author: Oxana Andriuc

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

import numpy as np
import sys
from operator import itemgetter
import seaborn as sns
sns.set(color_codes=True)
import ReadCube as rc
import CalculateCube as cc
from mayavi import mlab
#import time
from traits.api import HasTraits, Range, Instance, on_trait_change
from traitsui.api import View, Item, Group
from mayavi.core.api import PipelineBase
from mayavi.core.ui.api import MayaviScene, SceneEditor, MlabSceneModel



coldict = {  1:'#cecece',   2:'#efffff',   3:'#db5bff',   4:'#ebff7a',   5:'#ffb7be',   6:'#7c7c7c',   7:'#0032e8',   8:'#ff0000',   9:'#c4fcff',  10:'#addbdd',
            11:'#d749ff',  12:'#cee25a',  13:'#d8adb1',  14:'#7b918d',  15:'#ff9400',  16:'#ffe500',  17:'#30ff52',  18:'#8cc9cc',  19:'#d235ff',  20:'#b0c43c',
            21:'#e2e2e2',  22:'#d1d1d1',  23:'#b7b7b7',  24:'#afb0d8',  25:'#ad8dc9',  26:'#8875af',  27:'#6251ff',  28:'#737dc4',  29:'#cc7b4f',  30:'#7d739b',
            31:'#b28488',  32:'#648780',  33:'#ea91ff',  34:'#ffc300',  35:'#a51000',  36:'#6eb4b7',  37:'#b722e2',  38:'#917d1b',  39:'#d2f6f7',  40:'#abd9db',
            41:'#85bcbf',  42:'#66a2a5',  43:'#4e8b8e',  44:'#3e7e82',  45:'#2c6e72',  46:'#1d5f63',  47:'#9dc1e8',  48:'#e8d99d',  49:'#9b696d',  50:'#416860',
            51:'#b25dc6',  52:'#e59900',  53:'#a600b2',  54:'#4f9b9e',  55:'#9212b7',  56:'#72610b',  57:'#95d2e5',  58:'#fffabc',  59:'#ddffe1',  60:'#bbe8c1',
            61:'#a7e2af',  62:'#91e0ba',  63:'#81e2c5',  64:'#6bdbb9',  65:'#64d6b3',  66:'#59cca8',  67:'#4fc6a1',  68:'#4fc6a1',  69:'#4fc682',  70:'#33a041',
            71:'#21842d',  72:'#5cced6',  73:'#5cb7d6',  74:'#4299b7',  75:'#2e7f9b',  76:'#20708c',  77:'#166987',  78:'#0a5a77',  79:'#ffe100',  80:'#bab9b6',
            81:'#99575c',  82:'#4b514f',  83:'#9041a3',  84:'#a56e00',  85:'#6d5941',  86:'#307c7f',  87:'#760696',  88:'#564803',  89:'#7abedd',  90:'#62c0ef',
            91:'#53b9ed',  92:'#399cce',  93:'#268bbf',  94:'#2577bf',  95:'#4436c1',  96:'#6d49d8',  97:'#9248d8',  98:'#ac48d8',  99:'#bd48d8', 100:'#cc48d8',
           101:'#af36ab', 102:'#af3690', 103:'#b21c71', 104:'#ff9f7c', 105:'#f27f63', 106:'#cc5f47', 107:'#bc3e32', 108:'#b72a1d', 109:'#91180d'}

#@profile
def PlotSurface(fpath, factor=1, alpha=1, size=1, axes=False, colour='rainbow', bkg=(1, 1, 1), minmax=False, minmax_range=0.05, cb_lim=None, atoms=True, au=False, curvature=False, cv_lim=None, cv_pts=5, save=False, density=1, value_type='esp', mode='sphere', units=None, azimuth=None, elevation=None, animation=False, angle=10, cat=False, an=False):
    """
    arguments: path for a cube file (str), factor (float, opt), alpha (float, opt), size (float, opt), axes (bool, opt), colour (str, opt), background colour (tuple/list/str, opt), minmax (bool, opt), colour bar limits (list, opt), atoms (boolean, opt), au (boolean, opt), curvature (boolean, opt), curvature upper limit (float, opt), number of curvature points (int, opt)

    calls: ExtractData, Axes, ValuesAsDictionary, GetVdWPoints (from ReadCube), VdWLaplacian (from CalculateCube)

    returns: -

    this function creates a plot of all the points which reside within +/- half a distance from the surface defined in terms of the Van der Waals radii of the atoms in the molecule
    (half a distance represents half the distance between two diagonally adjacent points in the grid)

    fpath (string) = path for a cube file
    factor (float, default=1) = factor by which to multiply the VdW radii
    alpha (float, default=1) = transparency value (between 0 and 1)
    size (float, default=1) = marker size
    axes (boolean, default=False) = if True, show axes; if False, hide axes
    colour (string, default='rainbow') = colour scheme ('rainbow' by default, change it to 'bwr' for blue-white-red) - full list here
    bkg (tuple/list/string, default=(1, 1, 1)) = background colour (RGB tuple/list or name - each value between 0 and 1, white by default)
    minmax (boolean, default=False) = if True, show minima and maxima; if False, hide minima and maxima
    minmax_range (float, default=0.01) =
    cb_lim (list, default=None) = limits for the colour bar (list of two floats; if not specified, the limits will be the min and max of the values; if the limits do not cover the whole range of values, a warning is printed)
    atoms (boolean, default=True) = if True, show atoms; if False, hide atoms
    au (boolean, default=False) = if True, use atomic units (i.e. bohr for coordinates); if False, use Angstroms for coordinates and units for values
    curvature (boolean, default=False) = if True, plots curvature points (either based on cv_lim if given, or cv_pts)
    cv_lim (float, default=None) = the upper limit of the curvature for the low curvature points to be plotted (it overrides cv_pts if given)
    cv_pts (integer, default=5) = the number of low curvature points to be plotted (the points will be chosen in increasing order of their absolute curvature
    save (boolean, default=False) = if True, it saves a .txt file with a summary of the results (minimum, maximum and lowest curvature)    
    density (integer, default=1) = an integer number n which indicates that the number of points stored along each axis is reduced by the density value n (therefore the total number of points is reduced by n3)
    value_type (string, default='esp') = can be ‘esp’/’potential’ or ‘dens’/’density’ (or any uppercase version of the aforementioned); used in conjunction with the keywords au and units in order to convert the values to 
                                        the appropriate units
    units (string, default=None) = the units to be used for the value; currently supports ‘V’ for ESP and ‘A’/‘angstrom’, ’nm’, ’pm’, ’m’ for density (or any lowercase/uppercase version of these). In the case of the density, 
                                    the actual units are the specified units ^(-3). The keywords value_type and au override units. If au is True, the values will be in atomic units regardless of the value passed for units. 
                                    If value_type is ‘dens’, the values will be read as density values even if the units argument has a valid ESP value (such as ‘V’). The default unit for ESP is V and for density 1/Å3.    
    """
#    start_time = time.time()
    
    print('\n\n   “Patience is the companion of wisdom.”\n'+'(St. Augustine)'.rjust(41))
    
    #########################################
    # Extracting data in convenient formats #
    #########################################
    
    exdata = rc.ExtractData(fpath, au=au, value_type=value_type, density=density, units=units)
    origin=exdata[1]
    n_x = exdata[3]
    x_vector=exdata[4]
    n_y=exdata[5]
    y_vector=exdata[6]
    n_z=exdata[7]
    z_vector=exdata[8]
    atoms_list=exdata[9]
    val=exdata[11]
    vals=exdata[12]
    axes_list=rc.Axes(origin=origin, n_x = n_x, x_vector=x_vector, n_y=n_y, y_vector=y_vector,n_z=n_z, z_vector=z_vector)
    gridpts=rc.ValuesAsDictionary(vals=vals,axes_list=axes_list, origin=origin)[0]

    [o_x, o_y, o_z] = origin
    [x_x, y_x, z_x] = x_vector
    [x_y, y_y, z_y] = y_vector
    [x_z, y_z, z_z] = z_vector

    [coords,vdw_ind] = rc.GetVdWPoints(axes_list=axes_list,gridpts=gridpts,atoms_list=atoms_list, x_vector=x_vector, y_vector=y_vector, z_vector=z_vector, origin=origin, factor=factor, au=au, cat=cat, an=an)[0:2]
    coord_dict={}
    for i in range(len(coords)):
        coord_dict[(coords[i][0],coords[i][1],coords[i][2])]=coords[i][3]
    sorted_c = sorted(coords, key=itemgetter(3))  # sorted list of coordinates by value
    xs, ys, zs, vs = [list(x) for x in zip(*coords)]


    def get_abs_third_elem(iterable):
        return abs(iterable[3])

    cv_vdw = sorted(cc.VdWLaplacian(axes_list=axes_list, gridpts=gridpts, atoms_list=atoms_list, x_vector=x_vector,y_vector=y_vector, z_vector=z_vector, val=val, factor=factor, vdw_pts=coords,vdw_ind=vdw_ind, au=au)[0], key=get_abs_third_elem)  # sorted curvature for points on vdw surface

    vs2 = np.asarray(vs)

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

#    if au:
#        l_units = '$a_0$'
#    else:
#        l_units = r'$\AA$'

    if cb_lim is None:
        cb_lim = [None, None]
    elif cb_lim[0] > min(vs2) or cb_lim[1] < max(vs2):
        print('\nWarning: The colour bar limits you have selected do not cover the whole range of values\n')


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

    fig = mlab.figure(1,fgcolor=(0,0,0),bgcolor=bkg)
    
    ########################
    # Additional variables #
    ########################
    
    hl_x = None
    hl_y = None
    hl_z = None
    
    hl = None
    tx = None
    
    def ComputeAndPlot(alpha, size, axes, minmax, bkg, atoms, curvature):
        # plot atoms if True
        for at in atoms_list:
            x_origin = at[2]
            y_origin = at[3]
            z_origin = at[4]

            if au:
                rf = 0.529177  # lengths in atomic units (i.e. bohr)
            else:
                rf = 1  # lengths in angstroms

            r = rc.vdwdict[at[0]]/rf

            # hex to rgb
            h=coldict[at[0]].lstrip("#")
            rgb_col=tuple(int(h[i:i+2], 16)/255 for i in (0, 2 ,4))

            mlab.points3d(x_origin, y_origin, z_origin, color=rgb_col, scale_mode='none', scale_factor=r, mode='sphere', opacity=int(atoms), name="Atom "+str(at[0]))

        pts = mlab.points3d(xs, ys, zs, vs, scale_mode='none', scale_factor=0.2*size,colormap=colour,vmin=cb_lim[0],vmax=cb_lim[1],opacity=alpha,mode=mode, name="Data points")
        
        mlab.colorbar(orientation='vertical')

        if axes:
            mlab.axes()
        else:
            mlab.axes(x_axis_visibility=False, y_axis_visibility=False, z_axis_visibility=False)

#        if minmax:
#            mm_alpha=0.3
#        else:
#            mm_alpha=0
        i = 0
        min_x=[]
        min_y=[]
        min_z=[]
        max_x=[]
        max_y=[]
        max_z=[]
        while sorted_c[i][3] <= min(vs2)+minmax_range:
            min_x.append(sorted_c[i][0])
            min_y.append(sorted_c[i][1])
            min_z.append(sorted_c[i][2])
            i = i+1
            if i>=len(sorted_c[i]):
                break
        mlab.points3d(min_x, min_y, min_z, scale_mode='none', scale_factor=size*0.55, mode="sphere", opacity=0.3*int(minmax), color=(181/255,236/255,1), name="Minima")  # plot all points with minimum value
        i = len(sorted_c)-1
        while sorted_c[i][3] >= max(vs2)-minmax_range:
            max_x.append(sorted_c[i][0])
            max_y.append(sorted_c[i][1])
            max_z.append(sorted_c[i][2])
            i = i-1
            if i<0:
                break
        mlab.points3d(max_x, max_y, max_z, scale_mode='none', scale_factor=size*0.55, mode="sphere", opacity=0.3*int(minmax), color=(1,181/255,1), name="Maxima")  # plot all points with maximum value

        if curvature:
            cv_alpha=1
        else:
            cv_alpha=0
            
        cv_x=[]
        cv_y=[]
        cv_z=[]
        if cv_lim:
            i = 0
            while abs(cv_vdw[i][3]) <= cv_lim:
                cv_x.append(cv_vdw[i][0])
                cv_y.append(cv_vdw[i][1])
                cv_z.append(cv_vdw[i][2])
                i += 1
            mlab.points3d(cv_x, cv_y, cv_z, scale_mode='none', scale_factor=size*0.5, mode="2ddiamond", color=(0,0,0), name="Low curvature", opacity=cv_alpha)
        else:
            for i in range(cv_pts):
                cv_x.append(cv_vdw[i][0])
                cv_y.append(cv_vdw[i][1])
                cv_z.append(cv_vdw[i][2])
            mlab.points3d(cv_x, cv_y, cv_z, scale_mode='none', scale_factor=size*0.5, mode="2ddiamond", color=(0,0,0),name="Low curvature", opacity=cv_alpha)

        glyph_points = pts.glyph.glyph_source.glyph_source.output.points.to_array()
        mlab.text(0.05,0.05,fpath.split("/")[-1],name="File name",width=0.016*len(fpath.split("/")[-1]))
        
        def picker_callback(picker):
            nonlocal hl_x, hl_y, hl_z
            nonlocal hl,tx
            """ Picker callback: this get called when on pick events.
            """
            if picker.actor in pts.actor.actors:
                
                # Find which data point corresponds to the point picked:
                # we have to account for the fact that each data point is
                # represented by a glyph with several points
                point_id = picker.point_id/glyph_points.shape[0]
                # If the no points have been selected, we have '-1'
                if point_id != -1:
                    point_id=round(point_id)
                    # Retrieve the coordinnates coorresponding to that data
                    # point
                    hl_x = xs[point_id]
                    hl_y = ys[point_id]
                    hl_z = zs[point_id]
                    if hl:
                        hl.remove()
                    hl = mlab.points3d([hl_x],[hl_y],[hl_z],scale_mode='none', scale_factor=0.2*size,color=(1,1,0),mode=mode, name="Highlighted point", reset_zoom=False)
                    if tx:
                        tx.remove()
                    tx = mlab.text(0.8,0.8,"Highlighted:".center(15)+"\n"+('%.4E' % (vs[point_id])).center(15),name="Highlighted value text",width=0.12)
                    # Move the outline to the data point.
        
        
        picker = fig.on_mouse_pick(picker_callback)
        
        picker.tolerance = 0.01
#        print("--- %s seconds ---" % (time.time() - start_time))
        
    ComputeAndPlot(alpha, size, axes, minmax, bkg, atoms, curvature)    
    
    print('\n\n '+'-'*74+'\n|'+' Summary '.center(74)+'|\n '+'-'*74)
    print('Minimum: '+'%.4E' % (sorted_c[0][3])+' at (%.3f, %.3f, %.3f)' % (sorted_c[0][0],sorted_c[0][1],sorted_c[0][2]))
    print('Maximum: '+'%.4E' % (sorted_c[len(sorted_c)-1][3])+' at (%.3f, %.3f, %.3f)' % (sorted_c[len(sorted_c)-1][0],sorted_c[len(sorted_c)-1][1],sorted_c[len(sorted_c)-1][2]))
    print('Lowest curvature: '+'%.4E' % (cv_vdw[0][3])+' at (%.3f, %.3f, %.3f)' % (cv_vdw[0][0],cv_vdw[0][1],cv_vdw[0][2])+'\n')
    
    if save:
        with open(fpath[:-5]+'_summary.txt','w') as f:
            f.write('\n########## Summary ##########\n\n'+'Minimum: '+str(sorted_c[0][3])+' at (%.3f, %.3f, %.3f)' % (sorted_c[0][0],sorted_c[0][1],sorted_c[0][2])+
                    '\nMaximum: '+str(sorted_c[len(sorted_c)-1][3])+' at (%.3f, %.3f, %.3f)' % (sorted_c[len(sorted_c)-1][0],sorted_c[len(sorted_c)-1][1],sorted_c[len(sorted_c)-1][2])+
                    '\nLowest curvature: '+str(cv_vdw[0][3])+' at (%.3f, %.3f, %.3f)' % (cv_vdw[0][0],cv_vdw[0][1],cv_vdw[0][2]))

    mlab.view(azimuth=azimuth,elevation=elevation)
    if animation:
        input("When you are happy with the current settings press Enter to proceed with the animation:")
        (azim, elev, dist, fpoint)=mlab.view()
        for i in np.arange(0,360,angle):
            mlab.view(azimuth=(azim+i)%360, elevation=elev, distance=dist, focalpoint=fpoint)
            mlab.savefig(fpath[:-5]+"0"*(3-len(str(i)))+str(i)+".png")
#    print("--- %s seconds ---" % (time.time() - start_time))
    mlab.show()
    # Decrease the tolerance, so that we can more easily select a precise
    # point.
#    picker.tolerance = 0.01


def PlotIsoVdW(fpath, factor=1, alpha=1, size=4, axes=True, colour='rainbow', cb_lim=None, atoms=True, au=False, density=1, iso=None, iso_range=0.05, value_type='esp', units=None):
    """
    arguments: path for a cube file (str), factor (float, opt), alpha (float, opt), size (float, opt), axes (bool, opt), colour (str, opt), background colour (tuple/list/str, opt), minmax (bool, opt), colour bar limits (list, opt), atoms (boolean, opt), au (boolean, opt), curvature (boolean, opt), curvature upper limit (float, opt), number of curvature points (int, opt)

    calls: ExtractData, Axes, ValuesAsDictionary, GetVdWPoints (from ReadCube)

    returns: -

    this function creates an interactive plot of all the points which reside within +/- half a distance from the surface defined in terms of the Van der Waals radii of the atoms in the molecule
    (half a distance represents half the distance between two diagonally adjacent points in the grid)

    it has a slider for iso value

    fpath (string) = path for a cube file
    factor (float, default=1) = factor by which to multiply the VdW radii
    alpha (float, default=1) = transparency value (between 0 and 1)
    size (float, default=4) = marker size
    axes (boolean, default=False) = if True, show axes; if False, hide axes
    colour (string, default='rainbow') = colour scheme ('rainbow' by default, change it to 'bwr' for blue-white-red) - full list here
    cb_lim (list, default=None) = limits for the colour bar (list of two floats; if not specified, the limits will be the min and max of the values; if the limits do not cover the whole range of values, a warning is printed)
    atoms (boolean, default=True) = if True, show atoms; if False, hide atoms
    au (boolean, default=False) = if True, use atomic units (i.e. bohr for coordinates); if False, use Angstroms for coordinates and units for values
    density (integer, default=1) = an integer number n which indicates that the number of points stored along each axis is reduced by the density value n (therefore the total number of points is reduced by n3)
    iso (float, default=None) = the value of the points to be plotted in colour
    iso_range (float, default=0.05) = the value range to consider when selecting the points (i.e. iso±0.05 by default)
    value_type (string, default='esp') = can be ‘esp’/’potential’ or ‘dens’/’density’ (or any uppercase version of the aforementioned); used in conjunction with the keywords au and units in order to convert the values to 
                                        the appropriate units
    units (string, default=None) = the units to be used for the value; currently supports ‘V’ for ESP and ‘A’/‘angstrom’, ’nm’, ’pm’, ’m’ for density (or any lowercase/uppercase version of these). In the case of the density, the actual units are the specified units ^(-3). The keywords value_type and au override units. If au is True, the values will be in atomic units regardless of the value passed for units. If value_type is ‘dens’, the values will be read as density values even if the units argument has a valid ESP value (such as ‘V’). The default unit for ESP is V and for density 1/Å3.    """
    
    print('\n\n   “Patience is the companion of wisdom.”\n'+'(St. Augustine)'.rjust(41))
    
    #########################################
    # Extracting data in convenient formats #
    #########################################
    
    exdata = rc.ExtractData(fpath, au=au, value_type=value_type, density=density, units=units)
    origin=exdata[1]
    n_x = exdata[3]
    x_vector=exdata[4]
    n_y=exdata[5]
    y_vector=exdata[6]
    n_z=exdata[7]
    z_vector=exdata[8]
    atoms_list = exdata[9]
    vals=exdata[12]
    axes_list=rc.Axes(origin=origin, n_x = n_x, x_vector=x_vector, n_y=n_y, y_vector=y_vector,n_z=n_z, z_vector=z_vector)
    (gridpts,d_ind)=rc.ValuesAsDictionary(vals=vals,axes_list=axes_list, origin=origin)
    
    vdw_coords = rc.GetVdWPoints(axes_list=axes_list,gridpts=gridpts,atoms_list=atoms_list, x_vector=x_vector, y_vector=y_vector, z_vector=z_vector, origin=origin, factor=factor, au=au)[0]
    
    x_vdw, y_vdw, z_vdw, v_vdw = [list(j) for j in zip(*vdw_coords)]
    
    ##########################
    # Interpreting arguments #
    ##########################

    if iso is None:
        iso = (min(v_vdw)+max(v_vdw))/2
        
#    if au:
#        l_units = '$a_0$'
#    else:
#        l_units = r'$\AA$'

    if cb_lim is None:
        cb_lim = [min(v_vdw), max(v_vdw)]
    elif cb_lim[0] > min(v_vdw) or cb_lim[1] < max(v_vdw):
        print('\nWarning: The colour bar limits you have selected do not cover the whole range of values\n')

    ########################
    # Additional variables #
    ########################


    pts1 = None
    pts2 = None
    a = None
    
    once=True
    
    def select_pts(iso,iso_range): 
        x_p = []
        y_p = []
        z_p = []
        v_p = []
        
        x_g=[]
        y_g=[]
        z_g=[]
        v_g=[]
            
        for i in range(len(v_vdw)):
            if v_vdw[i]>iso+abs(iso_range) or v_vdw[i]<iso-abs(iso_range):
                x_g.append(x_vdw[i])
                y_g.append(y_vdw[i])
                z_g.append(z_vdw[i])
                v_g.append(v_vdw[i])
            else:
                x_p.append(x_vdw[i])
                y_p.append(y_vdw[i])
                z_p.append(z_vdw[i])
                v_p.append(v_vdw[i])
        return [x_p,y_p,z_p,v_p],[x_g,y_g,z_g,v_g]

    [x_p,y_p,z_p,v_p],[x_g,y_g,z_g,v_g]=select_pts(iso,iso_range)
        
    def ComputeAndPlot(alpha, size, axes, atoms, x_p,y_p,z_p,v_p,x_g,y_g,z_g,v_g):
        # plot atoms if True
        nonlocal pts1, pts2, a, once
        for at in atoms_list:
            x_origin = at[2]
            y_origin = at[3]
            z_origin = at[4]

            if au:
                rf = 0.529177  # lengths in atomic units (i.e. bohr)
            else:
                rf = 1  # lengths in angstroms

            r = rc.vdwdict[at[0]]/rf

            # hex to rgb
            h=coldict[at[0]].lstrip("#")
            rgb_col=tuple(int(h[i:i+2], 16)/255 for i in (0, 2 ,4))

            mlab.points3d(x_origin, y_origin, z_origin, color=rgb_col, scale_mode='none', scale_factor=r, mode='sphere',opacity=int(atoms), name="Atom "+str(at[0]))

        
        if pts1:
            pts1.remove()
        if pts2:
            pts2.remove()
        pts1 = mlab.points3d(x_p, y_p, z_p, v_p,scale_mode='none',scale_factor=0.1*size,colormap=colour,vmin=cb_lim[0],vmax=cb_lim[1],opacity=alpha,mode='sphere', name="Iso points")
        pts2 = mlab.points3d(x_g, y_g, z_g, v_g,scale_mode='none',scale_factor=0.1*size,color=(0.8,0.8,0.8),opacity=alpha,mode='sphere', name="Other points")
        mlab.colorbar(orientation='vertical')

        if once:
            mlab.text(0.05,0.05,fpath.split("/")[-1],name="File name",width=0.015*len(fpath.split("/")[-1]))
            once=False
        
    class MyModel(HasTraits):
        iso_slider = Range(round(min(v_vdw),3), round(max(v_vdw),3), value=round(iso,3))#mode='spinner')
        scene = Instance(MlabSceneModel, ())
    
        plot = Instance(PipelineBase)
    
    
        # When the scene is activated, or when the parameters are changed, we
        # update the plot.
        @on_trait_change('iso_slider,scene.activated')
        def update_plot(self):
            [x_p,y_p,z_p,v_p],[x_g,y_g,z_g,v_g]=select_pts(self.iso_slider,iso_range)
            ComputeAndPlot(alpha, size, axes, atoms, x_p,y_p,z_p,v_p,x_g,y_g,z_g,v_g)
    
        # The layout of the dialog created
        view = View(Item('scene', editor=SceneEditor(scene_class=MayaviScene),
                         height=250, width=300),
                    Group(
                            '_', 'iso_slider',
                         ),
                    resizable=True,
                    )
    
    my_model = MyModel()
    my_model.configure_traits()  

def PlotIsosurface(fpath, factor=1, alpha=1, size=4, axes=True, colour='rainbow', minmax=False, minmax_range=0, cb_lim=None, atoms=True, au=False, density=1, iso=None, iso_range=0.01, value_type='esp', units=None):
    """
    arguments: path for a cube file (str), factor (float, opt), alpha (float, opt), size (float, opt), axes (bool, opt), colour (str, opt), background colour (tuple/list/str, opt), minmax (bool, opt), colour bar limits (list, opt), atoms (boolean, opt), au (boolean, opt), curvature (boolean, opt), curvature upper limit (float, opt), number of curvature points (int, opt)

    calls: ExtractData, Axes, ValuesAsDictionary, GetVdWPoints (from ReadCube)

    returns: -

    this function creates an interactive plot of all the points which reside within +/- half a distance from the surface defined in terms of the Van der Waals radii of the atoms in the molecule
    (half a distance represents half the distance between two diagonally adjacent points in the grid)

    it has a slider for iso value
    
    fpath (string) = path for a cube file
    factor (float, default=1) = factor by which to multiply the VdW radii
    alpha (float, default=1) = transparency value (between 0 and 1)
    size (float, default=1) = marker size
    axes (boolean, default=False) = if True, show axes; if False, hide axes
    colour (string, default='rainbow') = colour scheme ('rainbow' by default, change it to 'bwr' for blue-white-red) - full list here
    bkg (tuple/list/string, default=(1, 1, 1)) = background colour (RGB tuple/list or name - each value between 0 and 1, white by default)
    minmax (boolean, default=False) = if True, show minima and maxima; if False, hide minima and maxima
    minmax_range (float, default=0.05) = the value range to consider for the minima and maxima (i.e. ±0.05 by default)
    cb_lim (list, default=None) = limits for the colour bar (list of two floats; if not specified, the limits will be the min and max of the values; if the limits do not cover the whole range of values, a warning is printed)
    atoms (boolean, default=True) = if True, show atoms; if False, hide atoms
    au (boolean, default=False) = if True, use atomic units (i.e. bohr for coordinates); if False, use Angstroms for coordinates and units for values
    curvature (boolean, default=False) = if True, plots curvature points (either based on cv_lim if given, or cv_pts)
    cv_lim (float, default=None) = the upper limit of the curvature for the low curvature points to be plotted (it overrides cv_pts if given)
    cv_pts (integer, default=5) = the number of low curvature points to be plotted (the points will be chosen in increasing order of their absolute curvature
    save (boolean, default=False) = if True, it saves a .txt file with a summary of the results (minimum, maximum and lowest curvature)
    density (integer, default=1) = an integer number n which indicates that the number of points stored along each axis is reduced by the density value n (therefore the total number of points is reduced by n3)
    value_type (string, default='esp') = can be ‘esp’/’potential’ or ‘dens’/’density’ (or any uppercase version of the aforementioned); used in conjunction with the keywords au and units in order to convert the values
                                        to the appropriate units
    units (string, default=None) = the units to be used for the value; currently supports ‘V’ for ESP and ‘A’/‘angstrom’, ’nm’, ’pm’, ’m’ for density (or any lowercase/uppercase version of these). In the case of the density, 
                                    the actual units are the specified units ^(-3). The keywords value_type and au override units. If au is True, the values will be in atomic units regardless of the value passed for units. 
                                    If value_type is ‘dens’, the values will be read as density values even if the units argument has a valid ESP value (such as ‘V’). The default unit for ESP is V and for density 1/Å3.    
    """
    
    print('\n\n   “Patience is the companion of wisdom.”\n'+'(St. Augustine)'.rjust(41))
    
    #########################################
    # Extracting data in convenient formats #
    #########################################
    
    exdata = rc.ExtractData(fpath, au=au, value_type=value_type, density=density, units=units)
    origin=exdata[1]
    n_x = exdata[3]
    x_vector=exdata[4]
    n_y=exdata[5]
    y_vector=exdata[6]
    n_z=exdata[7]
    z_vector=exdata[8]
    atoms_list = exdata[9]
    vals=exdata[12]
    axes_list=rc.Axes(origin=origin, n_x = n_x, x_vector=x_vector, n_y=n_y, y_vector=y_vector,n_z=n_z, z_vector=z_vector)
    (gridpts,d_ind)=rc.ValuesAsDictionary(vals=vals,axes_list=axes_list, origin=origin)
    
    vdw_coords = rc.GetVdWPoints(axes_list=axes_list,gridpts=gridpts,atoms_list=atoms_list, x_vector=x_vector, y_vector=y_vector, z_vector=z_vector, origin=origin, factor=factor, au=au)[0]
    
    coords = [list(i)+[gridpts[i]] for i in list(gridpts.keys())]
    sorted_c = sorted(coords, key=itemgetter(3))  # sorted list of coordinates by value
    xs, ys, zs, vs = [list(j) for j in zip(*coords)]
    vs2 = np.asarray(vs)

    x_vdw, y_vdw, z_vdw, v_vdw = [list(j) for j in zip(*vdw_coords)]
    
    ##########################
    # Interpreting arguments #
    ##########################

    if iso is None:
        iso = (min(v_vdw)+max(v_vdw))/2
        
#    if au:
#        l_units = '$a_0$'
#    else:
#        l_units = r'$\AA$'

    if cb_lim is None:
        cb_lim = [min(v_vdw), max(v_vdw)]
    elif cb_lim[0] > min(vs2) or cb_lim[1] < max(vs2):
        print('\nWarning: The colour bar limits you have selected do not cover the whole range of values\n')

    ########################
    # Additional variables #
    ########################


    pts = None
    a = None
    
    once=True
    
    def select_pts(iso,iso_range): 
        x_p = []
        y_p = []
        z_p = []
        v_p = []
            
        for i in range(len(sorted_c)):

            if sorted_c[i][3]>iso+abs(iso_range):
                break
            if sorted_c[i][3]>=iso-abs(iso_range):
                x_p.append(sorted_c[i][0])
                y_p.append(sorted_c[i][1])
                z_p.append(sorted_c[i][2])
                v_p.append(sorted_c[i][3])
        return [x_p,y_p,z_p,v_p]

    [x_p,y_p,z_p,v_p]=select_pts(iso,iso_range)
        
    def ComputeAndPlot(alpha, size, axes, minmax, atoms, x_p,y_p,z_p,v_p):
        # plot atoms if True
        nonlocal pts, a, once
        mlab.points3d([coords[0][0],coords[len(coords)-1][0]], [coords[0][1],coords[len(coords)-1][1]], [coords[0][2],coords[len(coords)-1][2]], scale_mode='none', mode='sphere', opacity=0, reset_zoom=False)
        for at in atoms_list:
            x_origin = at[2]
            y_origin = at[3]
            z_origin = at[4]

            if au:
                rf = 0.529177  # lengths in atomic units (i.e. bohr)
            else:
                rf = 1  # lengths in angstroms

            r = rc.vdwdict[at[0]]/rf

            # hex to rgb
            h=coldict[at[0]].lstrip("#")
            rgb_col=tuple(int(h[i:i+2], 16)/255 for i in (0, 2 ,4))

            mlab.points3d(x_origin, y_origin, z_origin, color=rgb_col, scale_mode='none', scale_factor=r, mode='sphere',opacity=int(atoms), name="Atom "+str(at[0]), reset_zoom=False)
        
        if pts:
            pts.remove()
        
        pts = mlab.points3d(x_p, y_p, z_p, v_p,scale_mode='none',scale_factor=0.1*size,colormap=colour,vmin=cb_lim[0],vmax=cb_lim[1],opacity=alpha,mode='sphere', name="Data points", reset_zoom=False)
        
        mlab.colorbar(orientation='vertical',label_fmt='%.2f')

            
        if minmax:
            i = 0
            min_x=[]
            min_y=[]
            min_z=[]
            max_x=[]
            max_y=[]
            max_z=[]
            while sorted_c[i][3] <= min(vs2)+minmax_range:
                min_x.append(sorted_c[i][0])
                min_y.append(sorted_c[i][1])
                min_z.append(sorted_c[i][2])
                i = i+1
            mlab.points3d(min_x, min_y, min_z, scale_mode='none',scale_factor=size*0.5, mode="2ddiamond", color=(181/255,236/255,1), name="Minima")  # plot all points with minimum value
            i = len(sorted_c)-1
            while sorted_c[i][3] >= max(vs2)-minmax_range:
                max_x.append(sorted_c[i][0])
                max_y.append(sorted_c[i][1])
                max_z.append(sorted_c[i][2])
                i = i-1
            mlab.points3d(max_x, max_y, max_z, scale_mode='none',scale_factor=size*0.5, mode="2ddiamond", color=(1,181/255,1), name="Maxima")  # plot all points with maximum value
                
        if once:
            mlab.text(0.05,0.05,fpath.split("/")[-1],name="File name",width=0.015*len(fpath.split("/")[-1]))
            once=False
        
    class MyModel(HasTraits):
        iso_slider = Range(round(min(v_vdw),3), round(max(v_vdw),3), value=round(iso,3))#mode='spinner')
        scene = Instance(MlabSceneModel, ())
    
        plot = Instance(PipelineBase)
    
    
        # When the scene is activated, or when the parameters are changed, we
        # update the plot.
        @on_trait_change('iso_slider,scene.activated')
        def update_plot(self):
            [x_p,y_p,z_p,v_p]=select_pts(self.iso_slider,iso_range)
            ComputeAndPlot(alpha, size, axes, minmax, atoms, x_p,y_p,z_p,v_p)
    
        # The layout of the dialog created
        view = View(Item('scene', editor=SceneEditor(scene_class=MayaviScene),
                         height=250, width=300),
                    Group(
                            '_', 'iso_slider',
                         ),
                    resizable=True,
                    )
    
    my_model = MyModel()
    my_model.configure_traits()  

    # Decrease the tolerance, so that we can more easily select a precise
    # point.
#    picker.tolerance = 0.01

def PlotSlice(fpath, factor=1, alpha=1, size=4, axes=True, colour='rainbow', minmax=False, minmax_range=0, cb_lim=None, atoms=False, au=False, a=0, b=0, c=1, x0=None, y0=None, z0=None, density=1, contour=True, value_type='esp', units=None):
    """
    arguments: path for a cube file (str), factor (float, opt), alpha (float, opt), size (float, opt), axes (bool, opt), colour (str, opt), background colour (tuple/list/str, opt), minmax (bool, opt), colour bar limits (list, opt), atoms (boolean, opt), au (boolean, opt), curvature (boolean, opt), curvature upper limit (float, opt), number of curvature points (int, opt)

    calls: ExtractData, Axes, ValuesAsDictionary, GetVdWPoints (from ReadCube)

    returns: -

    this function creates an interactive plot of all the points which reside within +/- half a distance from the surface defined in terms of the Van der Waals radii of the atoms in the molecule
    (half a distance represents half the distance between two diagonally adjacent points in the grid)

    it has sliders for a, b, c, x0, y0, z0

    fpath (string) = path for a cube file
    factor (float, default=1) = factor by which to multiply the VdW radii
    alpha (float, default=1) = transparency value (between 0 and 1)
    size (float, default=4) = marker size
    axes (boolean, default=True) = if True, show axes; if False, hide axes
    colour (string, default='rainbow') = colour scheme ('rainbow' by default, change it to 'bwr' for blue-white-red) - full list here
    minmax (boolean, default=False) = if True, show minima and maxima; if False, hide minima and maxima
    minmax_range (float, default=0) = the value range to consider for the minima and maxima
    cb_lim (list, default=None) = limits for the colour bar (list of two floats; if not specified, the limits will be the min and max of the values; if the limits do not cover the whole range of values, a warning is printed)
    atoms (boolean, default=False) = if True, show atoms; if False, hide atoms
    au (boolean, default=False) = if True, use atomic units (i.e. bohr for coordinates); if False, use Angstroms for coordinates and units for values
    a (float, default=0) = the x component of the normal vector defining the plane
    b (float, default=0) = the y component of the normal vector defining the plane
    c (float, default=1) = the z component of the normal vector defining the plane
    x0 (float, default=None) = the x coordinate of the origin for the normal vector defining the plane
    y0 (float, default=None) = the y coordinate of the origin for the normal vector defining the plane
    z0 (float, default=None) = the z coordinate of the origin for the normal vector defining the plane
    density (integer, default=1) = an integer number n which indicates that the number of points stored along each axis is reduced by the density value n (therefore the total number of points is reduced by n3)
    contour (boolean, default=True) = if True, show van der Waals contour; if False, hide van der Waals contour
    value_type (string, default='esp') = can be ‘esp’/’potential’ or ‘dens’/’density’ (or any uppercase version of the aforementioned); used in conjunction with the keywords au and units in order to convert the values to the
                                        appropriate units
    units (string, default=None) = the units to be used for the value; currently supports ‘V’ for ESP and ‘A’/‘angstrom’, ’nm’, ’pm’, ’m’ for density (or any lowercase/uppercase version of these). In the case of the density, 
                                    the actual units are the specified units ^(-3). The keywords value_type and au override units. If au is True, the values will be in atomic units regardless of the value passed for units. 
                                    If value_type is ‘dens’, the values will be read as density values even if the units argument has a valid ESP value (such as ‘V’). The default unit for ESP is V and for density 1/Å3.    
    """
    
    print('\n\n   “Patience is the companion of wisdom.”\n'+'(St. Augustine)'.rjust(41))
    
    #########################################
    # Extracting data in convenient formats #
    #########################################
    
    exdata = rc.ExtractData(fpath, au=au, value_type=value_type, density=density, units=units)
    origin=exdata[1]
    n_x = exdata[3]
    x_vector=exdata[4]
    n_y=exdata[5]
    y_vector=exdata[6]
    n_z=exdata[7]
    z_vector=exdata[8]
    atoms_list = exdata[9]
    vals=exdata[12]
    axes_list=rc.Axes(origin=origin, n_x = n_x, x_vector=x_vector, n_y=n_y, y_vector=y_vector,n_z=n_z, z_vector=z_vector)
    (gridpts,d_ind)=rc.ValuesAsDictionary(vals=vals,axes_list=axes_list, origin=origin)
    
    vdw_coords = rc.GetVdWPoints(axes_list=axes_list,gridpts=gridpts,atoms_list=atoms_list, x_vector=x_vector, y_vector=y_vector, z_vector=z_vector, origin=origin, factor=factor, au=au)[0]
    
    coords = [list(i)+[gridpts[i]] for i in list(gridpts.keys())]
    
    dist = ((x_vector[0]+y_vector[0]+z_vector[0])**2+(x_vector[1]+y_vector[1]+z_vector[1])**2+(x_vector[2]+y_vector[2]+z_vector[2])**2)**0.5
    half_dist = 0.3*dist
    
    
    x_vdw, y_vdw, z_vdw, v_vdw = [list(j) for j in zip(*vdw_coords)]
    
    sorted_c = sorted(coords, key=itemgetter(3))  # sorted list of coordinates by value
    
    xs, ys, zs, vs = [list(j) for j in zip(*coords)]
    vs2 = np.asarray(vs)
    
    ##########################
    # Interpreting arguments #
    ##########################
    
    if x0 is None:
        x0 = (min(xs)+max(xs))/2
    if y0 is None:
        y0 = (min(ys)+max(ys))/2
    if z0 is None:
        z0 = (min(zs)+max(zs))/2
        
#    if au:
#        l_units = '$a_0$'
#    else:
#        l_units = r'$\AA$'

    if cb_lim is None:
        cb_lim = [min(v_vdw), max(v_vdw)]
    elif cb_lim[0] > min(vs2) or cb_lim[1] < max(vs2):
        print('\nWarning: The colour bar limits you have selected do not cover the whole range of values\n')


    ########################
    # Additional variables #
    ########################

    
    
    pts = None
    
    cont = None
    
    once=True
    
    def select_pts(a,b,c,x0,y0,z0): 
        x_p = []
        y_p = []
        z_p = []
        v_p = []
        
        x_out = []
        y_out = []
        z_out = []
        
        d = -(a*x0+b*y0+c*z0)
        
        for i in range(len(coords)):
            d_to_plane = abs(a*coords[i][0]+b*coords[i][1]+c*coords[i][2]+d)/(a**2+b**2+c**2)**0.5
            if d_to_plane < half_dist:
                x_p.append(coords[i][0])
                y_p.append(coords[i][1])
                z_p.append(coords[i][2])
                v_p.append(coords[i][3])
                
        for i in range(len(x_vdw)):
            d_to_plane = abs(a*x_vdw[i]+b*y_vdw[i]+c*z_vdw[i]+d)/(a**2+b**2+c**2)**0.5
            if d_to_plane < half_dist:
                x_out.append(x_vdw[i])
                y_out.append(y_vdw[i])
                z_out.append(z_vdw[i])                        

        return [x_p,y_p,z_p,v_p],[x_out,y_out,z_out]

    [[x_p,y_p,z_p,v_p], [x_out,y_out,z_out]]=select_pts(a,b,c,x0,y0,z0)
        
    def ComputeAndPlot(alpha, size, axes, minmax, atoms, x_p,y_p,z_p,v_p, x_out,y_out,z_out):
        # plot atoms if True
        nonlocal pts, a, once, cont
        mlab.points3d([coords[0][0],coords[len(coords)-1][0]], [coords[0][1],coords[len(coords)-1][1]], [coords[0][2],coords[len(coords)-1][2]], scale_mode='none', mode='sphere', opacity=0, reset_zoom=False)
        
        a = mlab.axes()
        
        for at in atoms_list:
            x_origin = at[2]
            y_origin = at[3]
            z_origin = at[4]

            if au:
                rf = 0.529177  # lengths in atomic units (i.e. bohr)
            else:
                rf = 1  # lengths in angstroms

            r = rc.vdwdict[at[0]]/rf

            # hex to rgb
            h=coldict[at[0]].lstrip("#")
            rgb_col=tuple(int(h[i:i+2], 16)/255 for i in (0, 2 ,4))

            mlab.points3d(x_origin, y_origin, z_origin, color=rgb_col, scale_mode='none', scale_factor=r, mode='sphere',opacity=int(atoms), name="Atom "+str(at[0]), reset_zoom=False)

        
        if pts:
            pts.remove()
        
        pts = mlab.points3d(x_p, y_p, z_p, v_p,scale_mode='none',scale_factor=0.1*size,colormap=colour,vmin=cb_lim[0],vmax=cb_lim[1],opacity=alpha, name="Data points", reset_zoom=False)

        # plot contour
        
        if contour:
            if cont:
                cont.remove()

            cont=mlab.points3d(x_out, y_out, z_out, scale_mode='none',scale_factor=0.11*size,color=(0,0,0),opacity=0.3, mode='sphere',name="VdW outline", reset_zoom=False)
        
        mlab.colorbar(orientation='vertical')

            
        if minmax:
            i = 0
            min_x=[]
            min_y=[]
            min_z=[]
            max_x=[]
            max_y=[]
            max_z=[]
            
            while sorted_c[i][3] <= min(vs2)+minmax_range:
                min_x.append(sorted_c[i][0])
                min_y.append(sorted_c[i][1])
                min_z.append(sorted_c[i][2])
                i = i+1
            mlab.points3d(min_x, min_y, min_z, scale_mode='none',scale_factor=size*0.5, mode="2ddiamond", color=(181/255,236/255,1), name="Minima")  # plot all points with minimum value
            i = len(sorted_c)-1
            while sorted_c[i][3] >= max(vs2)-minmax_range:
                max_x.append(sorted_c[i][0])
                max_y.append(sorted_c[i][1])
                max_z.append(sorted_c[i][2])
                i = i-1
            mlab.points3d(max_x, max_y, max_z, scale_mode='none',scale_factor=size*0.5, mode="2ddiamond", color=(1,181/255,1), name="Maxima")  # plot all points with maximum value
        if once:
            mlab.text(0.05,0.05,fpath.split("/")[-1],name="File name",width=0.015*len(fpath.split("/")[-1]))
            once=False
        
    class MyModel(HasTraits):
        a_slider = Range(0., 1., value=a)#mode='spinner')
        b_slider = Range(0., 1., value=b)#mode='spinner')
        c_slider = Range(0., 1., value=c)#mode='spinner')
        x0_slider = Range(round(min(xs),2), round(max(xs),2), value=round(x0,2))#mode='spinner')
        y0_slider = Range(round(min(ys),2), round(max(ys),2), value=round(y0,2))#mode='spinner')
        z0_slider = Range(round(min(zs),2), round(max(zs),2), value=round(z0,2))#mode='spinner')
        
        
        scene = Instance(MlabSceneModel, ())
    
        plot = Instance(PipelineBase)
    
    
        # When the scene is activated, or when the parameters are changed, we
        # update the plot.
        @on_trait_change('a_slider,b_slider,c_slider,x0_slider,y0_slider,z0_slider,scene.activated')
        def update_plot(self):
            a = self.a_slider
            b = self.b_slider
            c = self.c_slider
            x0 = self.x0_slider
            y0 = self.y0_slider
            z0 = self.z0_slider
            [[x_p,y_p,z_p,v_p], [x_out,y_out,z_out]]=select_pts(a,b,c,x0,y0,z0)
            ComputeAndPlot(alpha, size, axes, minmax, atoms, x_p,y_p,z_p,v_p,x_out,y_out,z_out)
    
        # The layout of the dialog created
        view = View(Item('scene', editor=SceneEditor(scene_class=MayaviScene),
                         height=250, width=300),
                    Group(
                            '_', 'a_slider','b_slider','c_slider','x0_slider','y0_slider','z0_slider',
                         ),
                    resizable=True,
                    )
    
    my_model = MyModel()
    my_model.configure_traits()  

    print('\n\n '+'-'*74+'\n|'+' Summary '.center(74)+'|\n '+'-'*74)
    print('Minimum: '+'%.4E' % (sorted_c[0][3])+' at (%.3f, %.3f, %.3f)' % (sorted_c[0][0],sorted_c[0][1],sorted_c[0][2]))
    print('Maximum: '+'%.4E' % (sorted_c[len(sorted_c)-1][3])+' at (%.3f, %.3f, %.3f)' % (sorted_c[len(sorted_c)-1][0],sorted_c[len(sorted_c)-1][1],sorted_c[len(sorted_c)-1][2]))

    # Decrease the tolerance, so that we can more easily select a precise
    # point.
#    picker.tolerance = 0.01

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)
