"""
X線回折のブラッグ角計算スクリプト。

概要:
    結晶格子パラメータとサイト情報に基づいて、X線回折ピークの2θ角、格子面間隔d、ミラー指数(hkl)を計算し、表示します。

詳細説明:
    このスクリプトは、指定された格子パラメータと原子サイト情報を用いて、結晶構造のX線回折パターンをシミュレートします。
    直接格子および逆格子の各種パラメータ（格子ベクトル、メートルテンソル、単位胞体積）を計算し、
    その後、ブラッグの法則に基づいて各ミラー指数(hkl)に対する2θ回折角と格子面間隔(d)を算出します。
    計算された回折データは、2θ角でソートされて表示されます。
    `tkcrystalbase.py` モジュールに依存しており、結晶学的な計算のための基本関数を提供します。

関連リンク:
    :doc:`crystal_XRD_usage`
"""
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]

# Site information (atom name, site label, atomic number, atomic mass, charge, radius, color, position)
sites = [
         ['Na', 'Na1', 11, 22.98997, +1.0, 0.7, 'red',  np.array([0.0, 0.0, 0.0])]
        ,['Na', 'Na2', 11, 22.98997, +1.0, 0.7, 'red',  np.array([0.0, 0.5, 0.5])]
        ,['Na', 'Na3', 11, 22.98997, +1.0, 0.7, 'red',  np.array([0.5, 0.0, 0.5])]
        ,['Na', 'Na4', 11, 22.98997, +1.0, 0.7, 'red',  np.array([0.5, 0.5, 0.0])]
        ,['Cl', 'Cl1', 17, 35.4527,  -1.0, 1.4, 'blue', np.array([0.5, 0.0, 0.0])]
        ,['Cl', 'Cl2', 17, 35.4527,  -1.0, 1.4, 'blue', np.array([0.5, 0.5, 0.5])]
        ,['Cl', 'Cl3', 17, 35.4527,  -1.0, 1.4, 'blue', np.array([0.0, 0.0, 0.5])]
        ,['Cl', 'Cl4', 17, 35.4527,  -1.0, 1.4, 'blue', np.array([0.0, 0.5, 0.0])]
        ]

# X-ray wavelength
wl = 1.5405         # angstrom

# G min to remove the origin of reciprocal space
Gmin = 1.0e-5

# 2Theta max
Q2max = 150.0       # degree in 2Theta


def main():
    """
    メイン処理を実行し、結晶の基本情報とX線回折結果を表示します。

    概要:
        結晶の格子情報、逆格子情報、そしてX線回折ピークの2θ角とd値を計算し、コンソールに出力します。

    詳細説明:
        この関数は以下のステップを実行します。
        1. 直接格子の格子ベクトル、メートルテンソル、単位胞体積を計算し表示します。
        2. 逆格子の格子パラメータ、格子ベクトル、メートルテンソル、単位胞体積を計算し表示します。
        3. 指定されたX線波長と2θ最大値に基づいて、dminとhklの最大範囲を決定します。
        4. その範囲内の全ての(h, k, l)ミラー指数に対し、逆格子空間での距離G、格子面間隔d、
           そしてブラッグの法則に従って2θ回折角を計算します。
        5. 計算された(2θ, d, h, k, l)のリストを2θ角でソートし、コンソールに出力します。

    Returns:
        None
    """
    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))

    dmin = wl / 2.0 / sin(0.5 * Q2max * torad)
    hmax = int(lattice_parameters[0] / dmin)
    kmax = int(lattice_parameters[1] / dmin)
    lmax = int(lattice_parameters[2] / dmin)

    print("")
    print("hkl range:", hmax, kmax, lmax)

# Calculate diffraction angles and store them in qlist list variable
    org = np.array([0.0, 0.0, 0.0])
    qlist = []
    for l in range(-lmax, lmax+1):
      for k in range(-kmax, kmax+1):
        for h in range(-hmax, hmax+1):
# Calculate distance in reciprocal space between (0, 0, 0) and (h, k, l)
            G = distance(org, np.array([h, k, l]), Rgij)
            if G < Gmin:
                continue

# Calculate lattice space from G
            d = 1.0 / G

            sinQ = wl / 2.0 / d
            if sinQ >= 1.0:
                continue
            
# Calculate diffraction angle 2Theta
            Q2 = 2.0 * todeg * arcsin(sinQ)
            if Q2 > Q2max:
                continue

            qlist.append([Q2, d, h, k, l])
#            print("  2Q={:12.4g}  d={:12.6g}  ({:3d} {:3d} {:3d})".format(Q2, d, h, k, l))

# Sort rlist by 2Theta (x[0] priority)
    qlist.sort(key = lambda x: (x[0], x[2], x[3], x[4]))

    print("")
    print("Diffraction angle, d, h, k, l:")
    for qinf in qlist:
        Q2, d, h, k, l = qinf
        print("  2Q={:12.4g}  d={:12.6g}  ({:3d} {:3d} {:3d})".format(Q2, d, h, k, l))

    print("")
    exit()


if __name__ == '__main__':
    main()