import os
import sys
import shutil
import glob
import csv
from itertools import product
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


from tklib.tkfile import tkFile
from tklib.tkutils import IsDir, IsFile
from tklib.tkutils import terminate, pint, pfloat, getarg, getintarg, getfloatarg
from tklib.tkcrystal.tkcif import tkCIF
from tklib.tkcrystal.tkcrystal import tkCrystal
from tklib.tkcrystal.tkatomtype import tkAtomType


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

# mode: 'asf', 'anom', 'F', 'inf'
#    asf : atomic form factor for target_atom
#    anom: anomalous scattering factor for target_atom
#    F   : crystal structure factor for the CIF specified by ciffile
#    inf : show crystal structure information for the CIF specified by ciffile
mode = 'asf'
#mode = 'anom'
#mode = 'F'
#mode = 'inf'

# X-ray
target_atom = 'Ni'
xray_source = 'CuKa1'

wavelength = {
      'MoKa1': 0.070926 # nm
    , 'MoKa2': 0.071354
    , 'MoKa' : 0.071069
    , 'CuKa1': 0.15405
    , 'CuKa2': 0.15443
    , 'CuKa' : 0.15418
    , 'CuKb1': 0.13810
    , 'CuKb2': 0.13922
    , 'CuKb' : 0.13847
    }

# CIF file configuration
ciffile = 'test.cif'

single = 1
findvalidstructure = 1

# scattering vector (s = sin Q / lambda) plot range
smin =  0.0
smax = 20.0
ns   = 31

# s value to plot asf(s)
s_value = 0.1

# Wavelength plot range for  anomolous scattering factor
WLmin = None
WLmax = None

# wavelength to calculate atomic form factor and anomolous scattering factor
wl_value = None

# Miller index to calculate Fhkl
hkl = [5, 1, 3]
#hkl = [-9, -1, 2]

# graph configuration
figuresize = (8, 4)
scale  = 'linear'
#scale  = 'log'


#=============================
# Treat argments
#=============================
def usage():
    global mode
    global ciffile
    global target_atom, xray_source
    global smin, smax, ns
    global hkl
    global scale

    print("")
    print("Usage:")
    print("  (Ia)  python {} asf (target_atom xray_source smin smax ns s_value plot_scale".format(sys.argv[0]))
    print("      plot atomic form factor function for target_atom")
    print("      ex: python {} {} {} {} {} {} {} {} {}"
                    .format(sys.argv[0], 'asf', target_atom, xray_source, smin, smax, ns, s_value, scale))
    print("  (Ib)  python {} anom (target_atom wl_min wl_max wl_value wl scale)".format(sys.argv[0]))
    print("      plot anomalous scattering factor spectrum for target_atom")
    print("      ex: python {} {} {} {} {} {} {}"
                    .format(sys.argv[0], 'anom', target_atom, WLmin, WLmax, wl_value, scale))
    print("  (IIa) python {} F CIF_file (xray_source h k l)".format(sys.argv[0]))
    print("      calculate crystal structure factor for the CIF specified by ciffile")
    print("          plot_scale: 'linear', 'log'")
    print("     ex: python {} {} {} {} {} {}"
                .format(sys.argv[0], ciffile, xray_source, hkl[0], hkl[1], hkl[2]))
    print("  (IIb) python {} inf CIF_file".format(sys.argv[0]))
    print("      show crystal structure information for the CIF specified by ciffile")
    print("     ex: python {} {}".format(sys.argv[0], ciffile))

