import os.path
import re

from .const import __ATOM_LIST__

from . import mol





def str_atom(_atom):
    """
    str_atom(_atom: int) -> atom: str

    Convert integer atom to string atom.

    in:
    _atom: Integer with the atomic number.

    out:
    atom: String with the atomic element.

    """

    atom = __ATOM_LIST__[_atom - 1]
    return atom


def load(_path):
    """
    load(_path: str) -> molrepr: Mol

    Check input file extension and call appropriate load function. Return Mol Object 
    generated by the load function.

    in:
    _path: String with the path of the input file.

    out:
    molrepr: Mol Object containing the molecular information (coordinates and elements).

    """

    filetype = determine_filetype(_path)
    if filetype == 'xyz':
        molrepr = load_xyz(_path)

    if filetype == 'in':
        molrepr = load_gaussian(_path)
        
    return molrepr


def determine_filetype(_path):
    """
    determine_filetype(_path: str) -> filetype: str

    Return the extension of the input file.

    in:
    _path: String with the path of the input file.

    out:
    filetype: String containing the extentsion of the input file.

    """

    filetype = os.path.splitext(_path)[1][1:]
    return filetype


def load_xyz(_path):
    """
    load_xyz(_path: str) -> molrepr: Mol

    Load molecule from an XYZ input file and initialize a Mol Object for it.

    in:
    _path: String with the path of the input file.

    out:
    molrepr: Mol Object containing the molecular information (coordinates and elements).

    """

    molrepr = []
    with open(_path, "r") as file:
        for line_number, line in enumerate(file):
            if line_number > 1:
                atomic_symbol, x, y, z = line.split()
                if not atomic_symbol.isalpha():
                    atomic_symbol = int(atomic_symbol)
                    atomic_symbol = str_atom(atomic_symbol)
                molrepr.append([atomic_symbol.capitalize(), float(x), float(y), float(z)])

    molrepr = mol.Mol(molrepr)
    return molrepr


def load_gaussian(_path):
    """
    load_gaussian(_path: str) -> molrepr: Mol

    Load molecule from an gaussian input file and initialize a Mol Object for it.

    in:
    _path: String with the path of the input file.

    out:
    molrepr: Mol Object containing the molecular information (coordinates and elements).
    
    """
    
    molrepr = []
    data = ""
    blank_line_regex = r"(?:\r?\n){2,}" # match multiple (at least one) empty line
    
    with open(_path, "r") as file:
        for line in file:
            if line.startswith('#') or line.startswith('%'): # remove link0 and route section
                continue
            else:
                data += line
    data = re.split(blank_line_regex, data.strip()) # remove leading and trailing whitespace and then split along newlines

    # data[0] should now be the title field
    # data[1] should be the coordinates
    lines = data[1].splitlines()[1:]

    for line in lines:
        atomic_symbol, x, y, z = line.split()
        if not atomic_symbol.isalpha():
            atomic_symbol = int(atomic_symbol)
            atomic_symbol = str_atom(atomic_symbol)
        molrepr.append([atomic_symbol.capitalize(), float(x), float(y), float(z)])

    molrepr = mol.Mol(molrepr)
    return molrepr
