#!/usr/bin/env python3
# -*- python -*-

import sys
import numpy     as np
import argparse  as ap

#local libraries
sys.path.append('Lattices')
sys.path.append('Molecules')
sys.path.append('lib')
import rigid
from libgenice import generate_ice, generate_cages


def arrange_atoms(coord, cell, rotmatrices, intra, labels, name):
    atoms = []
    if len(intra) == 0:
        return atoms
    for node in range(len(coord)):
        abscom = np.dot(coord[node],cell)  # relative to absolute
        rotated = np.dot(intra,rotmatrices[node])
        for i in range(len(labels)):
            atoms.append([i,name,labels[i],rotated[i,:]+abscom])
    return atoms



def format_gro(atoms, cell, celltype):
    s = ""
    s += "#Generated by genice.py\n"
    s += "{0}\n".format(len(atoms))
    molorder = 0
    for i in range(len(atoms)):
        rank, molname, sitename, position = atoms[i]
        if rank == 0:
            molorder += 1
        s += "{0:5d}{1:5s}{2:>5s}{3:5d}{4:8.3f}{5:8.3f}{6:8.3f}\n".format(molorder,molname, sitename, i+1,position[0],position[1],position[2])
    if celltype == "rect":
        s += "    {0} {1} {2}\n".format(cell[0,0],cell[1,1],cell[2,2])
    return s



def format_mdv(atoms, cell, celltype):
    s = ""
    if celltype == "rect":
        s += "-length '({0}, {1}, {2})'\n".format(cell[0,0]*10,cell[1,1]*10,cell[2,2]*10)
    s += "-center 0 0 0\n"
    s += "-fold\n"
    s += "{0}\n".format(len(atoms))
    for i in range(len(atoms)):
        molorder, molname, sitename, position = atoms[i]
        s += "{0:5} {1:9.4f} {2:9.4f} {3:9.4f}\n".format(sitename,position[0]*10,position[1]*10,position[2]*10)
    return s



def format_euler(positions, cell, rotmatrices, celltype):
    s = ""
    if celltype == "rect":
        s += "@BOX3\n"
        s += "{0} {1} {2}\n".format(cell[0,0]*10,cell[1,1]*10,cell[2,2]*10)
    else:
        s += "@BOX9\n"
        for d in range(3):
            s += "{0} {1} {2}\n".format(cell[0,d]*10,cell[1,d]*10,cell[2,d]*10)
    s += "@NX3A\n"
    s += "{0}\n".format(len(positions))
    for i in range(len(positions)):
        position = np.dot(positions[i],cell)*10   #in Angstrom
        euler = rigid.quat2euler(rigid.rotmat2quat(rotmatrices[i]))
        s += "{0:9.4f} {1:9.4f} {2:9.4f}  {3:9.4f} {4:9.4f} {5:9.4f}\n".format(position[0],
                                                                               position[1],
                                                                               position[2],
                                                                               euler[0],
                                                                               euler[1],
                                                                               euler[2])
    return s


def format_quaternion(positions, cell, rotmatrices, celltype):
    s = ""
    if celltype == "rect":
        s += "@BOX3\n"
        s += "{0} {1} {2}\n".format(cell[0,0]*10,cell[1,1]*10,cell[2,2]*10)
    else:
        s += "@BOX9\n"
        for d in range(3):
            s += "{0} {1} {2}\n".format(cell[0,d]*10,cell[1,d]*10,cell[2,d]*10)
    s += "@NX4A\n"
    s += "{0}\n".format(len(positions))
    for i in range(len(positions)):
        position = np.dot(positions[i],cell)*10   #in Angstrom
        quat     = rigid.rotmat2quat(rotmatrices[i])
        s += "{0:9.4f} {1:9.4f} {2:9.4f}  {3:9.4f} {4:9.4f} {5:9.4f} {6:9.4f}\n".format(position[0],
                                                                               position[1],
                                                                               position[2],
                                                                               quat[0],
                                                                               quat[1],
                                                                               quat[2],
                                                                               quat[3])
    return s


def format_python(positions, cell, celltype, bondlen):
    s = ""
    s += "bondlen={0}\n".format(bondlen)
    s += "celltype='{0}'\n".format(celltype)
    s += "coord='absolute'\n"
    if celltype != "rect":
        sys.exit(1)
    NB   = 6.022e23
    density = 18.0 * len(positions) / (cell[0,0]*cell[1,1]*cell[2,2]*1e-24*NB) * 1e-3
    s += "density={0}\n".format(density)
    s += "cell='{0} {1} {2}'\n".format(cell[0,0],cell[1,1],cell[2,2])
    s += "waters=\"\"\"\n"
    for i in range(len(positions)):
        position = np.dot(positions[i],cell)
        s += "{0:9.4f} {1:9.4f} {2:9.4f}\n".format(position[0],position[1],position[2])
    s += "\"\"\"\n"
    return s


