#  Copyright (C) 2021 The authors of Py-ChemShell
#
#  This file is part of Py-ChemShell.
#
#  Py-ChemShell is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Lesser General Public License as
#  published by the Free Software Foundation, either version 3 of the
#  License, or (at your option) any later version.
#
#  Py-ChemShell is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public
#  License along with Py-ChemShell.  If not, see
#  <http://www.gnu.org/licenses/>.


__author__ = 'You Lu <you.lu@stfc.ukri.org>'


# local debug level passed in by platforms/__init__.py
CHEMSH_DEBUG = locals().get('debug', 0)

# stop Qt reporting debug information
if CHEMSH_DEBUG < 5:
    from os import environ
    environ['QT_LOGGING_RULES']='*=false'

def run(args, **kwargs):
    ''''''

    # YL 26/02/2021: moved all the importing stuff under the run() function so that
    #                the user only sees this error when plotting
    try:
        from matplotlib import pyplot
    except:
        print("\n >>> ChemShell ERROR: matplotlib is not available on this machine. Exiting...\n")
        exit(999)

    from numpy import isin, load

    # run from input script
    try:
        from ..data import unitconvert
    # run from command line
    except:
        from pathlib import Path
        from sys     import modules, path
        file = Path(__file__).resolve()
        parent, top = file.parent, file.parents[2]
        # add top and remove parent
        path.append(str(top))
        try:
             path.remove(str(parent))
        except ValueError:
             pass
#        import chemsh.tools
        __package__ = 'chemsh.tools'

    # we only support a single file
    input = args.plot
    data = load(input)

    keys = list(data.keys())

    if isin(['energies', 'potentials', 'kinetics', 'temperatures'], keys).sum() >= 3:

        plotMD(data, units=args.units, **kwargs)

    if isin(['energies', 'max_gradient', 'max_step', 'RMS_gradients', 'RMS_step'], keys).sum() >= 4:

        plotOpt(data, **kwargs)

    if isin(['rmsds'], keys).sum() >= 1:

        plotMDAnalysis(data, units=args.units, **kwargs)


def plotMD(data, units=['kcal/mol','ns'], grid=True, **kwargs):
    '''Plot information of an MD trajectory'''

    from matplotlib import pyplot, ticker
    from numpy      import all, arange, array, array2string, char, load, min
    from .          import getUnits
    from ..data     import unitconvert

    eunit, tunit, efactor, tfactor = getUnits(units)

    sizes = [ v.size for k, v in data.items() if k not in [ 'opt_energies' ] ]

    # points of snapshots
    npoints = min(sizes)
    indices = arange(npoints)

    # points of time
    times = data['times']*tfactor

    # energies
    fig_e = pyplot.figure('Energies')
    # against simulation time
    axis0_e = fig_e.add_subplot()
    axis1_e = axis0_e.twiny()

    # primary X-axis
    axis0_e.set_xlim(times.min(), times.max())
    axis0_e.set_xlabel(f'Simulation time ({tunit})')
    xticks = axis0_e.get_xticks(minor=False)
    xlabels = array2string(xticks, formatter={'float_kind':lambda x:f'{x:>.3f}'}).strip('[]').split()
    # to avoid "UserWarning: FixedFormatter should only be used together with FixedLocator"
    axis0_e.set_xticks(xticks)
    axis0_e.set_xticklabels(xlabels)

    # secondary X-axis
    axis1_e.set_xlim(0, indices.max())
    axis1_e.set_xlabel('Snapshot')

    # harmonise the Y-axis ticks and labels
    energies = [ 'energies', 'kinetics', 'potentials' ]
    min_energy = min([ data[e].min() for e in energies ])*efactor
    max_energy = max([ data[e].max() for e in energies ])*efactor
    fig_dummy = pyplot.figure('dummy')
    axis0_e_dummy = fig_dummy.add_subplot()
    p_dummy = axis0_e_dummy.plot([0,1], [min_energy, max_energy])
    yticks  = axis0_e_dummy.get_yticks(minor=False)
    ylabels = array2string(yticks, formatter={'float_kind':lambda x:f'{x:>.3f}'}).strip('[]').split()

    # primary Y-axis
