import os
import sys
import shutil
import glob
import csv
import numpy as np
from numpy import exp, log, sin, cos, tan, arcsin, arccos, arctan, pi
from scipy.interpolate import interp1d
from pprint import pprint
from matplotlib import pyplot as plt

sys.path.append("c:/Programs/python/lib")
sys.path.append("d:/Programs/python/lib")

from tkfile import tkFile
from tkutils import IsDir, IsFile, SplitFilePath
from tkutils import terminate, pint, pfloat, getarg, getintarg, getfloatarg
from tksci.tksci import Reduce01, Round
from tksci.tkmatrix import make_matrix1, make_matrix2, make_matrix3
from tkcrystal.tkcif import tkCIF, tkCIFData
from tkcrystal.tkcrystal import tkCrystal
from tkcrystal.tkatomtype import tkAtomType


#================================
# global parameters
#================================
debug = 0

atomtypes = ['O', 'Zn', 'Ga', 'In']
xlim, ylim, zlim = 1.0, 1.0, 1.0
#xlim, ylim, zlim = 0.7, 0.7, 0.7

cfgfile = 'STD0.cfg'
ciffile = 'STD0.cif'

freduce01 = 1


#=============================
# Treat argments
#=============================
def usage():
    global cfgfile
    global freduce01, xlim, ylim, zlim

    print("")
    print("Usage:")
    print("  (i) python {} cfgfile (reduce01 xilm ylim zlim)".format(sys.argv[0]))
    print("      reduce01: Reduce fractional coordinates to [0, 1} if reduce01 == 1")
    print("      ex: python {} {} {} {} {} {}"
                    .format(sys.argv[0], cfgfile, freduce01, xlim, ylim, zlim))

def updatevars():
    global cfgfile, ciffile
    global freduce01, xlim, ylim, zlim

    argv = sys.argv
#    if len(argv) == 1:
#        terminate(usage = usage)

    cfgfile   = getarg     (1, cfgfile)
    freduce01 = getintarg  (2, freduce01)
    xlim      = getfloatarg(3, xlim)
    ylim      = getfloatarg(4, ylim)
    zlim      = getfloatarg(5, zlim)

    header, ext = os.path.splitext(cfgfile)
    filebody    = os.path.basename(header)
    cfgfile     = filebody + '.cfg'
    ciffile     = filebody + '.cif'


def read_cfgfile(cfgfile):
    cfg = tkFile(cfgfile, 'r')
    if not cfg:
        terminate("Error in read_cfgfile: Can not read [{}]".format(cfgfile), usage = usage)

    line        = cfg.ReadLine()
    sample_name = cfg.ReadLine().strip()
    print("sample name: [{}]".format(sample_name))
    
    line = cfg.SkipTo(r"molecules\s+of\s+all")
    ret = cfg.Search(r"(\d+)", line)
    natoms = pint(ret[1])

    line = cfg.SkipTo("Defining vectors")
    aij = np.empty([3, 3])
    aij[0] = cfg.ReadLine().split()
#    aij[1] = cfg.ReadLine().split()
#    aij[2] = cfg.ReadLine().split()
    for i in range(3):
        aij[0][i] = pfloat(aij[0][i]) * 2.0
    aij[1] = [0.0, aij[0][0], 0.0]
    aij[2] = [0.0, 0.0, aij[0][0]]

    print("")
    print("Lattice vectors:")
    for i in range(3):
        print("  {:16.12f}  {:16.12f}  {:16.12f}".format(*aij[i]))

    ntypes = []
    ntot = 0
    while 1:
        line = cfg.SkipTo(r"molecules\s+of\s+type")
        if not line:
            break

        ret = cfg.Search(r"(\d+)", line)
        n = pint(ret[1])
        ntypes.append(n)
        
        ntot += n
        if ntot >= natoms:
            break

    print("")
    print("natoms=", natoms)
    print("  ntypes=", ntypes) 

# skip two lines after the "molecules of type" line
    cfg.ReadLine()
    cfg.ReadLine()

    site = []
    idx = 0
    idxtype = 0
    itype = 0
    while 1:
        line = cfg.ReadLine()
        if not line:
            break

        ret = line.split()
# skip blank line
        if ret is None or len(ret) < 3:
            continue

        name = atomtypes[itype]
        xc, yc, zc = ret
        xc = pfloat(xc)
        yc = pfloat(yc)
        zc = pfloat(zc)
        print("{:04d} {:2} #{:04d} ({:8.4f}, {:8.4f}, {:8.4f})".format(idx+1, name, idxtype+1, xc, yc, zc))
        site.append([name, xc, yc, zc])

        idx += 1
        idxtype += 1
        if idxtype >= ntypes[itype]:
            itype += 1
            idxtype = 0

    cfg.Close()

    print("")
    print("Build Crystal object:")
    cry = tkCrystal()
    cry.SetSampleName(sample_name)
    cry.SetCrystalName(sample_name)
    cry.SetLatticeVectors(aij)
    latt = cry.LatticeParameters()
    for i in range(natoms):
        if i % 100 == 0:
            print("{} ".format(i))
        x, y, z = site[i][1], site[i][2], site[i][3]
#        x, y, z = cry.CartesianToFractional(site[i][1], site[i][2], site[i][3])
        if freduce01:
            x, y, z = Reduce01(x), Reduce01(y), Reduce01(z)
        if x > xlim or y > ylim or z > zlim:
            continue
#        print("xyz {}:{}: {},{},{} => {},{},{}".format(i+1, site[i][0], site[i][1], site[i][2], site[i][3], x, y, z))
        cry.AddAtomSite(name = site[i][0], pos = [x, y, z])

    print("")
    
    cry.ExpandCoordinates()

#    cry.PrintInf()

    return cry


def cfg2cif():
    global cfgfile
    
    print("CFG file: {}".format(cfgfile))

    print("")
    print("Read [{}]".format(cfgfile))

#    if not tkutils.IsFile(infile):
#        terminate("Error: Invalid infile [{}]".format(infile), usage = usage)
    cry = read_cfgfile(cfgfile)

    if not cry:
        terminate("Error: Could not get crystal object from infile [{}]".format(cfgfile), usage = usage)

    print("")
    print("==============================================")
    print("    Crystal inf")
    print("==============================================")
    cry.PrintInf()

    print("")
    print("Save to [{}]".format(ciffile))
    cif = tkCIFData()
    cif.CreateCIFFileFromCCrystal(cry, ciffile)

    terminate(usage = usage)


def main():
    updatevars()

    print("")
    print("=============== Convert CFG files to CIF file ============")
    cfg2cif()


if __name__ == "__main__":
    main()