def format_digraph(graph, positions):
    s = ""
    s += "@NGPH\n"
    s += "{0}\n".format(len(positions))
    for i,j,k in graph.edges_iter(data=True):
        s += "{0} {1}\n".format(i,j)
    s += "-1 -1\n"
    return s



def getoptions():
    parser = ap.ArgumentParser(description='')
    parser.add_argument('--rep',  '-r', nargs = 3, type=int,   dest='rep',  default=(2,2,2),
                        help='Repeat the unit cell in x,y, and z directions. [2,2,2]')
    parser.add_argument('--dens', '-d', nargs = 1, type=float, dest='dens', default=(-1,),
                        help='Specify the ice density in g/cm3')
    parser.add_argument('--seed', '-s', nargs = 1, type=int,   dest='seed', default=(1000,),
                        help='Random seed [1000]')
    parser.add_argument('--format', '-f', nargs = 1,           dest='format', default=("gromacs",), metavar="gmeqdX",
                        help='Specify file format [g(romacs)|m(dview)|e(uler)|q(uaternion)|d(igraph)]')
    parser.add_argument('--water', '-w', nargs = 1,           dest='water', default=("tip3p",), metavar="model",
                        help='Specify water model. (tip3p, tip4p, etc.)')
    parser.add_argument('--g12', '-D', nargs = 1,           dest='g12', default=("g12",), metavar="model",
                        help='Specify guest in the 12-hedral cage. (empty, co2, uathf, etc.)')
    parser.add_argument('--g14', '-T', nargs = 1,           dest='g14', default=("g14",), metavar="model",
                        help='Specify guest in the 14-hedral cage. (empty, co2, uathf, etc.)')
    parser.add_argument('--g15', '-P', nargs = 1,           dest='g15', default=("g15",), metavar="model",
                        help='Specify guest in the 15-hedral cage. (empty, co2, uathf, etc.)')
    parser.add_argument('--g16', '-H', nargs = 1,           dest='g16', default=("g16",), metavar="model",
                        help='Specify guest in the 16-hedral cage. (empty, co2, uathf, etc.)')
    parser.add_argument('--cages', '-c', action='store_true', dest='cages',
                        help='Also output the cage positions. (g or m format only)')
    parser.add_argument('Type', nargs=1,
                       help='Crystal type (1c,1h,etc.)')
    return parser.parse_args()


def main():
    options = getoptions()
    cagepos = None
    if options.cages:
        cagepos, cagetype = generate_cages(options.Type[0], options.rep)
        
    water_type    = options.water[0]
    g12_type      = options.g12[0]
    g14_type      = options.g14[0]
    g15_type      = options.g15[0]
    g16_type      = options.g16[0]
    output_format = options.format[0][0]

    noGraph = False
    if output_format == "X":
        noGraph = True
    positions, rotmatrices, graph, cell, celltype, bondlen = generate_ice(options.Type[0], seed=options.seed[0], rep=options.rep, density=options.dens[0], noGraph = noGraph)

    #Final output formats
    #In case only the graph is wanted.
    if output_format == "d":
        s = format_digraph(graph, positions)
        print(s,end="")
        sys.exit(0)

    #For rigid rotors; no atomic information is required.
    if output_format == "e":          # NX3A
        s = format_euler(positions, cell, rotmatrices, celltype)
    elif output_format == "q":        # NX4A
        s = format_quaternion(positions, cell, rotmatrices, celltype)
    elif output_format == "X":        # python Lattice library
        s = format_python(positions, cell, celltype, bondlen)
    else:
        #arrange atoms
        water = __import__(water_type)
        atoms = arrange_atoms(positions, cell, rotmatrices, water.sites, water.labels, water.name)
        if cagepos is not None:
            g = dict()
            g["12"] = __import__(g12_type)
            g["14"] = __import__(g14_type)
            g["15"] = __import__(g15_type)
            g["16"] = __import__(g16_type)
            cagetype = np.array([cagetype[i%len(cagetype)] for i in range(cagepos.shape[0])])
            for ctype in ("12", "14", "15", "16"):
                cpos = cagepos[cagetype == ctype]
                cmat = np.array([np.identity(3) for i in range(cpos.shape[0])])  #guest molecules are not rotated.
                gmol = g[ctype]
                atoms += arrange_atoms(cpos, cell, cmat, gmol.sites, gmol.labels, gmol.name)
        if output_format == "m":      # MDView
            s = format_mdv(atoms, cell, celltype)
        elif output_format == "g":    # GROMACS
            s = format_gro(atoms, cell, celltype)

    print(s, end="")




if __name__ == "__main__":
    main()
