"""
格子面 (hkl) を3Dで可視化するスクリプト

このスクリプトは、指定された格子定数とミラー指数 (hkl) に基づいて、
結晶の単位格子とその中の (hkl) 面を3Dで描画します。
コマンドライン引数としてミラー指数h, k, lを受け入れ、それらの値に応じて面を可視化します。
描画されるのは単位格子に最も近い (hkl) 面の一部です。

:doc:`lattice_plane_usage`
"""
import sys
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

# ─── 格子定数とミラー指数 ───────────────────────
a = 1.0  # a 軸長さ
b = 1.2  # b 軸長さ
c = 0.8  # c 軸長さ
alpha, beta, gamma = 90, 100, 110  # α, β, γ （度）
h, k, l = 1, 2, 3  # ミラー指数
# ───────────────────────────────────────────

plane_color = 'cyan'
edge_color = 'blue'
plane_alpha = 0.1

# フォントを MS Gothic に設定
plt.rcParams['font.family'] = 'MS Gothic'


def init_frac_cell():
    """
    分数座標系における単位格子の頂点と辺を初期化します。

    この関数は、立方晶を想定した分数座標の8つの頂点と、
    それらの頂点を結ぶ12本の辺のインデックスのリストを生成します。
    これらの情報は、単位格子を表現するために使用されます。

    :returns: (fverts, edges)
    :rtype: (numpy.ndarray, list[tuple[int, int]])
    :returns fverts: 分数座標系における単位格子の8つの頂点の配列。
    :returns edges: 単位格子の辺を定義する頂点インデックスのペアのリスト。
    """
    fverts = np.array([
        [0,0,0], [1,0,0], [1,1,0], [0,1,0],
        [0,0,1], [1,0,1], [1,1,1], [0,1,1]
    ])
    edges = [
        (0,1),(1,2),(2,3),(3,0),
        (4,5),(5,6),(6,7),(7,4),
        (0,4),(1,5),(2,6),(3,7)
    ]
    return fverts, edges


def compute_frac_polygon(h: int, k: int, l: int, fverts: np.ndarray, edges: list[tuple[int, int]]) -> np.ndarray | None:
    """
    ミラー指数 (hkl) によって定義される平面と単位格子の辺との交点を計算し、ポリゴンとして返します。

    この関数は、平面の方程式 hx + ky + lz = 1 を使用して、
    単位格子を構成する各辺との交点を求めます。
    交点が辺の内部 (0 <= t <= 1) にある場合のみ採用し、
    それらの交点を多角形の頂点として収集します。
    最後に、これらの頂点を平面上の角度に基づいてソートし、
    凸多角形を形成するように並べ替えます。

    :param h: ミラー指数 h。
    :type h: int
    :param k: ミラー指数 k。
    :type k: int
    :param l: ミラー指数 l。
    :type l: int
    :param fverts: 分数座標系における単位格子の頂点配列。
    :type fverts: numpy.ndarray
    :param edges: 単位格子の辺を定義する頂点インデックスのペアのリスト。
    :type edges: list[tuple[int, int]]
    :returns: 計算された交点を角度順に並べたポリゴン頂点の配列。交点が存在しない場合はNone。
    :rtype: numpy.ndarray or None
    """
    pts = []
    for i0, i1 in edges:
        v0, v1 = fverts[i0], fverts[i1]
        d = v1 - v0
        denom = h*d[0] + k*d[1] + l*d[2]
        if abs(denom) < 1e-8: # 平面と辺が平行な場合
            continue
        t = (1 - (h*v0[0] + k*v0[1] + l*v0[2])) / denom
        if 0 <= t <= 1:
            pts.append(v0 + t*d)
    if not pts:
        return None
    
    # 浮動小数点数の誤差により重複する点を考慮し、ユニークな点を抽出
    # rtolとatolを調整して重複判定の閾値を設定することも可能だが、ここではデフォルトを使用
    pts = np.unique(np.array(pts), axis=0)

    # 交点をソートしてポリゴンを形成
    # 平面内の任意の中心点と、平面内の直交する2つのベクトルu, vを計算
    cen = pts.mean(axis=0)
    n = np.array([h, k, l], dtype=float)
    n /= np.linalg.norm(n) # 法線ベクトル

    # nに直交するベクトルuを計算
    # nが[1,0,0]に平行でない限り、[1,0,0]との外積が有効
    # nが[1,0,0]に平行な場合は[0,1,0]との外積を使用
    if abs(np.dot(n, [1,0,0])) < 0.99: # ほぼ平行でない場合
        u = np.cross(n, [1,0,0])
    else: # ほぼ平行な場合
        u = np.cross(n, [0,1,0])
    u /= np.linalg.norm(u) # 単位ベクトル化
    v = np.cross(n, u) # uとnに直交するvを計算
    
    # 各点を中心からのベクトルとu, vの内積を使って角度を計算しソート
    ang = [np.arctan2(np.dot(p-cen, v), np.dot(p-cen, u)) for p in pts]
    return pts[np.argsort(ang)]


