import numpy as np
from readvtk_sno import readvtk_sno
from mpi4py import MPI
from mpi4py_fft.mpifft import PFFT, Function
from mpi4py.MPI import Wtime
comm = MPI.COMM_WORLD


def get_local_wavenumbermesh(FFT, L):
  """Returns local wavenumber mesh."""

  s = FFT.local_slice()
  N = FFT.input_shape()

  # Set wavenumbers in grid
  k = [np.fft.fftfreq(n, 1./n).astype(int) for n in N[:-1]]
  k.append(np.fft.rfftfreq(N[-1], 1./N[-1]).astype(int))
  K = [ki[si] for ki, si in zip(k, s)]
  Ks = np.meshgrid(*K, indexing='ij', sparse=True)
  
  return [np.broadcast_to(k, FFT.local_shape(True)) for k in Ks]


def dshell(shnum1, shnum2, vk, dir):
  out = 1j * K[dir] * vk

  # zero out everything outside the range
  K_floor = np.floor(np.sqrt(K2))
  index = ~ (K_floor > shnum1) & (K_floor < shnum2)

  out[index] = 0.0

  return out


def shell(shnum1, shnum2, vk):

  out = 1.0*vk

  K_floor = np.floor(np.sqrt(K2))
  index = ~ (K_floor > shnum1) & (K_floor < shnum2)

  out[index] = 0.0

  return out


def step(Fk, Gk, Hr, dir):

  c1 = dshell(q1, q2, Fk, dir=dir)
  r1 = Function(FFT, False)
  r1 = FFT.backward(c1, r1)
  # Fix normalization
  r1 /= np.sqrt(np.product(N))

  c1 = shell(k1, k2, Gk)
  r2 = Function(FFT, False)
  r2 = FFT.backward(c1, r2)
  # Fix normalization
  r2 /= np.sqrt(np.product(N))

  return np.sum(r1.real * Hr * r2.real)


def shelltran(k1, k2, q1, q2,  Ak, Br, Ck):

  # Step 1a: A_x^{dK} (B_x dx) C_x^{dQ}
  stran = step(Ck[0], Ak[0], Br[0], dir=0)

  # Step 1b: A_x^{dK} (B_y dy) C_x^{dQ}
  stran += step(Ck[0], Ak[0], Br[1], dir=1)

  # Step 1c: A_x^{dK} (B_z dz) C_x^{dQ}
  stran += step(Ck[0], Ak[0], Br[2], dir=2)

  # Step 2a: A_y^{dK} (B_x dx) C_y^{dQ}
  stran += step(Ck[1], Ak[1], Br[0], dir=0)

  # Step 2b: A_y^{dK} (B_y dy) C_y^{dQ}
  stran += step(Ck[1], Ak[1], Br[1], dir=1)

  # Step 2c: A_y^{dK} (B_z dz) C_y^{dQ}
  stran += step(Ck[1], Ak[1], Br[2], dir=2)

  # Step 3a: A_z^{dK} (B_x dx) C_z^{dQ}
  stran += step(Ck[2], Ak[2], Br[0], dir=0)

  # Step 3b: A_z^{dK} (B_y dy) C_z^{dQ}
  stran += step(Ck[2], Ak[2], Br[1], dir=1)

  # Step 3c: A_z^{dK} (B_z dz) C_z^{dQ}
  stran += step(Ck[2], Ak[2], Br[2], dir=2)

  print('k1,k2,q1,q2,stran:', k1, k2, q1, q2, stran)

  return stran


def shelltran_concise(k1, k2, q1, q2,  Ak, Br, Ck):
  """
  Same as shelltran but using a double loop
  """
  stran = 0.0
  for i in range(3):
    for j in range(3):
      stran += step(Ck[i], Ak[i], Br[j], dir=j)
      msg = 'rank = {}, k1 = {}, k2 = {}, q1 = {}, q2 = {}, stran = {:2.4e},'
      print(msg.format(comm.rank, k1, k2, q1, q2, stran))

  return stran


# SNOOPY data reader
tt, vx, vy, vz, bx, by, bz, N, D = readvtk_sno('v0200.vtk')
vx = vx.astype(float)
vy = vy.astype(float)
vz = vz.astype(float)
bx = bx.astype(float)
by = by.astype(float)
bz = bz.astype(float)

print(np.mean(bx**2))

L = np.array([1.0, 1.0, 1.0], dtype=float)

FFT = PFFT(comm, N, collapse=False)

K = get_local_wavenumbermesh(FFT, L)
K = np.array(K).astype(float)
K2 = np.sum(K*K, 0, dtype=float)

# Velocity
U = Function(FFT, False, tensor=3)
U_hat = Function(FFT, tensor=3)

