import numpy as np
import os
import sys
import h5py
import itertools
from mpl_toolkits.mplot3d import Axes3D


def quant_fill():
  quant = np.empty((35,3))

  quant[0,0]=0;quant[0,1]=0;quant[0,2]=0;

  quant[1,0]=1;quant[1,1]=0;quant[1,2]=0;
  quant[2,0]=0;quant[2,1]=1;quant[2,2]=0;
  quant[3,0]=0;quant[3,1]=0;quant[3,2]=1;

  quant[4,0]=2;quant[4,1]=0;quant[4,2]=0;
  quant[5,0]=0;quant[5,1]=2;quant[5,2]=0;
  quant[6,0]=0;quant[6,1]=0;quant[6,2]=2;
  quant[7,0]=1;quant[7,1]=1;quant[7,2]=0;
  quant[8,0]=1;quant[8,1]=0;quant[8,2]=1;
  quant[9,0]=0;quant[9,1]=1;quant[9,2]=1;

  quant[10,0]=3;quant[10,1]=0;quant[10,2]=0;
  quant[11,0]=0;quant[11,1]=3;quant[11,2]=0;
  quant[12,0]=0;quant[12,1]=0;quant[12,2]=3;
  quant[13,0]=2;quant[13,1]=1;quant[13,2]=0;
  quant[14,0]=2;quant[14,1]=0;quant[14,2]=1;
  quant[15,0]=0;quant[15,1]=2;quant[15,2]=1;
  quant[16,0]=1;quant[16,1]=2;quant[16,2]=0;
  quant[17,0]=1;quant[17,1]=0;quant[17,2]=2;
  quant[18,0]=0;quant[18,1]=1;quant[18,2]=2;
  quant[19,0]=1;quant[19,1]=1;quant[19,2]=1;

  quant[20,0]=4;quant[20,1]=0;quant[20,2]=0;
  quant[21,0]=0;quant[21,1]=4;quant[21,2]=0;
  quant[22,0]=0;quant[22,1]=0;quant[22,2]=4;
  quant[23,0]=3;quant[23,1]=1;quant[23,2]=0;
  quant[24,0]=3;quant[24,1]=0;quant[24,2]=1;
  quant[25,0]=1;quant[25,1]=3;quant[25,2]=0;
  quant[26,0]=0;quant[26,1]=3;quant[26,2]=1;
  quant[27,0]=1;quant[27,1]=0;quant[27,2]=3;
  quant[28,0]=0;quant[28,1]=1;quant[28,2]=3;
  quant[29,0]=2;quant[29,1]=2;quant[29,2]=0;
  quant[30,0]=2;quant[30,1]=0;quant[30,2]=2;
  quant[31,0]=0;quant[31,1]=2;quant[31,2]=2;
  quant[32,0]=2;quant[32,1]=1;quant[32,2]=1;
  quant[33,0]=1;quant[33,1]=2;quant[33,2]=1;
  quant[34,0]=1;quant[34,1]=1;quant[34,2]=2;

  return quant

def calcu_prt_power(num_grid, dpos, max_nlm):
  # 3*nlm(power)*num_grid
  prt_power = np.ones(([3, max_nlm+1] + [num_grid]))
  for a in range(3):
    dpos_a = dpos[:,a]
    for nlm in range(1, max_nlm+1): # nlm=0 are all initiated as ones
      prt_power[a,nlm] = np.power(dpos_a, nlm)

  return prt_power

def calcu_prt_exp(num_grid, dpos, exp_c):
  # len(exp_c)*num_grid
  prt_exp = np.empty(([len(exp_c)] + [num_grid]))
  dpos_norm = np.linalg.norm(dpos, axis=-1)
  for ei in range(len(exp_c)):
    prt_exp[ei] = np.exp(-exp_c[ei]*np.power(dpos_norm,2))

  return prt_exp

