#!/usr/bin/env python
#
# SMARTREDIS Input Output
#
# Last rev: 24/10/2025
from __future__ import print_function, division

import numpy as np

# pyQvarsi imports
from ..cr             import cr
from ..utils.common   import raiseError
from ..utils.parallel import MPI_SIZE, MPI_RANK
from ..utils.parallel import mpi_barrier, mpi_send, mpi_recv, mpi_bcast


# Dict that relates how the variables are stored
# in the SMARTREDIS databased, hardcoded in the code
SMARTREDIS_VAR2IDX = {
	'u'     : [0,1,2],
	'pr'    : [3],
	'rho'   : [4],
	'avvel' : [0,1,2],
	'avpre' : [3],
	'avrho' : [4],
}

try:
	# SmartRedis imports
	from smartredis import Client

	def SMARTSIM_init_client(SSDB):
		return Client(address=SSDB,cluster=False)
except:
	def SMARTSIM_init_client(SSDB):
		raiseError('Import - Problems loading smartredis!',all=False)


@cr('pull_data')
def SMARTSIM_pull_data(client, rerank, key, nGPUs, dtype):
	'''
	Pulling data in this order: u, v, w, pr, rho
	'''
	for irank in range(nGPUs):
		if rerank == irank:
			data  = client.get_tensor(key).astype(dtype)
		mpi_barrier()
	if rerank < nGPUs:
		buff = np.vstack((np.array(data[0::5], dtype=dtype), np.array(data[1::5], dtype=dtype), np.array(data[2::5], dtype=dtype), np.array(data[3::5], dtype=dtype))).T ## ::5 because we can have up to 5 values.
	else:
		buff= np.empty((0,0), dtype=dtype)
	mpi_barrier()
	return buff


@cr('renummber_data')
def SMARTSIM_renumber_data(buff, rerank, CPUlist, nGPUs):
	if MPI_RANK != rerank:
		if MPI_RANK in CPUlist:
			mpi_send(buff,dest=np.argwhere(CPUlist==MPI_RANK))
		elif MPI_RANK < nGPUs:
			buff = mpi_recv(source=CPUlist[MPI_RANK])
	mpi_barrier()
	return buff


@cr('send_empty_partitions')
def SMARTSIM_repartition_data(buffdata, npoints, nGPUs, ranks, myranksG, mapping, who2send, partsend, varList, VAR2IDX):
	'''
	Send data to the empty CPUs when executing the postprocesing with more partitions than the running mesh
	'''
	varDict = {v:None for v in varList}
	if MPI_SIZE > nGPUs:
		for v in varDict.keys():
			varDict[v] = np.zeros((npoints,3) if v in ['u','avvel'] else (npoints,),dtype=buffdata.dtype)
		for irank in range(nGPUs):
			receivers = mpi_bcast(who2send, root=irank)
			if MPI_RANK == irank:
				for jj, jrank in enumerate(who2send):
					localmapping = mapping[partsend[jj,0]:partsend[jj,1]]
					localranks   = ranks[partsend[jj,0]:partsend[jj,1]]
					sendnodes    = localmapping[localranks==irank]
					if irank != jrank:
						mpi_send(buffdata[sendnodes],dest=jrank)
					else:
						for v in varDict.keys():
							idx = VAR2IDX[v]
							if len(idx) > 1:
								for j,i in enumerate(idx):
									varDict[v][myranksG==irank,j] = buffdata[sendnodes,idx[i]]
							else:
								varDict[v][myranksG==irank] = buffdata[sendnodes,idx[0]]
			else:
				for jrank in receivers:
					if MPI_RANK == jrank:
						buffrecv = mpi_recv(source=irank)
						for v in varDict.keys():
							idx = VAR2IDX[v]
							if len(idx) > 1:
								for j,i in enumerate(idx):
									varDict[v][myranksG==irank,j] = buffrecv[:,idx[i]]
							else:
								varDict[v][myranksG==irank] = buffrecv[:,idx[0]]
	else:
		for v in varDict.keys():
			idx = VAR2IDX[v]
			if len(idx) > 1:
				for j,i in enumerate(idx):
					varDict[v][myranksG==irank,j] = buffdata[:,idx[i]]
			else:
				varDict[v][myranksG==irank] = buffdata[:,idx[0]]
	mpi_barrier()
	return varDict
