#
# PAOFLOW
#
# Copyright 2016-2022 - Marco BUONGIORNO NARDELLI (mbn@unt.edu)
#
# Reference:
#
# F.T. Cerasoli, A.R. Supka, A. Jayaraj, I. Siloi, M. Costa, J. Slawinska, S. Curtarolo, M. Fornari, D. Ceresoli, and M. Buongiorno Nardelli,
# Advanced modeling of materials with PAOFLOW 2.0: New features and software design, Comp. Mat. Sci. 200, 110828 (2021).
#
# M. Buongiorno Nardelli, F. T. Cerasoli, M. Costa, S Curtarolo,R. De Gennaro, M. Fornari, L. Liyanage, A. Supka and H. Wang,
# PAOFLOW: A utility to construct and operate on ab initio Hamiltonians from the Projections of electronic wavefunctions on
# Atomic Orbital bases, including characterization of topological materials, Comp. Mat. Sci. vol. 143, 462 (2018).
#
# This file is distributed under the terms of the
# GNU General Public License. See the file `License'
# in the root directory of the present distribution,
# or http://www.gnu.org/copyleft/gpl.txt .

import numpy as np
import scipy.sparse as sprs
import scipy.linalg as LA
from scipy.special import factorial as fac
from tempfile import NamedTemporaryFile
import re
from .communication import scatter_full, gather_full, gather_scatter
from scipy.spatial.distance import cdist
from mpi4py import MPI
from .zero_pad import zero_pad
import time

comm = MPI.COMM_WORLD
rank = comm.Get_rank()


def check(Hksp_s, si_per_k, new_k_ind, orig_k_ind, phase_shifts, U, a_index, inv_flag, equiv_atom, kp, symop, fg, isl,
          sym_TR):
  # compares saved full grid hamiltonain to one generated
  # from wedge... only for debugging purposes
  nawf = Hksp_s.shape[1]
  # load for testing
  Hksp_f = np.load("ONLY_FOR_TESTING_kham.npy")
  Hksp_f = np.reshape(Hksp_f, (nawf, nawf, Hksp_s.shape[0]))
  Hksp_f = np.transpose(Hksp_f, axes=(2, 0, 1))

  print(np.allclose(Hksp_f, Hksp_s, atol=1.e-4, rtol=1.e-4))

  bad_symop = np.ones(symop.shape[0], dtype=bool)
  good_symop = np.ones(symop.shape[0], dtype=bool)
  good = []
  bad = []
  print(a_index)
  st = 0
  fn = 8
  #    for j in range(20):
  for j in range(Hksp_s.shape[0]):
    isym = si_per_k[j]
    nki = new_k_ind[j]
    oki = orig_k_ind[j]

    HP = Hksp_f[nki]
    THP = Hksp_s[nki]

    U_k = get_U_k(kp[oki], phase_shifts[isym], a_index, U[isym])

    good.append(isym)

    if np.all(np.isclose(HP, THP, rtol=1.e-4, atol=1.e-4)):
      bad.append(isym)
    else:

      good_symop[isym] = False
      continue
      print(j, isym)
      print('old_k=', kp[oki])
      print('new_k=', fg[nki])
      print(equiv_atom[isym])
      print(symop[isym])
      print()
      print("inv:", inv_flag[isym])
      print()
      print("*" * 50)
      print("REAL " * 10)
      print("*" * 50)
      print()
      #           print(U_k[st:fn,st:fn])
      print()
      print("CORRECT")
      print(HP[st:fn, st:fn].real)
      print("NEW")
      print(THP[st:fn, st:fn].real)
      print("EQUAL?")
      print(np.isclose(HP[st:fn, st:fn].real, THP[st:fn, st:fn].real,
                       rtol=1.e-4, atol=1.e-4))

      print("*" * 50)
      print("IMAG " * 10)
      print("*" * 50)
      print("CORRECT")
      print(HP[st:fn, st:fn].imag)
      print("NEW")
      print(THP[st:fn, st:fn].imag)
      print("EQUAL?")
      print(np.isclose(HP[st:fn, st:fn].imag, THP[st:fn, st:fn].imag,
                       rtol=1.e-4, atol=1.e-4))

      print("*" * 100)
      print("*" * 100)
      print("*" * 100)
  #            raise SystemExit
  print(len(good) - kp.shape[0], len(bad) - kp.shape[0])

  print()
  for i in range(symop.shape[0]):
    if i in si_per_k:
      if not good_symop[i]:
        print("%2d" % (i + 1), "BAD ", isl[i], sym_TR[i])
      else:
        print("%2d" % (i + 1), "GOOD", isl[i], sym_TR[i])


############################################################################################
############################################################################################
############################################################################################

def LPF(aux, nk1, nk2, nk3, nfft1, nfft2, nfft3):
  '''
  Pad frequency domain with zeroes, such that any relationship between
      aux[k] and aux[N-k] is preserved.

  Arguments:
      aux (ndarray): unpadded frequency domain data
      nk1 (int): current size of aux along axis 0
      nk2 (int): current size of aux along axis 1
      nk3 (int): current size of aux along axis 2
      nfft1 (int): number of zeroes to pad axis 0 by
      nfft1 (int): number of zeroes to pad axis 1 by
      nfft1 (int): number of zeroes to pad axis 2 by

  Returns:
      auxp3 (ndarray): padded frequency domain data
  '''

  # post-padding dimensions
  nk1p = nfft1 + nk1
  nk2p = nfft2 + nk2
  nk3p = nfft3 + nk3
  # halfway points
  sk1 = int((nk1 + 1) / 2)
  sk2 = int((nk2 + 1) / 2)
  sk3 = int((nk3 + 1) / 2)
  # parities (even <-> p==1)
  p1 = (nk1 & 1) ^ 1
  p2 = (nk2 & 1) ^ 1
  p3 = (nk3 & 1) ^ 1

  # accomodate nfft==0
  if nfft1 == 0:  p1 = 0
  if nfft2 == 0:  p2 = 0
  if nfft3 == 0:  p3 = 0

  # first dimension
  auxp1 = np.zeros((nk1, nk2, nk3p), dtype=complex)
  auxp1[:, :, :sk3 + p3] = aux[:, :, :sk3 + p3]
  auxp1[:, :, nfft3 + sk3:] = aux[:, :, sk3:]
  # second dimension
  auxp2 = np.zeros((nk1, nk2p, nk3p), dtype=complex)
  auxp2[:, :sk2 + p2, :] = auxp1[:, :sk2 + p2, :]
  auxp2[:, nfft2 + sk2:, :] = auxp1[:, sk2:, :]
  # third dimension
  auxp3 = np.zeros((nk1p, nk2p, nk3p), dtype=complex)
  auxp3[:sk1 + p1, :, :] = auxp2[:sk1 + p1, :, :]
  auxp3[nfft1 + sk1:, :, :] = auxp2[sk1:, :, :]

  if rank == 0:
    print(sk1)
  # halve Nyquist axes
  if p1:
    auxp3[sk1 - 1, :, :] = 0.0
    auxp3[-sk1 + 1, :, :] = 0.0
    auxp3[sk1, :, :] = 0.0
  if p2:
    auxp3[:, sk2 - 1, :] = 0.0
    auxp3[:, -sk2 + 1, :] = 0.0
    auxp3[:, -sk2, :] = 0.0
  if p3:
    auxp3[:, :, sk3 - 1] = 0.0
    auxp3[:, :, -sk3 + 1] = 0.0
    auxp3[:, :, -sk3] = 0.0
  return (auxp3)


############################################################################################
############################################################################################
############################################################################################

