"""Build O Image Depth Cutout."""
import bpy
from mathutils import Matrix
from anyimage.operators.cutout_tool.shape import CUTOUT_Z_TO_X, NORMAL_SCALE_ATTRIBUTE_NAME
from ..common.nodes import interface_socket, group_input, math_node, float_input, menu_switch, read_float_attribute, store_float_attribute, prepare_node_group, sample_field, evaluate_field
from ..common.smoothing import expand_smooth
from ..common.boundary_smoothing import smooth_cut_boundary
from ..common.outline_smoothing import smooth_outline_camera
from ..common.cutout import _island_cleanup, _position_field, _remove_zero_thickness, _shade_output
from ..common.depth_surface import _math, project_depth_surface, build_surface_camera, split_depth_surface, sample_face_camera, thicken_depth_surface
from ..common.smoothing import edge_boundary_field
from ..common.cutout_boundary import remove_boundary_triangles, taper_split_profile


NODE_GROUP_NAME = "O Image Depth Cutout"


def build_normal_weight(group, uniform):
    """Scale normal smoothing by the active mode's relative thickness."""
    control = group_input(group.nodes, {"Thickness", "Uniform Thickness", "Reference Depth"})
    amount = group.nodes.new("GeometryNodeSwitch")
    amount.input_type = "FLOAT"
    group.links.new(uniform, amount.inputs["Switch"])
    group.links.new(control.outputs["Thickness"], amount.inputs["False"])
    ratio = _math(group, "DIVIDE", control.outputs["Uniform Thickness"],
                  _math(group, "MAXIMUM", control.outputs["Reference Depth"], 1e-8))
    group.links.new(ratio, amount.inputs["True"])
    ramp = group.nodes.new("ShaderNodeMapRange")
    ramp.interpolation_type = "SMOOTHSTEP"
    ramp.clamp = True
    group.links.new(amount.outputs[0], ramp.inputs["Value"])
    return ramp.outputs["Result"]


def blend_normal(group, original, smoothed, weight):
    """Normalize the thickness-weighted blend on the consumer geometry."""
    mix = group.nodes.new("ShaderNodeMix")
    mix.data_type = "VECTOR"
    if isinstance(weight, (int, float)):
        mix.inputs[0].default_value = weight
    else:
        group.links.new(weight, mix.inputs[0])
    group.links.new(original, mix.inputs[4])
    group.links.new(smoothed, mix.inputs[5])
    normalize = group.nodes.new("ShaderNodeVectorMath")
    normalize.operation = "NORMALIZE"
    group.links.new(mix.outputs["Result"], normalize.inputs[0])
    return normalize.outputs[0]


def build_volume_fields(group):
    """Evaluate the center and thickness on the original front surface."""
    links = group.links
    rim = evaluate_field(group, edge_boundary_field(group.nodes, links), "FLOAT", "POINT")
    profile = read_float_attribute(group, "o_balloon")
    position = group.nodes.new("GeometryNodeInputPosition")
    axes = group.nodes.new("ShaderNodeSeparateXYZ")
    links.new(position.outputs[0], axes.inputs[0])
    seed = _math(group, "SUBTRACT", axes.outputs["Z"], profile)
    interior = _math(group, "SUBTRACT", 1, rim)
    blur = group.nodes.new("GeometryNodeBlurAttribute")
    blur.data_type = "FLOAT"
    blur.inputs["Iterations"].default_value = 256
    links.new(seed, blur.inputs["Value"])
    links.new(interior, blur.inputs["Weight"])
    center = blur.outputs[0]
    residual = _math(group, "SUBTRACT", axes.outputs["Z"], center)
    full = _math(group, "ADD", residual, profile)
    minimum = _math(group, "MULTIPLY", profile, 0.05)
    positive = _math(group, "MAXIMUM", full, minimum)
    control = group_input(group.nodes, {"Thickness"})
    thickness = _math(group, "MULTIPLY", positive, control.outputs["Thickness"])
    return {"rim": rim, "center": center, "thickness": thickness}


