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

This script adds writing and reading functions
"""

import os
import time
import numpy as np
import scipy as sp
import zipfile as zf
from rgmol.objects import *
import rgmol
import rgmol.grid


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_vector_for_cube(Lines,vector):
    """Writes a vector in a file"""
    for elem in vector:
        long = len(str(abs(int(elem))))
        Lines.append(" "*(4-long))
        Lines.append("{:9.6f}".format(elem))

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

def _write_voxel(file_name,grid_points,delta):
    """Writes the voxel parameters in a file"""

    with open(file_name,"w") as file:
        file.write("Grid Points :\n")
        file.write("{} {} {}\n".format(grid_points[0],grid_points[1],grid_points[2]))
        file.write("Delta :\n")
        file.write("{}".format(delta))


def _read_voxel(file_name):
    """Reads the voxel parameters from a file written using the write_voxel function"""

    count = 0
    voxel_matrix=[]
    for line in open(file_name,"r"):
        lsplit = line.split()
        if count == 1:
            grid_points = [int(lsplit[0]),int(lsplit[1]),int(lsplit[2])]
        elif count == 3:
            delta = float(lsplit[0])
        count+=1
    return grid_points,delta


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
    """
    Lines = []
    Lines.append("Cube data generated by rgmol\n")
    Lines.append(description+"\n")

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

    Lines.append(" "*(5-len(num_atoms)) +num_atoms)
    _write_vector_for_cube(Lines,voxel_origin)

    Lines.append("\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])
    Lines.append(" "*(5-len(grid_points_0)) + grid_points_0)
    _write_vector_for_cube(Lines,voxel_matrix[0])
    Lines.append("\n")
    grid_points_1 = str(grid_points[1])
    Lines.append(" "*(5-len(grid_points_1)) + grid_points_1)
    _write_vector_for_cube(Lines,voxel_matrix[1])
    Lines.append("\n")
    grid_points_2 = str(grid_points[2])
    Lines.append(" "*(5-len(grid_points_2)) + grid_points_2)
    _write_vector_for_cube(Lines,voxel_matrix[2])
    Lines.append("\n")

    #Write atom positions
    for atom in self.atoms:
        atomic_number = atom.atomic_number
        Lines.append(" "*(5-len(str(atomic_number)))+str(atomic_number))
        Lines.append(" "*(4-len(str(atomic_number)))+"{:9.6f}".format(atomic_number))
        _write_vector_for_cube(Lines,atom.pos)
        Lines.append("\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(Lines,vec_y[count:count+6])
                Lines.append('\n')
                count+=6

    with open(file_name,"w") as file:
        file.writelines(Lines)

# def write_txt(self,cube,file_name):
#     """
#     write_txt(cube,file_name)
#
#     This functions writes a txt file containing a 3D array.
#     This method is not the preferred one, just use mol.save()
#
#     Parameters
#     ----------
#         cube : ndarray
#             The cube file to save
#         file_name : str
#             The name of the file
#
#     Returns
#     -------
#         None
#             The file is written
#     """
#
#     grid_points = self.properties["grid_points"]
#     nx,ny,nz = grid_points
#     comments = ""
#     for coord in grid_points:
#         comments += str(coord) + " "
#
#     np.savetxt(file_name,cube.reshape((nx*ny,nz)),header=comments,fmt="%3.6e")


molecule.write_cube = write_cube
# molecule.write_txt = write_txt



#########################################
## Adding save functions for molecules ##
#########################################


def _save_kernel(mol,file_location,rgmol_folder,kernel_name,eigenvectors):
    """
    This functions saves the eigenmodes of a kernel
    """

    try: os.mkdir(file_location + rgmol_folder+"//"+kernel_name)
    except: pass
    for eigenvector in range(len(eigenvectors)):
        file_name = "eigenvector_{}.npy".format(eigenvector)
        np.save(file_location+rgmol_folder+"//"+kernel_name+"//"+file_name,eigenvectors[eigenvector])



def save(self,kernel,append_name=""):
    """
    save(kernel,append_name="")

    This function saves the eignemodes of the already computed kernels in a folder called rgmol.
    The rgmol folder is located in the same folder as the input file used for the molecule.

    Parameters
    ----------
        append_name : str, optional
            A string to append to the name of the rgmol folder.

    Returns
    -------
        None
            The files are written in the rgmol folder

    """

    file_location = _get_file_location(self.file)
    rgmol_folder = "rgmol" + append_name

    try: os.mkdir(file_location + rgmol_folder)
    except: pass
    try: os.mkdir(file_location + "temp")
    except: pass

    print("############################")
    print("# Writing the eigenvectors #")
    print("############################")
    time_before_writing = time.time()
    if kernel == "linear_response_function":
        if "Reconstructed_linear_response_eigenvectors" in self.properties:
            linear_response_eigenvalues = self.properties["linear_response_eigenvalues"]
            contribution_linear_response_eigenvectors = self.properties["contribution_linear_response_eigenvectors"]
            Reconstructed_linear_response_eigenvectors = self.properties["Reconstructed_linear_response_eigenvectors"]

            np.savetxt(file_location + rgmol_folder + "//linear_response_eigenvalues.txt",linear_response_eigenvalues,comments="#",header="Linear Response Eigenvalues")
            np.savetxt(file_location + rgmol_folder + "//contribution_linear_response.txt",contribution_linear_response_eigenvectors,comments="#",header="Contribution of transition densities on eigenmodes")
            _save_kernel(self,file_location,rgmol_folder,"linear_response_function",Reconstructed_linear_response_eigenvectors)
            grid_points = self.properties["grid_points"]
            delta = self.properties["delta"]
            _write_voxel(file_location + rgmol_folder +"//voxel_parameters.txt",grid_points,delta)


        if "linear_response_eigenvectors" in self.properties:
            linear_response_eigenvectors = self.properties["linear_response_eigenvectors"]
            linear_response_eigenvalues = self.properties["linear_response_eigenvalues"]
            contribution_linear_response_eigenvectors = self.properties["contribution_linear_response_eigenvectors"]
            np.savetxt(file_location + rgmol_folder + "//linear_response_eigenvalues.txt",linear_response_eigenvalues,comments="#",header="Linear Response Eigenvalues")
            np.savetxt(file_location + rgmol_folder + "//contribution_linear_response.txt",contribution_linear_response_eigenvectors,comments="#",header="Contribution of transition densities on eigenmodes")
            np.save(file_location + rgmol_folder + "//linear_response_eigenvectors",linear_response_eigenvectors)

    elif kernel == "softness_kernel":
        if "Reconstructed_softness_kernel_eigenvectors" in self.properties:
            softness_kernel_eigenvalues = self.properties["softness_kernel_eigenvalues"]
            contribution_softness_kernel_eigenvectors = self.properties["contribution_softness_kernel_eigenvectors"]
            Reconstructed_softness_kernel_eigenvectors = self.properties["Reconstructed_softness_kernel_eigenvectors"]

            np.savetxt(file_location + rgmol_folder + "//softness_kernel_eigenvalues.txt",softness_kernel_eigenvalues,comments="#",header="Linear Response Eigenvalues")
            np.savetxt(file_location + rgmol_folder + "//contribution_softness_kernel.txt",contribution_softness_kernel_eigenvectors,comments="#",header="Contribution of transition densities on eigenmodes")
            _save_kernel(self,file_location,rgmol_folder,"softness_kernel",Reconstructed_softness_kernel_eigenvectors)
            grid_points = self.properties["grid_points"]
            delta = self.properties["delta"]
            _write_voxel(file_location + rgmol_folder +"//voxel_parameters.txt",grid_points,delta)


        if "softness_kernel_eigenvectors" in self.properties:
            softness_kernel_eigenvectors = self.properties["softness_kernel_eigenvectors"]
            softness_kernel_eigenvalues = self.properties["softness_kernel_eigenvalues"]
            contribution_softness_kernel_eigenvectors = self.properties["contribution_softness_kernel_eigenvectors"]
            np.savetxt(file_location + rgmol_folder + "//softness_kernel_eigenvalues.txt",softness_kernel_eigenvalues,comments="#",header="Linear Response Eigenvalues")
            np.savetxt(file_location + rgmol_folder + "//contribution_softness_kernel.txt",contribution_softness_kernel_eigenvectors,comments="#",header="Contribution of transition densities on eigenmodes")
            np.save(file_location + rgmol_folder + "//softness_kernel_eigenvectors",softness_kernel_eigenvectors)

    else:
        raise ValueError("This kernel {} has not been implemented. Please choose linear_response_function or softness kernel".format(kernel))

    os.rmdir(file_location + "temp")

    print("#####################################")
    print("# Finished Writing the eigenvectors #")
    print("# in {:3.3f} #".format(time.time()-time_before_writing))
    print("#####################################")



def read(self,append_name="",nb_eigen=0):
    """
    read()

    This function reads the files located in the rgmol folder where the molecule is saved.

    Parameters
    ----------
        append_name : str, optional
            A string to append to the name of the rgmol folder.
        nb_eigen : int, optional
            How many eigenmodes will be extracted

    Returns
    -------
        None
            The properties will be put inside molecule.properties

    """

    file_location = _get_file_location(self.file)
    rgmol_folder = "rgmol" + append_name

    try: os.mkdir(file_location + "temp")
    except: pass
    listdir = os.listdir(file_location)
    if not rgmol_folder in listdir:
        raise TypeError("No {} folder were found in the location where the input file of the molecule is.".format(rgmol_folder))

    listdir_rgmol = os.listdir(file_location + rgmol_folder)

    if "voxel_parameters.txt" in listdir_rgmol:
        grid_points,delta = _read_voxel(file_location+rgmol_folder+"//voxel_parameters.txt")
        self.properties["grid_points"] = grid_points
        self.properties["delta"] = delta
        r,c = rgmol.grid.create_cubic_grid_from_molecule(self,grid_points=grid_points,delta=delta)
        self.properties["Reconstructed_coords"] = c

    else:
        print("No representation grid parameters were found, the extraction will proceed anyway")


    if "linear_response_function" in listdir_rgmol:
        print("#############################")
        print("# Extracting the eigenmodes #")
        print("#############################")
        time_before_extract = time.time()

        linear_response_eigenvectors = []

        listdir = os.listdir(file_location+rgmol_folder+"//linear_response_function")
        listdir_int = [int(files.split(".")[0][12:]) for files in listdir]

        arr_sort = np.argsort(listdir_int)
        listdir_sorted = []
        for index in arr_sort:
            listdir_sorted.append(listdir[index])

        if nb_eigen:
            listdir_sorted = listdir_sorted[:nb_eigen]

        for file in listdir_sorted:
            linear_response_eigenvectors.append(np.load(file_location+rgmol_folder+"//linear_response_function//"+file))
        self.properties["Reconstructed_linear_response_eigenvectors"] = linear_response_eigenvectors

        print("######################################")
        print("# Finished Extracting the eigenmodes #")
        print("# in {:3.3f} #".format(time.time()-time_before_extract))
        print("######################################")

    if "linear_response_eigenvectors.npy" in listdir_rgmol:
        print("############################################")
        print("# Extracting the eigenmodes on atomic grid #")
        print("############################################")
        time_before_extract = time.time()

        linear_response_eigenvectors = np.load(file_location+rgmol_folder+"//linear_response_eigenvectors.npy")
        self.properties["linear_response_eigenvectors"] = linear_response_eigenvectors

        print("#####################################################")
        print("# Finished Extracting the eigenmodes on atomic grid #")
        print("# in {:3.3f} #".format(time.time()-time_before_extract))
        print("#####################################################")

        #Recreate grid
        rgmol.grid.create_grid_from_mol(self,N_r=None,d_leb=None,zeta_list=None,alpha_list=None)



    if "softness_kernel" in listdir_rgmol:
        print("#############################")
        print("# Extracting the eigenmodes #")
        print("#############################")
        time_before_extract = time.time()

        softness_kernel_eigenvectors = []

        listdir = os.listdir(file_location+rgmol_folder+"//softness_kernel")
        listdir_int = [int(files.split(".")[0][12:]) for files in listdir]

        arr_sort = np.argsort(listdir_int)
        listdir_sorted = []
        for index in arr_sort:
            listdir_sorted.append(listdir[index])

        for file in listdir_sorted:
            softness_kernel_eigenvectors.append(np.load(file_location+rgmol_folder+"//softness_kernel//"+file))
        self.properties["Reconstructed_softness_kernel_eigenvectors"] = softness_kernel_eigenvectors

        print("######################################")
        print("# Finished Extracting the eigenmodes #")
        print("# in {:3.3f} #".format(time.time()-time_before_extract))
        print("######################################")

    if "softness_kernel_eigenvectors.npy" in listdir_rgmol:
        print("############################################")
        print("# Extracting the eigenmodes on atomic grid #")
        print("############################################")
        time_before_extract = time.time()

        softness_kernel_eigenvectors = np.load(file_location+rgmol_folder+"//softness_kernel_eigenvectors.npy")
        self.properties["softness_kernel_eigenvectors"] = softness_kernel_eigenvectors

        print("#####################################################")
        print("# Finished Extracting the eigenmodes on atomic grid #")
        print("# in {:3.3f} #".format(time.time()-time_before_extract))
        print("#####################################################")

        #Recreate grid
        rgmol.grid.create_grid_from_mol(self,N_r=None,d_leb=None,zeta_list=None,alpha_list=None)


    if "linear_response_eigenvalues.txt" in listdir_rgmol:
        linear_response_eigenvalues = np.loadtxt(file_location + rgmol_folder + "//linear_response_eigenvalues.txt")
        contribution_linear_response = np.real(np.loadtxt(file_location + rgmol_folder + "//contribution_linear_response.txt",dtype=np.complex_))

        self.properties["linear_response_eigenvalues"] = linear_response_eigenvalues
        self.properties["contribution_linear_response_eigenvectors"] = contribution_linear_response

    if "softness_kernel_eigenvalues.txt" in listdir_rgmol:
        softness_kernel_eigenvalues = np.loadtxt(file_location + rgmol_folder + "//softness_kernel_eigenvalues.txt")
        contribution_softness_kernel = np.loadtxt(file_location + rgmol_folder + "//contribution_softness_kernel.txt")

        self.properties["softness_kernel_eigenvalues"] = softness_kernel_eigenvalues
        self.properties["contribution_softness_kernel_eigenvectors"] = contribution_softness_kernel


molecule.save = save
molecule.read = read