
import numpy as np

import bpy
from bpy.props import BoolProperty, IntProperty

from sverchok.node_tree import SverchCustomTreeNode
from sverchok.data_structure import updateNode, zip_long_repeat
from sverchok.utils.sv_mesh_utils import mesh_join
from sverchok.utils.marching_squares import make_contours
from sverchok.dependencies import skimage

if skimage is not None:
    from skimage import measure


class SvExScalarFieldGraphNode(SverchCustomTreeNode, bpy.types.Node):
    """
    Triggers: Scalar Field Graph
    Tooltip: Generate a graphical representation of the scalar field
    """
    bl_idname = 'SvExScalarFieldGraphNode'
    bl_label = 'Scalar Field Graph'
    sv_dependencies = {'skimage'}

    samples_xy : IntProperty(
        name = "Samples X/Y",
        default = 50,
        min = 3,
        update = updateNode)

    samples_z : IntProperty(
        name = "Samples Z",
        default = 10,
        min = 2,
        update = updateNode)

    samples_value : IntProperty(
        name = "Value samples",
        default = 10,
        min = 2,
        update = updateNode)

    join : BoolProperty(
        name = "Join",
        default = True,
        update = updateNode)

    def update_sockets(self, context):
        self.outputs['Faces'].hide_safe = not self.make_faces
        updateNode(self, context)

    make_faces : BoolProperty(
            name = "Make faces",
            default = False,
            update = update_sockets)

    connect_bounds : BoolProperty(
            name = "Connect boundary",
            default = False,
            update = updateNode)

    def draw_buttons(self, context, layout):
        layout.prop(self, 'make_faces', toggle=True)
        layout.prop(self, 'connect_bounds', toggle=True)
        layout.prop(self, 'join', toggle=True)

    def sv_init(self, context):
        self.inputs.new('SvScalarFieldSocket', "Field")
        self.inputs.new('SvVerticesSocket', "Bounds")
        self.inputs.new('SvStringsSocket', "SamplesXY").prop_name = 'samples_xy'
        self.inputs.new('SvStringsSocket', "SamplesZ").prop_name = 'samples_z'
        self.inputs.new('SvStringsSocket', "ValueSamples").prop_name = 'samples_value'
        self.outputs.new('SvVerticesSocket', 'Vertices')
        self.outputs.new('SvStringsSocket', 'Edges')
        self.outputs.new('SvStringsSocket', "Faces")
        self.update_sockets(context)

    def get_bounds(self, vertices):
        vs = np.array(vertices)
        min = vs.min(axis=0)
        max = vs.max(axis=0)
        return min.tolist(), max.tolist()

    def process(self):
        if not any(socket.is_linked for socket in self.outputs):
            return

        field_s = self.inputs['Field'].sv_get()
        bounds_s = self.inputs['Bounds'].sv_get()
        samples_xy_s = self.inputs['SamplesXY'].sv_get()
        samples_z_s = self.inputs['SamplesZ'].sv_get()
        samples_value_s = self.inputs['ValueSamples'].sv_get()

        verts_out = []
        edges_out = []
        faces_out = []

        inputs = zip_long_repeat(field_s, bounds_s, samples_xy_s, samples_z_s, samples_value_s)
        for field, bounds, samples_xy, samples_z, samples_value in inputs:
            if isinstance(samples_xy, (list, tuple)):
                samples_xy = samples_xy[0]
            if isinstance(samples_z, (list, tuple)):
                samples_z = samples_z[0]
            if isinstance(samples_value, (list, tuple)):
                samples_value = samples_value[0]

            b1, b2 = self.get_bounds(bounds)
            min_x, max_x = b1[0], b2[0]
            min_y, max_y = b1[1], b2[1]
            min_z, max_z = b1[2], b2[2]

            x_size = (max_x - min_x)/samples_xy
            y_size = (max_y - min_y)/samples_xy

            x_range = np.linspace(min_x, max_x, num=samples_xy)
            y_range = np.linspace(min_y, max_y, num=samples_xy)
            z_range = np.linspace(min_z, max_z, num=samples_z)
            xs, ys, zs = np.meshgrid(x_range, y_range, z_range, indexing='ij')
            xs, ys, zs = xs.flatten(), ys.flatten(), zs.flatten()
            field_values = field.evaluate_grid(xs, ys, zs)
            min_field = field_values.min()
            max_field = field_values.max()
            field_values = field_values.reshape((samples_xy, samples_xy, samples_z))
            field_stops = np.linspace(min_field, max_field, num=samples_value)

            new_verts = []
            new_edges = []
            new_faces = []
            for i in range(samples_z):
                z_value = zs[i]
                z_verts = []
                z_edges = []
                z_faces = []
                for j in range(samples_value):
                    value = field_stops[j]
                    contours = measure.find_contours(field_values[:,:,i], level=value)

                    value_verts, value_edges, value_faces = make_contours(samples_xy, samples_xy,
                                    min_x, x_size, min_y, y_size, z_value, contours,
                                    make_faces=self.make_faces, connect_bounds = self.connect_bounds)
                    if value_verts:
                        z_verts.extend(value_verts)
                        z_edges.extend(value_edges)
                        z_faces.extend(value_faces)

                new_verts.extend(z_verts)
                new_edges.extend(z_edges)
                new_faces.extend(z_faces)

            if self.join:
                new_verts, new_edges, new_faces = mesh_join(new_verts, new_edges, new_faces)

                verts_out.append(new_verts)
                edges_out.append(new_edges)
                faces_out.append(new_faces)
            else:
                verts_out.extend(new_verts)
                edges_out.extend(new_edges)
                faces_out.extend(new_faces)

        self.outputs['Vertices'].sv_set(verts_out)
        self.outputs['Edges'].sv_set(edges_out)
        self.outputs['Faces'].sv_set(faces_out)


def register():
    bpy.utils.register_class(SvExScalarFieldGraphNode)


def unregister():
    bpy.utils.unregister_class(SvExScalarFieldGraphNode)