#    axis0_e.set_ylim()
    axis0_e.set_yticks(yticks)
    axis0_e.set_yticklabels(ylabels)
    axis0_e.set_ylabel(f'Energy ({eunit})')
    pyplot.close(fig_dummy)

    for k, v in data.items():

        if k not in energies:
            continue

        # get the data
        v = data[k]*unitconvert.au2unit[eunit]

        # primary X-axis
        pbuff_e = axis0_e.plot(times, v[:npoints], label=k)

    # legend and title
    axis0_e.set_title('Energies during the molecular dyanmics simulation trajectory')
    axis0_e.legend()
    axis0_e.grid(grid)

    # temperature
    fig_t = pyplot.figure('Temperature')
    # against simulation time
    axis0_t = fig_t.add_subplot()
    axis1_t = axis0_t.twiny()

    # primary X-axis
    axis0_t.set_xlim(times.min(), times.max())
    axis0_t.set_xlabel(f'Simulation time ({tunit})')
    xticks = axis0_t.get_xticks(minor=False)
    xlabels = array2string(xticks, formatter={'float_kind':lambda x:f'{x:>.3f}'}).strip('[]').split()
    # to avoid "UserWarning: FixedFormatter should only be used together with FixedLocator"
    axis0_t.set_xticks(xticks)
    axis0_t.set_xticklabels(xlabels)
    # get the data
    v = data['temperatures']
    pbuff_t = axis0_t.plot(times, v[:npoints], label='temperatures')
    
    # secondary X-axis
    axis1_t.set_xlim(0, indices.max())
    axis1_t.set_xlabel('Snapshot')

    # primary Y-axis
#    axis0.set_ylim()
    yticks  = axis0_t.get_yticks(minor=False)
    ylabels = array2string(yticks, formatter={'float_kind':lambda x:f'{x:>.1f}'}).strip('[]').split()
    axis0_t.set_yticks(yticks)
    axis0_t.set_yticklabels(ylabels)
    axis0_t.set_ylabel('Temperature (Kelvin)')

    # legend and title
    axis0_t.set_title('Temperature during the molecular dyanmics simulation trajectory')
    axis0_t.legend()
    axis0_t.grid(grid)

    # minimisation
    v = data.get('opt_energies', array([]))
    nminimisations = v.size
    if nminimisations:
        pyplot.figure('Minimisation')
        pyplot.plot(arange(nminimisations), v, label='opt_energies')
        pyplot.xlim(0, nminimisations-1)
        pyplot.xlabel('Minimisation cycle')
        pyplot.gca().yaxis.set_major_formatter(ticker.StrMethodFormatter('{x:,>.3f}'))
        pyplot.ylabel(f'Energy ({eunit})')
        pyplot.grid(grid)
        pyplot.title('Energy during the minimisation')
        pyplot.legend()
        strbuff_min = f' No. of minimisations: {nminimisations}'
    else:
        strbuff_min = ' No. of minimisations: N/A'

    print('\n ChemShell Plot Toolkit: Plotting MD trajectory profile...\n')
    print(f' Time length         : {times.max()} {tunit}')
    print(' No. of snapshots    :', npoints)
    print(strbuff_min)
    print(' Units requested     :', eunit, tunit)

    pyplot.show()
   

def plotOpt(data, unit='a.u.', grid=True, **kwargs):
    '''Plot convergence information of a (geometry) optimisation'''

    from matplotlib import pyplot, ticker
    from numpy      import all, arange, char, load, min
    from ..data     import unitconvert

    arrays = data.values()

    sizes = [ a.size for a in arrays ]

    npoints = min(sizes)
    indices = arange(npoints)
   
    # gradients
    pyplot.figure('Gradient')
    for k in [ 'max_gradient', 'RMS_gradient' ]:
        v = data[k]
        pbuff = pyplot.plot(indices, v[:npoints], label=k)
        pyplot.xlim(0, npoints-1)
        pyplot.xlabel('Cycle')
        pyplot.gca().yaxis.set_major_formatter(ticker.StrMethodFormatter('{x:,>.4f}'))
        pyplot.ylabel('Gradient (a.u.)')
        pyplot.grid(grid)
        pyplot.title('Maximum and RMS gradient during the geometry optimisation')
        pyplot.legend()

    # steps
    pyplot.figure('Step')
    for k in [ 'max_step', 'RMS_step' ]:
        v = data[k]
        pbuff = pyplot.plot(indices, v[:npoints], label=k)
        pyplot.xlim(0, npoints-1)
        pyplot.xlabel('Cycle')
        pyplot.gca().yaxis.set_major_formatter(ticker.StrMethodFormatter('{x:,>.4f}'))
        pyplot.ylabel('Step (a.u.)')
        pyplot.grid(grid)
        pyplot.title('Maximum and RMS step during the geometry optimisation')
        pyplot.legend()

    # temperature
    pyplot.figure('Energy')
    if 'energies' in data:
        v = data['energies']*unitconvert.au2unit[unit]
        pyplot.plot(indices, v[:npoints], label='energies')
        pyplot.xlim(0, npoints-1)
        pyplot.xlabel('Cycle')
        pyplot.gca().yaxis.set_major_formatter(ticker.StrMethodFormatter('{x:,>.6f}'))
        pyplot.ylabel('Energy ({})'.format(unit))
        pyplot.grid(grid)
        pyplot.title('Energy during the geometry optimisation')
        pyplot.legend()

    pyplot.show()
   