def updatevars():
    global mode
    global ciffile
    global target_atom, xray_source
    global smin, smax, ns, s_value
    global WLmin, WLmax, wl_value
    global scale

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

    mode        = getarg(1, mode)   
    if mode == 'asf':
        target_atom = getarg     (2, target_atom)
        xray_source = getarg     (3, xray_source)
        smin        = getfloatarg(4, smin)
        smax        = getfloatarg(5, smax)
        ns          = getintarg  (6, ns)
        s_value     = getfloatarg(7, s_value)
        scale       = getarg     (8, scale)
    elif mode == 'anom':
        target_atom = getarg     (2, target_atom)
        WLmin       = getfloatarg(3, WLmin)
        WLmax       = getfloatarg(4, WLmax)
        wl_value    = getfloatarg(5, wl_value)
        scale       = getarg     (6, scale)
    elif mode == 'F':
        ciffile     = getarg     (2, ciffile)
        xray_source = getarg     (3, xray_source)
        hkl[0]      = getintarg  (4, hkl[0])
        hkl[1]      = getintarg  (5, hkl[1])
        hkl[2]      = getintarg  (6, hkl[2])
    elif mode == 'inf':
        ciffile = getarg(2, ciffile)
    else:
        terminate("Error: Invalide mode [{}]".format(mode), usage = usage)

def cal_Fhkl(cry, xray_source):
    hmin, hmax = -4, 4
    kmin, kmax = -4, 4
    lmin, lmax = -4, 4
    
    print("X-ray source:", xray_source)
    wl = pfloat(xray_source, defval = None)
    if wl is None or wl == 0.0:
        wl = wavelength[xray_source]
    else:
        xray_source = wl
    print("  wavelength: {} nm".format(wl))

    AtomTypes = cry.AtomTypeList()
    AtomSites = cry.ExpandedAtomSiteList()

    inf = []
    print("")
    print("Crystal structure factors:")
    for h, k, l in product(range(hmin, hmax+1), range(kmin, kmax+1), range(lmin, lmax+1)):
        dhkl, sB, Q2 = cry.DiffractionAngle(wl * 10.0, h, k, l)    # angstrom in crystal object
        if dhkl is None or Q2 is None: continue

#        print("Diffraction for ({} {} {})".format(*hkl))
#        print("  d({} {} {}) = {:10.6g} A".format(*hkl, dhkl))
#        print("  s_B  = {:10.6g} nm^-1".format(sB))
#        print("  2Q_B = {:10.6g} rad = {:10.6g} deg".format(Q2, Q2 * 180.0 / pi))
#        print("  Atomic form factors at s_B:")
        for i in range(len(AtomTypes)):
            atom = AtomTypes[i]
            atomname  = atom.AtomTypeOnly()
            asf = atom.asf(sB) 
#            print("    {:2}: asf(sB) = {:10.6g}".format(atomname, asf))

        Fhkl = cry.Fhkl(sB, h, k, l)
        inf.append({
            'hkl' : (h, k, l),
            'dhkl': dhkl,
            'Q2'  : Q2 * 180.0 / pi,
            'Fhkl': Fhkl,
            })

    sorted_inf = sorted(inf, key=lambda x: x['Q2'])
    return sorted_inf

def cal_Fhkl_single(cry, hkl, xray_source):
    print("X-ray source:", xray_source)
    wl = pfloat(xray_source, defval = None)
    if wl is None or wl == 0.0:
        wl = wavelength[xray_source]
    else:
        xray_source = wl
    print("  wavelength: {} nm".format(wl))
    print("h k l = {} {} {}".format(hkl[0], hkl[1], hkl[2]))

    AtomTypes = cry.AtomTypeList()
    AtomSites = cry.ExpandedAtomSiteList()

    print("")
    dhkl, sB, Q2 = cry.DiffractionAngle(wl * 10.0, *hkl)    # angstrom in crystal object
    print("Diffraction for ({} {} {})".format(*hkl))
    print("  d({} {} {}) = {:10.6g} A".format(*hkl, dhkl))
    print("  s_B  = {:10.6g} nm^-1".format(sB))
    print("  2Q_B = {:10.6g} rad = {:10.6g} deg".format(Q2, Q2 * 180.0 / pi))
    print("  Atomic form factors at s_B:")
    for i in range(len(AtomTypes)):
        atom = AtomTypes[i]
        atomname  = atom.AtomTypeOnly()
        asf = atom.asf(sB) 
        print("    {:2}: asf(sB) = {:10.6g}".format(atomname, asf))

    Fhkl = cry.Fhkl(sB, *hkl)
    print("  F({} {} {}) = {:10.6g}".format(*hkl, Fhkl))

    print("")
    print("Plot atomic form factor:")
    sstep = (smax - smin) / (ns - 1)
    print("  plot range: s = sin(Q) / wl = {} to {} at {} nm^-1 step, {} points".format(smin, smax, sstep, ns))

    xs = [smin + i * sstep for i in range(ns)]
    yfs = []
    yfsB = []
    for i in range(len(AtomTypes)):
        at = AtomTypes[i]
        yfsB.append(at.asf(sB))
        print("")
        print("{:8}\t{:12}".format("s=sin(Q)/wl (nm^-1)", "f(s)"))
        _fs = []
        for i in range(ns):
            s = xs[i]
            fs = at.asf(s)
            _fs.append(fs)
            print("{:8.4f}\t{:12.6g}".format(s, fs))
        yfs.append(_fs)

    print("")
    print("Plot")

