#!/usr/bin/env python3
"""
NEB path generator for VASP
Usage: python neb.py -n 5 00/POSCAR 0N/POSCAR
"""

import numpy as np
import os
import sys
import argparse
from ase import Atoms
from ase.io import read, write
from ase.neb import NEB
from ase.calculators.vasp import Vasp

def create_interpolated_images(initial, final, n_images):
    """Create NEB images between initial and final structures"""
    images = [initial]
    
    # Create intermediate images
    for i in range(1, n_images):
        image = initial.copy()
        # Linear interpolation
        image.positions = initial.positions + i/(n_images+1) * (final.positions - initial.positions)
        images.append(image)
    
    images.append(final)
    
    # Create NEB object for interpolation
    neb = NEB(images)
    neb.interpolate()
    
    return images

def main():
    parser = argparse.ArgumentParser(description='Generate NEB images for VASP')
    parser.add_argument('-n', '--n_images', type=int, default=3,
                       help='Number of intermediate images')
    parser.add_argument('initial_poscar', help='Path to initial POSCAR')
    parser.add_argument('final_poscar', help='Path to final POSCAR')
    parser.add_argument('--method', default='idpp',
                       help='Interpolation method (idpp or linear)')
    
    args = parser.parse_args()
    
    # Read structures
    initial = read(args.initial_poscar, format='vasp')
    final = read(args.final_poscar, format='vasp')
    
    print(f"Initial structure: {len(initial)} atoms")
    print(f"Final structure: {len(final)} atoms")
    
    # Check atom count
    if len(initial) != len(final):
        print("ERROR: Initial and final structures must have same number of atoms!")
        sys.exit(1)
    
    # Create images
    images = create_interpolated_images(initial, final, args.n_images)
    
    # Create directories and write POSCAR files
    for i, image in enumerate(images):
        dir_name = f"{i:02d}"
        os.makedirs(dir_name, exist_ok=True)
        write(f"{dir_name}/POSCAR", image, format='vasp', direct=True)
        print(f"Created {dir_name}/POSCAR")
    
    # Write NEB information
    with open('NEB_INFO.txt', 'w') as f:
        f.write(f"NEB calculation with {len(images)} images\n")
        f.write(f"Initial: {args.initial_poscar}\n")
        f.write(f"Final: {args.final_poscar}\n")
        f.write(f"Intermediate images: {args.n_images}\n")
    
    print(f"\nCreated {len(images)} images in directories 00-{len(images)-1:02d}")
    print("Next steps:")
    print("1. Copy INCAR, KPOINTS, POTCAR to all image directories")
    print("2. Run VASP in the main directory")
    print("3. Check convergence in OUTCAR files")

if __name__ == '__main__':
    main()