import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

# ============================
# パラメータ設定
# ============================

def initialize():
    """
    1D 分子（Hückel）モデル
    """
    params = {}

    # 例：4 原子直鎖分子。s軌道同士の場合、移送積分は負であることに注意
    params["onsite"]  = [ 0.0,  0.1,  0.2, 0.1]
    params["hopping"] = [-1.0, -1.1, -0.9]  # 原子 i–(i+1)

    return params


# ============================
# Hamiltonian（有限鎖）
# ============================

def hamiltonian_molecule(onsite, hopping):
    """
    1D 分子の Hückel Hamiltonian
    （周期条件なし）
    """
    norb = len(onsite)
    H = np.zeros((norb, norb))

    # onsite
    for i in range(norb):
        H[i, i] = onsite[i]

    # nearest-neighbor hopping
    for i in range(norb - 1):
        H[i, i + 1] = hopping[i]
        H[i + 1, i] = hopping[i]

    return H


# ============================
# 計算
# ============================

def solve_molecule(onsite, hopping):
    """
    固有値・固有ベクトルを解く
    """
    H = hamiltonian_molecule(onsite, hopping)
    eigvals, eigvecs = np.linalg.eigh(H)
    return eigvals, eigvecs


# ============================
# 出力
# ============================

def save_results(filename, eigvals, eigvecs):
    """
    CSV 出力
    """
    df = pd.DataFrame(eigvecs, columns=[f"MO_{i}" for i in range(len(eigvals))])
    df.insert(0, "Energy", eigvals)
    df.to_csv(filename, index=False)

def plot_molecular_orbitals_vstack(eigvecs, eigvals=None):
    """
    分子軌道を縦に積んで表示（x軸共有・枠密着）
    """
    norb = eigvecs.shape[0]
    x = np.arange(norb)

    n_mo = eigvecs.shape[1]

    fig, axes = plt.subplots(
        n_mo, 1,
        figsize=(6, 1.2 * n_mo),
        sharex=True,
        gridspec_kw={"hspace": 0}
    )

    # MO が 1 本のときの保険
    if n_mo == 1:
        axes = [axes]

    # グラフの表示スケールを合わせるため、係数の最大値を取得
    ymax = np.max(np.abs(eigvecs))

    for i, ax in enumerate(axes):
        i_mo = n_mo - 1 - i

        if eigvals is not None:
            label = f"MO #{i_mo} E={eigvals[i_mo]:.2f}"
        else:
            label = f"MO #{i_mo}"
        ax.plot(x, eigvecs[:, i_mo], marker="o", label = label)
        ax.axhline(0, color="gray", lw=0.5)

        # y ラベル（エネルギーがあれば併記）
        ax.set_ylabel(f"$c_i$")

        # 下以外は x 軸ラベル・目盛り文字を消す
        if i != n_mo - 1:
            ax.tick_params(labelbottom=False)

        ax.set_ylim(-ymax, ymax)        
        ax.legend()

    axes[-1].set_xlabel("Atom index")

    plt.tight_layout()
    plt.show()

# ============================
# main
# ============================

def main():
    params = initialize()
    onsite = params["onsite"]
    hopping = params["hopping"]

    eigvals, eigvecs = solve_molecule(onsite, hopping)

    print("固有値（分子軌道エネルギー）:")
    for i, e in enumerate(eigvals):
        print(f"MO {i}: {e:.4f}")

    save_results("huckel_1d_molecule.csv", eigvals, eigvecs)
    print("huckel_1d_molecule.csv を出力しました")

    plot_molecular_orbitals_vstack(eigvecs, eigvals)


# ============================
# entry point
# ============================

if __name__ == "__main__":
    main()