import sys
import os
from numpy import sin, cos, tan, arcsin, arccos, arctan, exp, log, sqrt
import numpy as np
from numpy import linalg as la
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt

from tkcrystalbase import *

# Lattice parameters (angstrom and degree)
#lattice_parameters = [ 5.62, 5.62, 5.62, 60.0, 60.0, 60.0]
lattice_parameters = [ 5.62, 5.62, 5.62, 90.0, 90.0, 90.0]

rM = 30.0
rX = 50.0
# Site information (atom Game, site label, atomic number, atomic mass, charge, radius, color, position)
sites = [
         ['Ga', 'Ga1', 11, 22.98997, +1.0, rM, 'red',  np.array([0.0, 0.0, 0.0])]
        ,['Ga', 'Ga2', 11, 22.98997, +1.0, rM, 'red',  np.array([0.0, 0.5, 0.5])]
        ,['Ga', 'Ga3', 11, 22.98997, +1.0, rM, 'red',  np.array([0.5, 0.0, 0.5])]
        ,['Ga', 'Ga4', 11, 22.98997, +1.0, rM, 'red',  np.array([0.5, 0.5, 0.0])]
        ,['As', 'As1', 17, 35.4527,  -1.0, rX, 'blue', np.array([0.25, 0.25, 0.25])]
        ,['As', 'As2', 17, 35.4527,  -1.0, rX, 'blue', np.array([0.0,  0.75, 0.75])]
        ,['As', 'As3', 17, 35.4527,  -1.0, rX, 'blue', np.array([0.75, 0.25, 0.75])]
        ,['As', 'As4', 17, 35.4527,  -1.0, rX, 'blue', np.array([0.75, 0.75, 0.25])]
        ]

aprim = [
        [0.5 * lattice_parameters[0], 0.5 * lattice_parameters[1], 0.0 * lattice_parameters[1]],
        [0.5 * lattice_parameters[0], 0.0 * lattice_parameters[1], 0.5 * lattice_parameters[2]],
        [0.0 * lattice_parameters[0], 0.5 * lattice_parameters[1], 0.5 * lattice_parameters[2]],
        ]

sites_prim = [
        ['Ga', 'Ga1', 11, 22.98997, +1.0, rM, 'red',       np.array([0.0,  0.0,  0.0])],
        ['As', 'As1', 17, 35.4527,  -1.0, rX, 'darkgreen', np.array([0.25, 0.25, 0.25])]
        ]

# Coefficient for atomic size to plot
kr = 100.0

# Coefficient to plot reciprocal unit cell w.r.t. real space unit cell
kRUC = 0.8

# Range of unit cells to draw crystal structure
nrange = [[-0.1, 1.1], [-0.1, 1.1], [-0.1, 1.1]]

# Figure configuration
figsize = (8, 8)


def main():
    print("")
    print("Lattice parameters:", lattice_parameters)
    aij = cal_lattice_vectors(lattice_parameters)
    print("Lattice vectors:")
    print("  ax: ({:10.4g}, {:10.4g}, {:10.4g}) A".format(aij[0][0], aij[0][1], aij[0][2]))
    print("  ay: ({:10.4g}, {:10.4g}, {:10.4g}) A".format(aij[1][0], aij[1][1], aij[1][2]))
    print("  az: ({:10.4g}, {:10.4g}, {:10.4g}) A".format(aij[2][0], aij[2][1], aij[2][2]))
    inf = cal_metrics(lattice_parameters)
    gij = inf['gij']
    print("Metric tensor:")
    print("  gij: ({:10.4g}, {:10.4g}, {:10.4g}) A".format(*gij[0]))
    print("       ({:10.4g}, {:10.4g}, {:10.4g}) A".format(*gij[1]))
    print("       ({:10.4g}, {:10.4g}, {:10.4g}) A".format(*gij[2]))
    volume = cal_volume(aij)
    print("Volume: {:12.4g} A^3".format(volume))

    print("")
    print("Unit cell volume: {:12.4g} A^3".format(volume))
    Raij  = cal_reciprocal_lattice_vectors(aij)
    Rlatt = cal_reciprocal_lattice_parameters(Raij)
    Rinf  = cal_metrics(Rlatt)
    Rgij  = Rinf['gij']
    print("Reciprocal lattice parameters:", Rlatt)
    print("Reciprocal lattice vectors:")
    print("  Rax: ({:10.4g}, {:10.4g}, {:10.4g}) A^-1".format(*Raij[0]))
    print("  Ray: ({:10.4g}, {:10.4g}, {:10.4g}) A^-1".format(*Raij[1]))
    print("  Raz: ({:10.4g}, {:10.4g}, {:10.4g}) A^-1".format(*Raij[2]))
    print("Reciprocal lattice metric tensor:")
    print("  Rgij: ({:10.4g}, {:10.4g}, {:10.4g}) A^-1".format(*Rgij[0]))
    print("        ({:10.4g}, {:10.4g}, {:10.4g}) A^-1".format(*Rgij[1]))
    print("        ({:10.4g}, {:10.4g}, {:10.4g}) A^-1".format(*Rgij[2]))
    Rvolume = cal_volume(Raij)
    print("Reciprocal unit cell volume: {:12.4g} A^-3".format(Rvolume))


    fig = plt.figure(figsize = figsize)
    ax = fig.add_subplot(111, projection='3d')

# Real space unit cell
    draw_unitcell(ax, sites, aij, nrange, kr = kr)
# Primitive cell
    draw_unitcell(ax, sites_prim, aprim, nrange, color = 'blue', edgecolor = 'gray', kr = kr)

    a, b, c = lattice_parameters[:3]
    configure_axis_structure(ax, a * np.array([-0.2, 1.2]), 
                                 b * np.array([-0.2, 1.2]), 
                                 c * np.array([-0.2, 1.2]),
                                 fontsize = 12, legend_fontsize = 12)

    plt.show()

    print("")
    exit()


if __name__ == '__main__':
    main()
