#!/usr/bin/env python3
"""
Analyze NEB results and calculate reaction rate
"""

import numpy as np
import os
import matplotlib.pyplot as plt
from ase.io import read
from ase.neb import NEBTools

def read_neb_energies():
    """Read energies from OUTCAR files"""
    energies = []
    forces = []
    
    # Find all image directories
    dirs = sorted([d for d in os.listdir('.') if os.path.isdir(d) and d.isdigit()])
    
    for d in dirs:
        outcar = os.path.join(d, 'OUTCAR')
        if os.path.exists(outcar):
            with open(outcar, 'r') as f:
                lines = f.readlines()
                for line in lines:
                    if 'free  energy   TOTEN' in line:
                        energy = float(line.split()[-2])
                        energies.append(energy)
                    elif 'TOTAL-FORCE' in line:
                        # Read forces (simplified)
                        pass
    
    return np.array(energies)

def calculate_reaction_rate(E_a, T=300, A=1e13):
    """
    Calculate reaction rate using Arrhenius equation
    k = A * exp(-E_a / (k_B * T))
    
    Parameters:
    E_a: activation energy (eV)
    T: temperature (K)
    A: pre-exponential factor (s^-1)
    
    Returns:
    k: reaction rate (s^-1)
    """
    k_B = 8.617333262145e-5  # eV/K
    k = A * np.exp(-E_a / (k_B * T))
    return k

def main():
    # Read energies
    energies = read_neb_energies()
    
    if len(energies) == 0:
        print("No OUTCAR files found!")
        return
    
    print(f"Found {len(energies)} images")
    print(f"Energies (eV): {energies}")
    
    # Calculate activation energy
    E_a = energies.max() - energies[0]
    print(f"\nActivation energy E_a = {E_a:.3f} eV")
    
    # Calculate reaction rates at different temperatures
    temperatures = [250, 300, 350, 400, 450, 500]
    print("\nReaction rates (s^-1):")
    for T in temperatures:
        k = calculate_reaction_rate(E_a, T)
        print(f"T = {T} K: k = {k:.2e} s^-1")
    
    # Plot energy profile
    plt.figure(figsize=(10, 6))
    plt.plot(range(len(energies)), energies - energies[0], 'o-', linewidth=2)
    plt.xlabel('Reaction coordinate (image number)')
    plt.ylabel('Energy (eV)')
    plt.title(f'NEB Energy Profile (E_a = {E_a:.3f} eV)')
    plt.grid(True, alpha=0.3)
    plt.savefig('neb_energy_profile.png', dpi=300)
    plt.show()
    
    # Save results
    with open('neb_results.txt', 'w') as f:
        f.write(f"Activation energy: {E_a:.6f} eV\n")
        f.write(f"Energy barrier: {energies.max() - energies.min():.6f} eV\n")
        f.write(f"Reaction energy: {energies[-1] - energies[0]:.6f} eV\n")
        f.write("\nReaction rates:\n")
        for T in temperatures:
            k = calculate_reaction_rate(E_a, T)
            f.write(f"T = {T} K: k = {k:.6e} s^-1\n")

if __name__ == '__main__':
    main()