from glob import glob
import cclib
import periodictable as pt
from .geometries import xyz_molecule, xyz_atom
import os
import numpy as np


class g16_input:
    """ a base level g16 input file
    """

    def __init__(self, input_line:str, geometry:xyz_molecule, file_name:str,charge:int,spin_mult:int, nproc = 36,mem = 5):
        self.checkpoint = file_name[:-3] + "chk"
        self.file_name = file_name
        self.geometry = geometry
        self.mem = mem
        self.nproc = nproc
        self.input_line = " ".join(input_line.lower().split()[1:])
        self.extra = ""
        self.title_card = "Title Card"
        self.charge = charge
        self.spin_mult = spin_mult


    def write_file(self):
        """writes a file associated with the g16 input instance
        """
        out_string = f"""%chk={self.checkpoint}
%nproc={self.nproc}
%mem={self.mem}GB
#p {self.input_line}

{self.title_card}

{self.charge} {self.spin_mult}
{str(self.geometry)}

{self.extra}


"""
        with open(self.file_name,"w") as file:
            file.write(out_string)

    @classmethod
    def from_file(cls,file_name):
        collect_input = False
        input_line = ""
        nproc = 36
        mem = 5
        charge_collected = False
        charge = 0
        spin_mult = 1
        Title_Card = False
        blank_counter = 0 
        geometry = xyz_molecule([])

        with open(file_name,"r") as file:
            for line in file:
                true_line = line.strip(" \n")

                if "%mem" in true_line:
                    mem = int(true_line.split("=")[1][:-2])

                if "%nproc" in true_line:
                    nproc = int(true_line.split("=")[1])

                if collect_input == True and blank_counter == 0:
                    input_line = input_line + " " + line.strip(" \n")

                if len(true_line) > 0 and "#" == true_line[0]:
                    collect_input = True
                    input_line = input_line + line.strip("\n")

                if blank_counter == 2 and charge_collected and true_line != "":
                    geometry.add_atom(xyz_atom.from_string(true_line))

                if blank_counter == 2 and charge_collected == False:
                    true_line = line.strip(" \n\t")
                    charge_mult = true_line.split()
                    charge = charge_mult[0]
                    spin_mult = charge_mult[1]
                    charge_collected = True
                
                if true_line.strip("\t") == "":
                    blank_counter = blank_counter + 1

                                 

        return cls(input_line,geometry,file_name,charge,spin_mult,nproc = nproc,mem = mem)
    
    def __str__(self):
        out_string = f"""%chk={self.checkpoint}
%nproc={self.nproc}
%mem={self.mem}GB
#p {self.input_line}

{self.title_card}

{self.charge} {self.spin_mult}
{str(self.geometry)}

{self.extra}


"""
        return out_string



def stationary_point_count(output_file_pattern:str="*.log"):
    """Evaluates files for number of stationary points
Parameters
----------
output_file_pattern (str): a file pattern to match and evaluate stationary points.

Returns
-------
list[tuple[str,int]]: a list of tuples containing the names of the log files and the number of stationary points"""

    results ={}

    for output_file in glob(output_file_pattern):
        results[output_file] = 0
        with open(output_file) as current_file:
            for line in current_file:
                results[output_file] += "Stationary point" in line

    results = [(key,results[key]) for key in results.keys()]
    if len(results) == 0:
        raise FileExistsError(f"No files found to match pattern '{output_file_pattern}'.")
    results = sorted(results,key = lambda x: x[1])
    return results


def get_n_geometry(log_file_name:str,n_geometry:int = -1):
    """Grabs the nth geometry (zero indexed) from a log file. Defaults to the last geometry
Parameters
----------
log_file_name (str): the path to a log file

n_geometry (int): the index of the desired geometry

Returns
-------
xyz_molecule"""

    outfile = cclib.io.ccread(log_file_name)
    atom_nums = outfile.atomnos
    correct_geom = outfile.geometries[n_geometry]
    return xyz_molecule.from_cclib_vals(atom_nums,correct_geom)

def rerun_from_last_geom(file_to_rerun):
    """reruns a job from the last geometry found in the log file.
Parameters
----------
file_to_rerun (str): the path to a log file

Returns
-------
None: reformats and resubmits g16 input"""
    new_geometry = get_n_geometry(file_to_rerun)
    old_input = g16_input.from_file(file_to_rerun[:-3]+"gjf")
    old_input.geometry = new_geometry
    old_input.write_file()
    print(f"Rerunning {file_to_rerun[:-4]}")
    os.system(f"sbatch {file_to_rerun[:-3]+"s"}")

def quick_rerun_optimizations(log_file_pattern:str="*.log"):
    """Reruns any files with only one stationary point, and notifies about files with 0 stationary points
Parameters
----------

log_file_pattern (str): a pattern for the path of files to rerun.

Returns
-------
None: reruns files with only 1 stationary point and prints advice for files with 0."""
    statp_eval = stationary_point_count(log_file_pattern)

    for file,statp_count in statp_eval:
        if statp_count == 1:
            rerun_from_last_geom(file)
        elif statp_count == 0:
            print(f"{file} requires manual check, not Stationary point found")