#=============================
# Plot graphs
#=============================
    fig = plt.figure(figsize = figuresize)

    ax1 = fig.add_subplot(1, 2, 1)

    for i in range(len(AtomTypes)):
        at = AtomTypes[i]
        name = at.AtomType()
        ax1.plot(xs, np.array(yfs[i]).real,  label = name, linewidth = 0.5)
        ax1.plot(sB, yfsB[i].real, linestyle = 'none', marker = 'o')
    ax1.set_xlim([smin, smax])
    ylim = ax1.get_ylim()
    ax1.set_ylim([0.0, ylim[1]])
    ax1.set_xlabel("s = sin $ \Theta  / \lambda$ (nm$^{-1}$)")
    ax1.set_ylabel("real part of $f(s)$")
    ax1.set_title('atomic form factor')
    ax1.legend()
    if scale == 'log':
        ax1.set_yscale('log')
    plt.tight_layout()

    plt.pause(0.1)

    print("")
    print("Press ENTER to exit>>", end = '')
    input()
    
    terminate(usage = usage)

def exec_Fhkl():
    global xray_source
    
    print("X-ray source:", xray_source)
    wl = pfloat(xray_source, defval = None)
    if wl is None or wl == 0.0:
        wl = wavelength[xray_source]
    else:
        xray_source = wl
    print("  wavelength: {} nm".format(wl))

    print("CIF file: {}".format(ciffile))
    print("single: {}".format(single))
    print("findvalidstructure: {}".format(findvalidstructure))

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

#    if not tkutils.IsFile(ciffile):
#        terminate("Error: Invalid ciffile [{}]".format(ciffile), usage = usage)
    cif = tkCIF()
    cif.debug = debug
    
    cifdata = cif.ReadCIF(ciffile, find_valid_structure = findvalidstructure)
    cif.Close()

    if not cifdata:
        terminate("Error: Could not get cifdat from ciffile [{}]".format(ciffile), usage = usage)

    cry = cifdata.GetCrystal()
#    cry.PrintInf()

    print("")
    print("Crystal structure:")
    a, b, c, alpha, beta, gamm = cry.LatticeParameters()
    print("cell: {} {} {} A   {} {} {}".format(a, b, c, alpha, beta, gamm))

    print("")
    print("Atom types:")
    AtomTypes = cry.AtomTypeList()
    for i in range(len(AtomTypes)):
        t = AtomTypes[i]
        typea = t.AtomType()
        typeo = t.AtomTypeOnly()
        charge = t.Charge()
        AN     = t.AtomicNumber()
        M      = t.AtomicMass()
        print("  %3d: %4s (%2s) charge=%8.4f  [Z=%3d  M=%6.4f]" % (i, typea, typeo, charge, AN, M))
        asfparams = t.ReadASFParameters(XraySource = wl)
        print("     asf parameters: ", asfparams[0:10])
        print("                     ", asfparams[10:])

    print("")
    print("Expanded atom sites:")
    AtomSites = cry.ExpandedAtomSiteList()
    for atom in AtomSites:
        label     = atom.Label()
        atomname  = atom.AtomName()
        atomname0 = atom.AtomNameOnly()
        charge    = atom.Charge()
        pos       = atom.Position()
        occ       = atom.Occupancy()
        id        = atom.IdAsymmetricAtomSite()
        iAtomType = atom.iAtomType()
        atomtype  = atom.AtomType()
        m         = atom.Multiplicity()
        print("  %3d: (iAtomType=%d) %5s: %4s charge=%6.3f (%6.3f, %6.3f, %6.3f) occ=%8.4f  m=%d" 
                % (id, iAtomType, label, atomname, charge, pos[0], pos[1], pos[2], occ, m))

    cal_Fhkl_single(cry, hkl, xray_source = 'CuKa1')


