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.
ERCOFTAC KB Wiki for the AC-5/6 case

File Created: April 09, 2026
J.M. Duró (UPC), A. Miró (UPC), E. Mestres (UPC), O. Lehmkuhl (BSC), I. Rodriguez (UPC)
'''

VTKTYPE = np.bytes_('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)