def build_shell_fields(group, geometry, fields, uniform, normal_weight=1):
    """Evaluate rear directions on the center surface and select shell thickness."""
    nodes, links = group.nodes, group.links
    position = nodes.new("GeometryNodeInputPosition")
    axes = nodes.new("ShaderNodeSeparateXYZ")
    links.new(position.outputs[0], axes.inputs[0])
    center_position = nodes.new("ShaderNodeCombineXYZ")
    for axis in ("X", "Y"):
        links.new(axes.outputs[axis], center_position.inputs[axis])
    links.new(fields["center"], center_position.inputs["Z"])
    surface = nodes.new("GeometryNodeSetPosition")
    links.new(geometry, surface.inputs["Geometry"])
    links.new(center_position.outputs[0], surface.inputs["Position"])
    normal = nodes.new("GeometryNodeInputNormal")
    blur = nodes.new("GeometryNodeBlurAttribute")
    blur.data_type = "FLOAT_VECTOR"
    blur.inputs["Iterations"].default_value = 100
    links.new(normal.outputs[0], blur.inputs["Value"])
    normalized = blend_normal(group, normal.outputs[0], blur.outputs[0], normal_weight)
    rear_normal = sample_field(group, surface.outputs[0], normalized, "FLOAT_VECTOR")
    control = group_input(nodes, {"Uniform Thickness"})
    thickness = nodes.new("GeometryNodeSwitch")
    thickness.input_type = "FLOAT"
    links.new(uniform, thickness.inputs["Switch"])
    links.new(fields["thickness"], thickness.inputs["False"])
    links.new(control.outputs["Uniform Thickness"], thickness.inputs["True"])
    return {"thickness": thickness.outputs[0], "normal": rear_normal}


def build_front_fields(group, rim, front_normal, normal_weight=1):
    """Fit structural amplitude over neighbouring points before creating the shell."""
    links = group.links

    def blur(value, weight=None):
        result = group.nodes.new("GeometryNodeBlurAttribute")
        result.data_type = "FLOAT"
        result.inputs["Iterations"].default_value = 256
        links.new(value, result.inputs["Value"])
        if weight is not None:
            links.new(weight, result.inputs["Weight"])
        return result.outputs[0]

    profile_node = group.nodes.new("GeometryNodeInputNamedAttribute")
    profile_node.data_type = "FLOAT"
    profile_node.inputs["Name"].default_value = "o_balloon"
    profile = profile_node.outputs["Attribute"]
    interior = _math(group, "SUBTRACT", 1, rim)
    position = group.nodes.new("GeometryNodeInputPosition")
    axes = group.nodes.new("ShaderNodeSeparateXYZ")
    links.new(position.outputs[0], axes.inputs[0])
    depth = axes.outputs["Z"]
    depth_slow = blur(depth, interior)
    profile_slow = blur(profile, interior)
    depth_bump = _math(group, "SUBTRACT", depth, depth_slow)
    profile_bump = _math(group, "SUBTRACT", profile, profile_slow)
    numerator = _math(group, "MULTIPLY", depth_bump, profile_bump)
    denominator = _math(group, "MULTIPLY", profile_bump, profile_bump)
    numerator = blur(numerator)
    denominator = blur(denominator)
    denominator = _math(group, "MAXIMUM", denominator, 1e-10)
    structure = _math(group, "DIVIDE", numerator, denominator)
    structure = _math(group, "MAXIMUM", structure, 0)

    normal = front_normal
    smooth = group.nodes.new("GeometryNodeBlurAttribute")
    smooth.data_type = "FLOAT_VECTOR"
    links.new(normal, smooth.inputs["Value"])
    smooth.inputs["Iterations"].default_value = 512
    normalized = blend_normal(group, normal, smooth.outputs[0], normal_weight)
    return {"structure": structure, "inflation_normal": normalized}


