kinetics

import numpy as np
import math
from dataclasses import dataclass

# Physical constants
R = 0.25  # O2 N2 ratio
PRESSURE = 1  # pressure
KB = 1.380649E-23  # Boltzmann constant
H = 6.62607015E-34  # Planck constant
C = 3.0E+10  # speed of light, cm/s
DT = 1.0E-8  # T step for finite differences
NA = 6.02214076E+23  # Avogadro number
NU = 1.0  # number of moles
P_ATM = 101325.0  # atmospheric pressure
EV = 1.60218E-19  # eV -> J translation coefficient


@dataclass
class Species:
    constants: np.ndarray

def calculate_thermodynamic_properties(temperature: float, species_idx: int, species_data: np.ndarray) -> dict:
    """
    Calculate thermodynamic properties for a given species at specified temperature.

    Args:
        temperature: Temperature in Kelvin
        species_idx: Species index (0 for N, 1 for N2)
        species_data: Array containing species data

    Returns:
        Dictionary containing:
        - Q_in: Internal partition function
        - Q_tr: Translational partition function
        - Q_total: Total partition function
        - S: Entropy (J/mol/K)
        - H_H0: Enthalpy difference from 0K (kJ/mol)
    """
    # Create Species object for the selected species
    species = Species(constants=species_data[species_idx])

    # Calculate Q_in
    if species.constants[1] == 0:
        Q_rot = 1.0
    else:
        Q_rot = (8.0 * math.pi**2 * species.constants[1] * KB * temperature) / (H**2 * species.constants[0])

    if species.constants[2] == 0:
        Q_vib = 1.0
    else:
        Q_vib = 1.0 / (1.0 - math.exp((-H * C * species.constants[2]) / (KB * temperature)))

    if species.constants[3] == 0:
        Q_el = 1.0
    else:
        Q_el = sum(
            species.constants[3+i] * math.exp(-H * C * species.constants[22+i] / (KB * temperature))
            for i in range(19)
        )

    Q_in = Q_rot * Q_vib * Q_el

    # Calculate Q_tr
    Q_tr = (NU * KB * NA * temperature / (PRESSURE * P_ATM)) * \
           (2 * math.pi * species.constants[41] * KB * temperature / (H * H))**(1.5)

    # Calculate Q_total
    Q_total = Q_in * Q_tr

    # Calculate Entropy [J/(mol*K)]
    S = KB * NA * (
        math.log(Q_in) + math.log(Q_tr)
    ) + KB * NA * temperature * (
        (Q_tr_calc(species, temperature + DT) - Q_tr_calc(species, temperature - DT)) /
        (2.0 * DT * Q_tr) +
        (Q_in_calc(species, temperature + DT) - Q_in_calc(species, temperature - DT)) /
        (2.0 * DT * Q_in)
    ) - KB * NA * math.log(NA)

    # Calculate H-H0 [J/mol]
    H_H0 = KB * NA * temperature * (
        (temperature / Q_in) *
        (Q_in_calc(species, temperature + DT) - Q_in_calc(species, temperature - DT)) /
        (2.0 * DT) + 2.5
    )

    return {
        'Q_in': Q_in,
        'Q_tr': Q_tr,
        'Q_total': Q_total,
        'S': S,
        'H_H0': H_H0
    }

# Helper functions for finite difference calculations
def Q_in_calc(species: Species, T: float) -> float:
    """Calculate Q_in at a specific temperature (helper function for derivatives)"""
    if species.constants[1] == 0:
        Q_rot = 1.0
    else:
        Q_rot = (8.0 * math.pi**2 * species.constants[1] * KB * T) / (H**2 * species.constants[0])

    if species.constants[2] == 0:
        Q_vib = 1.0
    else:
        Q_vib = 1.0 / (1.0 - math.exp((-H * C * species.constants[2]) / (KB * T)))

    if species.constants[3] == 0:
        Q_el = 1.0
    else:
        Q_el = sum(
            species.constants[3+i] * math.exp(-H * C * species.constants[22+i] / (KB * T))
            for i in range(19)
        )

    return Q_rot * Q_vib * Q_el

def Q_tr_calc(species: Species, T: float) -> float:
    """Calculate Q_tr at a specific temperature (helper function for derivatives)"""
    return (NU * KB * NA * T / (PRESSURE * P_ATM)) * \
           (2 * math.pi * species.constants[41] * KB * T / (H * H))**(1.5)