#!/bin/bash

# ====================================================================================
#
#                                   CIFファイルからのVASP全自動計算・作図スクリプト
#
# ====================================================================================
#
# ## 概要
#
# このスクリプトは、'./cif'ディレクトリ内に置かれた複数のCIFファイルを起点として、
# VASPを用いた一連の計算（構造緩和、物性計算）から、結果の作図までを全自動で
# 実行します。研究や多数の物質探索における計算フローを大幅に効率化することを
# 目的としています。
#
# ---
#
# ## 主な機能
#
# - **一括処理**: 複数のCIFファイルをループ処理し、計算から作図までを一気通貫で実行。
# - **柔軟な作図設定**: 設定ファイルでバンド構造やDOSなど、作成したい図の種類を自由に選択可能。
# - **段階的実行**:「①構造緩和 → ②メイン計算 → ③作図」の順で実行。各ステップの成功を確認してから次に進む安全設計。
# - **高度な自動化**: CIFファイルから物質名と元素リストを自動で抽出し、ディレクトリ名や作図の引数として活用。
# - **堅牢な実行制御**: 既存ディレクトリのスキップ、計算のタイムアウト処理、Ctrl+C/Ctrl+Zによる安全な中断・スキップ機能。
# - **詳細ログ**: 誰がいつ実行しても状況を追跡できるよう、実行内容をタイムスタンプ付きでログファイルに記録。
#
# ---
#
# ## 使用方法
#
# ### 1. ディレクトリとファイルの準備
#
# 以下の構成でファイルとディレクトリを配置します。
#
# ```
# . (カレントディレクトリ)
# ├── run_automation.sh        (このスクリプト)
# ├── DoVASP.sh                (メイン計算用スクリプト)
# ├── DoVASP_relax.sh          (構造緩和用スクリプト)
# ├── hosts                    (MPI用ホストファイル)
# └── cif/
#     ├── Material_A.cif       <- CIF内に _chemical_formula_structural '...' の行が必要
#     ├── Material_B.cif
#     └── ...
#
# # Pythonスクリプトはホームディレクトリ等の指定場所に配置
# ~/python_scripts/
#  ├── plot_band_dos_all.py    (作図用Pythonスクリプト)
#  └── generate_summarize.py   (全結果の集計用Pythonスクリプト)
# ```
#
# ### 2. スクリプト設定の編集
#
# 下記の「▼▼▼ ユーザー設定 ▼▼▼」セクションを、ご自身の環境に合わせて編集してください。
#
# ### 3. 実行
#
# ターミナルで以下のコマンドを実行します。
#
# ```sh
# bash run_automation.sh
# ```
#
#2025/07/02 jtanaka
# ====================================================================================


# ▼▼▼ ユーザー設定 (ご自身の環境に合わせて編集してください) ▼▼▼
# ====================================================================================

# --- 1. パス設定 ---
# CIFファイルが格納されているディレクトリ
readonly CIF_DIR="./cif"
# Pythonスクリプト群が格納されているディレクトリ
readonly PYTHON_SCRIPT_DIR="$HOME/python_file"

# --- 2. 実行スクリプト設定 ---
# 構造緩和を実行するスクリプト名
readonly RELAX_SCRIPT="DoVASP_relax.sh"
# 物性計算（メイン計算）を実行するスクリプト名
readonly MAIN_CALC_SCRIPT="DoVASP.sh"
# 作図を実行するPythonスクリプト名
readonly PYTHON_PLOT_SCRIPT="${PYTHON_SCRIPT_DIR}/plot_band_dos_all.py"
# 全ての結果を集計するPythonスクリプト名
readonly PYTHON_SUMMARIZE_SCRIPT="${PYTHON_SCRIPT_DIR}/generate_summarize.py"

# --- 3. 計算パラメータ設定 ---
# 計算ディレクトリにコピーする共通ファイル（スクリプトやホストファイルなど）
readonly FILES_TO_COPY=("DoVASP.sh" "DoVASP_relax.sh" "hosts")
# 各計算ステップのタイムアウト時間 (例: 30m, 2h, 1d)
readonly CALCULATION_TIMEOUT="5m"
# PDOS（部分状態密度）の作図で対象とする軌道
readonly ORBITALS_FOR_PDOS="s p d"

# --- 4. 作図設定 ---
# trueに設定したグラフのみ作成されます。不要なものはfalseにしてください。
readonly PLOT_BAND_DOS=true   # バンド構造 + DOS のグラフ
readonly PLOT_BAND_ONLY=true  # バンド構造のみのグラフ
readonly PLOT_DOS_ONLY=true   # DOSのみのグラフ（合計DOSと元素別PDOS）

