"""
This input/test is to run MgO using ChemShel style input. DL_MONTE is
executed by command line (DL_MONTE is from the DL_MONTE package and not
embeded in ChemShel at present)

The file is modelled on the input file of butanol in the dl_poly folder 
"""

import subprocess
import numpy as np

from collections import OrderedDict

from chemsh       import *


#Import the dlmonte (htk) module with alias "dlmonte"
import dlmontepython.htk.sources.dlmonte as dlmonte
import dlmontepython.htk.sources.dlconfig as dlconfig
import dlmontepython.htk.sources.dlfieldspecies as dlfieldspecies
print("### dlmonte =", dlmonte)

##############################################################################
# Produce the FIELD file for the calculation
##############################################################################
def create_field_file():

    """ function to create and write out a FIELD file for use in DL_MONTE
    It does not take any arguments and does not return anything. TODO error checking in non existant
    """

    #if you want to read the field file from a file
    #filename = "FIELD.ORIGINAL"
    #field1 = dlmonte.dlfield.from_file(filename)
    #print(repr(field1))

    #else if you want an instance then the components of the FIELD have to added
    # manually
    field = dlmonte.dlfield.FIELD()

    #set the parameters
    field.description = "MgO interactions using Lewis-Catlow potentials"
    field.cutoff = 8.0

    #create a list of atomic types
    atype1 = dlfieldspecies.AtomType("Mg","core", 24.0, 2.0)
    field.atomtypes.append(atype1)
    atype2 = dlfieldspecies.AtomType("O","core", 16.0, -2.0)
    field.atomtypes.append(atype2)
    #print ("Atomtypes: ", field.atomtypes)

    #create a list of molecular types
    mol1 = dlfieldspecies.MolType("MgO", 512) #there are 512 atoms in the config
    field.moltypes.append(mol1)
    #print ("Moltypes: ", field.moltypes)


    #now for the tricky bit - building the VdW interactions
    field.vdw_ecap = 1.0e10
    atom1 = dlfieldspecies.Atom("Mg","core")
    atom2 = dlfieldspecies.Atom("O","core")
    int1 = dlmonte.dlinteraction.InteractionBuckingham(1007.4, 0.3262, 0.0)
    int2 = dlmonte.dlinteraction.InteractionBuckingham(22764.3, 0.149, 20.37)

    vdw1 = dlmonte.dlfield.VDW(atom1, atom2, int1)
    vdw2 = dlmonte.dlfield.VDW(atom2, atom2, int2)

    field.vdw.append(vdw1)
    field.vdw.append(vdw2)

    return field


##############################################################################
# Produce the CONTROL file for the calculation
##############################################################################
def create_control_file():

    """ creates and writes out a control file for use in DL_MONTE using the HTK.
    The HTK defines the CONTROL as three parts: title, use_block and main_block
    The function niether takes input nor returns anything. TODO again put error checking in
    """

    #if you want to directly read the CONTROL file uncomment the next two lines
    #filename = "CONTROL"
    #ctrl = dlmonte.dlcontrol.from_file(filename)
    #print(dir(ctrl))

    #start to create a CONTROL file for DL_MONTE
    title = "Control file for MgO NVT"

    # this is the 'use' keyword for DL_MONTE - we are not using any here
    use = dlmonte.dlcontrol.UseBlock()
    #print(repr(use))

    #the main control commands
    d = OrderedDict()
    d['steps'] = 1000000
    d['temperature'] = 800.0
    d['maxatmdist'] = 0.001
    d['maxvolchange'] = 0.01
    d['equilibration'] = 100000
    d['print'] = 1000
    #d['paratom'] = ""
    d['acceptatmmoveupdate'] = 100
    #d['distewald'] = ""
    d['acceptvolupdate'] = 100
    d['stack'] = 100

    main = dlmonte.dlcontrol.MainBlock(d)

    #the movers are the MC moves
    movers = []
    mover = dlmonte.dlmove.parse_atom("Mg core")
    movers.append(mover)
    mover = dlmonte.dlmove.parse_atom("O core")
    movers.append(mover)
    amove = dlmonte.dlmove.AtomMove(80, movers)
    vmove = dlmonte.dlmove.VolumeVectorMove(20) #adjust all lattice vectors
    main.moves.append(amove)
    main.moves.append(vmove)

    print("fabricated control")

    #create the instance of the control file
    cntrl = dlmonte.dlcontrol.CONTROL(title, use, main)

    #control file to stdout
    print(cntrl)


    #create a file
    cntrl.to_file()


