#!/usr/bin/env python3

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

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

OUTFILE = "caffeine_tddft.out"

# Gauss-kiszélesítés szórása (nm)
SIGMA = 10.0

# Hullámhossz-tartomány
WL_MIN = 120
WL_MAX = 300

# ============================================================
# TDDFT átmenetek kiolvasása az ORCA outputból
# ============================================================

transitions = []

with open(OUTFILE, "r", encoding="utf-8", errors="ignore") as f:
    lines = f.readlines()

inside_block = False

for line in lines:

    if "ABSORPTION SPECTRUM VIA TRANSITION ELECTRIC DIPOLE MOMENTS" in line:
        inside_block = True
        continue

    if inside_block:

        if re.match(r"\s*0-1A", line):

            parts = line.split()

            try:
                wavelength = float(parts[5])  # nm
                fosc = float(parts[6])

                if fosc > 0:
                    transitions.append((wavelength, fosc))

            except (ValueError, IndexError):
                pass

# ============================================================
# 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))
    )

# Normalizálás
if np.max(spectrum) > 0:
    spectrum /= np.max(spectrum)

# ============================================================
# Ábrázolás
# ============================================================

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

# Folytonos UV–Vis spektrum
plt.plot(
    x,
    spectrum,
    color="blue",
    linewidth=2.5,
    label="Kiszélesített UV–Vis spektrum"
)

# Diszkrét elektronikus átmenetek
max_fosc = max(f for _, f in transitions)

ymax = plt.ylim()[1]

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

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

plt.xlabel("Hullámhossz (nm)", fontsize=12)
plt.ylabel("Relatív abszorbancia", fontsize=12)
plt.title("TDDFT számított UV–Vis spektrum - Koffein", fontsize=14)

plt.xlim(WL_MIN, WL_MAX)
plt.ylim(0, 1.05)

plt.grid(True, alpha=0.3)

plt.legend(
    frameon=True,
    loc="upper right"
)

plt.tight_layout()

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

plt.show()

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