# --- 5. ログファイル設定 ---
# 生成されるログファイル名の接頭辞
readonly LOG_FILE_BASENAME="automation_main"

# ====================================================================================
# ▲▲▲ ユーザー設定ここまで ▲▲▲
# ====================================================================================


# --- 内部定数・グローバル変数 (通常は変更不要) ---
readonly EXIT_CODE_TIMEOUT=124
readonly REQUIRED_COMMANDS=("find" "timeout" "python" "cp" "mkdir" "basename" "tee" "pwd" "grep" "sed" "awk" "tr" "xargs")
MAIN_LOG_FILE=""
STOP_REQUESTED=0
SKIP_REQUESTED=0
CALC_PID=""
# グローバル変数: get_material_info関数で設定される
MATERIAL_NAME=""    # ファイル名用 (例: GeO2)
PLOT_TITLE=""       # グラフタイトル用 (例: GeO₂)
ELEMENTS_STR=""


# ====================================================================================
# ヘルパー関数群
# ====================================================================================

# ------------------------------------------------------------------------------
# ログメッセージをフォーマットして出力する内部関数
# Globals:
#   MAIN_LOG_FILE
# Arguments:
#   $1: ログレベル (例: INFO, ERROR)
#   $2: 色コード (例: 34 は青)
#   $3: メッセージ本文
# ------------------------------------------------------------------------------
log_message() {
    local type="$1" color="$2" msg="$3"
    local formatted_msg="[${type}] ${msg}"
    # ログファイルへの追記
    echo "$(date '+%Y-%m-%d %H:%M:%S') ${formatted_msg}" >> "$MAIN_LOG_FILE"
    # コンソールへのカラー出力
    echo -e "\n\e[${color}m${formatted_msg}\e[0m"
}

# --- ログ出力関数のラッパー ---
log_info()    { log_message "INFO"    "34" "$1"; } # 青色: 通常情報
log_success() { log_message "SUCCESS" "32" "$1"; } # 緑色: 成功
log_warning() { log_message "WARNING" "33" "$1"; } # 黄色: 警告
log_error() {                                     # 赤色: 致命的なエラー
    log_message "ERROR" "31" "$1"
    [[ -n "$2" ]] && log_message "ERROR" "31" "スクリプトの $2 行目付近でエラーが発生しました。"
    echo -e "\n\e[31mスクリプトを異常終了します。\e[0m" >&2
    exit 1
}

# ------------------------------------------------------------------------------
# Ctrl+C (SIGINT) を補足した際の処理
# ------------------------------------------------------------------------------
handle_sigint() {
    if [[ "$STOP_REQUESTED" -eq 1 && -n "$CALC_PID" ]]; then
        log_warning "\n二度目のCtrl+Cです。計算プロセス(PID: ${CALC_PID})を強制終了します..."
        kill -9 "$CALC_PID" 2>/dev/null
    fi
    STOP_REQUESTED=1
    log_warning "\n終了リクエスト(Ctrl+C)を検出しました。現在の処理が完了次第、安全に終了します。"
    log_warning "即時強制終了したい場合は、もう一度Ctrl+Cを押してください。"
}

# ------------------------------------------------------------------------------
# Ctrl+Z (SIGTSTP) を補足した際の処理
# ------------------------------------------------------------------------------
handle_sigtstp() {
    if [[ -n "$CALC_PID" ]]; then
        log_warning "\nスキップリクエスト(Ctrl+Z)を検出しました。現在の計算(PID: ${CALC_PID})を強制終了し、次の物質へ進みます..."
        SKIP_REQUESTED=1
        kill -9 "$CALC_PID" 2>/dev/null
    else
        log_warning "\nスキップリクエスト(Ctrl+Z)を検出しましたが、実行中の計算はありません。"
    fi
}


# ====================================================================================
# メイン処理関数群
# ====================================================================================

