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

# Hückel行列の定義（ベンゼンは6炭素環）
H = 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  # ブロードニング幅

print()
print("Hückel行列（ベンゼン）:")
pprint(H)

# 固有値・固有ベクトルの計算
eigenvalues, eigenvectors = np.linalg.eig(H)
eigenvalues = np.sort(eigenvalues)

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

# 縮重度の計算
degeneracies = Counter(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)

# 図を描画
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 Approximation for Benzene')
for energy, deg in degeneracies.items():
    for _ in range(deg):
        ax0.plot([0, 1], [energy, energy], 'k-')
ax0.set_ylabel('Energy')
ax0.set_xticks([])
ax0.set_xlabel('')

# DOSの描画
ax1 = fig.add_subplot(gs[1], sharey=ax0)
ax1.plot(dos, energy_range, 'k-')
ax1.set_xlabel('DOS')

plt.show()