def anom():
    global target_atom
    global WLmin, WLmax
    
    at = tkAtomType(target_atom)
    typea  = at.AtomType()
    typeo  = at.AtomTypeOnly()
    charge = at.Charge()
    Z      = at.AtomicNumber()
    M      = at.AtomicMass()
    print("Target atom:", target_atom)
    print("  %4s (%2s) charge=%8.4f  [Z=%3d  M=%6.4f]" % (typea, typeo, charge, Z, M))

    print("")
    print("Anomalous scattering factor")
    print("citation: http://skuld.bmsc.washington.edu/scatter/")

    xE, xwl, yf1, yf2, dbpath = at.ReadAnomalousScatteringFactor()
    print("DB path: ", dbpath)
    if WLmin is None:
        WLmin = 0.0
    if WLmax is None:
        WLmax = max(xwl)
    print("  plot range: wavelength = {} to {} nm".format(WLmin, WLmax))

    print("{:8}\t{:12}\t{:12}\t{:12}".format("E (eV)", "wl (nm)", "f1", "f2"))
    idx0 = None
    idx1 = None
    for i in range(len(xE)):
        if xwl[i] > WLmin:
            idx0 = i
        if xwl[i] < WLmax and idx1 is None:
            idx1 = i
        print("{:8.4f}\t{:12.6g}\t{:12.6g}\t{:12.6g}".format(xE[i], xwl[i], yf1[i], yf2[i]))

    print("")
    print("  plot index range: {} - {}".format(idx1, idx0))
    xE   = xE[idx1:idx0]
    xwl  = xwl[idx1:idx0]
    yf1  = yf1[idx1:idx0]
    yf2  = yf2[idx1:idx0]
    
    print("")
    f1, f2 = at.AnomalousScatteringFactor(wl_value)
    print("anormalous scattering factor at wl = {:8.4g} nm = {:10.6g} + j{:10.6g}".format(wl_value, f1, f2))
    print("")
    print("Plot")

#=============================
# Plot graphs
#=============================
    fig = plt.figure(figsize = figuresize)

    ax1 = fig.add_subplot(1, 2, 1)
    ax2 = fig.add_subplot(1, 2, 2)

    ax1.plot(xwl, yf1, label = target_atom, color = 'black', linewidth = 0.5)
    ax2.plot(xwl, yf2, label = target_atom, color = 'black', linewidth = 0.5)
    if wl_value is not None:
        ax1.plot([wl_value], [f1], linestyle = 'none', marker = 'o') #, markersize = 0.5)
        ax2.plot([wl_value], [f2], linestyle = 'none', marker = 'o') #, markersize = 0.5)

    ax1.set_xlim([WLmin, WLmax])
    ax1.set_ylim([min(yf1), max(yf1) * 1.1])
    ax1.set_xlabel("Wavelength (nm)")
    ax1.set_ylabel("f'")
    ax1.set_title('real part of anomalous scattering factor')
    ax2.set_xlim([WLmin, WLmax])
    ax2.set_ylim([0.0, max(yf2) * 1.1])
    ax2.set_xlabel("Wavelength (nm)")
    ax2.set_ylabel("f''")
    ax2.set_title('imaginary part of anomalous scattering factor')
    ax1.legend()
    ax2.legend()
    if scale == 'log':
        ax1.set_yscale('log')
        ax2.set_yscale('log')

    plt.tight_layout()

    plt.pause(0.1)

    print("")
    print("Press ENTER to exit>>", end = '')
    input()
    
    terminate(usage = usage)

