#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Notes
-----

This script contains the functions used to extract data from Cube files.
"""

import codecs
import numpy as np
from rgmol.objects import *
from rgmol.general_function import find_bonds


##########################
## Extraction functions ##
##########################




def _extract_cube(file):
    """
    extract_cube(file)

    Extracts the information from a cube file

    Parameters
    ----------
        file : str

    Outputs
    -------
        voxel_origin : list
            the origin of the voxel
        voxel_matrix : list
            the voxel matrix
        atomic_number : list
            the list of the atomic number
        atom_position : list
            the list of the position of the atoms
        cube : ndarray
            the cube
    """
    #flag value correspond to the number of line for each part on the file
    flag_header = 2
    flag_atoms = 1
    flag_voxel = 3
    flag_other = 1 #For the line between the atoms and the cube. Only on mo cube generated by orca ?
    number_points = 1
    list_number_points = []
    voxel_matrix = []

    count_line = 0
    count_square = 0
    line = []
    square = []

    atom_number = []
    atom_position = []
    cube = []

    for file_line in codecs.open(file, 'r',encoding="utf-8"):
        L = file_line.split()
        if flag_header:
            flag_header -= 1

        elif flag_atoms:
            number_atoms = abs(int(L[0]))
            voxel_origin = [float(L[1]),float(L[2]),float(L[3])]
            flag_pos = number_atoms
            flag_atoms -= 1

        elif flag_voxel:
            number_points *= int(L[0])
            list_number_points.append(int(L[0]))
            voxel_matrix.append([float(L[1]),float(L[2]),float(L[3])])
            flag_voxel -=1

        elif flag_pos:
            atom_number.append(int(L[0]))
            atom_position.append([float(L[2]),float(L[3]),float(L[4])])
            flag_pos -= 1

        elif flag_other:
            number_points_x = list_number_points[0]
            number_squares_y = list_number_points[1]

            flag_other -= 1
            if len(L) != 2:#There is not the strange line from the MO cube files
                for l in L:
                    line.append(float(l))
                count_line += len(L)

        else:
            for l in L:
                line.append(float(l))
            count_line += len(L)

            if count_line == number_points_x:
                square.append(line)
                line = []
                count_line = 0
                count_square += 1

                if count_square == number_squares_y:
                    cube.append(square)
                    square = []
                    count_square = 0
    return voxel_origin,voxel_matrix,atom_number,atom_position,np.array(cube)


def extract(file,do_find_bonds=0):
    """
    extract(file,do_find_bonds=0)

    Extracts a molecule from a cube file.
    As the bonds are not defined in a cube file, one can use do_find_bonds to use an algorithm that tries to find bonds. It is still in WIP.

    Parameters
    ----------
        file : str
        do_find_bonds : bool, optional
            if one wants to use an algorithm to find the bonds (WIP)

    Returns
    -------
        mol : molecule
            the molecule extracted form the cube file with the voxel_origin, voxel_matrix and cube as properties

    """

    voxel_origin,voxel_matrix,atom_number,atom_position,cube = _extract_cube(file)


    list_atoms = []
    nicknaming = 0
    for prop in zip(atom_number,atom_position):
        atom_x = atom(prop[0],prop[1],nickname=str(nicknaming))
        list_atoms.append(atom_x)
        nicknaming+=1



    mol = molecule(list_atoms,[],file=file,properties={"voxel_origin":voxel_origin,"voxel_matrix":voxel_matrix,"cube":cube})

    if do_find_bonds:
        mol.bonds = find_bonds(mol)

    return mol
