# pip install xrayutilities pymatgen matplotlib numpy

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

# ====== 層構造（入力は “分かりやすい単位” で書いてOK） ======
layer_stack = [
    {"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},
]
substrate = {"composition": "Si", "density_gcm3": 2.33, "roughness_A": 0.5}

# ====== pymatgenで分子量表示（任意） ======
print("各層の化学組成と分子量:")
for layer in layer_stack:
    comp = Composition(layer["composition"])
    print(f"  {layer['composition']}: 分子量 = {comp.weight:.3f} g/mol")

# ====== 単位変換 ======
def gcm3_to_kgm3(rho_gcm3: float) -> float:
    return rho_gcm3 * 1000.0  # 1 g/cm^3 = 1000 kg/m^3

def nm_to_A(t_nm: float) -> float:
    return t_nm * 10.0  # 1 nm = 10 Å

# ====== LayerStack を作成 ======
# substrate thickness: inf
sub_mat = xu.materials.Amorphous(substrate["composition"], gcm3_to_kgm3(substrate["density_gcm3"]))
sub = xu.simpack.Layer(sub_mat, float("inf"), roughness=substrate["roughness_A"])

layers = []
for d in layer_stack:
    mat = xu.materials.Amorphous(d["composition"], gcm3_to_kgm3(d["density_gcm3"]))
    lay = xu.simpack.Layer(mat, nm_to_A(d["thickness_nm"]), roughness=d["roughness_A"])
    layers.append(lay)

# “sub + lay1 + lay2 ...” で積める（ドキュメント例と同じ）:contentReference[oaicite:2]{index=2}
stack = sub
for lay in layers:
    stack = stack + lay

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

# ====== XRR（SpecularReflectivityModel） ======
# simulate() は “入射角 alphai [deg]” を渡す（qzではない）:contentReference[oaicite:3]{index=3}
model = xu.simpack.SpecularReflectivityModel(sample, energy="CuKa1")  # energyは文字列指定可:contentReference[oaicite:4]{index=4}

alphai = np.linspace(0.05, 2.5, 2000)     # 入射角 [deg]
R = model.simulate(alphai)

# 2θで描きたいなら 2θ=2*alphai（対称反射）
two_theta = 2.0 * alphai

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