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

This script adds writing functions
"""

import numpy as np
import scipy as sp
from rgmol.objects import *
import os

def _write_vector_for_cube(file,vector):
    """Writes a vector in a file"""
    for elem in vector:
        long = len(str(abs(int(elem))))
        file.write(" "*(4-long))
        file.write("{:9.6f}".format(elem))

def _write_data_for_cube(file,vector):
    """Writes a vector in a file"""
    for elem in vector:
        m = int(np.sign(elem)/2-1/2)
        file.write(" "*(3+m))
        file.write("{:9.5e}".format(elem))

def _get_file_location(file):
    """gets the location of the file of the output of the molecule"""
    file_splitted = file.split("//")[:-1]
    file_location = ""

    for f in file_splitted:
        file_location += f
        file_location +="//"

    return file_location



def write_cube(self,cube,file_name,description="cube"):
    """
    write_cube(cube,file_name,description="cube")

    Writes a cube file of a 3D vector inside a folder called rgmol which will be located inside the folder containing the molecule.

    Parameters
    ----------
        cube : ndarray
            The 3D vector to be stored in the cube file
        file_name : str
            The name of the file
        description : str
            Description of the file inside the cube file

    Returns
    -------
        None
            The cube file is written
    """

    file_location = _get_file_location(self.file)

    if not "rgmol" in os.listdir(file_location):
        os.mkdir(file_location+"rgmol")

    with open(file_location+"rgmol//"+file_name,"w") as file:
        file.write("Cube data generated by rgmol\n")
        file.write(description+"\n")

        #Number atoms and voxel origin
        num_atoms = str(len(self.atoms))
        voxel_origin = self.properties["voxel_origin"]

        file.write(" "*(5-len(num_atoms)) +num_atoms)
        _write_vector_for_cube(file,voxel_origin)

        file.write("\n")

        #Number of points and voxel matrix
        voxel_matrix = self.properties["voxel_matrix"]
        grid_points = self.properties["grid_points"]

        grid_points_0 = str(grid_points[0])
        file.write(" "*(5-len(grid_points_0)) + grid_points_0)
        _write_vector_for_cube(file,voxel_matrix[0])
        file.write("\n")
        grid_points_1 = str(grid_points[1])
        file.write(" "*(5-len(grid_points_1)) + grid_points_1)
        _write_vector_for_cube(file,voxel_matrix[1])
        file.write("\n")
        grid_points_2 = str(grid_points[2])
        file.write(" "*(5-len(grid_points_2)) + grid_points_2)
        _write_vector_for_cube(file,voxel_matrix[2])
        file.write("\n")

        #Write atom positions
        for atom in self.atoms:
            atomic_number = atom.atomic_number
            file.write(" "*(5-len(str(atomic_number)))+str(atomic_number))
            file.write(" "*(4-len(str(atomic_number)))+"{:9.6f}".format(atomic_number))
            _write_vector_for_cube(file,atom.pos)
            file.write("\n")

        #Write cube
        for vec_x in cube:
            for vec_y in vec_x:
                line_length = grid_points[2]
                count = 0
                while count < line_length:
                    _write_data_for_cube(file,vec_y[count:count+6])
                    file.write('\n')
                    count+=6




molecule.write_cube = write_cube