# ------------------------------------------------------------------------------
# 実行に必要なコマンドやファイルの存在をチェックする
# ------------------------------------------------------------------------------
check_dependencies() {
    log_info "依存関係をチェックしています..."
    local has_error=0
    for cmd in "${REQUIRED_COMMANDS[@]}"; do
        ! command -v "$cmd" &> /dev/null && { log_message "ERROR" "31" "必須コマンドが見つかりません: $cmd"; has_error=1; }
    done
    for file in "${FILES_TO_COPY[@]}"; do
        [[ ! -f "$file" ]] && { log_message "ERROR" "31" "必須ファイルが見つかりません: $file"; has_error=1; continue; }
        [[ ! -r "$file" ]] && { log_message "ERROR" "31" "必須ファイルが読み取れません: $file"; has_error=1; }
    done
    for py_script in "$PYTHON_PLOT_SCRIPT" "$PYTHON_SUMMARIZE_SCRIPT"; do
        [[ ! -f "$py_script" ]] && { log_message "ERROR" "31" "Pythonスクリプトが見つかりません: $py_script"; has_error=1; }
    done
    [[ ! -d "$CIF_DIR" ]] && { log_message "ERROR" "31" "CIFディレクトリが見つかりません: $CIF_DIR"; has_error=1; }
    [[ ! -w "./" ]] && { log_message "ERROR" "31" "カレントディレクトリへの書き込み権限がありません。"; has_error=1; }

    [[ "$has_error" -eq 1 ]] && log_error "依存関係エラーのため、処理を中断します。"
    log_info "依存関係は正常です。"
}

# ------------------------------------------------------------------------------
# CIFファイルから物質名と元素リストを抽出し、グローバル変数に格納する
# Arguments:
#   $1: ベース名 (ログ用)
#   $2: CIFファイルのパス
# Outputs:
#   - グローバル変数 MATERIAL_NAME, PLOT_TITLE, ELEMENTS_STR を設定
# Returns:
#   0: 成功 / 1: 失敗
# ------------------------------------------------------------------------------
get_material_info() {
    local base_name="$1"
    local source_cif_file="$2"
    local raw_formula=""

    # 変数の初期化
    MATERIAL_NAME=""
    ELEMENTS_STR=""
    PLOT_TITLE=""

    log_info "[${base_name}] 元ファイル '${source_cif_file}' から物質情報を取得します。"

    # 化学式抽出の信頼性を向上: スペースを含む式 ('Ge O2') や改行コードにも対応
    raw_formula=$(sed -n "s/.*_chemical_formula_structural[[:space:]]\+//p" "$source_cif_file" | tr -d "'\r")

    if [[ -z "$raw_formula" ]]; then
        log_warning "[${base_name}] '${source_cif_file}' 内に '_chemical_formula_structural' の定義が見つかりません。"
        return 1
    fi

    # 物質名（保存名）: 'Ge O2' -> 'GeO2'
    MATERIAL_NAME=$(echo "$raw_formula" | tr -d ' ')
    # 元素リスト: 'Ge O2' -> 'Ge O'
    ELEMENTS_STR=$(echo "$raw_formula" | sed 's/[0-9]//g' | xargs)

    # グラフタイトル生成
    if [[ "$MATERIAL_NAME" =~ [0-9] ]]; then
        # 数字を含む場合 (例: GeO2 -> GeO₂)、Unicodeの下付き文字に置換
        PLOT_TITLE=$(echo "$MATERIAL_NAME" | sed \
            -e 's/0/₀/g' -e 's/1/₁/g' -e 's/2/₂/g' \
            -e 's/3/₃/g' -e 's/4/₄/g' -e 's/5/₅/g' \
            -e 's/6/₆/g' -e 's/7/₇/g' -e 's/8/₈/g' -e 's/9/₉/g')
    else
        # 数字を含まない場合 (例: Si -> Si)
        PLOT_TITLE="$MATERIAL_NAME"
    fi

    if [[ -z "$MATERIAL_NAME" || -z "$ELEMENTS_STR" ]]; then
        log_warning "[${base_name}] 物質名または元素リストの生成に失敗しました。"
        return 1
    fi

    log_info "[${base_name}] 保存名: ${MATERIAL_NAME}, 元素: ${ELEMENTS_STR}, タイトル: ${PLOT_TITLE}"
    return 0
}

