# pip install xrayutilities pymatgen matplotlib numpy

import os
import numpy as np
import matplotlib.pyplot as plt
import xrayutilities as xu
from pymatgen.core import Composition, Structure


# ====== 単位変換 ======
def gcm3_to_kgm3(rho):
    return float(rho) * 1000.0


def nm_to_A(t_nm):
    return float(t_nm) * 10.0


def is_cif_path(s):
    return isinstance(s, str) and s.lower().endswith(".cif")


def load_from_cif(path):
    structure = Structure.from_file(path)
    formula = structure.composition.reduced_formula
    density = float(structure.density)
    return structure, formula, density


def main():

    # ====== 層構造 ======
    layer_stack = [
        {"composition": "NaCl.cif",  "density_gcm3": 0.0, "thickness_nm": 50.0,  "roughness_A": 0.5},
#        {"composition": "TiO2",  "density_gcm3": 4.23, "thickness_nm": 50.0,  "roughness_A": 0.5},
        {"composition": "Al2O3", "density_gcm3": 3.95, "thickness_nm": 30.0,  "roughness_A": 0.5},
        {"composition": "SiO2",  "density_gcm3": 2.20, "thickness_nm": 100.0, "roughness_A": 0.5},
        # {"composition": "sample.cif", "density_gcm3": 0.0, "thickness_nm": 10.0, "roughness_A": 0.5},
    ]

    substrate = {"composition": "Si", "density_gcm3": 2.33, "roughness_A": 0.5}

    # ====== 表示：分子量 ======
    print("各層の化学組成と分子量:")
    for layer in layer_stack:
        comp = layer["composition"]
        d    = layer["thickness_nm"]
        r    = layer["roughness_A"]
        if is_cif_path(comp):
            _, formula, _ = load_from_cif(comp)
            mw = Composition(formula).weight
            print(f"  {os.path.basename(comp)} ({formula}): {mw:.4f} g/mol: {d:.4f} nm: roughneess {r:.4f} A")
        else:
            mw = Composition(comp).weight
            print(f"  {comp}: {mw:.4f} g/mol: {d:.4f} nm: roughneess {r:.4f} A")

    # ====== 基板 ======
    sub_mat = xu.materials.Amorphous(substrate["composition"], gcm3_to_kgm3(substrate["density_gcm3"]))
    stack = xu.simpack.Layer(sub_mat, float("inf"), roughness=substrate["roughness_A"])

    # ====== 薄膜層 ======
    for layer in layer_stack:
        comp = layer["composition"]
        rho = float(layer["density_gcm3"])

        if is_cif_path(comp):
            _, formula, rho_cif = load_from_cif(comp)
            if rho <= 0.0:
                rho = rho_cif
            comp = formula

        mat = xu.materials.Amorphous(comp, gcm3_to_kgm3(rho))
        lay = xu.simpack.Layer(mat, nm_to_A(layer["thickness_nm"]), roughness=layer["roughness_A"])
        stack = stack + lay

    sample = xu.simpack.LayerStack("sample", stack)

    # ====== XRR ======
    model = xu.simpack.SpecularReflectivityModel(sample, energy="CuKa1")
    alphai = np.linspace(0.05, 2.5, 2000)
    R = model.simulate(alphai)

    two_theta = 2 * alphai

    plt.figure(figsize=(8, 5))
    plt.semilogy(two_theta, R)
    plt.xlabel("2θ (deg)")
    plt.ylabel("Reflectivity")
    plt.title("Multilayer XRR Simulation")
    plt.grid(True, which="both")
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    main()