class orca_cube_file:

    def __init__(self,n_atoms:int,n_points:int,volume_origin:np.array,x_spacing:np.array,y_spacing:np.array,z_spacing:np.array,atoms:np.array,nuclear_charges:np.array,atom_coords:np.array,data:np.array):
        self.n_atoms = n_atoms
        self.volume_origin = volume_origin
        self.n_points = n_points
        self.x_spacing = x_spacing
        self.y_spacing = y_spacing
        self.z_spacing = z_spacing
        self.atoms = atoms
        self.nuclear_charges =nuclear_charges
        self.atom_coords = atom_coords
        self.data = data

    @classmethod
    def read_file(cls,file_name:str):
        with open(file_name,"r") as file:
            file_lines=file.readlines()
        
        n_atoms = int(file_lines[2].strip().split()[0])
        volume_origin = np.array(file_lines[2].strip().split()[1:],float)
        n_points = int(file_lines[3].strip().split()[0])
        x_spacing = np.array(file_lines[3].strip().split()[1:],float)
        y_spacing = np.array(file_lines[4].strip().split()[1:],float)
        z_spacing = np.array(file_lines[5].strip().split()[1:],float)
        atoms = np.zeros(n_atoms,int)
        atom_coords = np.zeros(shape=(n_atoms,3),dtype=float)
        nuclear_charges = np.zeros(dtype=int,shape=n_atoms)

        for index,value in enumerate(file_lines[6:6+n_atoms]):
            clean_values = value.strip().split()
            atoms[index] = int(clean_values[0])
            nuclear_charges[index] = float(clean_values[1]) 
            atom_coords[index][0] = float(clean_values[2])
            atom_coords[index][1] = float(clean_values[3])
            atom_coords[index][2] = float(clean_values[4])
        
        tracker = 0
        final_data = n_points**3
        actual_data = np.zeros(dtype=float,shape=n_points**3)
        for line in file_lines[7+n_atoms:-2]:
            clean_line = [float(i) for i in line.strip().split()]

            for value in clean_line:
                actual_data[tracker] = value
                tracker +=1
        del file_lines

        return cls(n_atoms,n_points,volume_origin,x_spacing,y_spacing,z_spacing,atoms,nuclear_charges,atom_coords,actual_data)
    
    def gen_fukui(self,other,electron_gap =1):
        """Finds the fukui function of two cube files from orca (subtract more electrons from less.)"""
        
        if not (self.volume_origin == other.volume_origin).all():
            print("Volume arrays are not equivalent")
        elif self.n_atoms != other.n_atoms:
            print("Number of atoms not equal")
        elif self.n_points != other.n_points:
            print("Number of points not equal")
        elif not (self.atoms == other.atoms).all():
            print("Atoms are not the same")
        elif not (self.atom_coords == other.atom_coords).all():
            print("Atoms coordinates are not the same")

        new_data = (self.data - other.data)/electron_gap

        return orca_cube_file(self.n_atoms,
                   self.n_points,
                   self.volume_origin.copy(),
                   self.x_spacing.copy(),
                   self.y_spacing.copy(),
                   self.z_spacing.copy(),
                   self.atoms.copy(),
                   self.nuclear_charges.copy(),
                   self.atom_coords.copy(),
                   new_data)

    

    def write_file(self,file_name:str,description:str="Electron Density"):
        with open(file_name,"w") as cube_file:
            cube_file.write("Cube data generated by ORCA and modified by Ben python\n")
            cube_file.write(description+"\n")
            # Write number of atoms and volume vector
            section_one = f""" {self.n_atoms: >3}  {self.volume_origin[0]:.6f}  {self.volume_origin[1]:.6f}  {self.volume_origin[2]:.6f}
 {self.n_points: >3}  {self.x_spacing[0]:.6f}  {self.x_spacing[1]:.6f}  {self.x_spacing[2]:.6f}
 {self.n_points: >3}  {self.y_spacing[0]:.6f}  {self.y_spacing[1]:.6f}  {self.y_spacing[2]:.6f}
 {self.n_points: >3}  {self.z_spacing[0]:.6f}  {self.z_spacing[1]:.6f}  {self.z_spacing[2]:.6f}\n"""
            cube_file.write(section_one)

            # Now write atom data
            full_atoms = ""
            for index,coords in enumerate(self.atom_coords):
                atom_type = self.atoms[index]
                nuclear_charge = self.nuclear_charges[index]
                atom_data_template = f" {atom_type: >3} {nuclear_charge:.6f}  {coords[0]:.6f} {coords[1]:.6f} {coords[2]:.6f}\n"
                full_atoms += atom_data_template
            cube_file.write(full_atoms)

            for index,value in enumerate(self.data):
                cube_file.write(f" {value:.6f} ")
                if index % 6 == 5:
                    cube_file.write("\n")

            cube_file.write("\n\n")


    