def calcu_prt(grid, num_grid, quant, pos_c, typ_c, exp_c):
  dpos = np.subtract(grid, pos_c)
  dpos[dpos == 0] = 1e-8
  # calculate exponent term for pos_c
  prt_exp = calcu_prt_exp(num_grid, dpos, exp_c)

  # precalculate all possible polynomial terms for pos_c
  max_nlm = max([q for t in typ_c for q in quant[t]])
  prt_power  = calcu_prt_power(num_grid, dpos, max_nlm)
  
  # calculate primitives
  prt_c = np.empty(([len(typ_c)] + [num_grid]))
  for ti in range(len(typ_c)):
    t = typ_c[ti]
    prt_c[ti] = prt_power[0,quant[t,0]] * prt_power[1,quant[t,1]] * prt_power[2,quant[t,2]] * prt_exp[ti]

  # calculate gradients of primitives (grdrho(x)=[nx/(xi-xA)-2epi*(xi-xA)]*rho)
  grdprt_c = np.empty(([len(typ_c),3] + [num_grid]))
  grd = np.empty(([len(typ_c),3] + [num_grid]))
  for dim in range(3):
    for ti in range(len(typ_c)):
      t = typ_c[ti]
      tmp = np.divide(quant[t,dim], dpos[:,dim])
      grd[ti,dim] = np.subtract(tmp, 2 * exp_c[ti] * dpos[:,dim])
      grdprt_c[ti,dim] = grd[ti,dim] * prt_c[ti]

  '''
  # calculate sec gradients of primitives grd2rho(x)
  grd2prt_c = np.empty(([len(typ_c),3,3] + [num_grid]))
  for dim1 in range(3):
    for dim2 in range(3):
      for ti in range(len(typ_c)):
        t = typ_c[ti]
        tmp = -np.divide(quant[t,dim], np.power(dpos[:,dim],2))
        if dim1 == dim2:
          grd2prt_c[ti,dim1,dim2] = np.add(np.subtract(tmp, 2 * exp_c[ti]) * prt_c[ti], grd[ti,dim2] * grdprt_c[ti,dim1])
        else:
          grd2prt_c[ti,dim1,dim2] = grd[ti,dim2] * grdprt_c[ti,dim1]
  '''

  return prt_c, grdprt_c #, grd2prt_c

def sel_shell(unoccmo, hf_occ, num_orb, sel_orb, noe):
  if sel_orb == 'all':
    mat = np.ones((unoccmo))
  else:
    sel_orb = int(sel_orb)
    mat = np.zeros((unoccmo))
    mat[hf_occ-sel_orb:] = np.ones((sel_orb+unoccmo-hf_occ))
  '''  
  for i in range(unoccmo):
    c = 0
    for j in range(unoccmo, num_orb):
      c += 1 / (noe[j] - noe[i])
    mat[i] *= c
  '''
  #mat = mat / (noe[unoccmo] - noe[unoccmo-1])

  print(mat)
  return mat

