#!/usr/bin/env python3

import re
import numpy as np
import matplotlib.pyplot as plt

# ============================================================
# Beállítások
# ============================================================

OUTFILE = "caffeine_uvvis.log"

SIGMA = 10.0

WL_MIN = 150
WL_MAX = 350

# ============================================================
# TD-DFT átmenetek kiolvasása Gaussian logból
# ============================================================

transitions = []

pattern = re.compile(
    r"Excited State\s+\d+.*?([0-9]+\.[0-9]+)\s+nm\s+f=([0-9]+\.[0-9]+)"
)

with open(OUTFILE, "r", encoding="utf-8", errors="ignore") as f:
    for line in f:
        m = pattern.search(line)

        if m:
            wavelength = float(m.group(1))
            oscillator = float(m.group(2))

            transitions.append(
                (wavelength, oscillator)
            )

print(f"{len(transitions)} átmenet beolvasva.")

# ============================================================
# Spektrum generálása
# ============================================================

x = np.linspace(WL_MIN, WL_MAX, 5000)

spectrum = np.zeros_like(x)

for wl, fosc in transitions:

    spectrum += fosc * np.exp(
        -(x - wl) ** 2 / (2 * SIGMA**2)
    )

if np.max(spectrum) > 0:
    spectrum /= np.max(spectrum)

# ============================================================
# Ábra
# ============================================================

plt.figure(figsize=(9, 5))

plt.plot(
    x,
    spectrum,
    color="blue",
    linewidth=2.5,
    label="TD-DFT UV–Vis spektrum"
)

max_f = max(f for _, f in transitions)

# Diszkrét elektronikus átmenetek

for i, (wl, fosc) in enumerate(transitions):

    plt.axvline(
        wl,
        color="red",
        linestyle="--",
        linewidth=1.2,
        alpha=0.6,
        label="Diszkrét elektronikus átmenetek" if i == 0 else None
    )

plt.xlabel("Hullámhossz (nm)")
plt.ylabel("Relatív abszorbancia")
plt.title("Koffein UV–Vis spektrum (Gaussian TD-DFT)")
plt.xlim(WL_MIN, WL_MAX)
plt.ylim(0, 1.05)
plt.grid(alpha=0.3)
plt.legend()

plt.tight_layout()

plt.savefig(
    "caffeine_uvvis_spectrum.png",
    dpi=300,
    bbox_inches="tight"
)

plt.show()

print("Ábra mentve: caffeine_uvvis_spectrum.png")