##########################################################
# Created by Laura SCALFI, Thomas DUFILS and Alessandro CORETTI
# laura.scalfi(at)sorbonne-universite.fr
# alessandro.coretti(at)epfl.ch
#
# Script to compute electric field profiles and electric
# potential profiles from charge density profiles by
# integrating the Poisson equation.
#
##########################################################

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

me = 0

# Constants
Eh = 4.359744650e-18 #Hartree2J
k = 1.3806485279e-23 #J.K-1
e = 1.602176620898e-19 #C
bohr2ang = 0.52917721067
bohr2nm = bohr2ang * 0.1
ang2cm = 1e-8
ang2m = 1e-10
ua2V = 27.21138505
eps0 = 8.85418782e-12 #F.m-1

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

def read_info():

    # Read eta and potential difference
    potential1 = 0.0
    potential2 = 0.0
    name = []
    with open('runtime.inpt') as run:
        for line in run:
            if (line.lstrip()).startswith("charge") and line.split()[1] == "gaussian":
                eta = float(line.split()[2]) #bohr-1
            if (line.lstrip()).startswith("potential"):
                potential = float(line.split()[1])
                potential1 = potential2
                potential2 = potential
            if line.lstrip().startswith("species ") and len(line.split())>1:
                name.append(line.split()[1])

    potential_diff = abs(potential2-potential1) * ua2V #V

    # Read 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])

    # Read elec coordinates
    atom = 0
    coord = np.zeros(nelec)
    with open('data.inpt') as run:
        for line in run:
            if (line.startswith(name[0]) or line.startswith(name[1])) and atom < nelec:
                coord[atom] = float(line.split()[3])
                atom += 1

    planes = np.unique(coord)
    nplanes = np.shape(planes)[0]

    return eta, potential_diff, planes, nplanes

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

def gaussian_charge(z, q, eta, z0):
    return np.exp(-eta**2 * (z - z0)**2) * q * eta / sqrt(pi)

def integrated_gaussian(z1, z2, q, eta):
    return q / 2 * (erf(eta * z1) - erf(eta * z2))

# Compute electric potential profile
def compute_electric_potential(charge_density, eta, planes, nplanes):
    # Electrolyte contribution
    electrolyte_charge_density = np.sum(charge_density[:, 2:-2], axis=1)
    electrolyte_electric_field = 4 * pi * cumtrapz(electrolyte_charge_density, charge_density[:, 0], initial=0)
    electrolyte_electric_potential = - cumtrapz(electrolyte_electric_field, charge_density[:, 0], initial=0)

    # Electrode contribution
    # If the bin size is not small, reduce it
    dz = charge_density[1, 0] - charge_density[0, 0]
    npoints = np.shape(charge_density)[0]
    dzth = 0.1 #bohr
    scale = int(np.floor(dz/dzth))
    dzn = dz / scale
    print("Bin size is {} bohr (with {} points), using {} bohr for electrode integration".format(dz, npoints, dzn))

    electrode_charge_density = np.zeros((npoints * scale, 2))
    for ibin in range(npoints * scale):
       electrode_charge_density[ibin, 0] = (ibin+1/2) * dzn

    index = np.argwhere(charge_density[:, -2]+charge_density[:, -1])
    index = np.transpose(index)[0]
    gaussian_values = charge_density[index, -2]+charge_density[index, -1]
    gaussian_index = charge_density[index, 0]

    # The charge density profile is truncated, we can symetrize it
    if len(planes) != len(gaussian_values) and np.any(planes < 0):
       print('Did not find all values for the electrodes, the profile must be truncated.')
       answer = input('Should I try to guess the values assuming electrode symmetry or antisymmetry? s, a or n\n')

       separation = gaussian_index[1:] - gaussian_index[:-1]
       indmax1 = np.argmax(separation)

       reconstruct_gaussian = np.zeros(nplanes)
       istart = nplanes//2 - 1 - indmax1
       istop = istart + len(gaussian_values)
       reconstruct_gaussian[istart:istop] = gaussian_values[:]
       if answer == 's':
          for plan in range(nplanes):
             if reconstruct_gaussian[plan] == 0.0:
                reconstruct_gaussian[plan] = reconstruct_gaussian[nplanes - 1 - plan]
       if answer == 'a':
          for plan in range(nplanes):
             if reconstruct_gaussian[plan] == 0.0:
                reconstruct_gaussian[plan] = - reconstruct_gaussian[nplanes - 1 - plan]
       gaussian_values = reconstruct_gaussian

    # Some charges are simply equal to zero
    elif (len(planes) != len(gaussian_values) and np.all(planes >= 0)):
       separation = gaussian_index[1:] - gaussian_index[:-1]
       indmax1 = np.argmax(separation)

       reconstruct_gaussian = np.zeros(nplanes)
       istart = nplanes//2 - 1 - indmax1
       istop = istart + len(gaussian_values)
       reconstruct_gaussian[istart:istop] = gaussian_values[:]
       gaussian_values = reconstruct_gaussian

    for plan in range(nplanes):
       pel = gaussian_values[plan]
       z0 = planes[plan] # position of planes
       for ind, z in enumerate(electrode_charge_density[:, 0]):
          z1 = (z - z0)
          z2 = z1 - dz
          electrode_charge_density[ind, 1] += integrated_gaussian(z1, z2, pel, eta)

    # Integration
    electrode_electric_field = 4 * pi * cumtrapz(electrode_charge_density[:, 1], electrode_charge_density[:, 0], initial=0)
    electrode_electric_potential = - cumtrapz(electrode_electric_field, electrode_charge_density[:, 0], initial=0)

    # Sum contributions taking into account the different bin size eventually
    total_charge = np.zeros(npoints)
    total_charge[:] = electrolyte_charge_density[:]
    total_charge[:] += electrode_charge_density[::scale, 1]

    electric_field = np.zeros(npoints)
    electric_field[:] = electrolyte_electric_field[:]
    electric_field[:] += electrode_electric_field[::scale]

    electric_potential = np.zeros(npoints)
    electric_potential[:] = electrolyte_electric_potential[:]
    electric_potential[:] += electrode_electric_potential[::scale]

    # Convert from atomic units
    total_charge /= bohr2ang**3 #ang-3
    electric_field *= Eh / e / bohr2nm # V/nm
    electric_potential *= Eh / e # V

    return total_charge, electric_field, electric_potential

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