def down_samp(aux, nk1, nk2, nk3, nfft1, nfft2, nfft3):
  '''
  Pad frequency domain with zeroes, such that any relationship between
      aux[k] and aux[N-k] is preserved.

  Arguments:
      aux (ndarray): unpadded frequency domain data
      nk1 (int): current size of aux along axis 0
      nk2 (int): current size of aux along axis 1
      nk3 (int): current size of aux along axis 2
      nfft1 (int): number of zeroes to pad axis 0 by
      nfft1 (int): number of zeroes to pad axis 1 by
      nfft1 (int): number of zeroes to pad axis 2 by

  Returns:
      auxp3 (ndarray): padded frequency domain data
  '''

  nk1 += nfft1
  nk2 += nfft2
  nk3 += nfft3

  nfft1 *= -1
  nfft2 *= -1
  nfft3 *= -1

  # post-padding dimensions
  nk1p = nfft1 + nk1
  nk2p = nfft2 + nk2
  nk3p = nfft3 + nk3
  # halfway points
  sk1 = int((nk1 + 1) / 2)
  sk2 = int((nk2 + 1) / 2)
  sk3 = int((nk3 + 1) / 2)
  # parities (even <-> p==1)
  p1 = (nk1 & 1) ^ 1
  p2 = (nk2 & 1) ^ 1
  p3 = (nk3 & 1) ^ 1

  # accomodate nfft==0
  if nfft1 == 0:  p1 = 0
  if nfft2 == 0:  p2 = 0
  if nfft3 == 0:  p3 = 0

  # first dimension
  auxp1 = np.zeros((nk1p, nk2p, nk3), dtype=complex)
  auxp1[:, :, :sk3 + p3] = aux[:, :, :sk3 + p3]
  auxp1[:, :, sk3:] = aux[:, :, nfft3 + sk3:]
  # second dimension
  auxp2 = np.zeros((nk1p, nk2, nk3), dtype=complex)
  auxp2[:, :sk2 + p2, :] = auxp1[:, :sk2 + p2, :]
  auxp2[:, sk2:, :] = auxp1[:, nfft2 + sk2:, :]
  # third dimension
  auxp3 = np.zeros((nk1, nk2, nk3), dtype=complex)
  auxp3[:sk1 + p1, :, :] = auxp2[:sk1 + p1, :, :]
  auxp3[sk1:, :, :] = auxp2[nfft1 + sk1:, :, :]

  # double Nyquist axes
  if p1:
    auxp3[sk1, :, :] *= 2
  if p2:
    auxp3[:, sk2, :] *= 2
  if p3:
    auxp3[:, :, sk3] *= 2

  return auxp3


############################################################################################
############################################################################################
############################################################################################

def map_equiv_atoms(a_index, map_ind):
  # generate unitary tranformation for swapping
  # between atoms equivalent by symmetry op

  nawf = a_index.shape[0]
  U_wyc = np.zeros((map_ind.shape[0], nawf, nawf), dtype=complex)

  for isym in range(map_ind.shape[0]):
    map_U = np.zeros((nawf, nawf), dtype=complex)
    for i in range(map_ind[isym].shape[0]):
      fi = np.where(a_index == i)[0]
      si = np.where(a_index == map_ind[isym][i])[0]

      for j in range(si.shape[0]):
        map_U[fi[j], si[j]] = 1

    U_wyc[isym] = map_U

  return U_wyc


############################################################################################
############################################################################################
############################################################################################

def get_U_TR(jchia):
  # operator that swaps J for -J

  J = 0.5
  TR_J05 = np.fliplr(np.eye(int(2 * J + 1)) * ((-1) ** (J - np.arange(-J, J + 1))))
  J = 1.5
  TR_J15 = np.fliplr(np.eye(int(2 * J + 1)) * ((-1) ** (J - np.arange(-J, J + 1))))
  J = 2.5
  TR_J25 = np.fliplr(np.eye(int(2 * J + 1)) * ((-1) ** (J - np.arange(-J, J + 1))))
  J = 3.5
  TR_J35 = np.fliplr(np.eye(int(2 * J + 1)) * ((-1) ** (J - np.arange(-J, J + 1))))

  TR = [TR_J05, TR_J15, TR_J25, TR_J35]

  nawf = int(np.sum(2 * jchia + 1) * 2)

  U_TR = np.zeros((nawf, nawf), dtype=complex)

  blocks = []

  flr_jchia = np.array([int(np.floor(a)) for a in jchia])

  # block diagonal transformation matrix of T @ H(k) @ T^-1
  for s in flr_jchia:
    blocks.append(TR[s])
  U_TR = LA.block_diag(*blocks)

  return U_TR


############################################################################################
############################################################################################
############################################################################################

def add_U_TR(U, sym_TR, jchia):
  # applies time inversion to symops as needed
  U_TR = get_U_TR(jchia)

  for isym in range(U.shape[0]):
    if sym_TR[isym]:
      U[isym] = U_TR @ U[isym]

  return U


############################################################################################
############################################################################################
############################################################################################


def get_trans():
  # get transformation from angular momentum number to
  # chemistry orbital form for the angular momentum
  rsh = np.sqrt(0.5)
  ish = 1.0j * np.sqrt(0.5)

  trans_l0 = np.array([[1.0]], dtype=complex)

  # to convert from m=[-1,0,1] to px,py,pz for l=1
  trans_l1 = np.array([[0, 1, 0],
                       [-1, 0, 1],
                       [1, 0, 1], ], dtype=complex)

  trans_l1[1] *= rsh
  trans_l1[2] *= ish

  # to convert from m=[-2,-1,0,1,2] to dx2,dzx,dzy,dx2-y2,dxy for l=2
  trans_l2 = np.array([[0, 0, 1, 0, 0],
                       [0, -1, 0, 1, 0],
                       [0, 1, 0, 1, 0],
                       [1, 0, 0, 0, 1],
                       [-1, 0, 0, 0, 1], ], dtype=complex)

  trans_l2[1] *= rsh
  trans_l2[2] *= ish
  trans_l2[3] *= rsh
  trans_l2[4] *= ish

  # to convert from m=[-3,-2,-1,0,1,2,3] to real combinations for l=3
  trans_l3 = np.array([[0, 0, 0, 1, 0, 0, 0],
                       [0, 0, -1, 0, 1, 0, 0],
                       [0, 0, 1, 0, 1, 0, 0],
                       [0, 1, 0, 0, 0, 1, 0],
                       [0, -1, 0, 0, 0, 1, 0],
                       [-1, 0, 0, 0, 0, 0, 1],
                       [1, 0, 0, 0, 0, 0, 1], ], dtype=complex)

  trans_l3[1] *= rsh
  trans_l3[2] *= ish
  trans_l3[3] *= rsh
  trans_l3[4] *= ish
  trans_l3[5] *= rsh
  trans_l3[6] *= ish

  return trans_l0, trans_l1, trans_l2, trans_l3


############################################################################################
############################################################################################
############################################################################################


def d_mat_l(AL, BE, GA, l):
  # gets wigner_d matrix for a given alpha,beta,gamma
  # that rotates basis from m' to m

  # list of m and m'
  ms = np.arange(-l, l + 1)
  mps = ms

  d_mat = np.zeros((ms.shape[0], ms.shape[0]), dtype=complex)

  # find wigner_d matrix elements for each m and m'
  for m_i in range(len(ms)):
    for mp_i in range(len(mps)):
      m = ms[m_i]
      mp = mps[mp_i]
      w_max = int(l + mp + 1)
      dmm = 0.0

      out_sum = np.sqrt(fac(int(l + m)) * fac(int(l - m)) * fac(int(l + mp)) * fac(int(l - mp)))

      # loop over w for summation
      for w in range(0, w_max):
        # factorials in denominator must be positive
        df1 = int(l + mp - w)
        df2 = int(l - m - w)
        df3 = int(w)
        df4 = int(w + m - mp)
        if df1 < 0 or df2 < 0 or df3 < 0 or df4 < 0: continue

        # for each w in summation over w
        tmp = (-1) ** w
        tmp /= fac(df1) * fac(df2) * fac(df3) * fac(df4)
        tmp *= np.cos(BE / 2.0) ** (2 * l - m + mp - 2 * w)
        tmp *= np.sin(BE / 2.0) ** (2 * w - mp + m)

        dmm += tmp

      d_mat[mp_i, m_i] = dmm * np.exp(-1.0j * AL * mp) * np.exp(-1.0j * GA * m) * out_sum

  return d_mat


