##########################################################
# Created by Alessandro CORETTI and Laura SCALFI
# alessandro.coretti(at)epfl.ch
# laura.scalfi(at)sorbonne-universite.fr
#
# Script to check intial configuration for MaZe simulation
# at constant potential with or without charge neutrality
# or with polarizable force field.
# The script identifies the type of run to be performed
# reading the runtime.inpt file.
# The minimization algorithm for the electrode charges or
# the polarizable force field shoud be maze.
# It then checks that the minimum requirements for a MaZe
# run are satisfied by the data.inpt file, i.e. that there
# are a sufficient number of previous timesteps for the
# evolution of the additional variables for the Verlet
# algorithm.
# If this is not the case, the script launches a two-step
# run to meet the requirements of the data.inpt.
# In particular, it launches a two-step run with the same
# parameters of the MaZe simulation, but with a cg
# for the additional variables.
#
##########################################################

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

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

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

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

def read_info():

    timestep = 0.
    pimrun = False
    elecrun = False
    charge_neutrality = False

    if me == 0: print("\nChecking runtime.inpt for type of run...")

    with open('runtime.inpt') as run:
        for line in run:
            if (line.lstrip()).startswith('timestep'):
                timestep = float(line.split()[1])
            if (line.lstrip()).startswith('dipoles_minimization'):
                pimrun = True
                algorithm = line.split()[1]
                if algorithm != 'maze':
                    if me == 0: print('The "dipoles_minimization" algorithm should be "maze" and not "{}".\n'.format(algorithm))
                    quit()
            if (line.lstrip()).startswith('electrode_charges'):
                elecrun = True
                algorithm = line.split()[1]
                if algorithm != 'maze':
                    if me == 0: print('The "electrode_charges" algorithm should be "maze" and not "{}".\n'.format(algorithm))
                    quit()
            if (line.lstrip()).startswith('charge_neutrality'):
                if (line.split()[1] == "true"):
                    charge_neutrality = True

    if me == 0:
        print("Checking complete:")
        if pimrun:
            print("    Polarizable force field run detected.")
        if elecrun and charge_neutrality:
            print("    Constant potential run detected with charge neutrality enforced.")
        elif elecrun:
            print("    Constant potential run detected.")
        print("    The simulation timestep is {}.".format(timestep))


    return timestep, pimrun, elecrun, charge_neutrality

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

def read_sections(timestep, pimrun, elecrun, charge_neutrality):

    t0 = 0
    timestep_given = True
    timestep_read = 0
    same_timestep = False
    dipoles_count = 0
    charges_count = 0
    potential_shift_count = 0

    if me == 0: print("\nChecking data.inpt for correct number of sections...")

    with open('data.inpt') as sections:
        if pimrun:
            for line in sections:
                if (line.lstrip()).startswith("# dipoles"):
                    dipoles_count += 1
                    try:
                        t0 = float(line.split()[-1])
                        timestep_read += t0
                    except ValueError:
                        timestep_given = False
                # if dipoles_count == 2: break
        if elecrun and charge_neutrality:
            for line in sections:
                if (line.lstrip()).startswith("# charges"):
                    charges_count += 1
                    try:
                        t0 = float(line.split()[-1])
                        timestep_read += t0
                    except ValueError:
                        timestep_given = False
                elif (line.lstrip()).startswith("# potential_shift"):
                    potential_shift_count += 1
                # if charges_count >= 2 and potential_shift_count >= 2: break
        elif elecrun:
            for line in sections:
                if (line.lstrip()).startswith("# charges"):
                    charges_count += 1
                    try:
                        t0 = float(line.split()[-1])
                        timestep_read += t0
                    except ValueError:
                        timestep_given = False
                # if charges_count == 2: break

    if me == 0: print("Checking complete:")
    if pimrun:
        if me == 0: print("    Found {} sections for dipoles.".format(dipoles_count))
        if (dipoles_count > 1):
            if timestep_given:
                timestep_read -= dipoles_count*t0
                timestep_read /= dipoles_count*(charges_count-1)/2
                if me == 0: print("    The timestep read from the data.inpt files is {}.".format(timestep_read))
            else:
                if me == 0: print("    The timestep has not been found in the data file.")
    if elecrun and charge_neutrality:
        if me == 0: print("    Found {} sections for charges and {} sections for potential shift.".format(charges_count, potential_shift_count))
        if (charges_count > 1):
            if timestep_given:
                timestep_read -= charges_count*t0
                timestep_read /= charges_count*(charges_count-1)/2
                if me == 0: print("    The timestep read from the data.inpt files is {}.".format(timestep_read))
            else:
                if me == 0: print("    The timestep has not been found in the data file.")
    elif elecrun:
        if me == 0: print("    Found {} sections for charges.".format(charges_count))
        if (charges_count > 1):
            if timestep_given:
                timestep_read -= charges_count*t0
                timestep_read /= charges_count*(charges_count-1)/2
                if me == 0: print("    The timestep read from the data.inpt files is {}.".format(timestep_read))
            else:
                if me == 0: print("    The timestep has not been found in the data file.")

    if timestep_given:
        same_timestep = (np.abs(timestep_read - timestep) < 1.0e-3)

    return dipoles_count, charges_count, potential_shift_count, same_timestep

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