def calcu_grid(fwfn, grid, sel_orb, hf_occ):

  num_grid = len(grid)

  # load wave functions from h5py
  with h5py.File(fwfn, 'r') as hf:
    typ_orb, num_orb, num_prt, num_nuc = hf['BASIC'][:]
    pos = hf['POS'][:]
    cassign = hf['CASSIGN'][:]
    tassign = hf['TASSIGN'][:]
    exp = hf['EXP'][:]
    noon = hf['NOON'][:]
    noe = hf['NOE'][:]
    nocoef = hf['NOCOEF'][:]

  # assign positions to primitives -> a nested center-type list
  # e.g. [1,1,1,1,2,2,2] [1,1,2,3,1,1,2] -> [[1,1,2,3],[1,1,2]]
  # same with nested center-exponet list
  ct_list = [[] for _ in range(len(set(cassign)))]
  ce_list = [[] for _ in range(len(set(cassign)))]
  #print(len(cassign))
  #print(len(tassign)) 
  for i in range(len(cassign)):
    ct_list[cassign[i]-1].append(tassign[i]-1) # tassign[i]-1 to correspond to quant_fill()
    ce_list[cassign[i]-1].append(exp[i])
  
   
  # fill the type with quantum number -> 35 * [n,m,l]
  quant = quant_fill().astype(int)
  
  # check the nonzero occ
  if np.min(noon) == 0:
    unoccmo = next(list(noon).index(o) for o in list(noon) if o==0)
  else:
    unoccmo = num_orb
  
  # select only the shell
  shell = sel_shell(unoccmo, hf_occ, num_orb, sel_orb, noe)
  
  # calculate for every center(atom) the primitive (dpd: atom)
  prt = np.empty(([0] + [num_grid]), dtype=np.float32) # initiate a matrix with unknown dim of concatenate axis for center(atom)
  grdprt = np.empty(([0,3] + [num_grid]), dtype=np.float32)
  #grd2prt = np.empty(([0,3,3] + [num_grid]))
  for c in range(num_nuc):
    #prt_c, grdprt_c, grd2prt_c = calcu_prt(grid, num_grid, quant, pos[c], ct_list[c], ce_list[c]) # this is primitives per atom
    prt_c, grdprt_c = calcu_prt(grid, num_grid, quant, pos[c], ct_list[c], ce_list[c])
    #prt, grdprt, grd2prt = np.append(prt, prt_c, axis=0), np.append(grdprt, grdprt_c, axis=0), np.append(grd2prt, grd2prt_c, axis=0)
    prt, grdprt = np.append(prt, prt_c, axis=0), np.append(grdprt, grdprt_c, axis=0)

  np.save('prt.npy', prt)
  np.save('grdprt.npy', grdprt)

  print('Finish prt and grdprt calculations!')

  # summation of the primitives based on nocoef and no
  # prt in dimension num_prt * num_grid
  # grdprt in dimension num_prt * 3 * num_grid
  nowfn = np.empty([num_orb] + [num_grid])
  norho = np.empty([unoccmo] + [num_grid])
  nokin = np.empty([unoccmo] + [num_grid])
  nogrdrho = np.empty([unoccmo,3] + [num_grid])
  nogrd2rho = np.empty([unoccmo,3,3] + [num_grid])
  
  # for nonzero occupied no
  for n in range(unoccmo):
    # rho & wfn #
    prtwfn = np.multiply(nocoef[n][:,np.newaxis], prt)
    nowfn[n] = np.sum(prtwfn, axis=0)
    norho[n] = shell[n] * noon[n] * np.power(nowfn[n], 2)
    # grd  #
    grdprtwfn = np.multiply(nocoef[n][:,np.newaxis,np.newaxis], grdprt)
    nogrdrho[n] = 2 * shell[n] * noon[n] * (np.multiply(np.sum(prtwfn, axis=0), np.sum(grdprtwfn, axis=0)))
    # kin #
    nokin[n] = shell[n] * noon[n] * np.sum(np.power(np.sum(grdprtwfn, axis=0), 2), axis=0)
    '''
    # grd2 #
    grd2prtwfn = np.multiply(nocoef[n][:,np.newaxis,np.newaxis,np.newaxis], grd2prt)
    tmp1 = np.multiply(np.sum(prtwfn, axis=0), np.sum(grd2prtwfn, axis=0))
    grdwfn = np.sum(grdprtwfn, axis=0)
    tmp2 = np.empty([3,3] + [num_grid])
    for dim1 in range(3):
      for dim2 in range(3):
        tmp2[dim1,dim2] = grdwfn[dim1] * grdwfn[dim2]
    nogrd2rho[n] = 2 * shell[n] * noon[n] * np.add(tmp1, tmp2)
    '''

  # for zero occupied no
  for n in range(unoccmo, num_orb):
    # wfn #
    prtwfn = np.multiply(nocoef[n][:,np.newaxis], prt)
    nowfn[n] = np.sum(prtwfn, axis=0)
          
  del prtwfn
  del grdprtwfn
  #del grd2prtwfn
  
  rho = np.sum(norho, axis=0)
  del norho
  grdrho = np.sum(nogrdrho, axis=0)
  del nogrdrho
  kin = np.sum(nokin, axis=0)
  del nokin
  grd2rho = np.sum(nogrd2rho, axis=0)
  del nogrd2rho

  
  nrm_grdrho = np.linalg.norm(grdrho, axis=0)

  
  # calculate grd2rho eigen values
  dig_grd2rho = np.empty([num_grid] + [3])
  for i in range(num_grid):
    w, _ = np.linalg.eig(grd2rho[:,:,i])
    dig_grd2rho[i] = np.sort(w)
  
  '''
  rg_grdrho = []
  for i in range(num_grid):     
    rg_grdrho.append(grdrho[:,i])
  rg_grdrho = np.array(rg_grdrho)
  '''
  
  
  # save grid points and their rho, gradrho, grad2rho
  np.save('rho-' + sys.argv[1].split('/')[-1].split('.')[0] + '.npy', rho)
  np.save('nrmgrdrho-' + sys.argv[1].split('/')[-1].split('.')[0] + '.npy', nrm_grdrho)
  np.save('kin-' + sys.argv[1].split('/')[-1].split('.')[0] + '.npy', kin)
  np.save('diglaprho-' + sys.argv[1].split('/')[-1].split('.')[0] + '.npy', dig_grd2rho)

  np.savetxt('rho.dat', rho)
  np.savetxt('nrmgrdrho.dat', nrm_grdrho)
  np.savetxt('kin.dat', kin)
  np.savetxt('diglaprho.dat', dig_grd2rho)

  # plot 3d scattering of grid
  # plot3d(rg_grid)

  return rho, nowfn.T


def main():  
  fwfn = sys.argv[1]
  fgrid = sys.argv[2]
  fwgt = sys.argv[3]
  sel_orb = sys.argv[4]
  hf_occ = int(sys.argv[5]) # in casci cases, we need to know lumo from hf occupation

  with h5py.File(fwfn, 'r') as hf:
    _, _, _, num_nuc = hf['BASIC'][:]

  # load grids and weights
  grid = np.load(fgrid)
  wgt = np.load(fwgt)

  # calculate rho grad lap on each grid
  rho, wfn = calcu_grid(fwfn, grid, sel_orb, hf_occ)

  # sum up all of them based on weight
  wrho = np.sum(np.multiply(wgt, rho))
    
  print('Integrated rho of of system' + ': ' + str(wrho))

  # save grid points and their NO wfn
  np.savez(sys.argv[1].split('/')[-1].split('.')[0] + '.npz', GRID=grid, WFN=wfn, WGT=wgt)

if __name__ == "__main__":
  main()