def asf():
    global target_atom
    global xray_source, wavelength
    global smin, smax, ns

    print("X-ray source:", xray_source)
    wl = pfloat(xray_source)
    if wl is None or wl == 0.0:
        wl = wavelength[xray_source]
    else:
        xray_source = wl
    print("  wavelength: {} nm".format(wl))

    print("")
    print("Atomic form factor:")
    sstep = (smax - smin) / (ns - 1)
    print("  plot range: s = sin(Q) / wl = {} to {} at {} nm^-1 step, {} points".format(smin, smax, sstep, ns))

    at = tkAtomType(target_atom)
    typea  = at.AtomType()
    typeo  = at.AtomTypeOnly()
    charge = at.Charge()
    Z      = at.AtomicNumber()
    M      = at.AtomicMass()
    print("Target atom:", target_atom)
    print("  %4s (%2s) charge=%8.4f  [Z=%3d  M=%6.4f]" % (typea, typeo, charge, Z, M))
    asf_params = at.ReadASFParameters(target_atom)
    print("  ASF parameters:", asf_params)
    print("  asf(s) at s={:8.4g} = {:10.6g}".format(s_value, at.asf(s_value)))

    print("")
    print("Calculate atomic form factor function f(s)")
    xs  = []
    yfs = []
    print("{:8}\t{:12}".format("s=sin(Q)/wl (nm^-1)", "f(s)"))
    for i in range(ns):
        s = smin + i * sstep
        fs = at.asf(s)
        xs.append(s)
        yfs.append(fs)
        print("{:8.4f}\t{:12.6g}".format(s, fs))

    print("")
    print("Plot")

#=============================
# Plot graphs
#=============================
    fig = plt.figure(figsize = figuresize)

    ax1 = fig.add_subplot(1, 2, 1)

    ax1.plot(xs, yfs,  label = target_atom, color = 'black', linewidth = 0.5)
    ax1.plot([s_value], [at.asf(s_value)], linestyle = 'none', marker = 'o')
    ax1.set_xlim([smin, smax])
    ax1.set_ylim([0.0, max(yfs) * 1.1])
    ax1.set_xlabel("s = sin $\Theta / \lambda$ (nm$^{-1}$)")
    ax1.set_ylabel("$f(s)$")
    ax1.set_title('atomic form factor')
    ax1.legend()
    if scale == 'log':
        ax1.set_yscale('log')
    plt.tight_layout()

    plt.pause(0.1)

    print("")
    print("Press ENTER to exit>>", end = '')
    input()
    
    terminate(usage = usage)

def inf():
    global target_atom
    global xray_source, wavelength
    global smin, smax, ns

    print("CIF file: {}".format(ciffile))
    print("single: {}".format(single))
    print("findvalidstructure: {}".format(findvalidstructure))

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

#    if not tkutils.IsFile(ciffile):
#        terminate("Error: Invalid ciffile [{}]".format(ciffile), usage = usage)
    cif = tkCIF()
    cif.debug = debug
    
    cifdata = cif.ReadCIF(ciffile, find_valid_structure = findvalidstructure)
    cif.Close()

    if not cifdata:
        terminate("Error: Could not get cifdat from ciffile [{}]".format(ciffile), usage = usage)

    print("")
    print("==============================================")
    print("    CIF inf")
    print("==============================================")
    cifdata.Print()
    cry = cifdata.GetCrystal()
    print("")
    print("==============================================")
    print("    Crystal inf")
    print("==============================================")
    cry.PrintInf()
    print("")
            
#    d = cry.Density()
#    ad = cry.AtomDensity()

    terminate(usage = usage)


def main():
    global mode
    global xray_source, wavelength
    
    updatevars()

    print("")
    print("=============== Atom informat, form factor, crystal structure factor ============")
    print("")
    print("mode: ", mode)

    if mode == 'asf':
        asf()
    elif mode == 'anom':
        anom()
    elif mode == 'F':
        exec_Fhkl()
    elif mode == 'inf':
        inf()
    else:
        terminate("Error: Invalide mode [{}]".format(mode), usage = usage)


if __name__ == "__main__":
    main()