############################################################################################
############################################################################################
############################################################################################

def get_wigner(symop):
  # gets wigner_d matrix associated with each symop
  wigner_l0 = np.zeros((symop.shape[0], 1, 1), dtype=complex)
  wigner_l1 = np.zeros((symop.shape[0], 3, 3), dtype=complex)
  wigner_l2 = np.zeros((symop.shape[0], 5, 5), dtype=complex)
  wigner_l3 = np.zeros((symop.shape[0], 7, 7), dtype=complex)
  # inversion flag
  inv_flag = np.zeros((symop.shape[0]), dtype=bool)

  for i in range(symop.shape[0]):
    # get euler angles alpha,beta,gamma from the symop
    AL, BE, GA = mat2eul(symop[i])
    AL, BE, GA = np.deg2rad(np.around(np.rad2deg([AL, BE, GA]), decimals=0))

    # check if there is an inversion in the symop
    if not np.all(np.isclose(eul2mat(AL, BE, GA), symop[i])):
      inv_flag[i] = True
      AL, BE, GA = mat2eul(-symop[i])
      AL, BE, GA = np.deg2rad(np.around(np.rad2deg([AL, BE, GA]), decimals=0))
      if not np.all(np.isclose(eul2mat(AL, BE, GA), -symop[i])):
        print("ERROR IN MAT2EUL!")
        print(i + 1)
        print(symop[i])
        raise SystemExit

    # wigner_d matrix for l=0
    wigner_l0[i] = d_mat_l(AL, BE, GA, 0)
    # wigner_d matrix for l=1
    wigner_l1[i] = d_mat_l(AL, BE, GA, 1)
    # wigner_d matrix for l=2
    wigner_l2[i] = d_mat_l(AL, BE, GA, 2)
    # wigner_d matrix for l=3
    wigner_l3[i] = d_mat_l(AL, BE, GA, 3)

  return [wigner_l0, wigner_l1, wigner_l2, wigner_l3, inv_flag], inv_flag


############################################################################################
############################################################################################
############################################################################################

def get_wigner_so(symop):
  # gets wigner_d matrix associated with each symop
  wigner_j05 = np.zeros((symop.shape[0], 2, 2), dtype=complex)
  wigner_j15 = np.zeros((symop.shape[0], 4, 4), dtype=complex)
  wigner_j25 = np.zeros((symop.shape[0], 6, 6), dtype=complex)
  wigner_j35 = np.zeros((symop.shape[0], 8, 8), dtype=complex)
  # inversion flag
  inv_flag = np.zeros((symop.shape[0]), dtype=bool)

  for i in range(symop.shape[0]):
    # get euler angles alpha,beta,gamma from the symop
    AL, BE, GA = mat2eul(symop[i])
    AL, BE, GA = np.deg2rad(np.around(np.rad2deg([AL, BE, GA]), decimals=0))
    # check if there is an inversion in the symop
    if not np.all(np.isclose(correct_roundoff(eul2mat(AL, BE, GA)),
                             symop[i], atol=1.e-3, rtol=1.e-2)):
      inv_flag[i] = True

      AL, BE, GA = mat2eul(-symop[i])
      AL, BE, GA = np.deg2rad(np.around(np.rad2deg([AL, BE, GA]), decimals=0))
      if not np.all(np.isclose(correct_roundoff(eul2mat(AL, BE, GA)),
                               -symop[i], atol=1.e-3, rtol=1.e-2)):
        print("ERROR IN MAT2EUL!")
        print(i + 1)
        print('RESULT')
        print(-eul2mat(AL, BE, GA))
        print('CORRECT')
        print(symop[i])
        raise SystemExit

    # wigner_d for l=0
    wigner_j05[i] = d_mat_l(AL, BE, GA, 0.5)
    # wigner_d for l=1
    wigner_j15[i] = d_mat_l(AL, BE, GA, 1.5)
    # wigner_d for l=2
    wigner_j25[i] = d_mat_l(AL, BE, GA, 2.5)
    # wigner_d for l=3
    wigner_j35[i] = d_mat_l(AL, BE, GA, 3.5)

  return [wigner_j05, wigner_j15, wigner_j25, wigner_j35, inv_flag], inv_flag


############################################################################################
############################################################################################
############################################################################################

def convert_wigner_d(wigner):
  # get transformation from angular momentum number to
  # chemistry orbital form for the angular momentum

  trans_l0, trans_l1, trans_l2, trans_l3 = get_trans()

  c_wigner_l0 = np.zeros_like(wigner[0])
  c_wigner_l1 = np.zeros_like(wigner[1])
  c_wigner_l2 = np.zeros_like(wigner[2])
  c_wigner_l3 = np.zeros_like(wigner[3])

  for i in range(wigner[0].shape[0]):
    c_wigner_l0[i] = trans_l0 @ wigner[0][i] @ np.conj(trans_l0.T)
    c_wigner_l1[i] = trans_l1 @ wigner[1][i] @ np.conj(trans_l1.T)
    c_wigner_l2[i] = trans_l2 @ wigner[2][i] @ np.conj(trans_l2.T)
    c_wigner_l3[i] = trans_l3 @ wigner[3][i] @ np.conj(trans_l3.T)

  return [c_wigner_l0, c_wigner_l1, c_wigner_l2, c_wigner_l3]


############################################################################################
############################################################################################
############################################################################################

def eul2mat(ALPHA, BETA, GAMMA):
  # generates rotation matrix from euler angles
  C_1 = np.cos(ALPHA)
  C_2 = np.cos(BETA)
  C_3 = np.cos(GAMMA)
  S_1 = np.sin(ALPHA)
  S_2 = np.sin(BETA)
  S_3 = np.sin(GAMMA)

  euler = np.array([[C_1 * C_2 * C_3 - S_1 * S_3, -C_1 * C_2 * S_3 - C_3 * S_1, C_1 * S_2, ],
                    [C_2 * C_3 * S_1 + C_1 * S_3, -C_2 * S_1 * S_3 + C_1 * C_3, S_1 * S_2, ],
                    [-C_3 * S_2, S_2 * S_3, C_2, ]])

  return euler


############################################################################################
############################################################################################
############################################################################################

def mat2eul(R):
  # finds euler angles from rotation matrix in ZYZ convention
  if R[2, 2] < 1.0:
    if R[2, 2] > -1.0:
      ALPHA = np.arctan2(R[1, 2], R[0, 2])
      BETA = np.arccos(R[2, 2])
      GAMMA = np.arctan2(R[2, 1], -R[2, 0])
    else:
      ALPHA = -np.arctan2(R[1, 0], R[1, 1])
      BETA = np.pi
      GAMMA = 0.0
  else:
    ALPHA = np.arctan2(R[1, 0], R[1, 1])
    BETA = 0.0
    GAMMA = 0.0

  return np.array([ALPHA, BETA, GAMMA])


############################################################################################
############################################################################################
############################################################################################

def find_equiv_k(kp, symop, full_grid, sym_TR, check=True, include_self=False):
  # find indices and symops that generate full grid H from wedge H
  orig_k_ind = []
  new_k_ind = []
  si_per_k = []
  counter = 0
  kp_track = []
  kp = correct_roundoff(kp)

  full_grid_mask = np.copy(full_grid)

  for isym in range(symop.shape[0]):
    # transform k -> k' with the sym op
    if sym_TR[isym]:
      newk = ((((-symop[isym] @ (kp.T % 1.0)) % 1.0) + 0.5) % 1.0) - 0.5
    else:
      newk = ((((symop[isym] @ (kp.T % 1.0)) % 1.0) + 0.5) % 1.0) - 0.5
    newk = correct_roundoff(newk)
    newk[np.where(np.isclose(newk, 0.5))] = -0.5
    newk[np.where(np.isclose(newk, -1.0))] = 0.0
    newk[np.where(np.isclose(newk, 1.0))] = 0.0

    # find index in the full grid where this k -> k' with this sym op
    nw = np.where(np.isclose(cdist(newk.T, full_grid_mask), 0.0, atol=1.e-6, rtol=1.e-6, ))

    if not include_self:
      full_grid_mask[nw[1]] = np.ma.masked

    new_k_ind.extend(nw[1].tolist())
    si_per_k.extend([isym] * nw[1].shape[0])
    orig_k_ind.extend(nw[0].tolist())

  new_k_ind = np.array(new_k_ind)
  orig_k_ind = np.array(orig_k_ind)
  si_per_k = np.array(si_per_k)

  if not include_self:
    inds = np.unique(new_k_ind, return_index=True)

    new_k_ind = new_k_ind[inds[1]]
    orig_k_ind = orig_k_ind[inds[1]]
    si_per_k = si_per_k[inds[1]]
    counter = inds[1].shape[0]

  # check to make sure we have all the k points accounted for
  if counter != full_grid.shape[0] and check:

    print('NOT ALL KPOINTS ACCOUNTED FOR')
    print(counter, full_grid.shape[0])
    print('missing k:')
    print(full_grid[np.setxor1d(new_k_ind, np.arange(full_grid.shape[0]))])
    raise SystemExit

  else:
    pass

  return new_k_ind, orig_k_ind, si_per_k