# Magnetic field
B = Function(FFT, False, tensor=3)
B_hat = Function(FFT, tensor=3)

# local view into the global data is given by local_slice(spectral=False)
# in real space and by local_slice(spectral=True) in k-space
loc_slice = FFT.local_slice(spectral=False)
loc_slice_hat = FFT.local_slice(spectral=True)

# U now contains the three velocity components
U[0] = vx[loc_slice]
U[1] = vy[loc_slice]
U[2] = vz[loc_slice]

# B now contains the three magnetic field components
B[0] = bx[loc_slice]
B[1] = by[loc_slice]
B[2] = bz[loc_slice]

# Perform the FFT using mpi4py_fft.mpifft
for i in range(3):
  U_hat[i] = FFT.forward(U[i], U_hat[i])
  # Fix normalization to be the same as above
  U_hat[i] *= np.sqrt(np.product(N))

for i in range(3):
  B_hat[i] = FFT.forward(B[i], B_hat[i])
  # Fix normalization to be the same as above
  B_hat[i] *= np.sqrt(np.prod(N))

print(B.shape)
print(B_hat.shape)

binn = 'log'

if binn == 'log':
  ktrn  = 25 # for resolution = 256, wavemodes = 128; add 4 for double resolution
  k_num = np.linspace(1,ktrn,ktrn) - 1
  print('Binning is LOG with Grete+ 2017 formula and total wavemodes: ',ktrn+1)
  for ksh in range(2,ktrn-1):
    k_num[ksh] = 1.0 * (2**((ksh-3)/4.+ 2))
    print('ksh,k_num[ksh]: ',ksh,k_num[ksh])
#if binn == 'log':
#  r3 = 8.0
#  rmax = 341
#  ktrn  = 18 # for resolution = 128, wavemodes = 64; add 4 for double this
#  k_num = np.linspace(1,ktrn) - 1
#  print('Binning is LOG with Sundar+ 2017 formula and total wavemodes: ',ktrn+1)
#  for ksh in range(2,ktrn):
#    k_num[ksh] = r3 * ( (rmax/16.0) **((ksh+1)/(ktrn-4)))
#    print('ksh,k_num[ksh]: ',ksh,k_num[ksh])


#ktrn = 4
#k_num = np.array([i for i in range(ktrn)], dtype=np.float32)
if comm.rank == 0:
  print('Binning is log with total wavemodes: ', ktrn+1)
tran_uu = np.zeros((ktrn, ktrn), dtype=np.float32)
tran_ub = np.zeros((ktrn, ktrn), dtype=np.float32)
tran_bu = np.zeros((ktrn, ktrn), dtype=np.float32)
tran_bb = np.zeros((ktrn, ktrn), dtype=np.float32)

t1 = Wtime()
for ksh in range(ktrn-1):
  for qsh in range(ktrn-1):
    k1 = k_num[ksh]
    k2 = k_num[ksh+1]
    q1 = k_num[qsh]
    q2 = k_num[qsh+1]
    tran_uu[ksh, qsh] = shelltran_concise(k1, k2, q1, q2, U_hat, U, U_hat)
    tran_ub[ksh, qsh] = shelltran_concise(k1, k2, q1, q2, U_hat, B, B_hat)
    tran_bu[ksh, qsh] = shelltran_concise(k1, k2, q1, q2, B_hat, B, U_hat)
    tran_bb[ksh, qsh] = shelltran_concise(k1, k2, q1, q2, B_hat, U, B_hat)

tran_uu = comm.reduce(tran_uu, op=MPI.SUM)
tran_ub = comm.reduce(tran_ub, op=MPI.SUM)
tran_bu = comm.reduce(tran_bu, op=MPI.SUM)
tran_bb = comm.reduce(tran_bb, op=MPI.SUM)


#B_tot_r = comm.reduce(np.sum(B[0]*B[0]))
#B_tot_k = comm.reduce(2*np.sum(np.abs(B_hat[0])**2))


#if comm.rank == 0:
#  print(B_tot_r.shape)
#  print(B_tot_k.shape)
#  print('Bx^2:'   ,B_tot_r)
#  print('Bx_hat^2',B_tot_k)
#  print(np.mean(np.abs(B_hat[0])**2)/np.mean(B[0]**2))
#  print(1./np.prod(N))

t2 = Wtime()

if comm.rank == 0:
  print('Elapsed time was {:1.2f} seconds'.format(t2-t1))
  np.savez('tran-commsize-{}.npz'.format(comm.size),
        tran_uu=tran_uu, tran_ub=tran_ub,
        tran_bu=tran_bu, tran_bb=tran_bb)