def build_inflation_position(group, position, distance, uniform, fields):
    """Grow both modes continuously while leaving the rear cap in place."""
    links = group.links

    structure = fields["structure"]
    control = group.nodes.new("NodeGroupInput")
    a = _math(group, "MULTIPLY", control.outputs["Thickness"], control.outputs["Front Inflation"])
    # A shifted cubic blend preserves continuous growth and height order.
    shifted = _math(group, "ADD", a, 0.75)
    quadratic = _math(group, "MULTIPLY", a, shifted)
    quadratic = _math(group, "ADD", quadratic, 0.1875)
    q = _math(group, "MULTIPLY", a, quadratic)
    k = _math(group, "ADD", structure, 0.25)
    square = _math(group, "MULTIPLY", k, k)
    cube = _math(group, "MULTIPLY", square, k)
    total = _math(group, "ADD", cube, q)
    root = _math(group, "POWER", total, 1 / 3)
    root_square = _math(group, "MULTIPLY", root, root)
    product = _math(group, "MULTIPLY", root, k)
    denominator = _math(group, "ADD", root_square, product)
    denominator = _math(group, "ADD", denominator, square)
    response = _math(group, "DIVIDE", q, denominator)
    profile = group.nodes.new("GeometryNodeInputNamedAttribute")
    profile.data_type = "FLOAT"
    profile.inputs["Name"].default_value = "o_balloon"
    balloon = _math(group, "MULTIPLY", profile.outputs["Attribute"], response)
    uniform_control = group.nodes.new("NodeGroupInput")
    uniform_amount = _math(group, "MULTIPLY", uniform_control.outputs["Uniform Thickness"], uniform_control.outputs["Front Inflation"])
    selected = group.nodes.new("GeometryNodeSwitch")
    selected.input_type = "FLOAT"
    links.new(uniform, selected.inputs["Switch"])
    links.new(balloon, selected.inputs["False"])
    links.new(uniform_amount, selected.inputs["True"])
    displacement = selected.outputs[0]
    height = fields["thickness"]
    layer = _math(group, "DIVIDE", distance, height)
    layer = _math(group, "ADD", 1, layer)
    layer = _math(group, "MAXIMUM", layer, 0)
    layer = _math(group, "MINIMUM", layer, 1)
    displacement = _math(group, "MULTIPLY", displacement, layer)
    normal = fields["inflation_normal"]
    offset = group.nodes.new("ShaderNodeVectorMath")
    offset.operation = "SCALE"
    links.new(normal, offset.inputs[0])
    links.new(displacement, offset.inputs[3])
    result = group.nodes.new("ShaderNodeVectorMath")
    result.operation = "ADD"
    links.new(position, result.inputs[0])
    links.new(offset.outputs[0], result.inputs[1])
    enabled = group.nodes.new("GeometryNodeSwitch")
    enabled.input_type = "VECTOR"
    links.new(control.outputs["Front Inflation"], enabled.inputs["Switch"])
    links.new(position, enabled.inputs["False"])
    links.new(result.outputs[0], enabled.inputs["True"])
    return enabled.outputs[0]


def build_front_normal(group, normal_weight):
    links = group.links
    normal = group.nodes.new("GeometryNodeInputNormal")
    blur = group.nodes.new("GeometryNodeBlurAttribute")
    blur.data_type = "FLOAT_VECTOR"
    links.new(normal.outputs[0], blur.inputs["Value"])
    blur.inputs["Iterations"].default_value = 100
    normalized = blend_normal(group, normal.outputs[0], blur.outputs[0], normal_weight)
    return evaluate_field(group, normalized, "FLOAT_VECTOR", "POINT")


def build_shell_position(group, position, front, uniform, fields):
    links = group.links

    def connect(value, target):
        if hasattr(value, "node"):
            links.new(value, target)
        else:
            target.default_value = value

    def vector(operation, a, b):
        result = group.nodes.new("ShaderNodeVectorMath")
        result.operation = operation
        connect(a, result.inputs[0])
        connect(b, result.inputs[3 if operation == "SCALE" else 1])
        return result.outputs[0]

    def z(value):
        result = group.nodes.new("ShaderNodeSeparateXYZ")
        links.new(value, result.inputs[0])
        return result.outputs["Z"]

    # Uniform and the thin end of Balloon share the actual front normal.
    front_z = z(front)
    raw_z = z(position)
    distance = _math(group, "SUBTRACT", raw_z, front_z)
    normal = fields["front_normal"]
    thin_offset = vector("SCALE", normal, distance)
    uniform_position = vector("ADD", front, thin_offset)

    control = group_input(group.nodes, {"Thickness"})
    clamp = group.nodes.new("ShaderNodeClamp")
    links.new(control.outputs["Thickness"], clamp.inputs["Value"])
    t = clamp.outputs[0]
    radius = _math(group, "MULTIPLY", _math(group, "SUBTRACT", front_z, fields["center"]), t)
    rear = _math(group, "SUBTRACT", 1, read_float_attribute(group, "o_normal_scale"))
    rear_radius = _math(group, "MULTIPLY", rear, _math(group, "ADD", radius, distance))
    direction = vector("SUBTRACT", fields["normal"], (0, 0, 1))
    rear_shift = vector("SCALE", direction, rear_radius)
    axis_offset = group.nodes.new("ShaderNodeCombineXYZ")
    links.new(distance, axis_offset.inputs["Z"])
    full_offset = vector("ADD", axis_offset.outputs[0], rear_shift)

    # Smoothly reach the existing full-volume shape at Thickness >= 1.
    transition = group.nodes.new("ShaderNodeMapRange")
    transition.interpolation_type = "SMOOTHERSTEP"
    links.new(t, transition.inputs["Value"])
    mix = group.nodes.new("ShaderNodeMix")
    mix.data_type = "VECTOR"
    links.new(transition.outputs[0], mix.inputs[0])
    links.new(thin_offset, mix.inputs[4])
    links.new(full_offset, mix.inputs[5])
    balloon_position = vector("ADD", front, mix.outputs["Result"])
    choose = group.nodes.new("GeometryNodeSwitch")
    choose.input_type = "VECTOR"
    links.new(uniform, choose.inputs["Switch"])
    links.new(balloon_position, choose.inputs["False"])
    links.new(uniform_position, choose.inputs["True"])
    return build_inflation_position(group, choose.outputs[0], distance, uniform, fields)


