#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
1D coupled oscillator wave-packet collision / scattering
with argparse options and optional GIF save.

Site mass:
  m_site[i]

Bond spring constants:
  k2_bond[i] : spring between site i and i+1
  k4_bond[i] : nonlinear spring between site i and i+1

Potential:
  V = sum_i [ 1/2 k2_bond[i] (u[i+1]-u[i])^2
            + 1/4 k4_bond[i] (u[i+1]-u[i])^4 ]

Force:
  Let du[i] = u[i+1] - u[i]
  Let g[i]  = k2_bond[i] du[i] + k4_bond[i] du[i]^3

  Then for an interior site j:
      F[j] = g[j] - g[j-1]
"""
"""
.. _md_packets_FuncAnimation:

1次元結合振動子波束散乱シミュレーション

概要:
    1次元結合振動子系における波束の衝突・散乱をシミュレートし、アニメーションで可視化します。
    コマンドライン引数により、シミュレーションパラメータや散乱体の設定、GIF保存の有無などを細かく制御できます。

詳細説明:
    このスクリプトは、線形および4次の非線形ばねで結合された1次元の質点系をシミュレートします。
    初期条件としてガウス波束を設定し、その伝播と、システムに配置された散乱体（質量欠陥、結合欠陥、バリアなど）との相互作用を観測します。
    時間積分にはオイラー法またはヴェルレ法が選択可能です。
    結果はMatplotlibのFuncAnimationを使用してリアルタイムで表示され、オプションでGIFとして保存することもできます。
    さらに、シミュレーション終了後には、FFTによる波数空間のパワー分布と、全エネルギーの時間履歴がプロットされます。

サイト質量:
    m_site[i]: サイト i の質量

結合ばね定数:
    k2_bond[i]: サイト i とサイト i+1 間の線形ばね定数
    k4_bond[i]: サイト i とサイト i+1 間の非線形ばね定数

ポテンシャルエネルギー:
    V = sum_i [ 1/2 k2_bond[i] (u[i+1]-u[i])^2
              + 1/4 k4_bond[i] (u[i+1]-u[i])^4 ]

力:
    du[i] = u[i+1] - u[i]
    g[i]  = k2_bond[i] du[i] + k4_bond[i] du[i]^3

    内部サイト j における力:
        F[j] = g[j] - g[j-1]

関連リンク:
    :doc:`md_packets_FuncAnimation_usage`
