##########################################################
# Created by Laura SCALFI and Alessandro CORETTI
# laura.scalfi(at)sorbonne-universite.fr
# alessandro.coretti(at)epfl.ch
#
# Script to compute integral and electrolyte differential
# capacitances and related errors using correction for
# correlated data defined in Zwanzig et Ailawadi "Statistical
# Error Due to Finite Time Averaging in Computer Experiments"
# Physical Review 182, 280-283 (1969)
#
##########################################################

import argparse
import numpy as np
import os
import sys
import math
import matplotlib.pyplot as plt
from scipy.integrate import cumtrapz

me = 0

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

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

# Read simulation info from runtime and data
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

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

# Plot time evolution and define the equilibration time
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.savefig('charge_vs_time.png', bbox_inches='tight')
    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

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

# Compute the autocorrelation function and define the correlation times
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

    # Integration of normalized autocorrelation function
    tccum = cumtrapz(autocorr[:]/autocorr[0], dx = 1.0)
    plt.plot(np.array(range(nframes//2)), autocorr[:]/autocorr[0])
    #plt.plot(np.array(range(nframes//2-1)), tccum)
    plt.grid(True)
    plt.ylabel(r'Correlation function $<Q(0) Q(t)>/<Q^2>$')
    plt.xlabel('Time (frames)')
    plt.title("What is the correlation time tc?")
    plt.show()
    tc = int(input("Choose the correlation time tc (in frames) where the autocorrelation function goes to 0:"))
    print("==========================================")

    # Integration of the squared normalized autocorrelation function
    Tccum = cumtrapz((autocorr[:]/autocorr[0])**2, dx = 1.0)
    plt.plot(np.array(range(nframes//2)), (autocorr[:]/autocorr[0])**2)
    #plt.plot(np.array(range(nframes//2-1)), Tccum)
    plt.grid(True)
    plt.ylabel(r'Correlation function $<Q(0) Q(t)>^2/<Q^2>^2$')
    plt.xlabel('Time (frames)')
    plt.title("What is the correlation time Tc?")
    plt.show()
    Tc = int(input("Choose the correlation time Tc (in frames) where the autocorrelation function goes to 0:"))
    print("==========================================")

    return mean_charge, variance_charge, tccum, tc, Tccum, Tc, nframes

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

# Compute Integral Capacitance
def compute_integral_capacitance(potential_diff, section, mean_charge, variance_charge, tccum, tc, nframes):
    if potential_diff != 0.0:
        print("Integral capacitance calculation")
        integral_capacitance = abs(mean_charge) * e / (potential_diff * section) * 1e6 #microF/cm2
        variance_capacitance = variance_charge * (e / (potential_diff * section) * 1e6)**2
        error_integral_capacitance = np.sqrt(2 * tccum[tc] / nframes * variance_capacitance) #microF/cm2
        print("Cint = {} +/- {} microF/cm^2".format(integral_capacitance, error_integral_capacitance))
        print("================================")

    return

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

# Compute Differential Capacitance
def compute_differential_capacitance(section, variance_charge, Tccum, Tc, nframes):
    print("Differential capacitance calculation")
    differential_capacitance = variance_charge * e**2 / (k * temp * section) * 1e6 #microF/cm2
    error_differential_capacitance = np.sqrt(4 * Tccum[Tc] / nframes) * differential_capacitance #microF/cm2
    print("Cdiff,elec = {} +/- {} microF/cm^2".format(differential_capacitance, error_differential_capacitance))
    print("====================================")

    return

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

# Plot total charge distribution
def gaussian(x, mu, var):
    return 1/np.sqrt(2*math.pi*var) * np.exp(-(x - mu)**2 / (2 * var))

def plot_charge_distribution(equil_charges, mean_charge, variance_charge):

    plt.figure(figsize=(20,10))
    plt.hist(equil_charges, bins=200, histtype='step', density=True, linewidth=2, label='Total charge on left electrode')
    x = np.array(np.arange(mean_charge - 20*variance_charge, mean_charge + 20*variance_charge, 0.01))
    plt.plot(x, gaussian(x, mean_charge, variance_charge), label='Gaussian distribution of mean {} and variance {}'.format(mean_charge, variance_charge))
    plt.xlabel(r'Total_charge $Q=\sum_i q_i$')
    plt.legend()
    plt.savefig("charge_distribution.png", bbox_inches='tight')
    plt.show()

    return

def plot_log_charge_distribution(equil_charges, mean_charge, variance_charge):

    hist, bin_edge = np.histogram(equil_charges - mean_charge, bins=200, density=True)
    xvalues = (bin_edge[1:] + bin_edge[:-1]) / (2 * np.sqrt(variance_charge))
    yvalues = -((bin_edge[1:] + bin_edge[:-1])/2)**2 / (2 * variance_charge)

    plt.figure(figsize=(10,10))
    plt.semilogy(xvalues, hist)
    plt.semilogy(xvalues, 1/math.sqrt(2* math.pi * variance_charge) * np.exp(yvalues), label='Gaussian function')
    plt.ylabel(r'Total charge distribution $P(Q)$')
    plt.xlabel(r'$\delta Q / \sqrt{<\delta Q^2>}$')
    plt.legend()
    plt.savefig("log_charge_distribution.png", bbox_inches='tight')
    plt.show()

    return

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

def parseOptions(me):

    # create the top-level parser
    parser = argparse.ArgumentParser(
        prog='compute_capacitance', 
        formatter_class=argparse.RawTextHelpFormatter,
        description='''
        This script computes the integral capacitance and the contribution to the
        differential capacitance due to the charge fluctuations on the electrode
        atoms as defined in Scalfi et al. "Charge fluctuations from molecular
        simulations in the constant-potential ensemble" (PCCP 2020)).
        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.
        Along with capacitance values, the standard errors on the values are computed
        using corrections for correlated samples introduced in Zwanzig et Ailawadi,
        "Statistical Error Due to Finite Time Averaging in Computer Experiments"
        (Physical Review 1969).
        The calculation is done in several steps:
        1- Choice of the equilibration time
        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).
        2- Choice of the correlation time tc
        The plot of the normalized autocorrelation function is shown as a function of
        time and one is asked to enter the approximate number of frames after which
        the autocorrelation function goes to zero.
        3- Choice of the correlation time Tc
        The same procedure is done for the squared normalized autocorrelation function.

        The output is the integral capacitance Cint (if the potential difference is
        not zero) and the electrolyte differential capacitance Cdiff in microF/cm^2,
        with their associated standard error corrected for correlated samples.
        A figure 'charge_distribution.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 = plot_time_evolution(total_charge)
    mean_charge, variance_charge, tccum, tc, Tccum, Tc, nframes = compute_autocorr(equil_charges)
    compute_integral_capacitance(potential_diff, section, mean_charge, variance_charge, tccum, tc, nframes)
    compute_differential_capacitance(section, variance_charge, Tccum, Tc, nframes)
    plot_charge_distribution(equil_charges, mean_charge, variance_charge)

    # Uncomment to obtain logarithmic plot of the charge distribution
    # as in Limmer et al. Charge Fluctuations in Nanoscale Capacitors (PRL 2013)
    # plot_log_charge_distribution(equil_charges, mean_charge, variance_charge)

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