import json
import vtk
import numpy as np
import geojson
from pygltflib import GLTF2, Buffer, Accessor, Mesh, Primitive, BufferView
# # Create a grid
# grid_size = 50
# x = np.linspace(0, 10, grid_size)
# y = np.linspace(0, 10, grid_size)
# z = np.linspace(0, 10, grid_size)
# x, y, z = np.meshgrid(x, y, z)

# # Create scalar values based on a Gaussian distribution
# values = np.exp(-((x - 5)**2 + (y - 5)**2 + (z - 5)**2) / 10)

# # Convert grid points and values to VTK format
# points = vtk.vtkPoints()
# scalars = vtk.vtkDoubleArray()
# scalars.SetName("HeatValues")

# for i in range(grid_size):
#     for j in range(grid_size):
#         for k in range(grid_size):
#             points.InsertNextPoint(x[i, j, k], y[i, j, k], z[i, j, k])
#             scalars.InsertNextValue(values[i, j, k])

# # Create a structured grid
# structured_grid = vtk.vtkStructuredGrid()
# structured_grid.SetDimensions(grid_size, grid_size, grid_size)
# structured_grid.SetPoints(points)
# structured_grid.GetPointData().SetScalars(scalars)

# # Write the grid to a VTK file
# vtk_writer = vtk.vtkStructuredGridWriter()
# vtk_writer.SetFileName("heatmap.vtk")
# vtk_writer.SetInputData(structured_grid)
# vtk_writer.Write()



def load_vtk(vtk_file):
    # 使用 vtkStructuredGridReader 来读取 structured_grid 数据类型
    reader = vtk.vtkStructuredGridReader()
    reader.SetFileName(vtk_file)
    reader.Update()

    # 获取数据集
    dataset = reader.GetOutput()

    # 提取点数据
    points = dataset.GetPoints()
    point_data = points.GetData()

    # 由于结构化网格的数据是按规则网格组织的，我们直接获取网格尺寸
    dimensions = dataset.GetDimensions()

    return point_data, dimensions  # 返回维度而不是 cells


def vtk_to_geojson(point_data, dimensions):
    # 转换 VTK 点数据
    points = []
    for i in range(point_data.GetNumberOfTuples()):
        points.append(list(point_data.GetTuple3(i)))

    # 这里我们不直接使用 cells，而是基于维度构建网格的多边形
    features = []
    x_dim, y_dim, z_dim = dimensions
    for x in range(x_dim - 1):
        for y in range(y_dim - 1):
            for z in range(z_dim - 1):
                # 获取每个方格的四个角
                pts = [
                    points[x + y * x_dim + z * x_dim * y_dim],
                    points[(x + 1) + y * x_dim + z * x_dim * y_dim],
                    points[(x + 1) + (y + 1) * x_dim + z * x_dim * y_dim],
                    points[x + (y + 1) * x_dim + z * x_dim * y_dim]
                ]
                # 创建多边形面
                features.append({
                    "type": "Feature",
                    "geometry": {
                        "type": "Polygon",
                        "coordinates": [pts]
                    },
                    "properties": {}
                })

    # 构建 GeoJSON 数据
    geojson_data = geojson.FeatureCollection(features)
    return geojson_data


def vtk_to_gltf(point_data, dimensions):
    # 提取点数据
    vertices = []
    for i in range(point_data.GetNumberOfTuples()):
        vertices.append(list(point_data.GetTuple3(i)))

    # 创建 GLTF 数据结构
    gltf = GLTF2()

    # 将点数据转换为 numpy 数组
    vertices = np.array(vertices, dtype=np.float32).flatten()

    # 创建缓冲区
    buffer = Buffer(uri="data.bin", byteLength=len(vertices) * 4)
    gltf.buffers.append(buffer)

    # 创建缓冲区视图
    bufferViewVertices = BufferView(
        buffer=0,
        byteOffset=0,
        byteLength=len(vertices) * 4,
        target=34962
    )
    gltf.bufferViews.append(bufferViewVertices)

    # 创建访问器
    accessorVertices = Accessor(
        bufferView=0,
        byteOffset=0,
        componentType=5126,
        count=len(vertices) // 3,
        type="VEC3"
    )
    gltf.accessors.append(accessorVertices)

    # 返回 GLTF 数据
    return gltf


def save_geojson(geojson_data, filename):
    with open(filename, 'w') as f:
        json.dump(geojson_data, f)


def save_gltf(gltf, filename):
    gltf.save(filename)


def main():
    # 读取 VTK 文件
    vtk_file = 'heatmap.vtk'
    point_data, dimensions = load_vtk(vtk_file)

    # 转换为 GeoJSON
    geojson_data = vtk_to_geojson(point_data, dimensions)
    save_geojson(geojson_data, 'output.geojson')

    # 转换为 GLTF
    gltf = vtk_to_gltf(point_data, dimensions)
    save_gltf(gltf, 'output.gltf')


if __name__ == '__main__':
    main()