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

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

def initialize():
    """
    モデル・数値計算パラメータをまとめて返す
    """
    params = {}

    # 1原子/単位格子 
    # params["onsite"]  = [0.0]
    # params["hopping"] = [-1.0] 

    # 2原子/単位格子 
#    params["onsite"]  = [0.0, 1.0]
#    params["hopping"] = [-1.0, -1.2]

    # 4原子/単位格子
    params["onsite"]  = [0,  0.1,  0.2,  0.1]
    params["hopping"] = [-1, -1.1, -0.9, -1.2]

    params["a"] = 1.0
    params["nk"] = 201
    return params


# ============================
# Hamiltonian
# ============================
def hamiltonian_k(k, onsite, hopping, a=1.0):
    """
    Bloch Hamiltonian H(k)
    """
    norb = len(onsite)
    a_cell = norb * a  # 単位格子長 L = N * a

    H = np.zeros((norb, norb), dtype=complex)

    # onsite (対角成分)
    for i in range(norb):
        H[i, i] = onsite[i]

    # intra-cell hopping (セル内ホッピング)
    # 1原子の場合はここを通らない
    for i in range(norb - 1):
        H[i, i + 1] = hopping[i]
        H[i + 1, i] = np.conj(hopping[i])

    # inter-cell hopping (セル間ホッピング: Bloch phase)
    # 1原子、2原子の場合でも上書きされないよう "+=" を使用
    phase = np.exp(1j * k * a_cell)
    H[0, norb - 1] += hopping[-1] * np.conj(phase)
    H[norb - 1, 0] += hopping[-1] * phase

    return H

# ============================
# バンド計算
# ============================

def calculate_bands(klist, onsite, hopping, a=1.0):
    bands = []
    for k in klist:
        H = hamiltonian_k(k, onsite, hopping, a)
        eigvals = np.linalg.eigvalsh(H)
        bands.append(eigvals)
    return np.array(bands)


# ============================
# main
# ============================
def main():
    params = initialize()

    onsite = params["onsite"]
    hopping = params["hopping"]
    a = params["a"]
    nk = params["nk"]
    norb = len(onsite)

    # --- 単位格子長に基づいた第一ブリルアンゾーンの設定 ---
    # BZの範囲は [-π/L, π/L] (L = norb * a)
    a_cell = norb * a
    k_limit = np.pi / a_cell
    klist = np.linspace(-k_limit, k_limit, nk)

    bands = calculate_bands(klist, onsite, hopping, a)

    # --- プロット用に正規化 (k * a_cell / π) ---
    # これにより横軸が -1 から 1 (BZ端) になる
    k_normalized = klist / k_limit

    save_band_data("band_1d_tb.csv", k_normalized, bands)
    print(f"原子数 {norb} の計算が完了しました。")

    plot_bands(k_normalized, bands)


def save_band_data(filename, klist, bands):
    norb = bands.shape[1]
    df = pd.DataFrame(bands, columns=[f"band_{i}" for i in range(norb)])
    df.insert(0, "k_normalized", klist)
    df.to_csv(filename, index=False)


def plot_bands(klist, bands):
    plt.figure(figsize=(6, 4))
    for i in range(bands.shape[1]):
        plt.plot(klist, bands[:, i], label=f"Band {i}")

    plt.xlabel(r"$k / (\pi / L)$")
    plt.ylabel(r"$E(k)$")
    plt.title("1D Tight-Binding Band Structure")
    plt.axvline(0, color="gray", lw=0.5)
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    main()