#mission by nickoe
import bpy
import csv
import math
import bmesh


def gridGet(grid, u, v):
    # grid is a sparse rectangular array, a list of lists
    col = grid[u]
    if v<len(col):
        return col[v]
    return None

def gridPut(grid, u, v, val):
    while (len(grid) <= u):
        grid.append([])
    col = grid[u]
    while (len(col) <= v):
        col.append(None)
    col[v] = val

class MinMax:
    """This object keeps track of the minimum and maximum value in a data set"""
    def __init__(self):
        self.min = None
        self.max = None

    def load(self, v):
        if self.min is None:
            self.min = v
            self.max = v
        else:
            if (v < self.min):
                self.min = v
            if (v > self.max):
                self.max = v

def toUVCoord(x,y,z, du, dv, minMaxZ):
    """ return a 2-element list usable as a UV texture coordinate"""
    return [ (y/dv +x)/du, (z-minMaxZ.min)/(minMaxZ.max-minMaxZ.min) ]

def loadBathymetry(fname):
    """ load bathymetry data from fname.  
The format is space-separated floating point values (x y z) one triple per line.  
We assume that the X and Y coordinates are at intervals of 1 
and construct a mesh of squares wherever we have all four corners. """
    r2 = csv.reader(open(fname, "r"), delimiter=' ')

    minMaxX = MinMax()
    minMaxY = MinMax()
    minMaxZ = MinMax()

    # scan for min/max values
    for row in r2:
        x_,y_,z_ = row
        z = float(z_)
        x = float(x_)
        y = float(y_)
        minMaxX.load(x)
        minMaxY.load(y)
        minMaxZ.load(z)

    # make a 2nd pass
    r2 = csv.reader(open(fname, "r"), delimiter=' ')

    verts = []
    grid = []

    # build a list of vertices 
    # and memorize the vertex index of each cell
    # in the sparse matrix (grid)
    for row in r2:
        x_,y_,z_ = row
        z = float(z_)
        x = float(x_)-minMaxX.min
        y = float(y_)-minMaxY.min
        u = math.floor(x)
        v = math.floor(y)

        idx = len(verts)
        verts.append( [ x,y,z] )
#        print( "%d %d %f"%(u,v,z) )

        gridPut(grid, u, v, idx)

    faces = []
    uvStash = []

    du = math.ceil(minMaxX.max-minMaxX.min)
    dv = math.ceil(minMaxY.max-minMaxY.min)
    # iterate through all imagined squares in the sparse matrix 
    # and build faces for each one that has all four corners present
    for u in range(0, du):
        for v in range(0, dv):

            # get the vertex indices for the corners
            zi1 = gridGet(grid, u, v)
            zi2 = gridGet(grid, u+1, v)
            zi4 = gridGet(grid, u, v+1)
            zi3 = gridGet(grid, u+1, v+1)

            # if all four corners have data
            if not (zi1 is None or zi2 is None or zi3 is None or zi4 is None):

                # make a face for it
                faces.append( [zi1, zi2, zi3, zi4] )
                
                # memorize some UV coordinates for that face as well
                uvStash.append( [ 
                    toUVCoord(u, v, verts[zi1][2], du, dv, minMaxZ),
                    toUVCoord(u+1, v, verts[zi2][2], du, dv, minMaxZ),
                    toUVCoord(u+1, v+1, verts[zi3][2], du, dv, minMaxZ),
                    toUVCoord(u, v+1, verts[zi4][2], du, dv, minMaxZ),
                ] )

    # build a mesh from the vertex and face list
    mesh = bpy.data.meshes.new("bathymetry")
    mesh.from_pydata(verts, [], faces )

    # create a new UV layer
    mesh.uv_textures.new("depth")

    bm = bmesh.new()
    bm.from_mesh(mesh)
    bm.faces.ensure_lookup_table()
    uv_layer = bm.loops.layers.uv[0]

    # loop through the faces and set the UV coordinates
    # according to the values we stashed in the earlier loop.
    for fi in range(len(bm.faces)):
        uvs = uvStash[fi]
        for vi in range(len(uvs)):
            bm.faces[fi].loops[vi][uv_layer].uv = uvs[vi]
