#!/usr/bin/env python
#
# pyQvarsi, LAMINE.
#
# Live Analysis Multicore Integration for Numerical Enginering. 
# This module allows to pull and repartition data from a SmartRedis 
# database for insitu visualization
#
# Last rev: 24/10/2025
from __future__ import print_function, division

import os, numpy as np

from .cr                   import cr
from .inp_out.SmartRedisIO import SMARTREDIS_VAR2IDX, SMARTSIM_init_client, SMARTSIM_pull_data, SMARTSIM_renumber_data, SMARTSIM_repartition_data
from .inp_out.sod2dIO      import SOD2DJSON_read, SOD2DHDF_read_mapping
from .utils.parallel       import MPI_SIZE, MPI_RANK, core_renumbering, build_interaction_matrix


@cr('LAMINE.read_json')
def read_json(casestr, basedir='./'):
	filename = os.path.join(basedir,'%s.json'%casestr)
	data     = SOD2DJSON_read(filename)
	CASESTR  = data['mesh_h5_file_name']
	ntime    = data['save_restartFile_step']
	nsend    = data['nlamine']
	isrest   = data['loadRestartFile']
	iiload   = data['restartFile_to_load'] if isrest else 2
	return CASESTR, ntime, nsend, isrest, iiload

@cr('LAMINE.init_client')
def init_client(SSDB=os.getenv('SSDB')):
	return SMARTSIM_init_client(SSDB)

@cr('LAMINE.cpu_renumbering')
def cpu_renumbering(nCPUpn, nGPUpn, nGPUs, result_type='ins'):
	rerank, CPUlist = core_renumbering(nCPUpn, nGPUpn, nGPUs)
	return  rerank, CPUlist, "%s.x.%i"%(result_type, rerank)

@cr('LAMINE.repartition_mesh')
def repartition_mesh(casestr, nSource, targetPtable, basedir='./'):
	if MPI_SIZE > nSource:
		# Read mapping between postprocessing and running meshes
		start, end = targetPtable.partition_bounds(MPI_RANK, 'Points')
		ranks, map = SOD2DHDF_read_mapping(casestr, nSource, basedir=basedir)
		myranksG   = ranks[start:end]              
		myranks    = np.sort(np.unique(myranksG))  
		# Interaction matrix
		who2send, partsend = build_interaction_matrix(nSource,myranks,targetPtable)
	else:
		who2send, partsend, ranks, myranksG, map = None, None, None, None, None
	return who2send, partsend, ranks, myranksG, map

@cr('LAMINE.pull_data')
def pull_data(client, rerank, data_key, nGPUs, CPUlist, npoints, ranks, myranksG, map, who2send, partsend, *, varList=['u','pr'], dtype=np.float32, VAR2IDX=SMARTREDIS_VAR2IDX):
	buff = SMARTSIM_pull_data(client,rerank,data_key,nGPUs,dtype)
	buff = SMARTSIM_renumber_data(buff,rerank,CPUlist,nGPUs)
	return SMARTSIM_repartition_data(buff,npoints,nGPUs,ranks,myranksG,map,who2send,partsend,varList,VAR2IDX)