def _build_depth_surface(group, geometry, controls):
    """Project the mesh and sample front-surface fields while shaping the shell."""
    links = group.links
    split = controls.outputs["Split Threshold"]
    image = controls.outputs["Depth Image"]
    scale = controls.outputs["Uniform Scale"]
    original = geometry
    geometry, corner_camera, cut_vertex = split_depth_surface(
        group, geometry, sample_face_camera(group, image), split, controls.outputs["Reference Depth"], triangles=True,
    )
    geometry = remove_boundary_triangles(group, geometry)
    camera = build_surface_camera(group, _position_field(group, image), corner_camera, cut_vertex, split)
    camera = smooth_outline_camera(group, original, geometry, camera, scale)
    projection = project_depth_surface(
        group, geometry, camera, scale, controls.outputs["Reference Depth"], controls.outputs["Depth Scale"], split,
    )
    projection = smooth_cut_boundary(group, projection)
    projection = taper_split_profile(group, original, projection)
    # The mode field is shared by the geometry and thickness choices.
    mode = menu_switch(group, "Mode", ("Balloon", "Uniform"), "BOOLEAN")
    mode.inputs["Balloon"].default_value = False
    mode.inputs["Uniform"].default_value = True
    uniform = mode.outputs[0]
    volume = build_volume_fields(group)
    normal_weight = build_normal_weight(group, uniform)
    normal = build_front_normal(group, normal_weight)
    front = build_front_fields(group, volume["rim"], normal, normal_weight)
    shell_fields = build_shell_fields(group, projection, volume, uniform, normal_weight)
    thickness = sample_field(group, projection, shell_fields["thickness"], "FLOAT")
    shell = thicken_depth_surface(
        group, projection, thickness, back_mask=NORMAL_SCALE_ATTRIBUTE_NAME, weld_zero_thickness=True,
    )
    fields = {
        name: sample_field(group, projection, value, data_type, index=shell["source_index"])
        for name, value, data_type in (
            ("center", volume["center"], "FLOAT"), ("front_normal", normal, "FLOAT_VECTOR"),
            ("structure", front["structure"], "FLOAT"), ("inflation_normal", front["inflation_normal"], "FLOAT_VECTOR"),
        )
    }
    fields.update({
        name: sample_field(group, projection, value, "FLOAT_VECTOR" if name == "normal" else "FLOAT",
                           index=shell["source_index"])
        for name, value in shell_fields.items()
    })
    oriented = build_shell_position(group, shell["position"], shell["front"], uniform, fields)
    links.new(oriented, shell["position_input"])
    amount = group.nodes.new("GeometryNodeSwitch")
    amount.input_type = "FLOAT"
    links.new(uniform, amount.inputs["Switch"])
    links.new(controls.outputs["Thickness"], amount.inputs["False"])
    links.new(controls.outputs["Uniform Thickness"], amount.inputs["True"])
    return _remove_zero_thickness(group, projection, shell["geometry"], amount.outputs[0])


