import re


def extract_geometries(file_path, output_xyz_file, geometry_index=-1):
    with open(file_path, 'r') as file:
        content = file.read()

        # Define the start and end patterns
        start_pattern = r"Output coordinates in a\.u\. \(scale by  1\.000000000 to convert to a\.u\.\)"
        end_pattern = r"Effective nuclear repulsion energy \(a\.u\.\)"

        # Find all geometry sections in the content using regex
        geometry_patterns = re.finditer(
            f"{start_pattern}(.*?){end_pattern}", content, re.DOTALL)

        geometries = []

        for geometry_match in geometry_patterns:
            geometry_content = geometry_match.group(1)

            # Extract rows from the geometry section
            rows = re.findall(
                r"\s+(\d+)\s+(\S+)\s+([\d.-]+)\s+([\d.-]+)\s+([\d.-]+)\s+([\d.-]+)", geometry_content)

            # Convert rows to a list of dictionaries
            geometry_data = []
            for row in rows:
                geometry_data.append({
                    "Atom_No": int(row[0]),
                    "Atom_Tag": row[1],
                    "Charge": float(row[2]),
                    "X": float(row[3]),
                    "Y": float(row[4]),
                    "Z": float(row[5]),
                })

            geometries.append(geometry_data)

    if geometries:
        # Check if the specified geometry index is within range
        if geometry_index >= -len(geometries) and geometry_index < len(geometries):
            selected_geometry = geometries[geometry_index]

            # Write the selected geometry to an XYZ file
            with open(output_xyz_file, 'w') as file:
                # Write the number of atoms
                file.write(f"{len(selected_geometry)}\n")

                # Write a comment line
                file.write("Generated by NWChem Geometry Extractor\n")

                # Write atom coordinates
                for row in selected_geometry:
                    file.write(
                        f"{row['Atom_Tag']} {row['X']} {row['Y']} {row['Z']}\n")

            print(
                f"XYZ file created for Geometry {geometry_index + 1}: {output_xyz_file}")
        else:
            print("Invalid geometry index.")
    else:
        print("No geometries found in the NWChem output file.")


# Example usage:
if __name__ == "__main__":
    input_file_path = 'test1.out'
    output_xyz_file = 'extracted_geometry.xyz'
    geometry_index_to_extract = -1  # Specify the desired geometry index
    extract_geometries(input_file_path, output_xyz_file,
                       geometry_index_to_extract)