# Plot electric potential profile
def plot_electric_potential(charge_density, total_charge, electric_field, electric_potential):

    # Plot results
    plt.rcParams['mathtext.fontset'] = 'stix'
    plt.rcParams['font.family'] = 'STIXGeneral'
    plt.rcParams.update({'font.size': 20})
    plt.rc('axes', linewidth=2)

    plt.figure(figsize=(10, 15))
    plt.subplot(311)
    plt.plot(charge_density[:, 0] * bohr2ang, total_charge, label=r'Charge density $\rho$')
    plt.ylabel(r'Charge density $\rho (\AA ^{-3})$')

    plt.subplot(312)
    plt.plot(charge_density[:, 0] * bohr2ang, electric_field, label=r'Electric field $E$')
    plt.ylabel(r'Electric field $E (V.nm^{-1})$')

    plt.subplot(313)
    plt.plot(charge_density[:, 0] * bohr2ang, electric_potential, label=r'Potential $\psi$')
    plt.ylabel(r'Electric potential $\psi (V)$')
    plt.xlabel(r'Coordinate $z (\AA)$')

    #plt.legend()
    plt.tight_layout()
    plt.savefig('electric_potential.png')
    plt.show()

    output = np.zeros((np.shape(charge_density)[0], 4))
    output[:, 0] = charge_density[:, 0] * bohr2ang
    output[:, 1] = total_charge[:]
    output[:, 2] = electric_field[:]
    output[:, 3] = electric_potential[:]
    np.savetxt('electric_potential.out', output)

    return

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

def parseOptions(me):

    # create the top-level parser
    parser = argparse.ArgumentParser(
        prog='compute_electric_potential',
        formatter_class=argparse.RawTextHelpFormatter,
        description='''
        This script computes the electric field profiles and electric potential
        profiles from charge density profiles by integrating the Poisson equation.
        The required files are a runtime and a data file corresponding to the system
        of interest, and a file containing a histogram of the total charge density and
        the charge densities for each species which name should be given as input.
        This file should be formatted as the 'charge_density.out' file, ie it is necessary
        that the first column contains the z coordinate, the second column the total
        density and the remaining columns the ionic and electrode charge densities.
        This script only works for setups with 2 planar electrodes.
        The output 'electric_potential.out' contains 4 columns with the z coordinate,
        the charge density, the electric field and the electric potential in usual units.
        A figure 'electric_potential.png' is saved.
        ''',
        epilog='''
        Remember to call the script from the run directory.
        '''
        )

    # defining program arguments
    parser.add_argument("charge_density_filename", help='The string containing the filename of the charge density file')

    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
    chargename = args.charge_density_filename
    if os.path.isfile(chargename):
       charge_density = np.loadtxt(chargename)
    else:
       print("Didn't find file {}".format(chargename))
       quit()

    return charge_density

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

if __name__ == "__main__":

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

    # check directory for required files, reading runtype, reading sections and checking correct maze initialization
    eta, potential_diff, planes, nplanes = read_info()
    print("System parameters:\nGaussian Width {} bohr-1\nPotential difference {} V\n{} planes found".format(eta, potential_diff, nplanes))
    print("=========================================")
    total_charge, electric_field, electric_potential = compute_electric_potential(charge_density, eta, planes, nplanes)
    plot_electric_potential(charge_density, total_charge, electric_field, electric_potential)

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