def write_config_file(field = None, config = None):
    """
      writes out a CONFIG file : it's a bit rubbish at the moment as it should be
      done in a similar way to the FIELD/CONTROL
    """
    filename = 'CONFIG'
    try:
        outstream = open(filename, 'w')

    except:
        print('cant open CONFIG file')
        sys.exit(1)

    outstream.write(cfg.title)
    outstream.write("\n")

#    outstream.write("{:d}".format(cfg.levcfg))
#    outstream.write("{:8d}".format(cfg.pbc))
    outstream.write("\n")

#    outstream.write("{:15.8f}".format(cfg.cell.vectors[0][0] * 0.529))
#    outstream.write("{:15.8f}".format(cfg.cell.vectors[1][0] * 0.529))
#    outstream.write("{:15.8f}".format(cfg.cell.vectors[2][0] * 0.529))
#    outstream.write("\n")
#
#    outstream.write("{:15.8f}".format(cfg.cell.vectors[0][1] * 0.529))
#    outstream.write("{:15.8f}".format(cfg.cell.vectors[1][1] * 0.529))
#    outstream.write("{:15.8f}".format(cfg.cell.vectors[2][1] * 0.529))
#    outstream.write("\n")
#
#    outstream.write("{:15.8f}".format(cfg.cell.vectors[0][2] * 0.529))
#    outstream.write("{:15.8f}".format(cfg.cell.vectors[1][2] * 0.529))
#    outstream.write("{:15.8f}".format(cfg.cell.vectors[2][2] * 0.529))
    outstream.write("\n")
    #the tricky bit here requires that the number of molecule types is known
    #and the maximum number of molecules allowed for each type

    numtypes = len(field.moltypes)

    nummols_of_type = np.zeros(numtypes)
    total_mols = 0

    for i in range(numtypes):

        mol = field.moltypes[i]

        for j in range(len(cfg.molecules)):

            frag = cfg.molecules[j]

            if mol.name == frag.frag_name:

                nummols_of_type[i] += 1
                total_mols += 1

    string = "NUMMOLS " + str(total_mols)
    for i in range(numtypes):
        string = string + " " + str(int(nummols_of_type[:]))

    outstream.write(string)
    outstream.write("\n")

    for i in range(len(cfg.molecules)):

        frag = cfg.molecules[i]
        outstream.write("MOLECULE " + frag.frag_name + " " + str(frag.natoms))
        outstream.write("\n")

        for j in range(frag.natoms):

            outstream.write(frag.names[j].decode('UTF-8') + " core")
            outstream.write("\n")
            outstream.write("{:15.8f}".format(frag.coords[j][0] * 0.529))
            outstream.write("{:15.8f}".format(frag.coords[j][1] * 0.529))
            outstream.write("{:15.8f}".format(frag.coords[j][2] * 0.529))
            outstream.write("\n")




##############################################################################
# The 'main' program starts here
##############################################################################
field = create_field_file()
#it is possible to get HTK to spit out the FIELD file to stdout
#print(field)
#also in json format
#print(field.to_json())

#actually produce the FIELD file
field.to_file()


create_control_file()

# unit cell
mgo = Fragment(coords="mgo.cjson")
#print("### unit lattice consts =", mgo.cell.consts)
#print("### unit lattice natoms =", mgo.natoms)
#print("### unit coords =", mgo.coords)

# make a supercell trippling each dimension (if symmetric=True it will be expanded by 5-fold along each dimension)
supercell = mgo.getSuperCell(nlayers = 3, symmetric=False)
# in dl_monte you must give the fragment a name otherwise there is no molecule name in CONFIG
supercell.frag_name = "MgO"
#print("### supercell radius  =", supercell.radius)
#print("### supercell natoms  =", supercell.natoms, supercell.cell.natoms)
#print("### supercell consts  =", supercell.cell.consts)
#print("### supercell vectors =", supercell.cell.vectors[:,:])
#print("### supercell coords =", supercell.coords)

cfg = dlconfig.CONFIG()

cfg.title = "MgO supercell for test run"
#cfg.set_cell(supercell.cell)
#cfg.append_fragment(supercell)

write_config_file(field, cfg)

#set the path to the DL_MONTE executable
DL_MONTE_HOME = "/home/ylu/chemsh/issue87/chemsh/interfaces/dl_monte/gnu/src/dl_monte_release_v2.07/"

cmd = DL_MONTE_HOME + "DLMONTE-SRL.X"
print ('running DL_MONTE ', cmd)
status = subprocess.call(cmd, shell=True)