def main():
    """
    スクリプトのメインエントリポイント。単位格子と指定されたミラー指数面を描画します。

    この関数は、コマンドライン引数からミラー指数 (h, k, l) を読み込み、
    結晶の格子定数と角度に基づいて実空間の格子ベクトルを計算します。
    その後、`init_frac_cell` と `compute_frac_polygon` を呼び出して、
    分数座標系の単位格子頂点と (hkl) 面のポリゴンを取得します。
    これらの分数座標を実空間座標に変換し、Matplotlibの3Dプロットを用いて、
    単位格子、(hkl) 面、格子軸ラベル、およびミラーインターセプトの注釈を描画します。
    最終的に、生成されたプロットを画像ファイルとして保存し、画面に表示します。

    :returns: なし
    :rtype: None
    """
    global h, k, l

    # コマンドライン引数処理
    argv = sys.argv
    nargs = len(argv)
    if nargs > 1: h = int(argv[1])
    if nargs > 2: k = int(argv[2])
    if nargs > 3: l = int(argv[3])

    # ラジアン変換
    alpha_r, beta_r, gamma_r = np.deg2rad(alpha), np.deg2rad(beta), np.deg2rad(gamma)
    cos_alpha, cos_beta, cos_gamma = np.cos(alpha_r), np.cos(beta_r), np.cos(gamma_r)
    sin_gamma = np.sin(gamma_r)

    # 格子ベクトル (実座標系)
    # 結晶学の慣例に従い、a1をx軸に沿わせ、a2をxy平面に配置、a3は空間に配置
    a1 = np.array([a, 0.0, 0.0])
    a2 = np.array([b * cos_gamma, b * sin_gamma, 0.0])
    
    # a3の成分計算
    a3_x = c * cos_beta
    a3_y = c * (cos_alpha - cos_beta * cos_gamma) / sin_gamma
    # a3_zは三つのベクトルが作る体積が正になるように計算
    a3_z_squared = 1 - cos_beta**2 - ((cos_alpha - cos_beta * cos_gamma)/sin_gamma)**2
    a3_z = c * np.sqrt(max(0, a3_z_squared)) # 浮動小数点誤差で負になる可能性があるのでmax(0, ...)
    a3 = np.array([a3_x, a3_y, a3_z])
    
    # 分数座標から実座標への変換行列M
    M = np.column_stack((a1, a2, a3))

    # 分数格子
    fverts, edges = init_frac_cell()

    # 平面ポリゴン
    frac_poly = compute_frac_polygon(h, k, l, fverts, edges)
    # 分数座標ポリゴンを実座標に変換
    cart_poly = np.dot(M, frac_poly.T).T if frac_poly is not None else None
    # 分数座標単位格子頂点を実座標に変換
    cverts = np.dot(M, fverts.T).T

    # プロット設定
    fig = plt.figure(figsize=(6,6))
    ax = fig.add_subplot(projection='3d')
    ax.set_box_aspect((1,1,1)) # 3Dプロットのアスペクト比を揃える
    ax.grid(False) # グリッドを非表示に
    ax.set_axis_off() # 軸を非表示に
    # 軸ペインの背景色を白に設定
    ax.xaxis.set_pane_color((1,1,1,1))
    ax.yaxis.set_pane_color((1,1,1,1))
    ax.zaxis.set_pane_color((1,1,1,1))

    # 単位格子描画
    for i0, i1 in edges:
        xs, ys, zs = zip(cverts[i0], cverts[i1])
        ax.plot(xs, ys, zs, 'k-', lw=1)

    # (hkl) 面描画
    if cart_poly is not None and len(cart_poly) >= 3: # ポリゴンが有効かつ3点以上で構成されている場合
        ax.add_collection3d(
            Poly3DCollection(
                [cart_poly], facecolors=plane_color,
                edgecolors=edge_color, alpha=plane_alpha
            )
        )

    # 格子軸ベクトルラベル（少し外側にオフセット）
    offset_factor = 1.05
    label_pos_a = a1 * offset_factor
    label_pos_b = a2 * offset_factor
    label_pos_c = a3 * offset_factor

    def label_vec(vec: np.ndarray, label: str):
        """格子軸ラベルをプロットに配置するヘルパー関数。"""
        ax.text(*vec, label, fontsize=14,
                fontfamily='MS Gothic', ha='center', va='center')

    label_vec(label_pos_a, 'a')
    label_vec(label_pos_b, 'b')
    label_vec(label_pos_c, 'c')

    # ミラーインターセプト注記（オフセット適用）
    offset_alabel = -0.1
    offset_blabel = -0.2
    offset_clabel = -0.05

    if h != 0: # hが0でない場合、a軸との交点が存在
        pos_h = np.dot(M, np.array([1/h, 0, 0])) # 分数座標[1/h,0,0]を実座標に変換
        pos_h_offset = [pos_h[0], pos_h[1] + offset_alabel, pos_h[2]]
        ax.text(*pos_h_offset, f'$1/{{{h}}}$', fontsize=12,
                fontfamily='MS Gothic', ha='center', va='bottom')

    if k != 0: # kが0でない場合、b軸との交点が存在
        pos_k = np.dot(M, np.array([0, 1/k, 0])) # 分数座標[0,1/k,0]を実座標に変換
        pos_k_offset = [pos_k[0] + offset_blabel, pos_k[1], pos_k[2]]
        ax.text(*pos_k_offset, f'$1/{{{k}}}$', fontsize=12,
                fontfamily='MS Gothic', ha='right', va='center')

    if l != 0: # lが0でない場合、c軸との交点が存在
        pos_l = np.dot(M, np.array([0, 0, 1/l])) # 分数座標[0,0,1/l]を実座標に変換
        pos_l_offset = [pos_l[0], pos_l[1] + offset_clabel, pos_l[2]]
        ax.text(*pos_l_offset, f'$1/{{{l}}}$', fontsize=12,
                fontfamily='MS Gothic', ha='right', va='bottom')

    # 視点調整
    ax.view_init(elev=20, azim=30) # 仰角と方位角を設定
    ax.dist = 12 # カメラとオブジェクト間の距離を設定

    plt.tight_layout() # レイアウトを自動調整
    plt.savefig(f"lattice_plane{h}{k}{l}.png",
                dpi=300, bbox_inches='tight', transparent=True) # 画像として保存
    plt.show() # プロットを表示


# ───────────────────────────────────────────
# ここから実行
# ───────────────────────────────────────────
if __name__ == "__main__":
    main()