import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.lines import Line2D

# フォントを MS Gothic に、フォントサイズを16に設定
plt.rcParams['font.family'] = 'MS Gothic'
plt.rcParams['font.size'] = 16

# ローレンツ変換パラメータ
theta_deg = 60                  # θ（度）
theta = np.deg2rad(theta_deg)   # ラジアンに変換
v = np.cos(theta)               # v/c = cos(θ)、c=1単位
c = 1.0
gamma = 1.0 / np.sqrt(1 - v**2)

# 長方形の原点シフトとサイズ設定
origin_shift = np.array([1.0, 1.0, 1.0])
a, b = 0.5, 1/3  # 長方形の幅 (x方向), 高さ (y方向)

# 座標軸の長さを「長方形の最大辺の2倍」に設定
L = 2 * max(a, b)

# 静止系の座標軸ベクトル（長さ L）
x_axis    = np.array([L, 0, 0])
y_axis    = np.array([0, L, 0])
t_axis    = np.array([0, 0, L])

# 運動系の座標軸ベクトル（ローレンツ変換適用）
x_p_axis  = gamma * np.array([L,    0,   -v * L])
y_p_axis  = np.array([0,    L,    0])
t_p_axis  = gamma * np.array([-v*L, 0,    L])

# 長方形の頂点（相対座標）
rect_rel = [
    [0,   0, 0],
    [a,   0, 0],
    [a,   b, 0],
    [0,   b, 0],
    [0,   0, 0],
]
# 絶対座標にシフト
x_rect, y_rect, t_rect = zip(*[
    (origin_shift[0] + x, origin_shift[1] + y, origin_shift[2] + t)
    for x, y, t in rect_rel
])

# ローレンツ変換後の長方形頂点
x_p_rect = [gamma * (x - v * t) for x, t in zip(x_rect, t_rect)]
y_p_rect = list(y_rect)  # y は不変
t_p_rect = [gamma * (t - v * x) for x, t in zip(x_rect, t_rect)]

# プロット準備
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

# ── 座標軸描画 ──
# 静止系：青系
ax.plot([0, x_axis[0]], [0, x_axis[1]], [0, x_axis[2]],
        color='blue', linewidth=2)
ax.plot([0, y_axis[0]], [0, y_axis[1]], [0, y_axis[2]],
        color='skyblue', linewidth=2)
ax.plot([0, t_axis[0]], [0, t_axis[1]], [0, t_axis[2]],
        color='blue', linestyle='--', linewidth=2)

# 運動系：赤系
ax.plot([0, x_p_axis[0]], [0, x_p_axis[1]], [0, x_p_axis[2]],
        color='red', linewidth=2)
ax.plot([0, y_p_axis[0]], [0, y_p_axis[1]], [0, y_p_axis[2]],
        color='salmon', linewidth=2)
ax.plot([0, t_p_axis[0]], [0, t_p_axis[1]], [0, t_p_axis[2]],
        color='red', linestyle='--', linewidth=2)

# ── 長方形描画 ──
ax.plot(x_rect, y_rect, t_rect,
        color='blue', linewidth=2, label='静止系の矩形')
ax.plot(x_p_rect, y_p_rect, t_p_rect,
        color='red', linewidth=2, label='運動系の矩形')

# 軸ラベルと描画範囲
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('t')
ax.set_xlim([0, origin_shift[0] + a + 0.2])
ax.set_ylim([0, origin_shift[1] + b + 0.2])
ax.set_zlim([min(0, min(t_p_rect) - 0.2), max(origin_shift[2] + b + 0.2, t_p_axis[2] + 0.2)])

# 凡例
legend_elements = [
    Line2D([0], [0], color='blue', lw=2, label='静止系 x 軸'),
    Line2D([0], [0], color='skyblue', lw=2, label='静止系 y 軸'),
    Line2D([0], [0], color='blue', lw=2, linestyle='--', label='静止系 t 軸'),
    Line2D([0], [0], color='red', lw=2, label="運動系 x' 軸"),
    Line2D([0], [0], color='salmon', lw=2, label="運動系 y' 軸"),
    Line2D([0], [0], color='red', lw=2, linestyle='--', label="運動系 t' 軸"),
    Line2D([0], [0], color='blue', lw=2, label="静止系の矩形"),
    Line2D([0], [0], color='red', lw=2, label="運動系の矩形"),
]
ax.legend(handles=legend_elements, fontsize=14, loc='upper left')

# タイトル
plt.title(f"長方形を (1,1,1) にシフト＋座標軸表示 (θ={theta_deg}°, v/c={v:.2f})", pad=20)

plt.show()