def plotMDAnalysis(data, units=[], grid=True, **kwargs):
    '''Plot analysis information of an MD simulation trajectory'''

    from sys import modules
    from .          import getUnits
    from ..data     import unitconvert
    from matplotlib import pyplot, ticker
    from numpy      import arange, array2string, ceil

    # default units
    lunit = 'a.u.'
    tunit = 'fs'

    # because the unit names originally as str are converted by savez_compressed()
    # to arrays, we need to get npoints from RMSDs or other entries
    sizes = []
    for key, val in data.items():
        if key == 'lunit':
            lunit = val[()]
        elif key == 'tunit':
            tunit = val[()]
        # not in use
        elif key == 'lfactor':
            lfactor = val[()]
        # not in use
        elif key == 'tfactor':
            tfactor = val[()]
        else:
            sizes.append(val.size)

    # factors for converting from fs and a.u. to the user units
    lunit_u, tunit_u, lfactor_u, tfactor_u = getUnits(units)

    # factors for converting from fs and a.u. to the data units
    lunit_d, tunit_d, lfactor_d, tfactor_d = getUnits([tunit,lunit])

    # factors for converting from data units to the user units
    lfactor = lfactor_u/lfactor_d
    tfactor = tfactor_u/tfactor_d

    npoints = min(sizes)
    indices = arange(npoints)
    times   = data['times']*tfactor
   
    # RMSDs
    fig = pyplot.figure('RMSDs')
    for k in [ 'rmsds' ]:
        v = data[k]*lfactor

        # against simulation time
        if times.max():

            # 
            axis0 = fig.add_subplot()
            axis1 = axis0.twiny()
            
#            pbuff = pyplot.plot(times, v[:npoints], label=k)
#            pyplot.xlim(0.0, times.max())
#            pyplot.xlabel(f'Simulation time ({tunit})')
#            pyplot.gca().xaxis.set_major_formatter(ticker.StrMethodFormatter('{x:,>.3f}'))

            # primary X-axis
            pbuff = axis0.plot(times, v[:npoints], label=k)
            axis0.set_xlim(times.min(), times.max())
            axis0.set_ylim(0.0)
            axis0.set_xlabel(f'Simulation time ({tunit_u})')
            xticks = axis0.get_xticks(minor=False)
            xlabels = array2string(xticks, formatter={'float_kind':lambda x:f'{x:>.3f}'}).strip('[]').split()
            # to avoid "UserWarning: FixedFormatter should only be used together with FixedLocator"
            axis0.set_xticks(xticks)
            axis0.set_xticklabels(xlabels)
            axis0.set_ylabel(f'RMSD ({lunit_u})')
            axis0.set_title('RMSDs of snapshot structures aligned against the starting point')
            axis0.legend()
            axis0.grid(grid)
            
            # secondary X-axis
            axis1.set_xlim(0, indices.max())
            axis1.set_xlabel('Snapshot')

        # against snapshots
        else:
            pbuff = pyplot.plot(indices, v[:npoints], label=k)
            pyplot.xlim(0, npoints-1)
            pyplot.xlabel('Snapshot')
#            pyplot.gca().xaxis.set_major_formatter(ticker.StrMethodFormatter('{x:f}'))
            pyplot.gca().xaxis.set_major_locator(ticker.MaxNLocator(integer=True))
            pyplot.gca().yaxis.set_major_formatter(ticker.StrMethodFormatter('{x:,>.4f}'))
            pyplot.ylabel(f'RMSD ({lunit})')
            pyplot.grid(grid)
            pyplot.title('RMSDs of snapshot structures aligned against the starting point')
            pyplot.legend()

    print('\n ChemShell Plot Toolkit: Plotting MD analysis...\n')
    print(' No. of snapshots:', npoints)
    print(' Units requested :', *units)

    pyplot.show()
   