"""

import argparse
import re
import sys
from pathlib import Path

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation


# ============================================================
# argparse
# ============================================================
def build_argparser():
    """
    概要: コマンドライン引数を解析するためのArgumentParserオブジェクトを構築します。

    詳細説明:
        1次元結合振動子波束散乱シミュレーションの様々なパラメータを定義します。
        これには、基本的なシミュレーション設定、波束パラメータ、散乱体パラメータ、
        および出力/表示に関するオプションが含まれます。

    戻り値:
        :returns: argparse.ArgumentParser: 構築されたArgumentParserオブジェクト。
    """
    p = argparse.ArgumentParser(
        description=(
            "1D coupled oscillator wave-packet scattering simulation "
            "with linear + quartic springs."
        ),
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    # --------------------------
    # basic simulation parameters
    # --------------------------
    p.add_argument("--n", type=int, default=200, help="number of masses")
    p.add_argument("--a", type=float, default=1.0, help="lattice spacing")

    p.add_argument("--m0", type=float, default=1.0, help="background mass")
    p.add_argument("--k2", type=float, default=5.0, help="background linear spring constant")
    p.add_argument("--k4", type=float, default=0.0, help="background nonlinear spring constant")

    p.add_argument(
        "--method",
        type=str,
        default="verlet",
        choices=["euler", "verlet"],
        help="time-integration method",
    )
    p.add_argument("--dt", type=float, default=0.05, help="time step")
    p.add_argument("--nstep", type=int, default=1000, help="number of MD/animation steps")

    p.add_argument("--interval", type=int, default=1, help="animation interval [ms]")
    p.add_argument(
        "--save_interval",
        type=int,
        default=10,
        help="interval for FFT and energy snapshots",
    )
    p.add_argument(
        "--boundary",
        type=str,
        default="fixed",
        choices=["fixed"],
        help="boundary condition",
    )

    # --------------------------
    # wave packet parameters
    # --------------------------
    p.add_argument("--A", type=float, default=0.20, help="packet amplitude")
    p.add_argument("--sigma", type=float, default=10.0, help="packet width: should be sigma*k0 >> 1")
    p.add_argument("--k0", type=float, default=0.15, help="central wave number: should be sigma*k0 >> 1")

    p.add_argument(
        "--initial_mode",
        type=str,
        default="left_only",
        choices=["left_only", "two_packets"],
        help="initial packet configuration",
    )
    p.add_argument(
        "--x0_left",
        type=float,
        default=None,
        help="left packet center. If omitted, n*a*0.25 is used.",
    )
    p.add_argument(
        "--x0_right",
        type=float,
        default=None,
        help="right packet center. If omitted, n*a*0.75 is used.",
    )

    # --------------------------
    # scatterer parameters
    # --------------------------
    p.add_argument(
        "--scatterer_mode",
        type=str,
        default="heavy_mass",
        choices=[
            "none",
            "heavy_mass",
            "light_mass",
            "soft_bond",
            "hard_bond",
            "nonlinear_bond",
            "barrier",
        ],
        help="scatterer type",
    )
    p.add_argument(
        "--scatterer_center",
        type=int,
        default=None,
        help="scatterer center index. If omitted, n//2 is used.",
    )

    # The following values reproduce the current hard-coded settings by default.
    p.add_argument("--mass_width", type=int, default=3, help="mass-defect width")
    p.add_argument("--heavy_mass_factor", type=float, default=10.0, help="heavy mass factor")
    p.add_argument("--light_mass_factor", type=float, default=0.2, help="light mass factor")

    p.add_argument("--bond_width", type=int, default=4, help="bond-defect width")
    p.add_argument("--soft_bond_k2_factor", type=float, default=0.2, help="soft-bond k2 factor")
    p.add_argument("--soft_bond_k4_factor", type=float, default=1.0, help="soft-bond k4 factor")
    p.add_argument("--hard_bond_k2_factor", type=float, default=5.0, help="hard-bond k2 factor")
    p.add_argument("--hard_bond_k4_factor", type=float, default=1.0, help="hard-bond k4 factor")
    p.add_argument("--nonlinear_bond_k2_factor", type=float, default=1.0, help="nonlinear-bond k2 factor")
    p.add_argument("--nonlinear_bond_k4_factor", type=float, default=10.0, help="nonlinear-bond k4 factor")

    p.add_argument("--barrier_half_width", type=int, default=3, help="barrier half width")
    p.add_argument("--barrier_mass", type=float, default=8.0, help="barrier mass value")
    p.add_argument("--barrier_k2", type=float, default=20.0, help="barrier k2 value")
    p.add_argument(
        "--barrier_k4",
        type=float,
        default=None,
        help="barrier k4 value. If omitted, background k4 is used.",
    )

    # --------------------------
    # output / display
    # --------------------------
    p.add_argument(
        "--save",
        type=int,
        default=0,
        choices=[0, 1],
        help="save GIF with ani.save after the animation window is closed",
    )
    p.add_argument(
        "--gif_fps",
        type=int,
        default=30,
        help="frames per second for GIF save",
    )
    p.add_argument(
        "--gif_dpi",
        type=int,
        default=100,
        help="DPI for GIF save",
    )

    return p


def fmt_param(x):
    """
    概要: パラメータ値をファイル名に適した短い文字列に変換します。

    詳細説明:
        浮動小数点数などの値を、`.`を`p`に、`-`を`m`に置換することで、
        ファイル名に安全な形式に変換します。

    引数:
        :param x: Any: 変換するパラメータ値。

    戻り値:
        :returns: str: ファイル名に適した短い文字列。
    """
    s = f"{x:g}"
    s = s.replace("-", "m").replace("+", "")
    s = s.replace(".", "p")
    return s


def safe_stem(s):
    """
    概要: スクリプトのステムをWindowsで安全な出力ファイル名にサニタイズします。

    詳細説明:
        ファイル名に使用できない文字（数字、英字、アンダースコア、ハイフン以外）を
        アンダースコアに置換し、前後のアンダースコアを削除します。
        これにより、異なるOS環境でも安全なファイル名が生成されます。

    引数:
        :param s: str: サニタイズする文字列。

    戻り値:
        :returns: str: サニタイズされた文字列。
    """
    return re.sub(r"[^0-9A-Za-z_\-]+", "_", s).strip("_") or "md_packets"


def make_gif_filename(args):
    """
    概要: GIF保存用のファイル名を生成します。

    詳細説明:
        スクリプト名、シミュレーションメソッド、初期モード、散乱体モード、
        ばね定数k2、k4の値に基づいてファイル名を構築します。
        生成されるファイル名は、シミュレーション設定を反映したものになります。

    引数:
        :param args: argparse.Namespace: コマンドライン引数を格納するオブジェクト。

    戻り値:
        :returns: str: 生成されたGIFファイル名。
    """
    stem = safe_stem(Path(sys.argv[0]).stem)
    return (
        f"{stem}_{args.method}_{args.initial_mode}_{args.scatterer_mode}_"
        f"{fmt_param(args.k2)}_{fmt_param(args.k4)}.gif"
    )




def main(argv=None):
    """
    概要: 1次元結合振動子波束散乱シミュレーションのメインエントリポイントです。

    詳細説明:
        この関数は、コマンドライン引数を解析し、シミュレーションを初期化、実行し、
        結果を可視化します。具体的には、以下の処理を行います。
        1. 引数の解析とパラメータの初期設定。
        2. サイト質量と結合ばね定数配列の準備。
        3. `scatterer_mode`に基づいた散乱体の設定。
        4. 初期波束（変位と速度）の生成。
        5. シミュレーション中の状態（FFTスナップショット、エネルギー履歴）の記録。
        6. Matplotlib FuncAnimationによるシミュレーションのリアルタイム可視化。
        7. オプションでシミュレーション結果のアニメーションGIF保存。
        8. FFTおよびエネルギーの時間履歴のプロットによる詳細分析。
        9. シミュレーション終了時のエネルギー保存のレポート。

    引数:
        :param argv: list, optional: コマンドライン引数のリスト。
                     デフォルトはNoneで、`sys.argv`を使用します。

    戻り値:
        :returns: None
    """
    args = build_argparser().parse_args(argv)

    # ============================================================
    # parameters
    # ============================================================
    n = args.n
    if n < 3:
        raise ValueError("--n must be >= 3")

    a = args.a
    m0 = args.m0
    k2_0 = args.k2
    k4_0 = args.k4

    method = args.method

    dt = args.dt
    nstep = args.nstep

    interval = args.interval
    save_interval = args.save_interval

    boundary = args.boundary

    A = args.A
    sigma = args.sigma
    k0 = args.k0

    initial_mode = args.initial_mode

    x0_left = args.x0_left if args.x0_left is not None else n * a * 0.25
    x0_right = args.x0_right if args.x0_right is not None else n * a * 0.75

    scatterer_mode = args.scatterer_mode
    scatterer_center = args.scatterer_center if args.scatterer_center is not None else n // 2
    barrier_k4 = args.barrier_k4 if args.barrier_k4 is not None else k4_0


    # ============================================================
    # equilibrium positions
    # ============================================================
    x_eq = np.arange(n) * a


    # ============================================================
    # site and bond arrays
    # ============================================================
    m_site = np.full(n, m0, dtype=float)

    # k2_bond[i], k4_bond[i] are for the bond between site i and i+1
    k2_bond = np.full(n - 1, k2_0, dtype=float)
    k4_bond = np.full(n - 1, k4_0, dtype=float)

    scatterer_spans = []


    # ============================================================
    # scatterer helper functions
    # ============================================================
    def add_mass_defect(center, width=1, m_factor=5.0):
        """
        概要: 質量欠陥を追加します。

        詳細説明:
            指定された中心サイトと幅に基づいて、サイトの質量を乗算係数で変更します。
            質量欠陥が適用されるサイトの範囲は、`[center - width // 2, center - width // 2 + width)` です。
            計算された範囲は、配列の有効なインデックス内に収まるようにクリップされます。
            欠陥領域は `scatterer_spans` リストに記録され、プロットに表示されます。

        引数:
            :param center: int: 欠陥領域の中心サイトインデックス。
            :param width: int, optional: 欠陥領域のサイト数。デフォルトは1。
            :param m_factor: float, optional: 質量に乗算する係数。例: `m_factor=10` は質量を10倍にします。
                             デフォルトは5.0。

        戻り値:
            :returns: None
        """
        i1 = max(0, int(center - width // 2))
        i2 = min(n, int(i1 + width))

        if i2 <= i1:
            return

        m_site[i1:i2] *= m_factor

        x1 = x_eq[i1]
        x2 = x_eq[i2 - 1]
        scatterer_spans.append((x1, x2, f"mass x{m_factor:g}"))


    def set_mass_region(i1, i2, m_value):
        """
        概要: 指定された範囲のサイト質量を設定します。

        詳細説明:
            サイトのインデックス範囲`[i1, i2)` (i2は排他的) の`m_site`の値を`m_value`に設定します。
            計算された範囲は、配列の有効なインデックス内に収まるようにクリップされます。
            欠陥領域は `scatterer_spans` リストに記録され、プロットに表示されます。

        引数:
            :param i1: int: 質量を設定する開始サイトインデックス (含む)。
            :param i2: int: 質量を設定する終了サイトインデックス (排他的)。
            :param m_value: float: 設定する質量値。

        戻り値:
            :returns: None
        """
        i1 = max(0, int(i1))
        i2 = min(n, int(i2))

        if i2 <= i1:
            return

        m_site[i1:i2] = m_value

        x1 = x_eq[i1]
        x2 = x_eq[i2 - 1]
        scatterer_spans.append((x1, x2, f"m={m_value:g}"))


    def add_bond_defect(center, width=1, k2_factor=1.0, k4_factor=1.0):
        """
        概要: 結合欠陥を追加します。

        詳細説明:
            これはサイトではなく結合を変更します。
            `k2_bond[i]`と`k4_bond[i]`はサイト`i`とサイト`i+1`間のばね定数です。
            指定された中心結合と幅に基づいて、結合のばね定数`k2_bond`と`k4_bond`を乗算係数で変更します。
            結合欠陥が適用される結合の範囲は、`[center - width // 2, center - width // 2 + width)` です。
            計算された範囲は、配列の有効なインデックス内に収まるようにクリップされます。
            欠陥領域は `scatterer_spans` リストに記録され、プロットに表示されます。

        引数:
            :param center: int: 欠陥領域の中心結合インデックス。
            :param width: int, optional: 欠陥領域の結合数。デフォルトは1。
            :param k2_factor: float, optional: k2に乗算する係数。デフォルトは1.0。
            :param k4_factor: float, optional: k4に乗算する係数。デフォルトは1.0。

        戻り値:
            :returns: None
        """
        i1 = max(0, int(center - width // 2))
        i2 = min(n - 1, int(i1 + width))

        if i2 <= i1:
            return

        k2_bond[i1:i2] *= k2_factor
        k4_bond[i1:i2] *= k4_factor

        x1 = x_eq[i1]
        x2 = x_eq[i2]
        scatterer_spans.append(
            (x1, x2, f"k2 x{k2_factor:g}, k4 x{k4_factor:g}")
        )


    def set_bond_region(i1, i2, k2_value=None, k4_value=None):
        """
        概要: 指定された結合領域のばね定数を設定します。

        詳細説明:
            結合インデックス範囲`[i1, i2)` (i2は排他的) のばね定数`k2_bond`と`k4_bond`の値を変更します。
            `k2_value`または`k4_value`がNoneでない場合、対応するばね定数が設定値に置き換えられます。
            計算された範囲は、配列の有効なインデックス内に収まるようにクリップされます。
            欠陥領域は `scatterer_spans` リストに記録され、プロットに表示されます。

        引数:
            :param i1: int: ばね定数を設定する開始結合インデックス (含む)。
            :param i2: int: ばね定数を設定する終了結合インデックス (排他的)。
            :param k2_value: float or None, optional: 設定するk2値。Noneの場合、変更されません。
            :param k4_value: float or None, optional: 設定するk4値。Noneの場合、変更されません。

        戻り値:
            :returns: None
        """
        i1 = max(0, int(i1))
        i2 = min(n - 1, int(i2))

        if i2 <= i1:
            return

        label_parts = []

        if k2_value is not None:
            k2_bond[i1:i2] = k2_value
            label_parts.append(f"k2={k2_value:g}")

        if k4_value is not None:
            k4_bond[i1:i2] = k4_value
            label_parts.append(f"k4={k4_value:g}")

        if label_parts:
            x1 = x_eq[i1]
            x2 = x_eq[i2]
            scatterer_spans.append((x1, x2, ", ".join(label_parts)))


    # ============================================================
    # scatterer settings
    # ============================================================
    def setup_scatterer():
        """
        概要: コマンドライン引数に基づいて散乱体を設定します。

        詳細説明:
            `scatterer_mode`引数の値に応じて、異なる種類の散乱体（質量欠陥、結合欠陥、またはバリア）を
            シミュレーションシステムに適用します。
            各モードは、対応するヘルパー関数（`add_mass_defect`, `add_bond_defect`, `set_mass_region`, `set_bond_region`）を呼び出します。
            `--scatterer_center`, `--mass_width`, `--bond_width` などの引数で散乱体の特性を制御できます。

        戻り値:
            :returns: None

        例外:
            ValueError: 未知の`scatterer_mode`が指定された場合。
        """
        if scatterer_mode == "none":
            pass

        elif scatterer_mode == "heavy_mass":
            # A local heavy-mass defect near the center.
            # This gives reflection and transmission.
            add_mass_defect(
                center=scatterer_center,
                width=args.mass_width,
                m_factor=args.heavy_mass_factor,
            )

        elif scatterer_mode == "light_mass":
            # A local light-mass defect.
            add_mass_defect(
                center=scatterer_center,
                width=args.mass_width,
                m_factor=args.light_mass_factor,
            )

        elif scatterer_mode == "soft_bond":
            # A soft spring region.
            add_bond_defect(
                center=scatterer_center,
                width=args.bond_width,
                k2_factor=args.soft_bond_k2_factor,
                k4_factor=args.soft_bond_k4_factor,
            )

        elif scatterer_mode == "hard_bond":
            # A hard spring region.
            add_bond_defect(
                center=scatterer_center,
                width=args.bond_width,
                k2_factor=args.hard_bond_k2_factor,
                k4_factor=args.hard_bond_k4_factor,
            )

        elif scatterer_mode == "nonlinear_bond":
            # Only nonlinear spring constant is enhanced.
            add_bond_defect(
                center=scatterer_center,
                width=args.bond_width,
                k2_factor=args.nonlinear_bond_k2_factor,
                k4_factor=args.nonlinear_bond_k4_factor,
            )

        elif scatterer_mode == "barrier":
            # A wider barrier-like region.
            hw = args.barrier_half_width
            set_mass_region(scatterer_center - hw, scatterer_center + hw + 1, m_value=args.barrier_mass)
            set_bond_region(scatterer_center - hw, scatterer_center + hw, k2_value=args.barrier_k2, k4_value=barrier_k4)

        else:
            raise ValueError(f"Unknown scatterer_mode: {scatterer_mode}")


    setup_scatterer()


    # ============================================================
    # reference dispersion relation for the homogeneous incident region
    # ============================================================
    def omega_k(k):
        """
        概要: 均質な入射領域における線形鎖の分散関係を計算します。

        詳細説明:
            `omega(k) = 2 sqrt(k2_0/m0) |sin(k a / 2)|` の式を用いて、
            波数`k`に対応する角振動数`omega`を計算します。
            この関数は主に、初期入射波束の構築に使用されます。

        引数:
            :param k: float: 波数を表す値。

        戻り値:
            :returns: float: 波長kに対応する角振動数。
        """
        return 2.0 * np.sqrt(k2_0 / m0) * np.abs(np.sin(0.5 * k * a))


    def vg_k(k):
        """
        概要: 均質な入射領域における群速度 `d omega / d k` を計算します。

        詳細説明:
            `vg = sqrt(k2_0/m0) a cos(k a / 2) sign(k)` の式を用いて、
            波数`k`に対応する群速度を計算します。
            `k=0`の場合は群速度を0とします。
            この関数は初期入射波束の構築にのみ使用されます。

        引数:
            :param k: float: 波数を表す値。

        戻り値:
            :returns: float: 波長kに対応する群速度。
        """
        if k == 0:
            return 0.0

        return (
            np.sqrt(k2_0 / m0)
            * a
            * np.cos(0.5 * np.abs(k) * a)
            * np.sign(k)
        )


    # ============================================================
    # traveling wave packet
    # ============================================================
    def packet_u(x, x0, k):
        """
        概要: ガウス波束の初期変位を計算します。

        詳細説明:
            `u(x,0) = A exp[-(x-x0)^2/(2 sigma^2)] cos[k(x-x0)]` の式に基づいて、
            時刻`t=0`における各サイトの変位を計算します。
            ここで、`A`は振幅、`sigma`は波束の幅、`k`は中心波数、`x0`は波束の中心位置です。

        引数:
            :param x: numpy.ndarray: サイトの平衡位置の配列。
            :param x0: float: 波束の中心位置。
            :param k: float: 中心波数。

        戻り値:
            :returns: numpy.ndarray: 各サイトの初期変位の配列。
        """
        dx = x - x0
        env = np.exp(-dx**2 / (2.0 * sigma**2))
        return A * env * np.cos(k * dx)


    def packet_v(x, x0, k):
        """
        概要: 進行するガウス波束の初期速度を計算します。

        詳細説明:
            近似的に `u(x,t) = A exp[-(x-x0-vg t)^2/(2 sigma^2)] cos[k(x-x0) - omega t]`
            と仮定し、時刻`t=0`における変位の時間微分 `v(x,0) = du/dt` を計算します。
            これにより、波束が初期に`vg0`の群速度で進行するような速度場が設定されます。
            `k > 0` の場合、波束は正のx方向へ移動し、`k < 0` の場合は負のx方向へ移動します。

        引数:
            :param x: numpy.ndarray: サイトの平衡位置の配列。
            :param x0: float: 波束の中心位置。
            :param k: float: 中心波数。

        戻り値:
            :returns: numpy.ndarray: 各サイトの初期速度の配列。
        """
        dx = x - x0
        env = np.exp(-dx**2 / (2.0 * sigma**2))

        omega0 = omega_k(k)
        vg0 = vg_k(k)

        return A * env * (
            vg0 * dx / sigma**2 * np.cos(k * dx)
            + omega0 * np.sin(k * dx)
        )


    # ============================================================
    # force calculation
    # ============================================================
    def force(u):
        """
        概要: 最近接結合の線形および4次非線形ばねによる力を計算します。

        詳細説明:
            各サイトに作用する力は、隣接するサイトとの変位差と、
            線形および4次非線形ばね定数`k2_bond`、`k4_bond`に基づいて計算されます。
            `k2_bond[i]`と`k4_bond[i]`はサイト`i`と`i+1`間のばね定数です。
            固定端境界条件が適用される場合、両端のサイトの力は0に設定されます。

        引数:
            :param u: numpy.ndarray: 各サイトの変位の配列。

        戻り値:
            :returns: numpy.ndarray: 各サイトに作用する力の配列。
        """
        f = np.zeros_like(u)

        du = u[1:] - u[:-1]

        # bond force-like quantity
        # g[i] = dV_i / d(u[i+1]-u[i])
        g = k2_bond * du + k4_bond * du**3

        # interior sites
        f[1:-1] = g[1:] - g[:-1]

        # endpoints are fixed, so forces there are not used
        f[0] = 0.0
        f[-1] = 0.0

        return f


    def apply_boundary(u, v):
        """
        概要: 固定境界条件を適用します。

        詳細説明:
            `boundary`変数が`"fixed"`の場合、両端のサイト（インデックス0と-1）の変位`u`と
            速度`v`をゼロに設定します。これにより、両端が壁に固定された効果をシミュレートします。
            他の境界条件が指定された場合はエラーを発生させます。

        引数:
            :param u: numpy.ndarray: 各サイトの変位の配列。
            :param v: numpy.ndarray: 各サイトの速度の配列。

        戻り値:
            :returns: tuple[numpy.ndarray, numpy.ndarray]: 境界条件が適用された変位と速度の配列。

        例外:
            ValueError: 未知の境界条件が指定された場合。
        """
        if boundary == "fixed":
            u[0] = 0.0
            u[-1] = 0.0
            v[0] = 0.0
            v[-1] = 0.0
        else:
            raise ValueError(f"Unknown boundary condition: {boundary}")

        return u, v


    # ============================================================
    # energy calculation
    # ============================================================
    def total_energy(u, v):
        """
        概要: システムの全エネルギー（運動エネルギー + ポテンシャルエネルギー）を計算します。

        詳細説明:
            各サイトの運動エネルギーと、隣接するサイト間の線形および4次非線形ばねによる
            ポテンシャルエネルギーの合計を計算します。

        引数:
            :param u: numpy.ndarray: 各サイトの変位の配列。
            :param v: numpy.ndarray: 各サイトの速度の配列。

        戻り値:
            :returns: float: システムの総エネルギー。
        """
        kinetic = 0.5 * np.sum(m_site * v**2)

        du = u[1:] - u[:-1]
        potential2 = 0.5 * np.sum(k2_bond * du**2)
        potential4 = 0.25 * np.sum(k4_bond * du**4)

        return kinetic + potential2 + potential4


    # ============================================================
    # initial condition
    # ============================================================
    def make_initial_condition():
        """
        概要: 初期条件（変位と速度）を生成します。

        詳細説明:
            `initial_mode`引数の値に応じて、初期波束の構成を選択します。
            - `"left_only"`: 左側から右方向へ進行する単一のガウス波束を生成します。
            - `"two_packets"`: 左から右へ、右から左へ進行する2つのガウス波束を生成します。
            波束生成後、`apply_boundary`関数を呼び出して境界条件を適用します。

        戻り値:
            :returns: tuple[numpy.ndarray, numpy.ndarray]: 初期変位と初期速度の配列。

        例外:
            ValueError: 未知の`initial_mode`が指定された場合。
        """
        if initial_mode == "left_only":
            # One packet from left to right.
            # This is useful for seeing scattering by a defect.
            u0 = packet_u(x_eq, x0_left, +k0)
            v0 = packet_v(x_eq, x0_left, +k0)

        elif initial_mode == "two_packets":
            # Left packet moves right:  k = +k0
            # Right packet moves left: k = -k0
            u0 = (
                packet_u(x_eq, x0_left, +k0)
                + packet_u(x_eq, x0_right, -k0)
            )

            v0 = (
                packet_v(x_eq, x0_left, +k0)
                + packet_v(x_eq, x0_right, -k0)
            )

        else:
            raise ValueError(f"Unknown initial_mode: {initial_mode}")

        u0, v0 = apply_boundary(u0, v0)
        return u0, v0


    u_initial, v_initial = make_initial_condition()
    u = u_initial.copy()
    v = v_initial.copy()

    E0 = total_energy(u, v)


    # ============================================================
    # console output
    # ============================================================
    print("=== Simulation settings ===")
    print(f"n              = {n}")
    print(f"a              = {a}")
    print(f"m0             = {m0}")
    print(f"k2_0           = {k2_0}")
    print(f"k4_0           = {k4_0}")
    print(f"dt             = {dt}")
    print(f"nstep          = {nstep}")
    print(f"method         = {method}")
    print(f"boundary       = {boundary}")
    print(f"scatterer_mode = {scatterer_mode}")
    print(f"initial_mode   = {initial_mode}")
    print(f"save           = {args.save}")
    print()
    print("=== Array summary ===")
    print(f"m_site   min/max = {m_site.min():.6g} / {m_site.max():.6g}")
    print(f"k2_bond  min/max = {k2_bond.min():.6g} / {k2_bond.max():.6g}")
    print(f"k4_bond  min/max = {k4_bond.min():.6g} / {k4_bond.max():.6g}")
    print()
    print("=== Wave packet settings ===")
    print(f"A        = {A}")
    print(f"sigma    = {sigma}")
    print(f"k0       = {k0}")
    print(f"x0_left  = {x0_left}")
    print(f"x0_right = {x0_right}")
    print(f"omega0   = {omega_k(k0):.8f}")
    print(f"vg0      = {vg_k(k0):.8f}")
    print()
    print(f"Initial energy = {E0:.12e}")
    print()


    # ============================================================
    # storage for FFT and energy
    # ============================================================
    snapshots = []
    snapshot_times = []

    energy_list = []
    energy_times = []

    animation_phase = "display"


    def clear_observables():
        """
        概要: FFTとエネルギーのスナップショットデータをクリアします。

        詳細説明:
            `snapshots`, `snapshot_times`, `energy_list`, `energy_times`のリストを空にします。
            これは、シミュレーションのリセット時や、GIF保存前など、
            新しいデータ収集サイクルを開始する際に使用されます。

        戻り値:
            :returns: None
        """
        snapshots.clear()
        snapshot_times.clear()
        energy_list.clear()
        energy_times.clear()


    def reset_state():
        """
        概要: シミュレーションの状態を初期状態にリセットします。

        詳細説明:
            変位`u`と速度`v`を、`make_initial_condition`で生成された初期値に戻します。
            また、アニメーションのプロットデータ（`line.set_ydata`）と
            タイトルテキストを初期状態に設定し、すべてのオブザーバブルデータをクリアします。
            主にGIF保存前に、アニメーションの再生を初期状態から開始するために使用されます。

        戻り値:
            :returns: None
        """
        nonlocal u, v
        u = u_initial.copy()
        v = v_initial.copy()
        line.set_ydata(u)
        title_text.set_text("1D wave packet scattering")
        clear_observables()


    # ============================================================
    # one MD step
    # ============================================================
    def step_euler(u, v):
        """
        概要: オイラー法を用いて1MDステップを進めます。

        詳細説明:
            現在の変位`u`と速度`v`から、現在の力`f`を計算します。
            その後、オイラー法（速度の更新に現在の力、変位の更新に新しい速度を使用）に
            よって次の時刻の変位`u_new`と速度`v_new`を更新します。
            最後に、境界条件を適用します。

        引数:
            :param u: numpy.ndarray: 現在の各サイトの変位の配列。
            :param v: numpy.ndarray: 現在の各サイトの速度の配列。

        戻り値:
            :returns: tuple[numpy.ndarray, numpy.ndarray]: 更新された変位と速度の配列。
        """
        f = force(u)

        v_new = v + dt * f / m_site
        u_new = u + dt * v_new

        u_new, v_new = apply_boundary(u_new, v_new)

        return u_new, v_new


    def step_verlet(u, v):
        """
        概要: ヴェルレ法を用いて1MDステップを進めます。

        詳細説明:
            現在の変位`u`と速度`v`から、現在の力`f`を計算します。
            ヴェルレ法（速度の半分更新、変位の更新、新しい変位に基づく力の再計算、
            速度の残り半分更新）によって次の時刻の変位`u_new`と速度`v_new`を更新します。
            変位更新後に境界条件を適用し、新しい力を計算します。

        引数:
            :param u: numpy.ndarray: 現在の各サイトの変位の配列。
            :param v: numpy.ndarray: 現在の各サイトの速度の配列。

        戻り値:
            :returns: tuple[numpy.ndarray, numpy.ndarray]: 更新された変位と速度の配列。
        """
        f = force(u)

        v_half = v + 0.5 * dt * f / m_site
        u_new = u + dt * v_half

        # Apply boundary before calculating new force
        if boundary == "fixed":
            u_new[0] = 0.0
            u_new[-1] = 0.0

        f_new = force(u_new)
        v_new = v_half + 0.5 * dt * f_new / m_site

        u_new, v_new = apply_boundary(u_new, v_new)

        return u_new, v_new


    def md_step(u, v):
        """
        概要: 指定された時間積分メソッドで1MDステップを進めます。

        詳細説明:
            `method`引数の値（`"euler"`または`"verlet"`）に基づいて、
            対応する時間積分関数（`step_euler`または`step_verlet`）を呼び出します。
            これにより、シミュレーションの次の状態が計算されます。

        引数:
            :param u: numpy.ndarray: 現在の各サイトの変位の配列。
            :param v: numpy.ndarray: 現在の各サイトの速度の配列。

        戻り値:
            :returns: tuple[numpy.ndarray, numpy.ndarray]: 更新された変位と速度の配列。

        例外:
            ValueError: 未知の`method`が指定された場合。
        """
        if method == "euler":
            return step_euler(u, v)
        elif method == "verlet":
            return step_verlet(u, v)
        else:
            raise ValueError(f"Unknown method: {method}")


    # ============================================================
    # animation update
    # ============================================================
    def update(frame):
        """
        概要: MatplotlibのFuncAnimationの各フレームを更新します。

        詳細説明:
            この関数はFuncAnimationによってフレームごとに呼び出されます。
            1. `md_step`関数を呼び出してシミュレーションを1ステップ進めます。
            2. `save_interval`で指定された間隔で、現在の変位`u`と時刻`t`を`snapshots`と`snapshot_times`に、
               現在の全エネルギー`E`を`energy_list`と`energy_times`に記録します。
            3. プロットのラインデータ（`line.set_ydata(u)`）を更新し、変位を可視化します。
            4. `step_no`が特定の間隔（20ステップごと、または1ステップ目）で、
               現在の全エネルギーと初期エネルギーからの相対ドリフトを計算し、タイトルテキストを更新します。
            5. `step_no`が特定の間隔（50ステップごと、または最終ステップ）で、
               現在のシミュレーション進捗とエネルギー情報がコンソールに出力されます。

        引数:
            :param frame: int: 現在のアニメーションフレーム番号（0から`nstep-1`まで）。

        戻り値:
            :returns: tuple: 更新されたmatplotlibオブジェクト（line, title_text）のタプル。
        """
        nonlocal u, v

        u, v = md_step(u, v)

        step_no = frame + 1
        t = step_no * dt

        if save_interval > 0 and step_no % save_interval == 0:
            snapshots.append(u.copy())
            snapshot_times.append(t)

            E = total_energy(u, v)
            energy_list.append(E)
            energy_times.append(t)

        line.set_ydata(u)

        if step_no % 20 == 0 or step_no == 1:
            E = total_energy(u, v)
            drift = (E - E0) / E0
            title_text.set_text(
                f"1D wave packet scattering  step={step_no},  t={t:.2f},  "
                f"energy drift={drift:.2e}"
            )

        if step_no % 50 == 0 or step_no == nstep:
            E = total_energy(u, v)
            drift = (E - E0) / E0
            print(
                f"[{animation_phase}] step={step_no:6d}/{nstep}, "
                f"t={t:.4g}, energy drift={drift:.3e}",
                flush=True,
            )

        return line, title_text


    # ============================================================
    # plot setup
    # ============================================================
    fig, ax = plt.subplots(figsize=(10, 4))

    line, = ax.plot(x_eq, u, lw=2)

    title_text = ax.set_title("1D wave packet scattering")

    ax.set_xlim(x_eq[0], x_eq[-1])
    ax.set_ylim(-0.5, 0.5)

    ax.set_xlabel("x")
    ax.set_ylabel("u")
    ax.grid(True)

    # mark scatterer regions
    for x1, x2, label in scatterer_spans:
        ax.axvspan(x1, x2, alpha=0.20)
        ax.text(
            0.5 * (x1 + x2),
            0.43,
            label,
            ha="center",
            va="center",
            fontsize=9,
            rotation=90,
        )

    ani = FuncAnimation(
        fig,
        update,
        frames=nstep,
        interval=interval,
        blit=False,
    )

    plt.show()


    # ============================================================
    # optional GIF save
    # ============================================================
    if args.save:
        # FuncAnimation with a stateful update function replays frames during ani.save.
        # Reset first so the saved GIF starts from the same initial condition.
        animation_phase = "save"
        reset_state()
        gif_name = make_gif_filename(args)
        print(f"Saving GIF: {gif_name}")
        try:
            ani.save(gif_name, writer="pillow", fps=args.gif_fps, dpi=args.gif_dpi)
        except Exception:
            print("ERROR: failed to save GIF.")
            print("Please check that pillow is installed:")
            print("  pip install pillow")
            raise
        print(f"Saved GIF: {gif_name}")


    # ============================================================
    # FFT visualization
    # ============================================================
    if len(snapshots) > 1:
        snapshots_arr = np.array(snapshots)
        snapshot_times_arr = np.array(snapshot_times)

        # Remove spatial average before FFT
        snapshots0 = snapshots_arr - snapshots_arr.mean(axis=1, keepdims=True)

        Uk = np.fft.rfft(snapshots0, axis=1)
        P = np.abs(Uk)**2

        k_axis = 2.0 * np.pi * np.fft.rfftfreq(n, d=a)

        plt.figure(figsize=(10, 5))
        plt.imshow(
            P.T,
            aspect="auto",
            origin="lower",
            extent=[
                snapshot_times_arr[0],
                snapshot_times_arr[-1],
                k_axis[0],
                k_axis[-1],
            ],
        )
        plt.xlabel("time")
        plt.ylabel("k")
        plt.ylim([0, min(np.pi / a, k0 * 5)])
        plt.title(r"Mode power $|u_k(t)|^2$")
        plt.colorbar(label="Power")
        plt.tight_layout()
        plt.show()


    # ============================================================
    # energy drift visualization
    # ============================================================
    if len(energy_list) > 1:
        energy_arr = np.array(energy_list)
        energy_times_arr = np.array(energy_times)

        plt.figure(figsize=(10, 4))
        plt.plot(energy_times_arr, (energy_arr - E0) / E0)
        plt.xlabel("time")
        plt.ylabel("relative energy drift")
        plt.title("Energy conservation check")
        plt.grid(True)
        plt.tight_layout()
        plt.show()


    # ============================================================
    # final report
    # ============================================================
    E1 = total_energy(u, v)
    print()
    print("=== Final energy check ===")
    print(f"Initial energy = {E0:.12e}")
    print(f"Final energy   = {E1:.12e}")
    print(f"Relative drift = {(E1 - E0) / E0:.12e}")


if __name__ == "__main__":
    main()