import numpy as np
import h5py

'''
VTKHDF Mesh Handler

This script provides a 'Dataset' class designed to interface
with the VTKHDF (Unstructured Grid) file format of the data
uploaded in the KBwiki, allowing to reconstruct the global nodes,
and connectivity arrays.

File created: March 11, 2026 
R. Montalà (UPC), O. Lehmkuhl (BSC) and I. Rodriguez (UPC)
'''

VTKTYPE = np.string_('UnstructuredGrid')
VTKVERS = np.array([1,0],np.int32)

def _vtkh5_connectivity_and_offsets(lnods):
	'''
	Build the offsets array (starting point per each element)
	'''
	# Compute the number of points per cell
	ppcell = np.sum(lnods >= 0,axis=1)

	# Compute the number of zeros per cell
	zpcell = np.sum(lnods < 0,axis=1)

	# Flatten the connectivity array
	lnodsf = lnods.flatten('c')

	# Now build the offsets vector
	offset = np.zeros((ppcell.shape[0]+1,),np.int32)
	offset[1:] = np.cumsum(ppcell) + np.cumsum(zpcell)
	return lnodsf, offset


class Dataset(object):
	def __init__(self, points, connec, eltype, varlist, data, time):
		if len(varlist) != len(data):
			raise ValueError('The variable list has a different size than the data provided!')
		self.points  = points
		self.connec  = connec
		self.eltype  = eltype
		self.varlist = varlist
		self.data    = data     ## Data is a list of np.arrays containing the data with the same order as the varlist
		self.time    = time

	@property
	def xyz(self):
		return self.points

	@property
	def npoints(self):
		return self.points.shape[0]
	
	@property
	def nelem(self):
		return self.connec.shape[0]

	@property
	def connectivity(self):
		return self.connec

	def get_variable(self, varname):
		if varname in self.varlist:
			return self.data[self.varlist.index(varname)]
		else:
			raise ValueError(f"Variable '{varname}' not found in varlist.")
	
	def save(self, filename):
		with h5py.File(filename, 'w') as file:

			# Create main group
			main = file.create_group('VTKHDF')
			main.attrs['Type']    = VTKTYPE
			main.attrs['Version'] = VTKVERS

			# Create cell data group
			main.create_group('CellData')
			main.create_group('PointData')
			main.create_group('FieldData')

			# Create dataset for number of points
			npoints, ndim = self.xyz.shape
			main.create_dataset('NumberOfPoints',(1,),dtype=int,data=npoints)
			main.create_dataset('Points',(npoints,ndim),dtype=np.double,data=self.xyz)

			# Create dataset for number of cells
			lnods, offsets = _vtkh5_connectivity_and_offsets(self.connec)
			ncells = self.eltype.shape[0]
			ncsize = lnods.shape[0]
			main.create_dataset('NumberOfCells',(1,),dtype=int,data=ncells)
			main.create_dataset('NumberOfConnectivityIds',(1,),dtype=int,data=ncsize)
			main.create_dataset('Connectivity',(ncsize,),dtype=int,data=lnods)
			main.create_dataset('Offsets',(ncells+1,),dtype=int,data=offsets)
			main.create_dataset('Types',(ncells,),dtype=np.uint8,data=self.eltype)
			main['FieldData'].create_dataset('TimeValue',(1,),dtype=float,data=self.time)

			# Save arrays as point data
			for varname, array in zip(self.varlist, self.data):
				main['PointData'].create_dataset(varname,array.shape,dtype=array.dtype,data=array)

	@classmethod
	def load(cls, filename):
		with h5py.File(filename, 'r') as file:
			main = file['VTKHDF']

			# Load mesh information
			points = main['Points'][:]
			npoints_x_rank = main['NumberOfPoints'][:]
			nel_x_rank = main['NumberOfCells'][:]
			nconnec_x_rank = main['NumberOfConnectivityIds'][:]
			connec = main['Connectivity'][:] 	# Local indices in each rank
			offsets = main['Offsets'][:] 		# Local indices in each rank
			eltype = main['Types'][:]

			# Reconstruct global connectivity
			glob_connec = connec.copy()
			connec_start_indices = np.concatenate(([0], np.cumsum(nconnec_x_rank)))
			point_offsets = np.concatenate(([0], np.cumsum(npoints_x_rank[:-1])))
        
			for i in range(len(npoints_x_rank)):
				start = connec_start_indices[i]
				end = connec_start_indices[i+1]
				glob_connec[start:end] += point_offsets[i] # Add the global point offset to the connectivity IDs of this rank

			# Reconstruct global offsets
			connec_cumulative_offsets = np.concatenate(([0], np.cumsum(nconnec_x_rank[:-1])))
			glob_offsets = []

			cell_start_idx = 0
			for i in range(len(nel_x_rank)):
				num_cells_in_rank = nel_x_rank[i]
				local_offset_block = offsets[cell_start_idx : cell_start_idx + num_cells_in_rank + 1]
				global_block = local_offset_block + connec_cumulative_offsets[i]
    
				if i < len(nel_x_rank) - 1:
					glob_offsets.append(global_block[:-1])
				else:
					glob_offsets.append(global_block)
				cell_start_idx += num_cells_in_rank + 1
			
			glob_offsets = np.concatenate(glob_offsets)

			# Reshape the connectivity array as [elements, points_x_elem]
			# (Valid for a fixed elem_type)
			reconstructed_connec = []
			for i in range(len(glob_offsets) - 1):
				reconstructed_connec.append(glob_connec[glob_offsets[i]:glob_offsets[i+1]])
			glob_connec = np.array(reconstructed_connec)

			# Load time array from the FieldData group
			time = main['FieldData']['TimeValue'][:]

			# Load variable list and data from the PointData group
			arrays_group = main['PointData']
			varlist = list(arrays_group.keys())
			data = [arrays_group[varname][:] for varname in varlist]

		return cls(points, glob_connec, eltype, varlist, data, time)