##########################################################
# Created by Laura SCALFI and Alessandro CORETTI
# laura.scalfi(at)sorbonne-universite.fr
# alessandro.coretti(at)epfl.ch
#
# Script to compute normalized autocorrelation function of
# the total charge using Fourier transforms
#
##########################################################

import argparse
import numpy as np
import os
import sys
import math
import matplotlib.pyplot as plt

me = 0

# Constants
k = 1.3806485279e-23 #J.K-1
e = 1.602176620898e-19 #C
bohr2ang = 0.52917721067
ang2cm = 1e-8
ua2V = 27.21138505

##########################################

def read_info():
    # Read temperature and potential difference
    potential1 = 0.0
    potential2 = 0.0
    with open('runtime.inpt') as run:
        for line in run:
            if (line.lstrip()).startswith("temperature"):
                temp = float(line.split()[1])
            if (line.lstrip()).startswith("potential"):
                potential = float(line.split()[1])
                potential1 = potential2
                potential2 = potential
    potential_diff = abs(potential2-potential1) * ua2V #V

    # Read box parameters and num elec
    with open('data.inpt') as run:
        for line in run:
            if (line.lstrip()).startswith("num_electrode_atoms") or (line.lstrip()).startswith("num_electrode_species"): # for retrocompatibility
                nelec = int(line.split()[1])
            if (line.lstrip()).startswith("# box"):
                line2 = run.readline()
                cellx = float(line2.split()[0])
                celly = float(line2.split()[1])
    section = cellx * celly * (bohr2ang * ang2cm)**2 #cm^2

    return temp, potential_diff, nelec, section

##########################################

def plot_time_evolution(total_charge):

    nframes = np.shape(total_charge)[0]
    stepssep = int(total_charge[1, 0] - total_charge[0, 0])

    # Time evolution of total charge
    fig, ax1 = plt.subplots(figsize=(20,10))
    ax2 = ax1.twiny()
    ax1.plot(stepssep*np.array(range(nframes)), total_charge[:, 1], '-', lw=2, label='Total charge on electrode 1')
    ax2.plot(np.array(range(nframes)), total_charge[:, 1], '-', lw=2, label='Total charge on electrode 1')
    plt.ylabel(r'Total_charge $Q=\sum_i q_i$')
    ax1.set_xlabel('Time (steps)')
    ax2.set_xlabel('Time (frames)')
    plt.legend()
    plt.show()

    # Set neq
    neq = 0
    neq = int(input("Choose the equilibration time (in frames):"))
    print("Using {} frames, ie {} steps".format(nframes - neq, (nframes - neq)*(total_charge[1, 0] - total_charge[0,0])))
    print("=========================================")

    # Charge distribution: total electrode charge
    equil_charges = total_charge[neq:, 1]

    return equil_charges, stepssep

##########################################

def compute_autocorr(equil_charges):

    nframes = np.shape(equil_charges)[0]
    mean_charge = np.mean(equil_charges)
    variance_charge = np.var(equil_charges)

    # Autocorrelation function
    print("Computing autocorrelation function")
    print("==================================")
    Cqq = np.fft.fftn(equil_charges - mean_charge)
    CC = Cqq[:] * np.conjugate(Cqq[:])
    CC[:] = np.fft.ifftn(CC[:])
    autocorr = (CC[:len(CC)//2]).real

    return autocorr

##########################################

def save_and_plot_autocorr(autocorr, stepssep):

    nframes = np.shape(equil_charges)[0]
    xdata = stepssep*np.array(range(nframes//2))
    ydata = autocorr[:]/autocorr[0]

    np.savetxt('charge_autocorrelation.out', np.stack((xdata, ydata), axis=1), header='# Timestep   Normalized_autocorrelation')

    fig, ax1 = plt.subplots(figsize=(20,10))
    ax2 = ax1.twiny()
    ax1.plot(xdata, ydata, marker='o', linestyle='None')
    ax2.plot(xdata/stepssep, ydata, marker='o', linestyle='None')
    plt.grid(True)
    ax1.set_ylabel(r'Charge autocorrelation function $<\delta Q(0) \delta Q(t)>/<\delta Q^2>$')
    ax1.set_xlabel('Time (steps)')
    ax2.set_xlabel('Time (frames)')
    plt.savefig('charge_autocorrelation.png')
    plt.show()

    return

##########################################

def parseOptions(me):

    # create the top-level parser
    parser = argparse.ArgumentParser(
        prog='compute_charge_charge_autocorrelation_function',
        formatter_class=argparse.RawTextHelpFormatter,
        description='''
        This script computes the normalized autocorrelation function of the total charge.
        The required files are a runtime and a data file corresponding to the system
        of interest, and a file containing the time series of the total charge on
        the electrodes which name should be given as input.
        This file should be formatted as the 'total_charge.out' file, ie it is necessary
        that the first column contains the time and the second column the total charge
        on either electrode.
        Upon running, the plot of the total charge as a function of the frame number
        is shown and one is asked to enter the equilibration time (in number of frames).

        The output is the autocorrelation function saved as 'charge_autocorrelation.out'
        A figure 'charge_autocorrelation.png' is saved.
        ''',
        epilog='''
        Remember to call the script from the run directory.
        '''
        )

    # defining program arguments
    parser.add_argument("total_charge_filename", nargs='+', help='The string(s) containing the filename(s) of the total_charges file(s)')

    args = None
    args = parser.parse_args(sys.argv[1:])

    if args is None:
        exit(0)

    if not(os.path.isfile('runtime.inpt')) or not(os.path.isfile('data.inpt')):
        parser.error('''
The script requires a runtime and a data file corresponding to the system of interest in the working directory.''')
        exit(0)

    return args

##########################################################

def check_args(args):
    # checking arguments for problems in formatting
    return

##########################################################

def process_args(args):
    # processing parsed arguments

    for chargename in args.total_charge_filename:
        if os.path.isfile(chargename):
            print("Reading charges from {}".format(chargename))
            try:
                total_charge = np.concatenate((total_charge, np.loadtxt(chargename)), axis=0)
            except NameError:
                total_charge = np.loadtxt(chargename)
        else:
            print("Didn't find file {}".format(chargename))
            quit()

    return total_charge

##########################################################

if __name__ == "__main__":

    # building parser, parsing arguments and checking formats
    args = parseOptions(me)
    check_args(args)
    total_charge = process_args(args)

    # check directory for required files, reading runtype, reading sections and checking correct maze initialization
    temp, potential_diff, nelec, section = read_info()
    print("System parameters:\nTemperature {} K\nPotential difference {} V\nN electrodes {}\nSection {} cm^2".format(temp, potential_diff, nelec, section))
    print("=========================================")
    equil_charges, stepssep = plot_time_evolution(total_charge)
    autocorr = compute_autocorr(equil_charges)
    save_and_plot_autocorr(autocorr, stepssep)

##########################################################
