import numpy as np
import ase
import ase.io
import sys
import os
import dft_utils as du
import random
import argparse

parser=argparse.ArgumentParser()
parser.add_argument('-p','--path', help='path to the cubefile',action='store',nargs=1)
parser.add_argument('-n','--n2', help='number of Fourier functions to use',action='store',nargs=1)
if len(sys.argv)==1:
    parser.print_help(sys.stderr)
    sys.exit(1)
args=parser.parse_args()
data_dir = args.path[0]
n2 = [int(args.n2[0]),int(args.n2[0]),int(args.n2[0])]

def to_basis3d(X, n2, d=None):
    # Input:
    # X (n, d) matrix where the 3d axis are flattend into last axis
    # n2 is n/2, half the number of basis functions to use along one axis
    dx = d[0]
    dy = d[1]
    dz = d[2]
    n2x = n2[0]
    n2y = n2[1]
    n2z = n2[2]
    X2 = X.reshape(-1, dx, dy, dz)
    X2 = np.fft.rfft(X2, axis=-1)[:, :, :, :n2z]
    X2 = np.concatenate((X2.real, X2.imag), -1)
    X2 = np.fft.rfft(X2, axis=-2)[:, :, :n2y, :]
    X2 = np.concatenate((X2.real, X2.imag), -2)
    X2 = np.fft.rfft(X2, axis=-3)[:, :n2x, :, :]
    X2 = np.concatenate((X2.real, X2.imag), -3)
    return X2.reshape(-1)

def read_cube_file(fname):
    #read a cube file generated from cpmd
    with open(fname) as f:
        lines = f.readlines()
    n_atoms = int(lines[2].strip().split(' ')[0])
    dx = int(lines[3].strip().split(' ')[0])
    voxel_x = float(lines[3].strip().split()[1])
    dy = int(lines[4].strip().split(' ')[0])
    voxel_y = float(lines[4].strip().split()[2])
    dz = int(lines[5].strip().split(' ')[0])
    voxel_z = float(lines[5].strip().split()[3])
    voxel=[voxel_x,voxel_y,voxel_z]
    # Fourier transfrom of cube files with nonzero offdiagonal volume elements are not supported 
    if  float(lines[3].strip().split()[2]) != 0 or float(lines[3].strip().split()[3]) != 0 :
        print('Error! Cube file not supported!')
        sys.exit(1)
    if  float(lines[4].strip().split()[1]) != 0 or float(lines[4].strip().split()[3]) != 0 :
        print('Error! Cube file not supported!')
        sys.exit(1)
    if  float(lines[5].strip().split()[1]) != 0 or float(lines[5].strip().split()[2]) != 0 :
        print('Error! Cube file not supported!')
        sys.exit(1)
    # get atom positions

    atom_p = []
    atom_t = []
    for i in range(n_atoms):
        atom_t.append(int(lines[6 + i].strip().split()[0]))
        atom_p.append([float(p) for p in lines[6 + i].strip().split()[2:]])

    volumetric_data = np.fromstring(' '.join(lines[6 + n_atoms:]), dtype=np.float, sep=' ')
    return volumetric_data.reshape(dx, dy, dz), np.array(atom_p), np.array(atom_t)

densityCPMD_ma, atom_posCPMD_ma, atom_typesCPMD_ma = read_cube_file(data_dir)
if densityCPMD_ma.shape[0] < n2[0]:
    print('Error! Number of Fourier functions larger than number of grid points')
    sys.exit(1)
if densityCPMD_ma.shape[1] < n2[1]:
    print('Error! Number of Fourier functions larger than number of grid points')
    sys.exit(1)
if densityCPMD_ma.shape[2] < n2[2]:
    print('Error! Number of Fourier functions larger than number of grid points')
    sys.exit(1)
densityFCPMD_ma = to_basis3d(densityCPMD_ma.reshape(1,-1),n2=n2,d=densityCPMD_ma.shape)
densityFCPMD_ma = densityFCPMD_ma.reshape(-1)
np.save('./densities_example.npy',densityFCPMD_ma)