"""
nist2mw.py

This script is used to convert nist input data into mw data files
"""
import numpy as np

angstrom2bohr = 1.88972612546
avogadro = 6.02214129e23
boltzmann = 1.806488e-23
K2KJpermol = boltzmann * avogadro * 1.0e-3

def read_config(config_file):
    """Read spce_sample_config_periodic#.txt file

    Line  1       :  x, y, and z dimensions of periodic volume
    Line  2       :  Number of SPC/E Molecules in the rest of file (N)
                     The number of atoms in the file will be 3*N (two Hydrogen + one oxygen per molecule)
    Lines 3->3*N+2 : 
          Column 1   : Identification number of current atom (will equal [Line # -2])
          Columns 2-4: x, y, and z coordinates of current atom
          Column 5   : Atom type (either O for Oxygen or H for Hydrogen)

    All coordinates and dimensions are given in Angstroms (1.0e-10 meters)

    """
    
    with open(config_file, 'r') as f:
        # line 1
        line = f.readline()
        box_length = [ angstrom2bohr*float(word) for word in line.split()]
        # line 2
        line = f.readline()
        num_molecules = int(line)

    xyz_by_molecule = np.loadtxt(config_file, skiprows=2, usecols=(1,2,3,))

    # reorder data by atom types
    xyz_O = xyz_by_molecule[0::3,:]
    xyz_H1 = xyz_by_molecule[1::3,:]
    xyz_H2 = xyz_by_molecule[2::3,:]
    xyz = np.concatenate((xyz_O, xyz_H1, xyz_H2,), axis=0)

    # scale and shift
    xyz[:,0] = xyz[:,0]*angstrom2bohr + 0.5*box_length[0]
    xyz[:,1] = xyz[:,1]*angstrom2bohr + 0.5*box_length[1]
    xyz[:,2] = xyz[:,2]*angstrom2bohr + 0.5*box_length[2]

    return box_length, num_molecules, xyz

def write_data(data_file, num_ions, xyz, box_length):
    """Write data.inpt input file for mw2

    Line  1       :  istep
    Line  2       :  num_ions num_xyz_step num_velocity_step num_forces_step
    Line  3       :  num_atoms num_xyz_step num_qelec_step
    line  5->...  :  data

    All coordinates and dimensions are given in au (bohr)
    """

    header = []
    header.append('{:6d}'.format(0))
    header.append('{:6d} {:6d} {:6d} {:6d}'.format(num_ions, 1, 0, 0))
    header.append('{:6d} {:6d} {:6d}'.format(0, 0, 0))
    header.append('{:6d} {:6d}'.format(0, 0))

    xyz_data = ['{0:22.15e} {1:22.15e} {2:22.15e}'.format(*xyz_ion) for xyz_ion in xyz]

    with open(data_file, 'w') as f:
        f.write('\n'.join(header))
        f.write('\n')
        f.write('\n'.join(xyz_data))
        f.write('\n')
        f.write('{0:22.15e} {1:22.15e} {2:22.15e}\n'.format(box_length[0],box_length[1],box_length[2]))


class ion:
    def __init__(self, name, count, charge, mass, mobile=True):
        self.name = name
        self.count = count
        self.charge = charge
        self.mass = mass
        self.mobile = mobile

    
    def to_string(self):
        """write ion_type parameters in configuration file

        ion_type
          name   self.name
          count  self.count
          charge self.charge
          mass   self.mass
          mobile self.mobile
        """
        config = []
        config.append("ion_type")
        config.append("  name   {:s}".format(self.name))
        config.append("  count  {:6d}".format(self.count))
        config.append("  charge {:.5f}".format(self.charge))
        config.append("  mass   {:.4f}".format(self.mass))
        if (self.mobile):
            config.append("  mobile {:s}".format('True'))
        else:
            config.append("  mobile {:s}".format('False'))

        return '\n'.join(config)


def write_parameters(parameter_file, configID, num_molecules, box_length):
    """Write parameters data
    
    parameters_file (str): name of parameters file (configuration.inpt)
    configID (int): index of NIST SCP/E water configuration
    num_molecules (int): number of molecules in the configuration
    box_length (floats): box length parameters
    """

    
    # model parameters
    q_H = 0.42380 # e
    m_O = 15.9994 # amu
    m_H = 1.008   # amu

    num_steps = 0
    timestep = 1.0
    temperature = 298.0 # K
    num_pbc = 3

    lj_OO_sigma   = 0.316555789 * 10.0 # angstrom
    lj_OO_epsilon = 78.19743111 * K2KJpermol # KJpermol

    rcut = 10.0 * angstrom2bohr # angstrom

    alpha = 5.6 / min(box_length)
    kmax = 5
    rkmaxsq = 26.5 * (2 * np.pi / min(box_length))**2

    coulomb_rtol = np.exp(-(alpha*alpha*rcut*rcut)) / rcut
    coulomb_ktol = np.exp(-rkmaxsq / (4.0*alpha*alpha))

    config = """# Check run
# =========
# SPC/E water -- NIST configuration 1

# Global parameters
num_steps 0
timestep  0.000001
temperature     298

# Box
# ---
num_pbc  3
box
  length {a:.14e} {b:.14e} {c:.14e}

# Ions definition
# ---------------

ions

  ion_type
    name        O
    count   {num_molecules:6d}
    charge  -0.8476
    mass   15.9994

  ion_type
    name H1
    count  {num_molecules:6d}
    charge 0.4238
    mass 1.008

  ion_type
    name H2
    count  {num_molecules:6d}
    charge 0.4238
    mass 1.008

molecules

  molecule_type
    name SPC
    count {num_molecules:6d}
    sites O H1 H2

    constraint O   H1   1.88972
    constraint O   H2   1.88972
    constraint H1  H2   3.08587

    constraints_algorithm rattle 1.0e-7 100

interactions
  coulomb
    coulomb_rtol       {rtol:.5e}
    coulomb_rcut       18.89726124565
    coulomb_ktol       {ktol:.5e}

  lennard-jones
     lj_rcut 18.8972612456506
     lj_pair   O       O  0.6501655     3.16555789  

output
  default 1

""".format(a=box_length[0], b=box_length[1], c=box_length[2], 
           num_molecules=num_molecules, rtol=coulomb_rtol, ktol=coulomb_ktol)

    with open(parameter_file,'w') as f:
        f.write(config)

def main(configID):
    config_file = "spce_sample_config_periodic{:d}.txt".format(configID)
    box_length, num_molecules, xyz = read_config(config_file)
    
    write_parameters("runtime.inpt", configID, num_molecules, box_length)
    write_data("data.inpt", 3*num_molecules, xyz, box_length)

if __name__ == "__main__":
    import sys
    configID = int(sys.argv[1])
    main(configID)

    



    
