"""
概要:
    ベンゼン分子のHückel近似による電子構造計算と状態密度（DOS）の可視化を行うスクリプト。
詳細説明:
    このスクリプトは、Hückel近似法を用いてベンゼン分子（6炭素環）の電子構造を解析します。
    Hückel行列を定義し、その固有値問題を解くことで、分子軌道のエネルギー準位と波動関数を計算します。
    計算されたエネルギー準位はソートされ、その縮重度が特定されます。
    その後、ガウス型ブロードニングを適用して状態密度（DOS）を計算します。
    最後に、matplotlibライブラリを使用して、計算された電子準位と状態密度の両方を一つの図にプロットし、
    ベンゼン分子の電子状態を視覚的に表現します。
関連リンク:
    benzene_usage
"""

import numpy as np
from pprint import pprint
import matplotlib.pyplot as plt
from scipy.ndimage import gaussian_filter1d
from collections import Counter
from matplotlib.gridspec import GridSpec

#===================
# Global parameters
#===================

# Hückel行列の定義（ベンゼンは6炭素環）
# α=0, β=1 とした隣接行列形式
H_BENZENE = np.array([[0, 1, 0, 0, 0, 1],
                      [1, 0, 1, 0, 0, 0],
                      [0, 1, 0, 1, 0, 0],
                      [0, 0, 1, 0, 1, 0],
                      [0, 0, 0, 1, 0, 1],
                      [1, 0, 0, 0, 1, 0]])

# 状態密度（DOS）パラメータ
NUM_POINTS = 1000
SIGMA = 0.1  # ブロードニング幅

def main():
    """
    概要:
        Hückel行列の対角化を行い、エネルギー準位と状態密度をプロットします。
    詳細説明:
        この関数は、ベンゼン分子のHückel行列を用いて固有値問題を解き、
        分子軌道のエネルギー準位を計算します。
        計算されたエネルギー準位はソートされ、その縮重度が特定されます。
        その後、ガウス関数を用いたブロードニングによって状態密度（DOS）を計算します。
        最後に、matplotlibの機能を用いて電子準位とDOSを一つの図に描画し、表示します。
    """
    print("\nHückel行列（ベンゼン）:")
    pprint(H_BENZENE)

    # 固有値・固有ベクトルの計算
    eigenvalues, eigenvectors = np.linalg.eig(H_BENZENE)
    # 実数部のみ取得しソート
    eigenvalues = np.sort(eigenvalues.real)

    # エネルギー準位のソートと出力
    print("\nエネルギー準位（ソート後）:")
    print(eigenvalues)
    print()

    # 縮重度の計算
    # 浮動小数点の微小な差を丸めてカウント
    rounded_eigenvalues = np.round(eigenvalues, 8)
    degeneracies = Counter(rounded_eigenvalues)

    # 状態密度（DOS）の計算
    energy_range = np.linspace(min(eigenvalues) - 1, max(eigenvalues) + 1, NUM_POINTS)
    dos = np.zeros(NUM_POINTS)

    for energy, deg in degeneracies.items():
        # ガウス関数によるブロードニング
        dos += deg * np.exp(-((energy_range - energy) ** 2) / (2 * SIGMA ** 2))

    # 状態密度のさらなる平滑化（オプション）
    dos = gaussian_filter1d(dos, SIGMA * (NUM_POINTS / (max(energy_range) - min(energy_range))))

    # 図を描画
    fig = plt.figure(figsize=(8, 6))
    gs = GridSpec(1, 2, width_ratios=[1, 3], wspace=0.4)

    # 電子準位図の描画
    ax0 = fig.add_subplot(gs[0])
    ax0.set_title('Hückel Levels')
    for energy, deg in degeneracies.items():
        # 縮重している場合は少しずらすか、線の太さを変える等の表現も可能ですが
        # オリジナルのプロット形式を維持します
        ax0.hlines(energy, 0, 1, colors='k', linestyles='solid')
        
    ax0.set_ylabel('Energy (units of β)')
    ax0.set_xticks([])
    ax0.set_xlabel('')

    # DOSの描画
    ax1 = fig.add_subplot(gs[1], sharey=ax0)
    ax1.plot(dos, energy_range, 'k-')
    ax1.fill_betweenx(energy_range, 0, dos, color='gray', alpha=0.3)
    ax1.set_title('Density of States (DOS)')
    ax1.set_xlabel('DOS (Arbitrary Units)')

    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    main()