##########################################################
# Created by Laura SCALFI and Alessandro CORETTI
# laura.scalfi(at)sorbonne-universite.fr
# alessandro.coretti(at)epfl.ch
#
# Script to compute the condition number of the MaZe matrix
# as introduced in Coretti et al. "Mass-zero constrained
# molecular dynamics for electrode charges in simulations
# of electrochemical systems" (JCP 2020)
#
##########################################################

import argparse
import mw
import numpy as np
import metalwalls
import sys
import os
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
potential_diff = 0.018374661240427328 * 2 #ua = 1 V
eps0 = 8.85418782e-12 #F.m-1
Eh = 4.359744650e-18 #Hartree2J

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

def read_info():
   # Read temperature and potential difference
   with open('runtime.inpt') as run:
      for line in run:
         if (line.lstrip()).startswith('electrode_charges'):
            algorithm = line.split()[1]
            if len(line.split()) > 2:
               algo_type = line.split()[2]
            else:
               if me==0: print('The "electrode_charges" algorithm should be "maze inversion" and not {}'.format(algorithm))
               quit()

            if algorithm != 'maze' or algo_type != 'inversion':
               if me==0: print('The "electrode_charges" algorithm should be "maze inversion" and not {} {}'.format(algorithm, algo_type))
               quit()
   return

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

def compute_matrix_condition_number():

    if me==0: print("Condition number 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)

    cnumber = np.linalg.cond(my_algorithms.maze.ainv)

    if me==0: print("CN = {}".format(cnumber))
    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_matrix_condition_number',
        formatter_class=argparse.RawTextHelpFormatter,
        description='''
        This script computes the condition number of the MaZe matrix, as introduced
        in Coretti et al. "Mass-zero constrained molecular dynamics for electrode
        charges in simulations of electrochemical systems" (JCP 2020).
        The required files are a runtime and a data file corresponding to the system of
        interest.
        The algorithm used to solve the electrode charges should be 'maze inversion'.
        The output is the condition number CN.
        ''',
        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
    read_info()
    compute_matrix_condition_number()

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