# ------------------------------------------------------------------------------
# 個別のVASP計算（構造緩和 or メイン計算）を実行する
# Arguments:
#   $1: 実行する計算スクリプト (DoVASP.sh など)
#   $2: ベース名 (ログ用)
#   $3: 計算の種類名 (ログ用、例: "構造緩和計算")
#   $4: 計算ログのファイル名
# Returns:
#   0: 成功 / 1: 失敗、タイムアウト、またはユーザーによる中断
# ------------------------------------------------------------------------------
run_calculation() {
    local calc_script="$1" base_name="$2" calc_name="$3" log_filename="$4"
    local exit_code=0

    log_info "[${base_name}] ${calc_name}の初期化を実行します..."
    if ! (echo | bash "${calc_script}" init &> /dev/null); then
        log_warning "[${base_name}] ${calc_name}の初期化に失敗しました。この物質をスキップします。"
        return 1
    fi

    log_info "[${base_name}] ${calc_name}を実行します (タイムアウト: ${CALCULATION_TIMEOUT})..."
    (set -o pipefail; echo | timeout "${CALCULATION_TIMEOUT}" bash "${calc_script}" 2>&1 | tee "${log_filename}") &
    CALC_PID=$!
    wait "$CALC_PID" || exit_code=$?
    CALC_PID=""

    if [[ "$SKIP_REQUESTED" -eq 1 ]]; then
        log_warning "[${base_name}] ユーザーリクエストにより${calc_name}をスキップしました。"
        SKIP_REQUESTED=0
        return 1
    elif [[ "$STOP_REQUESTED" -eq 1 ]]; then
        log_warning "[${base_name}] ユーザーリクエストにより処理を中断します。"
        return 1
    elif [[ "$exit_code" -eq "$EXIT_CODE_TIMEOUT" ]]; then
        log_warning "[${base_name}] ${calc_name}でタイムアウトが発生しました。"
        return 1
    elif [[ "$exit_code" -ne 0 ]]; then
        log_warning "[${base_name}] ${calc_name}がエラーで失敗しました(終了コード: $exit_code)。"
        return 1
    fi

    log_success "[${base_name}] ${calc_name}が正常に完了しました。"
    return 0
}

# ------------------------------------------------------------------------------
# Pythonスクリプトを呼び出して単一のプロットを実行する内部ヘルパー関数
# Arguments:
#   $1: ベース名 (ログ用)
#   $2: プロットの説明 (ログ用)
#   $@: 以降の引数は全てpythonスクリプトに渡される
# ------------------------------------------------------------------------------
_execute_plot() {
    local base_name="$1"
    local plot_description="$2"
    shift 2
    local python_options=("$@")
    local save_filename=""

    # 引数から --save の値を取得
    for i in "${!python_options[@]}"; do
        if [[ "${python_options[$i]}" == "--save" ]]; then
            save_filename="${python_options[$((i+1))]}"
            break
        fi
    done

    log_info "[${base_name}] ${plot_description}を作成します..."
    if python "${PYTHON_PLOT_SCRIPT}" "${python_options[@]}"; then
        log_success "[${base_name}] '${save_filename}' の作成に成功しました。"
    else
        log_warning "[${base_name}] '${save_filename}' の作成に失敗しました。"
    fi
}

# ------------------------------------------------------------------------------
# 設定に基づいて各種グラフ作成タスクを実行する
# Arguments:
#   $1: ベース名 (ログ用)
#   $2: 元のCIFファイルのパス
# ------------------------------------------------------------------------------
run_plotting_tasks() {
    local base_name="$1"
    local cif_file="$2"

    if ! $PLOT_BAND_DOS && ! $PLOT_BAND_ONLY && ! $PLOT_DOS_ONLY; then
        log_info "[${base_name}] 全ての作図設定が無効のため、作図処理をスキップします。"
        return 0
    fi

    log_info "[${base_name}] 作図処理を開始します..."

    if ! get_material_info "$base_name" "$cif_file"; then
        log_warning "[${base_name}] 物質情報が取得できなかったため、作図をスキップします。"
        return 1
    fi

    # --- バンド構造 + DOS プロット ---
    if [[ "$PLOT_BAND_DOS" == "true" ]]; then
        _execute_plot "$base_name" "バンド構造+DOSグラフ" \
            --mode band_dos --title "${PLOT_TITLE}" \
            --elements ${ELEMENTS_STR} --save "${MATERIAL_NAME}_band_dos.png"
    fi

    # --- バンド構造のみプロット ---
    if [[ "$PLOT_BAND_ONLY" == "true" ]]; then
        _execute_plot "$base_name" "バンド構造のみのグラフ" \
            --mode band --title "${PLOT_TITLE}" --save "${MATERIAL_NAME}_band.png"
    fi

    # --- DOSのみプロット (合計と元素別) ---
    if [[ "$PLOT_DOS_ONLY" == "true" ]]; then
        _execute_plot "$base_name" "合計DOSグラフ" \
            --mode dos --title "${PLOT_TITLE} Total DOS" \
            --elements ${ELEMENTS_STR} --save "${MATERIAL_NAME}_dos_total.png" --hide_gap_lines

        log_info "[${base_name}] 元素別の部分DOSグラフを作成します..."
        for element in $ELEMENTS_STR; do
             _execute_plot "$base_name" "'${element}' のPDOSグラフ" \
                 --mode dos --title "${element} PDOS in ${PLOT_TITLE}" \
                 --elements "${element}" --orbitals ${ORBITALS_FOR_PDOS} \
                 --save "${MATERIAL_NAME}_dos_${element}.png" --hide_gap_lines
        done
    fi
    return 0
}


