import bpy
from bpy.props import StringProperty
from bpy_extras.io_utils import ImportHelper
import rasterio

class GridVertexGenerator():
    def __init__(self, *, grid_size_x, grid_size_y, cell_size_x, cell_size_y, dem_array, map_scale, vertical_exaggeration):
        self.cell_size_x = cell_size_x
        if cell_size_y is not None:
            self.cell_size_y = cell_size_y
        else:
            self.cell_size_y = cell_size_x
        self.grid_size_x = grid_size_x
        self.grid_size_y = grid_size_y
        self.dem_array = dem_array
        self.map_scale = map_scale
        self.vertical_exaggeration = vertical_exaggeration

    def __len__(self):
        return self.grid_size_x * self.grid_size_y

    def __iter__(self):
        return self._generator()

    def _generator(self):
        for j in range(self.grid_size_y):
            for i in range(self.grid_size_x):
                yield (i * self.cell_size_x * self.map_scale,
                       j * self.cell_size_y * self.map_scale,
                       self.dem_array[j][i] * self.map_scale * self.vertical_exaggeration)

class GridQuadGenerator():
    def __init__(self, grid_size_x, grid_size_y):
        self.grid_size_x = grid_size_x
        self.grid_size_y = grid_size_y

    def __len__(self):
        return (self.grid_size_x - 1) * (self.grid_size_y - 1)

    def __iter__(self):
        return self._generator()

    def _generator(self):
        for i in range(self.grid_size_x - 1):
            for j in  range(self.grid_size_y - 1):
                v00 = j * self.grid_size_x + i
                v01 = (j + 1) * self.grid_size_x + i
                v10 = j * self.grid_size_x + i + 1
                v11 = (j + 1) * self.grid_size_x + i + 1
                yield (v00, v10, v11, v01)


class OBJECT_OT_create_mesh_from_geotiff(bpy.types.Operator):
    """Create a new object from a GeoTIFF"""
    bl_idname = "object.create_mesh_from_geotiff"
    bl_label = "Mesh from GeoTIFF"

    @classmethod
    def poll(cls, context):
        return context.scene.mapping_extension_geotiff_filename

    def execute(self, context):
        map_origin = context.scene.mapping_extension_map_origin
        map_scale = 1.0 / context.scene.mapping_extension_map_scale
        vertical_exaggeration = context.scene.mapping_extension_vertical_exaggeration
        filename = context.scene.mapping_extension_geotiff_filename
        with rasterio.open(filename) as geotiff:
            grid_size_x = geotiff.width
            grid_size_y = geotiff.height
            print(f"GeoTIFF is {grid_size_x}x{grid_size_y}")
            bounds = geotiff.bounds
            print(f"GeoTIFF bounds: {bounds}")
            cell_size_x = (bounds.right - bounds.left) / geotiff.width
            cell_size_y = (bounds.bottom - bounds.top) / geotiff.height
            print(f"Calculated cell size: {cell_size_x} x {cell_size_y}")
            print("Reading GeoTIFF...")
            dem_array = geotiff.read(1)
            print("Processing values...")
            print("Done.")
        mesh = bpy.data.meshes.new("mesh")
        dem_obj = bpy.data.objects.new("DEM", mesh)
        collection = bpy.data.collections["Collection"]
        collection.objects.link(dem_obj)
        context.view_layer.objects.active = dem_obj
        print("Creating mesh data...")
        (vertices, edges, faces) = self.construct_mesh(grid_size_x=grid_size_x,
                                                       grid_size_y=grid_size_y,
                                                       cell_size_x=cell_size_x,
                                                       cell_size_y=cell_size_y,
                                                       dem_array=dem_array,
                                                       map_scale=map_scale,
                                                       vertical_exaggeration=vertical_exaggeration)
        print("Constructing mesh...")
        mesh.from_pydata(vertices, edges, faces, shade_flat=False)
        print("Done.")
        mesh.validate(verbose=True)
        return {'FINISHED'}

    def construct_mesh(self, *, grid_size_x, grid_size_y, cell_size_x, cell_size_y, dem_array, map_scale, vertical_exaggeration):
        return (GridVertexGenerator(grid_size_x=grid_size_x,
                                    grid_size_y=grid_size_y,
                                    cell_size_x=cell_size_x,
                                    cell_size_y=cell_size_y,
                                    dem_array=dem_array,
                                    map_scale=map_scale,
                                    vertical_exaggeration=vertical_exaggeration),
                [],
                GridQuadGenerator(grid_size_x, grid_size_y))

def menu_func(self, context):
    self.layout.operator(CreateMeshFromGeotiffOperator.bl_idname,
                         text=CreateMeshFromGeotiffOperator.bl_label)