############################################################################################
############################################################################################
############################################################################################


def build_U_matrix(wigner, shells):
  # builds U from blocks

  nawf = int(np.sum(2 * shells + 1))

  U = np.zeros((wigner[0].shape[0], nawf, nawf), dtype=complex)
  # for the so case...
  flr_shells = np.array([int(np.floor(a)) for a in shells])

  for i in range(wigner[0].shape[0]):
    blocks = []
    # block diagonal transformation matrix of H(k)->H(k')
    for s in flr_shells:
      blocks.append(wigner[s][i])
    U[i] = LA.block_diag(*blocks)
  return U


############################################################################################
############################################################################################
############################################################################################

def get_phase_shifts(atom_pos, symop, equiv_atom):
  # calculate phase shifts for U
  phase_shift = np.zeros((symop.shape[0], atom_pos.shape[0], 3), dtype=float)
  for isym in range(symop.shape[0]):
    for p in range(atom_pos.shape[0]):
      p1 = equiv_atom[isym, p]
      phase_shift[isym, p1] = (symop[isym].T @ atom_pos[p]) - atom_pos[p1]

  return correct_roundoff(np.around(phase_shift, decimals=2))


############################################################################################
############################################################################################
############################################################################################

def get_U_k(k, shift, a_index, U):
  # add phase shift to U
  U_k = U * np.exp(2.0j * np.pi * (shift[a_index] @ k))
  return U_k


############################################################################################
############################################################################################
############################################################################################

def correct_roundoff(arr, incl_hex=False, atol=1.e-8):
  # correct for round off
  arr[np.where(np.isclose(arr, 0.0, atol=atol))] = 0.0
  arr[np.where(np.isclose(arr, 1.0, atol=atol))] = 1.0
  arr[np.where(np.isclose(arr, -1.0, atol=atol))] = -1.0

  if incl_hex:
    sq3o2 = np.sqrt(3) / 2.0
    arr[np.where(np.isclose(arr, sq3o2, atol=atol))] = sq3o2
    arr[np.where(np.isclose(arr, -sq3o2, atol=atol))] = -sq3o2
    arr[np.where(np.isclose(arr, 0.5, atol=atol))] = 0.5
    arr[np.where(np.isclose(arr, -0.5, atol=atol))] = -0.5

  return arr


############################################################################################
############################################################################################
############################################################################################

def get_full_grid(nk1, nk2, nk3, o1, o2, o3):
  # generates full k grid in crystal fractional coords
  nktot = nk1 * nk2 * nk3
  b_vectors = np.eye(3)
  Kint = np.zeros((nktot, 3), dtype=float)
  offset = np.array([o1, o2, o3]) / (2 * np.array([nk1, nk2, nk3]))

  for i in range(nk1):
    for j in range(nk2):
      for k in range(nk3):
        n = k + j * nk3 + i * nk2 * nk3
        Rx = float(i) / float(nk1)
        Ry = float(j) / float(nk2)
        Rz = float(k) / float(nk3)
        if Rx >= 0.5: Rx = Rx - 1.0
        if Ry >= 0.5: Ry = Ry - 1.0
        if Rz >= 0.5: Rz = Rz - 1.0
        Rx -= int(Rx)
        Ry -= int(Ry)
        Rz -= int(Rz)

        Kint[n] = Rx * b_vectors[0, :] + Ry * b_vectors[1, :] + Rz * b_vectors[2, :] + offset

  return Kint


############################################################################################
############################################################################################
############################################################################################

def get_inv_op(shells):
  # returns the inversion operator
  orb_index = np.hstack([[l] * (2 * l + 1) for l in shells])
  sign_inv = (-np.ones(orb_index.shape[0])) ** orb_index

  return np.outer(sign_inv, sign_inv)


############################################################################################
############################################################################################
############################################################################################

def enforce_t_rev(Hksp_s, nk1, nk2, nk3, spin_orb, U_inv, jchia):
  # enforce time reversal symmetry on H(k)
  nawf = Hksp_s.shape[1]

  Hksp_s = np.reshape(Hksp_s, (nk1, nk2, nk3, nawf, nawf), order="C")

  if spin_orb:
    U_TR = get_U_TR(jchia)

  for i in range(int(nk1 / 2) + 1):
    for j in range(int(nk2 / 2) + 1):
      for k in range(int(nk3 / 2) + 1):
        iv = (nk1 - i) % nk1
        jv = (nk2 - j) % nk2
        kv = (nk3 - k) % nk3
        if not spin_orb:
          temp1 = np.conj(Hksp_s[i, j, k])
          temp2 = np.conj(Hksp_s[iv, jv, kv])
          Hksp_s[iv, jv, kv] = (Hksp_s[iv, jv, kv] + temp1) / 2.0
          Hksp_s[i, j, k] = (Hksp_s[i, j, k] + temp2) / 2.0
        else:
          temp1 = np.conj(U_inv * (U_TR @ Hksp_s[i, j, k] @ np.conj(U_TR.T)))
          #                    temp2=np.conj(U_inv*(U_TR @ Hksp_s[iv,jv,kv] @ np.conj(U_TR.T)))
          #                    Hksp_s[iv,jv,kv] = (Hksp_s[iv,jv,kv] + temp1)/2.0
          #                    Hksp_s[i,j,k]    = (Hksp_s[i,j,k]    + temp2)/2.0
          Hksp_s[iv, jv, kv] = temp1

  Hksp_s = np.reshape(Hksp_s, (nk1 * nk2 * nk3, nawf, nawf), order="C")

  return Hksp_s


############################################################################################
############################################################################################
############################################################################################

def apply_t_rev(Hksp, kp, spin_orb, U_inv, jchia):
  # apply time reversal operator to get H(-k) from H(k)

  # gamma only case
  if Hksp.shape[0] == 1:
    return Hksp, kp
  new_kp_list = []
  new_Hk_list = []

  if spin_orb:
    U_TR = get_U_TR(jchia)

  for i in range(Hksp.shape[0]):
    new_kp = -kp[i]
    if not np.any(np.all(np.isclose(new_kp, kp), axis=1)):
      new_kp_list.append(new_kp)
      if not spin_orb:
        new_Hk_list.append(np.conj(Hksp[i]))
      else:
        #                print(U_inv.shape,U_TR.shape,Hksp.shape)
        new_Hk_list.append(np.conj(U_inv * (U_TR @ Hksp[i] @ np.conj(U_TR.T))))

  if len(new_kp_list) == 0:
    return Hksp, kp

  kp = np.vstack([kp, np.array(new_kp_list)])
  Hksp = np.vstack([Hksp, np.array(new_Hk_list)])

  return Hksp, kp


############################################################################################
############################################################################################
############################################################################################

def enforce_hermaticity(Hksp):
  # enforce H(k) to be hermitian (it should be already)
  for k in range(Hksp.shape[0]):
    Hksp[k] = (Hksp[k] + np.conj(Hksp[k].T)) / 2.0

  return Hksp


