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():
    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, k, l, fverts, edges):
    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
    pts = np.unique(np.array(pts), axis=0)
    cen = pts.mean(axis=0)
    n = np.array([h, k, l], dtype=float)
    n /= np.linalg.norm(n)
    if abs(n[0]) < abs(n[1]):
        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)
    ang = [np.arctan2(np.dot(p-cen, v), np.dot(p-cen, u)) for p in pts]
    return pts[np.argsort(ang)]


# メイン処理
def main():
    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 = np.array([a, 0.0, 0.0])
    a2 = np.array([b * cos_gamma, b * sin_gamma, 0.0])
    a3_x = c * cos_beta
    a3_y = c * (cos_alpha - cos_beta * cos_gamma) / sin_gamma
    a3_z = c * np.sqrt(max(0, 1 - cos_beta**2 - ((cos_alpha - cos_beta * cos_gamma)/sin_gamma)**2))
    a3 = np.array([a3_x, a3_y, a3_z])
    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))
    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:
        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, label):
        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:
        pos_h = np.dot(M, np.array([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:
        pos_k = np.dot(M, np.array([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:
        pos_l = np.dot(M, np.array([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()