#            print("face[%d] .uv[%d] = %s" % (fi, vi, uvs[vi]) )

    bm.to_mesh(mesh)

    return mesh


def getGradientImage():
    """rainbow gradient image"""
    rval = bpy.data.images.get("gradient")
#    print(rval)
    if rval is None:
        rval = bpy.data.images.new("gradient", 1, 6)

        rval.pixels = [
            #R,G,B,A,
            1,0,1,1, # magenta
            0,0,1,1, # blue
            0,1,1,1, # cyan
            0,1,0,1, # green
            1,1,0,1, # yellow
            1,0,0,1, # red
        ]
        # if we don't pack the image into the .blend, it will be lost when you re-load the project
        rval.pack(True)
    return rval

def getGradientTexture():
    """rainbow gradient texture"""
    rval = bpy.data.textures.get("gradient")
#    print(rval)
    if rval is None:
        rval = bpy.data.textures.new("gradient", 'IMAGE')
        rval.image = getGradientImage()
        print(rval.image)
    return rval

def getGradientMaterial():
    """rainbow gradient material from UV layer 'depth' """
    rval = bpy.data.materials.get("gradient")
    print(rval)
    if rval is None:
        rval = bpy.data.materials.new("gradient")
        rval.texture_slots.add()
        ts = rval.texture_slots[0]
        ts.texture = getGradientTexture()
        ts.texture_coords='UV'
        ts.uv_layer = "depth"

    return rval

def set_UV_editor_texture(mesh):
    """ set the image for the face.tex layer on all the faces 
    so we have a rough idea of what the mesh will look like 
    in the 3D view's Texture render mode"""

    # load the mesh data into a bmesh object
    bm = bmesh.new()
    bm.from_mesh(mesh)

    bm.faces.ensure_lookup_table()

    # Get the "tex" layer for the first UV map
    # If you don't already have a UV map, why are you even calling this function?
    tex_layer = bm.faces.layers.tex[mesh.uv_layers[0].name]

    for i in range(len(bm.faces)):

        # figure out which material this face uses
        mi = bm.faces[i].material_index
        mat = mesh.materials[mi]

        # Assume that we want to use the image from the first texture slot;
        # and assume that the material has a texture in that first slot;
        # and assume that the texture is an image texture instead of a procedural texture.
        # if any of several assumptions are wrong, this will explode
        img = mat.texture_slots[0].texture.image

        bm.faces[i][tex_layer].image = img

    # copy the modified data into the mesh
    bm.to_mesh(mesh)

#
#

if False:
    # debugging code to blow away old data blocks
    try:
        bpy.context.scene.objects.unlink(bpy.data.objects.get("bathymetry"))
    except:
        pass
    try:
        bpy.data.objects.remove(bpy.data.objects.get("bathymetry"))
    except:
        pass
    try:
        bpy.data.meshes.remove(bpy.data.meshes.get("bathymetry"))
    except:
        pass
    try:
        bpy.data.materials.remove(bpy.data.materials.get("gradient"))
    except:
        pass
    try:
        bpy.data.textures.remove(bpy.data.textures.get("gradient"))
    except:
        pass
    try:
        bpy.data.images.remove(bpy.data.images.get("gradient"))
    except:
        pass

# build the mesh and UVs from the CSV data
mesh = loadBathymetry("/var/tmp/testdata.csv")
# put the gradient material on the mesh
mesh.materials.append(getGradientMaterial())
# create the object for the mesh
obj = bpy.data.objects.new("bathymetry", mesh)

# put an image on the UV layer so we can see the 
# colors in texture mode of the 3d view
set_UV_editor_texture(obj.data)

# link the object to the current scene
bpy.context.scene.objects.link(obj)
obj.select = True
bpy.context.scene.objects.active = obj

print("done")