############################################################################################
############################################################################################
############################################################################################

def add_U_wyc(U, U_wyc):
  # add operator that shifts orbital from one equivalent site to another
  for isym in range(U.shape[0]):
    U[isym] = U_wyc[isym] @ U[isym]
    U[isym] = correct_roundoff(U[isym], incl_hex=True)

  return U


############################################################################################
############################################################################################
############################################################################################

def wedge_to_grid(Hksp, U, a_index, phase_shifts, kp, new_k_ind, orig_k_ind, si_per_k, inv_flag, U_inv, sym_TR,
                  spin_orb, npool):
  # generates full grid from k points in IBZ
  nawf = Hksp.shape[1]

  fgm = scatter_full(np.arange(si_per_k.shape[0], dtype=int), npool)
  new_k_ind = scatter_full(new_k_ind, npool)
  orig_k_ind = scatter_full(orig_k_ind, npool)
  si_per_k = scatter_full(si_per_k, npool)

  Hksp_s = np.zeros((new_k_ind.shape[0], nawf, nawf), dtype=complex)

  for j in range(new_k_ind.shape[0]):
    isym = si_per_k[j]
    oki = orig_k_ind[j]
    nki = np.where(new_k_ind[j] == fgm)[0]

    # if symop is identity
    if isym == 0:
      Hksp_s[nki] = Hksp[oki]
      continue

    # other cases
    H = Hksp[oki]

    # get k dependent U
    U_k = get_U_k(kp[oki], phase_shifts[isym], a_index, U[isym])

    # transformated H(k)
    THP = U_k @ H @ np.conj(U_k.T)

    # apply inversion operator if needed
    if inv_flag[isym]:
      THP *= U_inv

    # time inversion is anti-unitary
    # old version
    # if sym_TR[isym]:
    #    THP *= U_inv
    #    THP = np.conj(THP)
    if sym_TR[isym]:
      if spin_orb:
        THP *= U_inv
      THP = np.conj(THP)

    Hksp_s[nki] = THP

  # make sure of hermiticity of each H(k)
  Hksp_s = enforce_hermaticity(Hksp_s)

  Hksp = None
  Hksp_s = gather_full(Hksp_s, npool)

  return Hksp_s


############################################################################################
############################################################################################
############################################################################################

