#!/usr/bin/env python
# ./cube2mayavi_edge_corner.py -i 050Li3PS4-050LiI_PS4_I_zero.cube 050Li3PS4-050LiI_PS4_I_corner.cube 050Li3PS4-050LiI_PS4_I_edge.cube 050Li3PS4-050LiI_PS4_I_three.cube -iso 1.2e-09
import numpy as np
import os
from mayavi import mlab

ang_borr = 0.5291772109217


def read_cube(filename):
    """
    .cubeファイルを読み込み、原子座標と3次元ボクセルデータを返す関数
    """
    with open(filename, 'r') as f:
        # ヘッダーの読み飛ばし（最初の2行はコメント）
        f.readline()
        f.readline()

        # 3行目: 原子の数 と 原点座標
        line = f.readline().split()
        natoms = abs(int(line[0]))  # 原子の数
        origin = np.array([float(x) for x in line[1:4]])  # 原点

        # 4-6行目: グリッドの分割数(N)とベクトル(X, Y, Z軸)
        line = f.readline().split()
        nx, x_vec = int(line[0]), np.array([float(x) for x in line[1:4]])
        line = f.readline().split()
        ny, y_vec = int(line[0]), np.array([float(x) for x in line[1:4]])
        line = f.readline().split()
        nz, z_vec = int(line[0]), np.array([float(x) for x in line[1:4]])

        # 原子座標の読み込み
        atoms = []
        for _ in range(natoms):
            line = f.readline().split()
            atoms.append([float(x) for x in line[2:5]])
        atoms = np.array(atoms)

        # ボクセルデータの読み込み
        data = []
        for line in f:
            data.extend([float(x) for x in line.split()])

    vol_data = np.array(data).reshape(nx, ny, nz)

    return vol_data, atoms, origin, (x_vec, y_vec, z_vec)


def draw_radial_scale(max_dist=10.0, step=1.0, axis='x'):
    """
    原点から動径方向に目盛り（定規）を描画する
    """
    if axis == 'y':
        vec = np.array([0, 1, 0])
        tick_dir = np.array([0.2, 0, 0])
    elif axis == 'z':
        vec = np.array([0, 0, 1])
        tick_dir = np.array([0.2, 0, 0])
    else:  # default x
        vec = np.array([1, 0, 0])
        tick_dir = np.array([0, 0, 0.2])

    # 1. メインの直線を引く
    end_point = vec * max_dist
    mlab.plot3d([0, end_point[0]], [0, end_point[1]], [0, end_point[2]],
                tube_radius=0.02, color=(0, 0, 0))

    # 2. 目盛りと数字を配置
    for r in np.arange(step, max_dist + step, step):
        pos = vec * r
        t_start = pos - tick_dir
        t_end = pos + tick_dir

        mlab.plot3d([t_start[0], t_end[0]],
                    [t_start[1], t_end[1]],
                    [t_start[2], t_end[2]],
                    tube_radius=0.02, color=(0, 0, 0))

        text_pos = pos + tick_dir * 1.5
        mlab.text3d(text_pos[0], text_pos[1], text_pos[2],
                    f"{int(r)}", scale=0.25, color=(0, 0, 0))


def change_view(azimuth=30, elevation=60):
    mlab.view(azimuth=azimuth, elevation=elevation)


def get_color_from_filename(filename):
    """
    ファイル名から Edge, Corner などの判別を行い、RGBカラーを返す
    """
    fname_lower = os.path.basename(filename).lower()
    if 'zero' in fname_lower:
        return (0.227, 0.373, 0.804)  # 青色
    elif 'corner' in fname_lower:
        return (0.000, 0.733, 0.000)  # 緑色
    elif 'edge' in fname_lower:
        return (0.850, 0.058, 0.058)  # 赤色
    elif 'three' in fname_lower:
        return (0.943, 0.754, 0.000)  # 黄色
    else:
        return (0.5, 0.5, 0.5)        # 該当しない場合はグレー


def visualize_cubes(filenames, iso_val=None, cutoff=8):
    """
    複数のcubeファイルを読み込んで可視化する
    """
    mlab.figure(bgcolor=(1, 1, 1), size=(800, 600))
    atoms_drawn = False

    for filename in filenames:
        print(f"Reading {filename}...")
        vol_data, atoms, origin, vectors = read_cube(filename)
        atoms = atoms * ang_borr
        origin_ang = origin * ang_borr

        dx = np.linalg.norm(vectors[0]) * ang_borr
        dy = np.linalg.norm(vectors[1]) * ang_borr
        dz = np.linalg.norm(vectors[2]) * ang_borr

        # iso_val が未指定の場合はデータから自動計算
        current_iso = iso_val
        if current_iso is None:
            current_iso = np.std(vol_data) * 2.0

        print(f"  -> Plotting isosurface at value: +/- {current_iso:.4e}")

        # --- カットオフによるデータの球状マスク処理 ---
        if cutoff is not None:
            nx, ny, nz = vol_data.shape
            x = origin_ang[0] + np.arange(nx) * dx
            y = origin_ang[1] + np.arange(ny) * dy
            z = origin_ang[2] + np.arange(nz) * dz
            X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
            R = np.sqrt(X**2 + Y**2 + Z**2)
            vol_data[R > cutoff] = 0.0

        print(f"  -> max_value: {np.max(vol_data):.4e}, min_value: {np.min(vol_data):.4e}")

        # --- 確率密度 ---
        src_in = mlab.pipeline.scalar_field(vol_data)
        src_in.spacing = [dx, dy, dz]
        src_in.origin = origin_ang

        # ファイル名から色を取得
        color = get_color_from_filename(filename)

        # 等値面の描画
        mlab.pipeline.iso_surface(src_in, contours=[current_iso], opacity=0.3, color=color)

        # --- 原子の描画（最初の1回のみ実行） ---
        if not atoms_drawn:
            mlab.points3d(atoms[0, 0], atoms[0, 1], atoms[0, 2],
                          scale_factor=0.8, color=(0.576, 0.439, 0.8), resolution=20)
            mlab.points3d(atoms[1:, 0], atoms[1:, 1], atoms[1:, 2],
                          scale_factor=0.8, color=(1, 1, 0), resolution=20)

            # --- 結合の描画 ---
            p_coord = atoms[0]
            s_coords = atoms[1:]
            for s_coord in s_coords:
                mlab.plot3d([p_coord[0], s_coord[0]],
                            [p_coord[1], s_coord[1]],
                            [p_coord[2], s_coord[2]],
                            tube_radius=0.1, color=(0.6, 0.6, 0.6), opacity=1.0)

            # 半径の基準となるうすい球の描画
            mlab.points3d(0, 0, 0, scale_factor=cutoff * 2,
                          mode='sphere', color=(0, 0, 0), opacity=0.05, resolution=50)

            atoms_drawn = True


# --- 実行部分 ---
if __name__ == "__main__":
    import argparse

    description = """Visualizes multiple cube files with different colors based on their filenames."""
    par = argparse.ArgumentParser(description=description)

    # 複数ファイルを受け取れるように nargs="+" を設定（既存のまま）
    par.add_argument('-i', '--infiles', default=[], required=True, nargs="+",
                     help='input cube files (e.g., *_I_*.cube)')
    par.add_argument('-iso', '--iso_val', required=False, type=float,
                     help='Threshold for isosurface (しきい値)')
    args = par.parse_args()

    # 複数ファイルをリストとして渡す
    visualize_cubes(args.infiles, iso_val=args.iso_val)
    mlab.show()