# ------------------------------------------------------------------------------
# 単一のCIFファイルに対する全処理フロー（ディレクトリ作成、計算、作図）
# Arguments:
#   $1: CIFファイルのパス
# Returns:
#   0: 成功 / 1: 失敗
# ------------------------------------------------------------------------------
process_cif_file() {
    local cif_file="$1"
    local base_name
    base_name=$(basename "$cif_file" .cif)

    log_info "================================================="
    log_info "処理を開始: ${base_name}"

    mkdir -p "$base_name" || { log_warning "ディレクトリ作成に失敗: ${base_name}"; return 1; }
    cp "${FILES_TO_COPY[@]}" "$cif_file" "$base_name/" || { log_warning "必須ファイルのコピーに失敗: -> ${base_name}/"; return 1; }

    # ディレクトリを移動して処理を実行
    pushd "$base_name" > /dev/null

    # --- ステップ1: 構造緩和 ---
    log_info "--- [ステップ1/3] 構造緩和 ---"
    run_calculation "$RELAX_SCRIPT" "$base_name" "構造緩和計算" "DoVASP_relax.log" || { popd > /dev/null; return 1; }

    # --- ステップ2: メイン計算 ---
    log_info "--- [ステップ2/3] メイン計算 ---"
    run_calculation "$MAIN_CALC_SCRIPT" "$base_name" "メイン計算" "DoVASP_main.log" || { popd > /dev/null; return 1; }

    # --- ステップ3: 作図処理 ---
    log_info "--- [ステップ3/3] 作図処理 ---"
    # 作図の成否は全体の成否に影響させない
    run_plotting_tasks "$base_name" "../${cif_file}"

    popd > /dev/null
    log_success "全処理が正常に完了しました: ${base_name}"
    return 0
}


# ====================================================================================
# スクリプト実行部 (エントリーポイント)
# ====================================================================================
main() {
    MAIN_LOG_FILE="$(pwd)/${LOG_FILE_BASENAME}_$(date '+%Y%m%d_%H%M').log"
    touch "$MAIN_LOG_FILE"

    # シグナルハンドラを設定
    trap handle_sigint INT
    trap handle_sigtstp TSTP

    echo "スクリプト実行開始: $(date)" | tee -a "$MAIN_LOG_FILE"
    log_info "ログは ${MAIN_LOG_FILE} に記録されます。"
    log_info "スクリプト制御: [Ctrl+C] 全体終了 | [Ctrl+Z] 現在の計算をスキップして次へ"

    check_dependencies

    local cif_files=()
    while IFS= read -r -d $'\0' file; do
        cif_files+=("$file")
    done < <(find "$CIF_DIR" -maxdepth 1 -type f -name '*.cif' -print0)

    if [[ ${#cif_files[@]} -eq 0 ]]; then
        log_error "'${CIF_DIR}' ディレクトリに *.cif ファイルが見つかりません。"
    fi

    log_info "合計 ${#cif_files[@]} 個のCIFファイルを対象に処理を開始します。"
    local calculated_count=0 skipped_count=0 failed_count=0

    # CIFファイルごとにメインループを処理
    for cif_file in "${cif_files[@]}"; do
        [[ "$STOP_REQUESTED" -eq 1 ]] && { log_warning "終了リクエストにより、ループを中断します。"; break; }

        local base_name
        base_name=$(basename "$cif_file" .cif)
        log_info "-------------------------------------------------"
        log_info "ファイルをチェック中: ${base_name}.cif"

        if [[ -d "$base_name" ]]; then
            log_warning "ディレクトリ '${base_name}' が既に存在するため、スキップします。"
            ((skipped_count++))
        else
            if process_cif_file "$cif_file"; then
                ((calculated_count++))
            else
                ((failed_count++))
            fi
        fi
    done

    # --- 最終サマリー ---
    log_info "================================================="
    log_success "全てのキューが処理されました。"
    log_info "最終結果 -> 成功: ${calculated_count}件, スキップ(重複): ${skipped_count}件, 失敗/中断: ${failed_count}件"

    if [[ "$calculated_count" -gt 0 || "$failed_count" -gt 0 ]]; then
        log_info "結果の集約スクリプトを実行します..."
        if ! python "${PYTHON_SUMMARIZE_SCRIPT}"; then
            log_warning "結果の集約スクリプトの実行に失敗しました。"
        fi
    else
        log_info "新たに処理された物質がないため、集約スクリプトは実行しません。"
    fi
    log_success "全工程が完了しました。詳細は ${MAIN_LOG_FILE} をご確認ください。"
}

# スクリプトの実行を開始
main