def open_grid(Hksp, full_grid, kp, symop, symop_cart, atom_pos, shells, a_index, equiv_atom,
              nk1, nk2, nk3, o1, o2, o3, spin_orb, sym_TR, jchia, mag_calc, symm_grid, thresh, max_iter, verbose, npool):
  # calculates full H(k) grid from wedge

  nawf = Hksp.shape[1]
  # get inversion operator
  U_inv = get_inv_op(shells)

  # apply time reversal symmetry H(k) = H(-k)* where appropriate
  if not (spin_orb and mag_calc):
    Hksp, kp = apply_t_rev(Hksp, kp, spin_orb, U_inv, jchia)

  # get array with wigner_d rotation matrix for each symop
  # for each of the orbital angular momentum l=[0,1,2,3]
  if spin_orb:
    wigner, inv_flag = get_wigner_so(symop_cart)
  else:
    wigner, inv_flag = get_wigner(symop_cart)
    # convert the wigner_d into chemistry form for each symop
    wigner = convert_wigner_d(wigner)

  # get phase shifts from rotation symop
  phase_shifts = get_phase_shifts(atom_pos, symop, equiv_atom)

  # build U and U_inv from blocks
  if spin_orb:
    U = build_U_matrix(wigner, jchia)
  else:
    U = build_U_matrix(wigner, shells)

  # if any symop involve time inversion add the TR to the symop
  if np.any(sym_TR) and spin_orb:
    U = add_U_TR(U, sym_TR, jchia)

  # adds transformation to U that maps orbitals from
  # atom A to equivalent atom B atoms after symop
  U_wyc = map_equiv_atoms(a_index, equiv_atom)

  # combine U_wyc and U
  U = add_U_wyc(U, U_wyc)

  # get index of k in wedge, index in full grid,
  # and index of symop that transforms k to k'
  new_k_ind, orig_k_ind, si_per_k = find_equiv_k(kp, symop, full_grid, sym_TR, check=True)

  # transform H(k) -> H(k')
  Hksp = wedge_to_grid(Hksp, U, a_index, phase_shifts, kp,
                       new_k_ind, orig_k_ind, si_per_k, inv_flag, U_inv, sym_TR, spin_orb, npool)

  if rank == 0:
    # enforce time reversion where appropriate
    if not (spin_orb and mag_calc):
      Hksp = enforce_t_rev(Hksp, nk1, nk2, nk3, spin_orb, U_inv, jchia)

  else:
    Hksp = np.zeros((full_grid.shape[0], nawf, nawf), dtype=complex)

  comm.Bcast(Hksp)

  if symm_grid:

    symop_inv = np.zeros_like(symop)
    for i in range(symop.shape[0]):
      symop_inv[i] = LA.inv(symop[i])

    nkl = []
    partial_grid = scatter_full(full_grid, npool)
    for i in range(partial_grid.shape[0]):
      nkl.append(find_equiv_k(partial_grid[i][None], symop_inv, full_grid, sym_TR, check=False, include_self=True))
    nkl_no_interp = np.array(nkl)
    Hksp, tmax = symmetrize_grid(Hksp, U, a_index, phase_shifts, inv_flag, U_inv, sym_TR,
                                 full_grid, jchia, spin_orb, mag_calc, nk1, nk2, nk3,
                                 nkl_no_interp, partial_grid, npool)
    upscale1 = int(0.25 * nk1)
    upscale2 = int(0.25 * nk2)
    upscale3 = int(0.25 * nk3)

    nfft1 = nk1 + upscale1
    nfft2 = nk2 + upscale2
    nfft3 = nk3 + upscale3

    full_grid_interp = get_full_grid(nfft1, nfft2, nfft3, o1, o2, o3)
    nkl = []
    partial_grid_interp = scatter_full(full_grid_interp, npool)
    for i in range(partial_grid_interp.shape[0]):
      nkl.append(
        find_equiv_k(partial_grid_interp[i][None], symop_inv, full_grid_interp, sym_TR, check=False, include_self=True))
    nkl_interp = np.array(nkl)
    # max difference bewtween H(k) and H(k*)
    tmax = 999999

    for i in range(int(max_iter * 2)):
      st = time.time()
      add1 = upscale1 * ((-1) ** i)
      add2 = upscale2 * ((-1) ** i)
      add3 = upscale3 * ((-1) ** i)

      nfft1 = nk1 + add1
      nfft2 = nk2 + add2
      nfft3 = nk3 + add3

      if rank == 0:
        Hksp = np.reshape(Hksp, (nk1 * nk2 * nk3, nawf * nawf))
        Hksp = np.ascontiguousarray(Hksp.T)

      Hksp = scatter_full(Hksp, npool)

      Hksp = np.reshape(Hksp, (Hksp.shape[0], nk1, nk2, nk3))
      HRs = np.fft.ifftn(Hksp, axes=(1, 2, 3))

      switch = True
      if switch == True:

        Hksp = None
        Hksp = np.zeros((HRs.shape[0], nfft1, nfft2, nfft3), dtype=complex)

        for m in range(Hksp.shape[0]):
          Hksp[m, :, :, :] = np.fft.fftn(zero_pad(HRs[m, :, :, :], nk1, nk2, nk3, add1, add2, add3))
        #                     if not i%2:
        #                         Hksp[m,:,:,:]=np.fft.fftn(zero_pad(HRs[m,:,:,:],nk1,nk2,nk3,add1,add2,add3))
        #                     else:
        # #                        Hksp[m,:,:,:]=np.fft.fftn(LPF(HRs[m,:,:,:],nk1,nk2,nk3,add1,add2,add3))
        #                         Hksp[m,:,:,:]=np.fft.fftn(down_samp(HRs[m,:,:,:],nk1,nk2,nk3,add1,add2,add3))

        HRs = None
        Hksp = np.reshape(Hksp, (Hksp.shape[0], nfft1 * nfft2 * nfft3))
        Hksp = gather_full(Hksp, npool)

        if rank == 0:
          Hksp = np.ascontiguousarray(Hksp.T)
          Hksp = np.reshape(Hksp, (nfft1 * nfft2 * nfft3, nawf, nawf))
        else:
          Hksp = np.zeros((nfft1 * nfft2 * nfft3, nawf, nawf), dtype=complex)

        comm.Bcast(Hksp)

        # Hksp=None
        # Hksp=np.zeros((HRs.shape[0],nfft1,nfft2,nfft3),dtype=complex)

        # for m in range(Hksp.shape[0]):
        #     Hksp[m,:,:,:]=np.fft.fftn(zero_pad(HRs[m,:,:,:],nk1,nk2,nk3,add1,add2,add3))

        # HRs  = None
        # Hksp = np.reshape(Hksp,(Hksp.shape[0],nfft1*nfft2*nfft3))
        # Hksp = gather_full(Hksp,npool)
        # if rank==0:
        #     Hksp = np.ascontiguousarray(Hksp.T)
        #     Hksp = np.reshape(Hksp,(nfft1*nfft2*nfft3,nawf,nawf))
        # else:
        #     Hksp=np.zeros((nfft1*nfft2*nfft3,nawf,nawf),dtype=complex)
        # comm.Bcast(Hksp)

      else:
        HRs = np.zeros((Hksp.shape[0], nfft1, nfft2, nfft3), dtype=complex)
        for m in range(Hksp.shape[0]):
          HRs[m, :, :, :] = zero_pad(Hksp[m, :, :, :], nk1, nk2, nk3, add1, add2, add3)
        Hksp = None
        Hksp = np.fft.fftn(HRs, axes=(1, 2, 3))
        HRs = None
        Hksp = np.reshape(Hksp, (Hksp.shape[0], nfft1 * nfft2 * nfft3))
        Hksp = gather_scatter(Hksp, 1, npool)
        Hksp = np.ascontiguousarray(Hksp.T)
        Hksp = gather_full(Hksp, npool)
        if rank == 0:
          Hksp = np.reshape(Hksp, (nfft1 * nfft2 * nfft3, nawf, nawf))
          print(Hksp[0, :8, :8].real)

        Hksp = np.reshape(Hksp, (nfft1 * nfft2 * nfft3, nawf, nawf))

      # if it's the non interpolated grid
      if i % 2:
        Hksp, _ = symmetrize_grid(Hksp, U, a_index, phase_shifts, inv_flag, U_inv,
                                  sym_TR, full_grid, jchia, spin_orb, mag_calc,
                                  nfft1, nfft2, nfft3, nkl_no_interp, partial_grid, npool)

      # if it's the interpolated grid
      else:
        Hksp, tmax = symmetrize_grid(Hksp, U, a_index, phase_shifts, inv_flag, U_inv,
                                     sym_TR, full_grid_interp, jchia, spin_orb,
                                     mag_calc, nfft1, nfft2, nfft3, nkl_interp, partial_grid_interp, npool)

      nk1 += add1
      nk2 += add2
      nk3 += add3

      if rank == 0 and i % 2 and verbose:
        print("Sym iter #%2d: %6.4e" % ((i // 2) + 1, tmax[0]))

      # stop if we hit threshold
      if tmax < thresh:
        if i % 2 and i >= 3:
          break

  # for debugging purposes
  # try:
  #    if rank==0:
  #        check(Hksp,si_per_k,new_k_ind,orig_k_ind,phase_shifts,U,
  #              a_index,inv_flag,equiv_atom,kp,symop,full_grid,sym_info,sym_TR)
  # except Exception as e:
  #    print(('Exception in pao_sym, check.', e), flush=True)

  return Hksp

def open_grid_nspin2(Hksp, full_grid, kp, symop, symop_cart, atom_pos, shells, a_index, equiv_atom,
              nk1, nk2, nk3, o1, o2, o3, sym_TR, symm_grid, thresh, max_iter, verbose, npool):
  # calculates full H(k) grid from wedge
  nawf = Hksp.shape[1]

  H0 = Hksp[...,0]
  H1 = Hksp[...,1]
  # get inversion operator
  U_inv = get_inv_op(shells)

  # apply time reversal symmetry H(k) = H(-k)* where appropriate
  H0, _ = apply_t_rev(H0, kp, False, U_inv, None)
  H1, kp = apply_t_rev(H1, kp, False, U_inv, None)

  # get array with wigner_d rotation matrix for each symop
  # for each of the orbital angular momentum l=[0,1,2,3]
  wigner, inv_flag = get_wigner(symop_cart)
  # convert the wigner_d into chemistry form for each symop
  wigner = convert_wigner_d(wigner)

  # get phase shifts from rotation symop
  phase_shifts = get_phase_shifts(atom_pos, symop, equiv_atom)

  # build U and U_inv from blocks
  U = build_U_matrix(wigner, shells)

  # adds transformation to U that maps orbitals from
  # atom A to equivalent atom B atoms after symop
  U_wyc = map_equiv_atoms(a_index, equiv_atom)

  # combine U_wyc and U
  U = add_U_wyc(U, U_wyc)

  # get index of k in wedge, index in full grid,
  # and index of symop that transforms k to k'
  new_k_ind, orig_k_ind, si_per_k = find_equiv_k(kp, symop, full_grid, sym_TR, check=True)


  # transform H(k) -> H(k')

  H0 = wedge_to_grid(H0, U, a_index, phase_shifts, kp,
                       new_k_ind, orig_k_ind, si_per_k, inv_flag, U_inv, sym_TR, False, npool)
  H1 = wedge_to_grid(H1, U, a_index, phase_shifts, kp,
                        new_k_ind, orig_k_ind, si_per_k, inv_flag, U_inv, sym_TR, False, npool)

  if rank == 0:
    # enforce time reversion where appropriate

    H0 = enforce_t_rev(H0, nk1, nk2, nk3, False, U_inv, None)
    H1 = enforce_t_rev(H1, nk1, nk2, nk3, False, U_inv, None)

    if np.any(sym_TR) and (not np.all(sym_TR)):
      # determine which mapping of k involves TR
      k_TR = sym_TR[si_per_k]
      Hksp_up = np.zeros_like(H0)
      Hksp_down = np.zeros_like(H1)
      Hksp_up[~k_TR] = H0[~k_TR]
      Hksp_up[k_TR] = H1[k_TR]
      Hksp_down[~k_TR] = H1[~k_TR]
      Hksp_down[k_TR] = H0[k_TR]
    else:
      Hksp_up = H0
      Hksp_down = H1

    Hksp = np.stack((Hksp_up,Hksp_down),axis=-1)

  else:
    Hksp = np.zeros((full_grid.shape[0], nawf, nawf,2), dtype=complex)

  comm.Bcast(Hksp)

  if symm_grid:

    symop_inv = np.zeros_like(symop)
    for i in range(symop.shape[0]):
      symop_inv[i] = LA.inv(symop[i])

    nkl = []
    partial_grid = scatter_full(full_grid, npool)
    for i in range(partial_grid.shape[0]):
      nkl.append(find_equiv_k(partial_grid[i][None], symop_inv, full_grid, sym_TR, check=False, include_self=True))
    nkl_no_interp = np.array(nkl)

    Hksp, tmax = symmetrize_grid_nspin2(Hksp, U, a_index, phase_shifts, inv_flag, U_inv, sym_TR,
                                   full_grid, nk1, nk2, nk3,
                                   nkl_no_interp, partial_grid, npool)
    upscale1 = int(0.25 * nk1)
    upscale2 = int(0.25 * nk2)
    upscale3 = int(0.25 * nk3)

    nfft1 = nk1 + upscale1
    nfft2 = nk2 + upscale2
    nfft3 = nk3 + upscale3

    full_grid_interp = get_full_grid(nfft1, nfft2, nfft3, o1, o2, o3)
    nkl = []
    partial_grid_interp = scatter_full(full_grid_interp, npool)
    for i in range(partial_grid_interp.shape[0]):
      nkl.append(
        find_equiv_k(partial_grid_interp[i][None], symop_inv, full_grid_interp, sym_TR, check=False, include_self=True))
    nkl_interp = np.array(nkl)
    # max difference bewtween H(k) and H(k*)
    tmax = [999999,999999]

    for i in range(int(max_iter * 2)):
      st = time.time()
      add1 = upscale1 * ((-1) ** i)
      add2 = upscale2 * ((-1) ** i)
      add3 = upscale3 * ((-1) ** i)

      nfft1 = nk1 + add1
      nfft2 = nk2 + add2
      nfft3 = nk3 + add3

      if rank == 0:
        Hksp = np.reshape(Hksp, (nk1 * nk2 * nk3, nawf * nawf, 2))
        Hksp = np.ascontiguousarray(np.transpose(Hksp,axes=(1,0,2)))

      Hksp = scatter_full(Hksp, npool)

      Hksp = np.reshape(Hksp, (Hksp.shape[0], nk1, nk2, nk3, 2))
      HRs = np.fft.ifftn(Hksp, axes=(1, 2, 3))

      Hksp = None
      Hksp = np.zeros((HRs.shape[0], nfft1, nfft2, nfft3, 2), dtype=complex)

      for m in range(Hksp.shape[0]):
        Hksp[m, :, :, :, 0] = np.fft.fftn(zero_pad(HRs[m, :, :, :, 0], nk1, nk2, nk3, add1, add2, add3))
        Hksp[m, :, :, :, 1] = np.fft.fftn(zero_pad(HRs[m, :, :, :, 1], nk1, nk2, nk3, add1, add2, add3))

      HRs = None
      Hksp = np.reshape(Hksp, (Hksp.shape[0], nfft1 * nfft2 * nfft3, 2))
      Hksp = gather_full(Hksp, npool)

      if rank == 0:
        Hksp = np.ascontiguousarray(np.transpose(Hksp,axes=(1,0,2)))
        Hksp = np.reshape(Hksp, (nfft1 * nfft2 * nfft3, nawf, nawf, 2))
      else:
        Hksp = np.zeros((nfft1 * nfft2 * nfft3, nawf, nawf, 2), dtype=complex)

      comm.Bcast(Hksp)

      # if it's the non interpolated grid
      if i % 2:
        Hksp, _ = symmetrize_grid_nspin2(Hksp, U, a_index, phase_shifts, inv_flag, U_inv,
                  sym_TR, full_grid, nfft1, nfft2, nfft3, nkl_no_interp, partial_grid, npool)


      # if it's the interpolated grid
      else:
        Hksp, tmax = symmetrize_grid_nspin2(Hksp, U, a_index, phase_shifts, inv_flag, U_inv,
                    sym_TR, full_grid_interp, nfft1, nfft2, nfft3, nkl_interp, partial_grid_interp, npool)


      nk1 += add1
      nk2 += add2
      nk3 += add3

      if rank == 0 and i % 2 and verbose:
        print("Sym iter #%2d: up %6.4e, down %6.4e" % ((i // 2) + 1, tmax[0], tmax[1]))

      # stop if we hit threshold
      if max(tmax) < thresh:
        if i % 2 and i >= 3:
          break

  return Hksp


############################################################################################
############################################################################################
############################################################################################

def open_grid_wrapper(data_controller):
  #    np.set_printoptions(precision=3,suppress=True,linewidth=220)

  # wrapper function to unload everything from
  # data controller and do a few conversions

  data_arrays = data_controller.data_arrays
  data_attr = data_controller.data_attributes
  alat = data_attr['alat']
  # nelec       = data_attr['nelec']
  nk1 = data_attr['nk1']
  nk2 = data_attr['nk2']
  nk3 = data_attr['nk3']
  o1 = data_attr['ok1']
  o2 = data_attr['ok2']
  o3 = data_attr['ok3']
  spin_orb = data_attr['dftSO']
  mag_calc = data_attr['dftMAG']
  symm_grid = data_attr['symmetrize']
  thresh = data_attr['symm_thresh']
  max_iter = data_attr['symm_max_iter']
  verbose = data_attr['verbose']
  Hks = data_arrays['Hks']
  atom_pos = data_arrays['tau'] / alat
  atom_lab = data_arrays['atoms']
  equiv_atom = data_arrays['equiv_atom']
  kp_red = data_arrays['kpnts']
  b_vectors = data_arrays['b_vectors']
  a_vectors = data_arrays['a_vectors']
  # sym_info    = data_arrays['sym_info']
  # sym_shift   = data_arrays['sym_shift']
  symop = data_arrays['sym_rot']
  sym_TR = data_arrays['sym_TR']
  npool = data_attr['npool']

  a_vectors = a_vectors
  b_vectors = b_vectors
  nspin = Hks.shape[3]
  nawf = Hks.shape[0]

  # convert atomic positions to crystal fractional coords
  conv = LA.inv(a_vectors)
  atom_pos = atom_pos @ conv
  atom_pos = correct_roundoff(atom_pos)
  atom_pos = np.around(atom_pos, decimals=6)

  # get symop crystal -> cartesian
  symop = correct_roundoff(symop)
  symop_cart = np.zeros_like(symop, dtype=float)
  inv_a_vectors = LA.inv(a_vectors)
  for isym in range(symop.shape[0]):
    symop_cart[isym] = (inv_a_vectors @ symop[isym] @ a_vectors)

  symop_cart = correct_roundoff(symop_cart, incl_hex=True, atol=1.e-6)

  # convert k points from cartesian to crystal fractional
  conv = LA.inv(b_vectors)
  conv = correct_roundoff(conv)
  kp_red = kp_red @ conv
  kp_red = correct_roundoff(kp_red)

  # get full grid in crystal fractional coords
  full_grid = get_full_grid(nk1, nk2, nk3, o1, o2, o3)

  jchia = []
  shells = []
  a_index = []
  sh = data_arrays['shells']
  for i, a in enumerate(atom_lab):
    ash = []
    for v in sh[a]:
      ash += [v, v] if v == 0 and spin_orb else [v]
    shells += ash
    a_index += [i] * np.sum([2 * n + 1 for n in ash])
    if spin_orb:
      jchia += data_arrays['jchia'][a]

  jchia = np.array(jchia)
  shells = np.array(shells)
  a_index = np.array(a_index)

  # correct small differences due to conversion
  kp_red = correct_roundoff_kp(kp_red, full_grid)

  # we wont need this for now
  if rank == 0:
    if nspin == 2:
      data_arrays['Hks'] = np.zeros((nawf, nawf, nk1 * nk2 * nk3, nspin), dtype=complex)
    else:
      data_arrays['Hks'] = None
  else:
    data_arrays['Hks'] = None

  # expand grid from wedge (older version that works for QE)
  # for ispin in range(nspin):
  #     Hksp = np.ascontiguousarray(np.transpose(Hks, axes=(2, 0, 1, 3))[:, :, :, ispin])
  #
  #     Hksp = open_grid(Hksp, full_grid, kp_red, symop, symop_cart, atom_pos,
  #                      shells, a_index, equiv_atom, sym_info, sym_shift,
  #                      nk1, nk2, nk3, spin_orb, sym_TR, jchia, mag_calc,
  #                      symm_grid, thresh, max_iter, nelec, verbose)
  #
  #     if rank == 0:
  #         if nspin == 2:
  #             data_arrays['Hks'][:, :, :, ispin] = np.ascontiguousarray(np.transpose(Hksp, axes=(1, 2, 0)))
  #         else:
  #             data_arrays['Hks'] = np.ascontiguousarray(np.transpose(Hksp, axes=(1, 2, 0))[..., None])
  #     #            np.save("kham.npy",np.ascontiguousarray(np.transpose(Hksp,axes=(1,2,0))))
  #     else:
  #         Hksp = None

  Hksp = np.ascontiguousarray(np.transpose(Hks, axes=(2, 0, 1, 3)))
  if nspin == 1:
    Hks_full = open_grid(Hksp[...,0], full_grid, kp_red, symop, symop_cart, atom_pos,
                     shells, a_index, equiv_atom,
                     nk1, nk2, nk3, o1, o2, o3, spin_orb, sym_TR, jchia, mag_calc,
                     symm_grid, thresh, max_iter, verbose, npool)
    Hks_full = Hks_full[...,np.newaxis]
  else:
    Hks_full = open_grid_nspin2(Hksp, full_grid, kp_red, symop, symop_cart, atom_pos,
                         shells, a_index, equiv_atom,
                         nk1, nk2, nk3, o1, o2, o3, sym_TR,
                         symm_grid, thresh, max_iter, verbose, npool)
  if rank == 0:
    # if nspin==1:
    #     data_arrays['Hks'] = np.ascontiguousarray(np.transpose(Hksp, axes=(1, 2, 0))[..., None])
    #     #  np.save("kham.npy",np.ascontiguousarray(np.transpose(Hksp,axes=(1,2,0))))
    # else:
    #     data_arrays['Hks'][:, :, :, 0] = np.ascontiguousarray(np.transpose(Hksp_up, axes=(1, 2, 0)))
    #     data_arrays['Hks'][:, :, :, 1] = np.ascontiguousarray(np.transpose(Hksp_down, axes=(1, 2, 0)))
    data_arrays['Hks'] = np.ascontiguousarray(np.transpose(Hks_full, axes=(1, 2, 0, 3)))
  # else:
  #   Hks_full = None


############################################################################################
############################################################################################
############################################################################################

def symmetrize(Hksp, U, a_index, phase_shifts, new_k_ind, orig_k_ind, si_per_k, inv_flag, U_inv, sym_TR, spin_orb,
               full_grid, reverse=False):
  # generates full grid from k points in IBZ
  nawf = Hksp.shape[1]
  Hksp_s = np.zeros((new_k_ind.shape[0], nawf, nawf), dtype=complex)

  for j in range(new_k_ind.shape[0]):
    isym = si_per_k[j]
    oki = orig_k_ind[j]
    nki = new_k_ind[j]

    # if symop is identity
    if isym == 0:
      Hksp_s[j] = Hksp[nki]
      continue

    # other cases
    H = Hksp[nki]

    # get k dependent U
    U_k = get_U_k(full_grid[nki], phase_shifts[isym], a_index, U[isym])

    # transformated H(k)
    if not reverse:
      THP = U_k @ H @ np.conj(U_k.T)
    else:
      THP = np.conj(U_k.T) @ H @ U_k

    # apply inversion operator if needed
    if inv_flag[isym]:
      THP *= U_inv

    # time inversion is anti-unitary
    # old version
    # if sym_TR[isym]:
    #    THP *= U_inv
    #    THP = np.conj(THP)
    if sym_TR[isym]:
      if spin_orb:
        THP *= U_inv
      THP = np.conj(THP)

    Hksp_s[j] = THP

  return Hksp_s


############################################################################################
############################################################################################
############################################################################################

def symmetrize_grid(Hksp, U, a_index, phase_shifts, inv_flag, U_inv, sym_TR, full_grid, jchia, spin_orb,
                    mag_calc, nk1, nk2, nk3, nkl, partial_grid, npool):
  max_iter = 1
  tmax = []
  Hksp_d = np.zeros((partial_grid.shape[0], Hksp.shape[1], Hksp.shape[2]), dtype=complex)

  for i in range(partial_grid.shape[0]):
    new_k_ind, orig_k_ind, si_per_k = nkl[i]

    temp = symmetrize(Hksp, U, a_index, phase_shifts, new_k_ind,
                      orig_k_ind, si_per_k, inv_flag, U_inv, sym_TR, spin_orb, full_grid)

    tmax.append(np.amax(np.abs(temp[0][None] - temp)))
    Hksp_d[i] = np.sum(temp, axis=0) / (temp.shape[0])

  # make sure of hermiticity of each H(k)
  Hksp_d = enforce_hermaticity(Hksp_d)

  if rank == 0:
    Hksp = gather_full(Hksp_d, npool)
    if not (spin_orb and mag_calc):
      Hksp = enforce_t_rev(Hksp, nk1, nk2, nk3, spin_orb, U_inv, jchia)

  else:
    gather_full(Hksp_d, npool)

  tm = np.array([np.amax(tmax)])
  tmax = np.copy(tm)
  comm.Reduce(tm, tmax, op=MPI.MAX)
  comm.Bcast(tmax)

  return Hksp, tmax

def symmetrize_grid_nspin2(Hksp, U, a_index, phase_shifts, inv_flag, U_inv, sym_TR, full_grid,
                     nk1, nk2, nk3, nkl, partial_grid, npool):
  nawf = Hksp.shape[1]

  tmax_up = []; tmax_down = []
  H0 = Hksp[..., 0]
  H1 = Hksp[..., 1]
  Hksp_d_up = np.zeros((partial_grid.shape[0], nawf, nawf), dtype=complex)
  Hksp_d_down = np.zeros((partial_grid.shape[0], nawf, nawf), dtype=complex)
  for i in range(partial_grid.shape[0]):
    new_k_ind, orig_k_ind, si_per_k = nkl[i]
    temp0 = symmetrize(H0, U, a_index, phase_shifts, new_k_ind,
                      orig_k_ind, si_per_k, inv_flag, U_inv, sym_TR, False, full_grid)
    temp1 = symmetrize(H1, U, a_index, phase_shifts, new_k_ind,
                       orig_k_ind, si_per_k, inv_flag, U_inv, sym_TR, False, full_grid)
    if np.any(sym_TR) and (not np.all(sym_TR)):
      temp_up = np.zeros_like(temp0)
      temp_down = np.zeros_like(temp1)
      temp_up[~sym_TR] = temp0[~sym_TR]
      temp_up[sym_TR] = temp1[sym_TR]
      temp_down[~sym_TR] = temp1[~sym_TR]
      temp_down[sym_TR] = temp0[sym_TR]
    else:
      temp_up = temp0
      temp_down = temp1

    tmax_up.append(np.amax(np.abs(temp_up[0][None] - temp_up)))
    tmax_down.append(np.amax(np.abs(temp_down[0][None] - temp_down)))
    Hksp_d_up[i] = np.sum(temp_up, axis=0) / (temp_up.shape[0])
    Hksp_d_down[i] = np.sum(temp_down, axis=0) / (temp_down.shape[0])

  # make sure of hermiticity of each H(k)
  Hksp_d_up = enforce_hermaticity(Hksp_d_up)
  Hksp_d_down = enforce_hermaticity(Hksp_d_down)


  if rank == 0:

    Hksp_up = gather_full(Hksp_d_up, npool)
    Hksp_down = gather_full(Hksp_d_down, npool)

    Hksp_up = enforce_t_rev(Hksp_up, nk1, nk2, nk3, False, U_inv, None)
    Hksp_down = enforce_t_rev(Hksp_down, nk1, nk2, nk3, False, U_inv, None)

    Hksp = np.stack((Hksp_up,Hksp_down),axis=-1)

  else:
    gather_full(Hksp_d_up, npool)
    gather_full(Hksp_d_down, npool)


  tm_up = np.array([np.amax(tmax_up)])
  tm_down = np.array([np.amax(tmax_down)])
  tmax_up = np.copy(tm_up)
  tmax_down = np.copy(tm_down)
  comm.Reduce(tm_up, tmax_up, op=MPI.MAX)
  comm.Reduce(tm_down, tmax_down, op=MPI.MAX)
  comm.Bcast(tmax_up)
  comm.Bcast(tmax_down)
  tmax = [tmax_up[0],tmax_down[0]]

  return Hksp, tmax


############################################################################################
############################################################################################
############################################################################################

def correct_roundoff_kp(kp, full_grid):
  kp_c = np.copy(kp)
  nw = np.where(np.isclose(cdist(kp_c, full_grid), 0.0, atol=1.e-5, rtol=1.e-5, ))
  kp_c[nw[0]] = full_grid[nw[1]]

  return kp_c

############################################################################################
############################################################################################
############################################################################################
