import bpy
from mathutils import noise
import bmesh
from math import *

factor = 3.5 # How far up and down the landscape goes
snow_level = 1.0 #How high for snow
water_level = -.2 # What level is the water
frequency = 0.3 # affects gradient
height = 0.0 # Height offset of the landscape
subdivisions = 80


def make_vertex_color_material():
    mat = bpy.data.materials.new(name="VertexColorMaterial")
    mat.use_nodes = True
    bsdf = mat.node_tree.nodes.get("Principled BSDF")
    color_node = mat.node_tree.nodes.new(type='ShaderNodeVertexColor')
    color_node.layer_name = "Col"
    mat.node_tree.links.new(color_node.outputs['Color'], bsdf.inputs['Base Color'])
    return mat

def delete_all_objects():
    bpy.ops.object.select_all(action='DESELECT')
    bpy.ops.object.select_by_type(type='MESH')
    bpy.ops.object.delete()

def make_key(co):
    return (round(co[0], 3), round(co[1], 3))

delete_all_objects()

#bpy.ops.mesh.primitive_grid_add(size=20.0, x_subdivisions=100, y_subdivisions=100, location=(0, 0, water_level))
#obj = bpy.context.active_object
#
#obj.data.materials.append(make_vertex_color_material())
#
#bpy.ops.object.mode_set(mode='EDIT')
#bm = bmesh.from_edit_mesh(obj.data)
#bm.faces.ensure_lookup_table()
#
#color_layer = bm.loops.layers.color.new("Col")
#for i, face in enumerate(bm.faces):
#    for loop in face.loops:
#        loop[color_layer] = (0, 0, 1, .3)
#
#bmesh.update_edit_mesh(obj.data)

# bpy.ops.object.mode_set(mode='OBJECT')
bpy.ops.mesh.primitive_grid_add(size=20.0, x_subdivisions=subdivisions, y_subdivisions=subdivisions)
step = 20.0 / subdivisions

obj = bpy.context.active_object

v_heights = {}
for v in obj.data.vertices:
    v.co.z += height + (noise.noise(frequency * v.co) * factor)
    v_heights[make_key((v.co.x, v.co.y))] = v.co.z

#for v in obj.data.vertices:
#    v.co.xy *= 1.0 + noise.noise(20*v.co)/5

obj.data.materials.append(make_vertex_color_material())

bpy.ops.object.mode_set(mode='EDIT')
bm = bmesh.from_edit_mesh(obj.data)
bm.faces.ensure_lookup_table()

color_layer = bm.loops.layers.color.new("Col")
face_color = (1, 1, 1, 1)
for i, face in enumerate(bm.faces):
    for loop in face.loops:
        elev = loop.vert.co.z
        slope = 0
        water_direction = 0
        if abs(loop.vert.co.y) < 9.75 and abs(loop.vert.co.x) < 9.75 :
            e1 = v_heights[make_key((loop.vert.co.x - step, loop.vert.co.y + step))]
            e2 = v_heights[make_key((loop.vert.co.x + step, loop.vert.co.y - step))]
            s1 = e1 - e2
            e1 = v_heights[make_key((loop.vert.co.x, loop.vert.co.y + step))]
            e2 = v_heights[make_key((loop.vert.co.x, loop.vert.co.y - step))]
            s2 = e1 - e2
            e1 = v_heights[make_key((loop.vert.co.x + step, loop.vert.co.y + step))]
            e2 = v_heights[make_key((loop.vert.co.x - step, loop.vert.co.y - step))]
            s3 = e1 - e2
            e1 = v_heights[make_key((loop.vert.co.x + step, loop.vert.co.y))]
            e2 = v_heights[make_key((loop.vert.co.x - step, loop.vert.co.y))]
            s4 = e1 - e2
            slopes = [s1, s2, s3, s4]
            slope_values = [abs(s) for s in slopes]
            slope = max(slope_values)
            max_index = slope_values.index(slope)
            water_direction = max_index * pi/4
            if(slopes[max_index] < 0):
                water_direction += pi
# For the adjustment:
#   We want to move the lower vertices of the face under the upper vertices
#   All the way would make this face vertical (might not want that)
#   Or we could make the lower vertices lower and the higher ones higher
#   Either way, we need to figure out which vertices the lower vertices are!
        if abs(elev - water_level) < 0.4 and slope < 0.5:
            loop[color_layer] = (0, .4, .1, 1)
        elif abs(elev - water_level) < 0.4:
            cval = water_direction / (2*pi)
            loop[color_layer] = (cval, 0, 1-cval, 1)
        else:
            loop[color_layer] = (.35, .35, .35, 1) 
        continue
        if slope > 0.5:
            loop[color_layer] = (.35, .35, .35, 1) 
        elif elev < water_level: # underwater
            loop[color_layer] = (0, 0, 1, 1)
        elif elev < snow_level:
            loop[color_layer] = (0, .4, .1, 1)
        else:
            loop[color_layer] = (1, 1, 1, 1)

bmesh.update_edit_mesh(obj.data)

bpy.ops.object.mode_set(mode='OBJECT')