def build_image_depth_cutout_group():
    name = NODE_GROUP_NAME
    existing = bpy.data.node_groups.get(name)
    if existing is not None:
        return existing
    group = bpy.data.node_groups.new(name, 'GeometryNodeTree')
    group.is_modifier, group.use_fake_user = (True, True)
    group.color_tag = 'GEOMETRY'
    group['anyimage_asset_revision'] = 23
    interface_socket(group, 'Geometry', 'INPUT', 'NodeSocketGeometry')
    interface_socket(group, 'Mode', 'INPUT', 'NodeSocketMenu')
    float_input(group, 'Thickness', 1.0).description = 'Balloon profile multiplier; zero returns the sheet'
    float_input(group, 'Uniform Thickness', 0.0, subtype='DISTANCE').description = 'Distance used only by Uniform thickness mode'
    float_input(group, 'Depth Scale', 1.0)
    float_input(group, 'Split Threshold', 0.1, maximum=1, subtype='FACTOR')
    iterations = interface_socket(group, 'Smooth', 'INPUT', 'NodeSocketInt')
    iterations.default_value = 0
    iterations.min_value, iterations.max_value = (0, 20)
    options = group.interface.new_panel(name='Options', default_closed=True)
    float_input(group, 'Cleanup Threshold', 0.0, parent=options)
    float_input(group, 'Smooth Weight', 1.0, maximum=1, subtype='FACTOR', parent=options)
    float_input(group, 'Reference Depth', 1.0, subtype='DISTANCE', parent=options)
    boundary_smooth = interface_socket(group, 'Boundary Smooth', 'INPUT', 'NodeSocketInt', options)
    boundary_smooth.default_value, boundary_smooth.min_value, boundary_smooth.max_value = 2, 1, 20
    boundary_smooth.description = 'Boundary geometry iterations and depth blur steps per outline smoothing round'
    outline_smooth = interface_socket(group, 'Outline Depth Fix', 'INPUT', 'NodeSocketInt', options)
    outline_smooth.default_value, outline_smooth.min_value, outline_smooth.max_value = 16, 0, 200
    outline_smooth.description = 'Adaptive ray-distance smoothing rounds around the original outline; preserves split boundary distance'
    float_input(group, 'Front Inflation', 0.0, maximum=1, subtype='FACTOR', parent=options)
    data = group.interface.new_panel(name='Data', default_closed=True)
    interface_socket(group, 'Inward Axis', 'INPUT', 'NodeSocketMenu', data).hide_in_modifier = True
    float_input(group, 'Uniform Scale', 1.0, parent=data).hide_in_modifier = True
    interface_socket(group, 'Depth Image', 'INPUT', 'NodeSocketImage', data).hide_in_modifier = True
    interface_socket(group, 'Geometry', 'OUTPUT', 'NodeSocketGeometry')
    base = group_input(group.nodes, {'Geometry', 'Cleanup Threshold'})
    area = math_node(group.nodes, 'MULTIPLY')
    area.inputs[1].default_value = 0.01
    group.links.new(base.outputs['Cleanup Threshold'], area.inputs[0])
    cleaned = _island_cleanup(group, base.outputs['Geometry'], area.outputs[0])
    cleaned = store_float_attribute(group, cleaned, NORMAL_SCALE_ATTRIBUTE_NAME, 1.0, domain='FACE')
    controls = group_input(group.nodes, set())
    geometry = _build_depth_surface(group, cleaned, controls)
    geometry = expand_smooth(group, geometry, iterations_name='Smooth', weight_name='Smooth Weight')
    orient = group.nodes.new('GeometryNodeTransform')
    orient.inputs['Rotation'].default_value = Matrix(CUTOUT_Z_TO_X).to_euler()
    group.links.new(geometry, orient.inputs['Geometry'])
    axis = menu_switch(group, 'Inward Axis', ('-Z', '+X'), 'GEOMETRY')
    group.links.new(geometry, axis.inputs['-Z'])
    group.links.new(orient.outputs[0], axis.inputs['+X'])
    _shade_output(group, axis.outputs[0])
    prepare_node_group(group)
    for item in group.interface.items_tree:
        if item.item_type != 'SOCKET' or item.in_out != 'INPUT':
            continue
        if item.name == 'Mode':
            item.default_value = 'Balloon'
        elif item.name == 'Inward Axis':
            item.default_value = '-Z'
        elif item.name == 'Uniform Thickness':
            item.name = 'Thickness'
    return group