##########################################################
# Created by Laura SCALFI and Alessandro CORETTI
# laura.scalfi(at)sorbonne-universite.fr
# alessandro.coretti(at)epfl.ch
#
# Script to compute the empty capacitance introduced in
# Scalfi et al. "Charge fluctuations from molecular
# simulations in the constant-potential ensemble" (PCCP 2020).
#
##########################################################

import argparse
import mw
import numpy as np
import metalwalls
import os
import sys
from mpi4py import MPI

me = 0
comm = MPI.COMM_WORLD
fcomm = comm.py2f()
me = comm.Get_rank()
nprocs = comm.Get_size()

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

# Define some constants
e = 1.602176620898e-19 #C
bohr2ang = 0.52917721067
ang2cm = 1e-8
Eh = 4.359744650e-18 #Hartree2J
ua2V = 27.21138505

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

def read_info():
    # Read temperature and potential difference
    with open('runtime.inpt') as run:
        names = []
        counts = []
        for line in run:
            if (line.lstrip()).startswith('electrode_charges'):
                algorithm = line.split()[1]
                if algorithm != 'matrix_inversion':
                    if me==0: print('The "electrode_charges" algorithm should be "matrix_inversion" and not {}'.format(algorithm))
                    quit()
            if (line.lstrip()).startswith('name'):
                names.append(line.split()[1])
            if (line.lstrip()).startswith('count'):
                counts.append(line.split()[1])

    ntypes = len(counts)
    nelec1 = 0
    nelec2 = 0

    if len(names) != ntypes + 2:
        if me==0: print("Something is wrong with the input reading")
        quit()

    for types in range(ntypes):
        if names[types] == names[ntypes]:
            nelec1 = int(counts[types])
        if names[types] == names[ntypes+1]:
            nelec2 = int(counts[types])

    # 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

    if nelec != nelec1 + nelec2:
        if me==0: print("Electrodes numbers are inconsistent: Nelec {} Nelec1 {} Nelec2 {}".format(nelec, nelec1, nelec2))
        quit()

    return section, nelec1, nelec2

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

def compute_empty_capacitance(section, nelec1, nelec2):

    if me==0: print("Empty differential capacitance calculation")

    # Read inputs, setup the simulation and the parallel environment and precompute the matrix A
    # Initialize system, parallel, algorithms
    my_parallel = metalwalls.mw_parallel.MW_parallel_t()
    my_system = metalwalls.mw_system.MW_system_t()
    my_algorithms = metalwalls.mw_algorithms.MW_algorithms_t()

    # Start_parallel
    metalwalls.mw_tools.start_parallel_python(my_parallel, fcomm)

    do_output = metalwalls.mw_tools.initialize(my_system, my_parallel, my_algorithms)

    # Setup simulation
    step_output_frequency = metalwalls.mw_tools.setup(my_system, my_parallel, my_algorithms, do_output)

    # Invert the matrix
    invA = np.copy(my_algorithms.matrix_inversion.ainv)

    # Compare with tDSD matrix multiplication
    D = np.zeros((nelec1 + nelec2, 1))
    D[:nelec1, 0] = nelec2 / (nelec1 + nelec2)
    D[nelec1:, 0] = -nelec1 / (nelec1 + nelec2)
    tD = np.transpose(D)

    E = np.ones((nelec1 + nelec2, 1))
    tE = np.transpose(E)

    tEinvA = np.matmul(tE, invA)
    tEinvAE = np.matmul(tEinvA, E)
    EtEinvA = np.matmul(E, tEinvA)
    invAEtEinvA = np.matmul(invA, EtEinvA)

    S = invA - invAEtEinvA / tEinvAE
    tDSD = np.matmul(tD, np.matmul(S, D))

    empty_capacitance = tDSD[0,0] * e / (Eh / e * section) * 1e6 # microF / cm^2

    if me==0: print("Cempty = tDSD = {} microF/cm^2".format(empty_capacitance))
    if me==0: print("=================================")

    # Finalize the simulation
    mw.simulation.finalize_simulation(my_system, my_parallel, my_algorithms)

    return

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

def parseOptions(me):

    # create the top-level parser
    parser = argparse.ArgumentParser(
        prog='compute_empty_capacitance',
        formatter_class=argparse.RawTextHelpFormatter,
        description='''
        This script computes the contribution to the capacitance due to neglected charge
        fluctuations 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, with or without electrolyte atoms.
        The algorithm used to solve the electrode charges should be 'matrix_inversion'.
        The output is Cempty = tDSD in microF/cm^2.
        ''',
        epilog='''
        Remember to call the script from the run directory.
        '''
        )

    # defining program arguments

    # parsing arguments with MPI
    args = None
    try:
        if me == 0:
            args = parser.parse_args(sys.argv[1:])
    finally:
        args = comm.bcast(args, root=0)

    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
    return

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

if __name__ == "__main__":

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

    # check directory for required files, reading runtype, reading sections and checking correct maze initialization
    section, nelec1, nelec2 = read_info()
    if me==0: print("System parameters: Section {} cm^2\nNelec1 {}\nNelec2 {}".format(section, nelec1, nelec2))
    if me==0: print("=========================================")
    compute_empty_capacitance(section, nelec1, nelec2)

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