def check_maze_initialization(pimrun, dipoles_sections, elecrun, charge_neutrality, charges_sections, potential_shift_sections, same_dt):

    if me == 0: print("\nChecking MaZe initialization...")

    if pimrun:
        if dipoles_sections >= 2:
            if (same_dt):
                if me == 0:
                    print("Checking complete:")
                    print("    Initialization is good. Simulation with MaZe can be safely launched.\n")
                return True
            else:
                if me == 0:
                    print("Checking complete:")
                    print("    The data.inpt file has been obtained with a different timestep or the timestep has not been found in the data file.")
                return False
        else:
            if me == 0:
                print("Checking complete:")
                print("    Not enough dipoles sections for MaZe initialization.")
            return False


    if elecrun and charge_neutrality:
        if charges_sections >= 2 and potential_shift_sections >= 2:
            if (same_dt):
                if me == 0:
                    print("Checking complete:")
                    print("    Initialization is good. Simulation with MaZe can be safely launched.\n")
                return True
            else:
                if me == 0:
                    print("Checking complete:")
                    print("    The data.inpt file has been obtained with a different timestep or the timestep has not been found in the data file.")
                return False
        else:
            if me == 0:
                print("Checking complete:")
                if (charges_sections < 2 and potential_shift_sections < 2):
                    print("    Not enough charges and potential shift sections for MaZe initialization.")
                elif (charges_sections < 2):
                    print("    Not enough charges sections for MaZe initialization.")
                elif (potential_shift_sections < 2):
                    print("    Not enough potential shift sections for MaZe initialization.")
            return False

    elif elecrun:
        if charges_sections >= 2:
            if (same_dt):
                if me == 0:
                    print("Checking complete:")
                    print("    Initialization is good. Simulation with MaZe can be safely launched.\n")
                return True
            else:
                if me == 0:
                    print("Checking complete:")
                    print("    The data.inpt file has been obtained with a different timestep or the timestep has not been found in the data file.")
                return False
        else:
            if me == 0:
                print("Checking complete:")
                print("    Not enough charges sections for MaZe initialization.")
            return False

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

def generate_datafile(pimrun, elecrun):

    if me == 0: print("\nGenerating initial configuration for MaZe simulation...")

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

    # Modify runtime to run CG calculation
    if me==0:
        shutil.copyfile('data.inpt', 'data.inpt.bck')
        shutil.copyfile('runtime.inpt', 'runtime.inpt.bck')
        with open('runtime.inpt.bck') as bck:
            lines = bck.readlines()
            with open('runtime.inpt', 'w') as runtime:
                if pimrun:
                    for i, line in enumerate(lines):
                        if len(line.split())>0 and line.split()[0] == 'dipoles_minimization':
                            line = 'dipoles_minimization cg 1.0e-12 200\n'
                        runtime.write(line)
                if elecrun:
                    for i, line in enumerate(lines):
                        if len(line.split())>0 and line.split()[0] == 'electrode_charges':
                            line = 'electrode_charges cg 1.0e-12 200\n'
                        runtime.write(line)
    comm.Barrier()

    # Initialize simulation, read input files
    do_output = metalwalls.mw_tools.initialize(my_system, my_parallel, my_algorithms)

    if me==0:
        shutil.copyfile('runtime.inpt.bck', 'runtime.inpt')
        os.remove('runtime.inpt.bck')
    comm.Barrier()

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

    # Run 2 steps
    metalwalls.mw_tools.run_step(my_system, my_parallel, my_algorithms, 0, 2, do_output, step_output_frequency, True, True)

    # Make sure to write the restart using
    metalwalls.mw_system.write_data(my_system, 2, 'new_data')
    if me==0:
        for fil in os.listdir('./'):
            if fil.startswith('new_data'):
                os.rename(fil, 'data.inpt')

    # Finalize
    metalwalls.mw_tools.finalize(my_system, my_parallel, my_algorithms)
    metalwalls.mw_tools.finalize_parallel()

    return

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

def parseOptions(me):

    # create the top-level parser
    parser = argparse.ArgumentParser(
        prog='check_MaZe_initialization',
        formatter_class=argparse.RawTextHelpFormatter,
        description='''
        Script to check intial configuration for MaZe simulation
        at constant potential with or without charge neutrality
        or with polarizable force field.
        The script identifies the type of run to be performed
        reading the runtime.inpt file.
        The minimization algorithm for the electrode charges or
        the polarizable force field shoud be 'maze'.
        It then checks that the minimum requirements for a MaZe
        run are satisfied by the data.inpt file, i.e. that there
        are a sufficient number of previous timesteps for the
        evolution of the additional variables for the Verlet
        algorithm.
        If this is not the case, the script launches a two-step
        run to meet the requirements of the data.inpt.
        In particular, it launches a two-step run with the same
        parameters of the MaZe simulation, but with a cg
        for the additional variables.
        ''',
        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

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

if __name__ == "__main__":

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

    # check directory for required files, reading runtype, reading sections and checking correct maze initialization
    dt, pimrun, elecrun, charge_neutrality = read_info()
    dipoles_sections, charges_sections, potential_shift_sections, same_dt = read_sections(dt, pimrun, elecrun, charge_neutrality)
    check_initialization = check_maze_initialization(pimrun, dipoles_sections, elecrun, charge_neutrality, charges_sections, potential_shift_sections, same_dt)

    # if initialization is wrong, running 2-step simulation using cg to obtain 0 and -dt timesteps
    if not check_initialization:
        generate_datafile(pimrun, elecrun)

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