"""Mirror completed depth geometry and fill matching boundaries."""
import bpy
from mathutils import Matrix

from anyimage.operators.cutout_tool.shape import NORMAL_REDUCTION_ATTRIBUTE_NAME

from ..common.nodes import (
    interface_socket, group_input, float_input, bool_input, vector_input,
    compare_node, evaluate_field, store_float_attribute, read_float_attribute,
    prepare_node_group,
)
from ..common.depth_surface import _math
from ..common.normal_map import ROTATION_ATTRIBUTE, FACE_ATTRIBUTE, AXIS_ATTRIBUTE
from ..common.direction import symmetry_legacy_direction_to_canonical
from ..common.smoothing import edge_boundary_field, pinned_smooth

NODE_GROUP_NAME = "O Image Cutout Symmetry"

# 衔接带从轮廓向内覆盖的圈数。
FILL_SMOOTH_RINGS = 4

# 对称面附近按中位边长的倍数降低法线贴图强度。
SEAM_NORMAL_BAND = 4


def _align_front(group, geometry, controls):
    """Align existing geometry with the fixed XY symmetry plane."""
    nodes, links = group.nodes, group.links
    direction = symmetry_legacy_direction_to_canonical(group, controls['Direction'])
    align = nodes.new('FunctionNodeAlignEulerToVector')
    align.axis = 'Y'
    links.new(direction, align.inputs['Vector'])
    position = nodes.new('GeometryNodeInputPosition')
    rotate = nodes.new('ShaderNodeVectorRotate')
    rotate.rotation_type = 'EULER_XYZ'
    rotate.invert = True
    links.new(position.outputs[0], rotate.inputs['Vector'])
    links.new(align.outputs['Rotation'], rotate.inputs['Rotation'])
    offset = nodes.new('ShaderNodeCombineXYZ')
    offset_value = nodes.new('ShaderNodeMath')
    offset_value.operation = 'MULTIPLY'
    offset_value.inputs[1].default_value = -1.0
    links.new(controls['Offset'], offset_value.inputs[0])
    links.new(offset_value.outputs[0], offset.inputs['Y'])
    offset_position = nodes.new('ShaderNodeVectorMath')
    offset_position.operation = 'ADD'
    links.new(rotate.outputs[0], offset_position.inputs[0])
    links.new(offset.outputs[0], offset_position.inputs[1])
    scale_vector = nodes.new('ShaderNodeCombineXYZ')
    scale_vector.inputs['X'].default_value = 1
    scale_vector.inputs['Z'].default_value = 1
    links.new(controls['Scale'], scale_vector.inputs['Y'])
    scale = nodes.new('ShaderNodeVectorMath')
    scale.operation = 'MULTIPLY'
    links.new(offset_position.outputs[0], scale.inputs[0])
    links.new(scale_vector.outputs[0], scale.inputs[1])
    place = nodes.new('GeometryNodeSetPosition')
    links.new(geometry, place.inputs['Geometry'])
    links.new(scale.outputs[0], place.inputs['Position'])
    store = nodes.new('GeometryNodeStoreNamedAttribute')
    store.data_type, store.domain = 'FLOAT_VECTOR', 'FACE'
    store.inputs['Name'].default_value = ROTATION_ATTRIBUTE
    links.new(place.outputs[0], store.inputs['Geometry'])
    links.new(align.outputs['Rotation'], store.inputs['Value'])
    return store.outputs['Geometry']


def _boundary_edge_field(group):
    """Report one on edges that carry a single face."""
    neighbors = group.nodes.new('GeometryNodeInputMeshEdgeNeighbors')
    return compare_node(group, 'EQUAL', neighbors.outputs['Face Count'], 1, data_type='INT')


def _fold_to_plane(group, geometry, tolerance):
    """Move over-plane points onto the plane and drop the faces that collapse."""
    nodes, links = group.nodes, group.links
    position = nodes.new('GeometryNodeInputPosition')
    components = nodes.new('ShaderNodeSeparateXYZ')
    links.new(position.outputs[0], components.inputs[0])
    beyond = compare_node(group, 'GREATER_THAN', components.outputs['Y'], tolerance)
    target = nodes.new('ShaderNodeCombineXYZ')
    links.new(components.outputs['X'], target.inputs['X'])
    links.new(components.outputs['Z'], target.inputs['Z'])
    place = nodes.new('GeometryNodeSetPosition')
    links.new(geometry, place.inputs['Geometry'])
    links.new(beyond, place.inputs['Selection'])
    links.new(target.outputs[0], place.inputs['Position'])
    # A face lying on the plane would be mirrored onto itself.
    center = evaluate_field(group, _math(group, 'ABSOLUTE', components.outputs['Y']), 'FLOAT', 'FACE')
    normal = nodes.new('GeometryNodeInputNormal')
    normal_axes = nodes.new('ShaderNodeSeparateXYZ')
    links.new(normal.outputs[0], normal_axes.inputs[0])
    flat = nodes.new('FunctionNodeBooleanMath')
    flat.operation = 'AND'
    links.new(compare_node(group, 'LESS_THAN', center, tolerance), flat.inputs[0])
    links.new(compare_node(group, 'GREATER_THAN', _math(group, 'ABSOLUTE', normal_axes.outputs['Y']), 1 - 1e-3), flat.inputs[1])
    delete = nodes.new('GeometryNodeDeleteGeometry')
    delete.domain, delete.mode = 'FACE', 'ALL'
    links.new(place.outputs[0], delete.inputs['Geometry'])
    links.new(flat.outputs[0], delete.inputs['Selection'])
    return _clean_crease(group, delete.outputs[0], tolerance)


def _clean_crease(group, geometry, tolerance):
    """Open interior in-plane edges left by folding a closed surface."""
    nodes, links = group.nodes, group.links
    heights = _edge_abs_y(group)
    edge_plane = nodes.new('FunctionNodeBooleanMath')
    edge_plane.operation = 'AND'
    links.new(compare_node(group, 'LESS_THAN', heights[0], tolerance), edge_plane.inputs[0])
    links.new(compare_node(group, 'LESS_THAN', heights[1], tolerance), edge_plane.inputs[1])
    neighbors = nodes.new('GeometryNodeInputMeshEdgeNeighbors')
    interior = compare_node(group, 'EQUAL', neighbors.outputs['Face Count'], 2, data_type='INT')
    fold = nodes.new('FunctionNodeBooleanMath')
    fold.operation = 'AND'
    links.new(edge_plane.outputs[0], fold.inputs[0])
    links.new(interior, fold.inputs[1])

    face_index = nodes.new('GeometryNodeInputIndex')
    corners = nodes.new('GeometryNodeCornersOfFace')
    links.new(face_index.outputs[0], corners.inputs['Face Index'])
    face_fold = None
    for offset in range(4):
        corner = nodes.new('GeometryNodeOffsetCornerInFace')
        links.new(corners.outputs['Corner Index'], corner.inputs['Corner Index'])
        corner.inputs['Offset'].default_value = offset
        edge_of_corner = nodes.new('GeometryNodeEdgesOfCorner')
        links.new(corner.outputs['Corner Index'], edge_of_corner.inputs['Corner Index'])
        sample = nodes.new('GeometryNodeFieldAtIndex')
        sample.data_type, sample.domain = 'BOOLEAN', 'EDGE'
        links.new(fold.outputs[0], sample.inputs['Value'])
        links.new(edge_of_corner.outputs['Next Edge Index'], sample.inputs['Index'])
        if face_fold is None:
            face_fold = sample.outputs['Value']
        else:
            combined = nodes.new('FunctionNodeBooleanMath')
            combined.operation = 'OR'
            links.new(face_fold, combined.inputs[0])
            links.new(sample.outputs['Value'], combined.inputs[1])
            face_fold = combined.outputs['Boolean']
    cleanup = nodes.new('GeometryNodeDeleteGeometry')
    cleanup.domain, cleanup.mode = 'FACE', 'ALL'
    links.new(geometry, cleanup.inputs['Geometry'])
    links.new(face_fold, cleanup.inputs['Selection'])
    return cleanup.outputs[0]


def _off_plane_edge_field(group, tolerance):
    """Report one on edges that keep an endpoint off the symmetry plane."""
    heights = _edge_abs_y(group)
    return compare_node(group, 'GREATER_THAN', _math(group, 'MAXIMUM', heights[0], heights[1]), tolerance)


def _edge_abs_y(group):
    """Return both endpoint heights on the symmetry axis for an edge."""
    edge = group.nodes.new('GeometryNodeInputMeshEdgeVertices')
    heights = []
    for socket in ('Position 1', 'Position 2'):
        axis = group.nodes.new('ShaderNodeSeparateXYZ')
        group.links.new(edge.outputs[socket], axis.inputs[0])
        heights.append(_math(group, 'ABSOLUTE', axis.outputs['Y']))
    return heights


def _mirror(group, geometry):
    mirror = group.nodes.new('GeometryNodeTransform')
    mirror.inputs['Scale'].default_value = (1, -1, 1)
    group.links.new(geometry, mirror.inputs['Geometry'])
    flip = group.nodes.new('GeometryNodeFlipFaces')
    group.links.new(mirror.outputs[0], flip.inputs['Mesh'])
    return flip.outputs[0]


def _boundary_fields(group, tolerance):
    """Return the folded front's boundary field and its on-plane seam."""
    nodes, links = group.nodes, group.links
    boundary = evaluate_field(group, edge_boundary_field(nodes, links), 'FLOAT', 'POINT')
    position = nodes.new('GeometryNodeInputPosition')
    components = nodes.new('ShaderNodeSeparateXYZ')
    links.new(position.outputs[0], components.inputs[0])
    seam = nodes.new('FunctionNodeBooleanMath')
    seam.operation = 'AND'
    links.new(compare_node(group, 'GREATER_THAN', boundary, 0), seam.inputs[0])
    links.new(
        compare_node(group, 'LESS_THAN', _math(group, 'ABSOLUTE', components.outputs['Y']), tolerance),
        seam.inputs[1],
    )
    return boundary, seam.outputs[0]


def _falloff(group, center, rings=2):
    """Blur a 0/1 center into a falloff over the requested rings."""
    blur = group.nodes.new('GeometryNodeBlurAttribute')
    blur.data_type = 'FLOAT'
    blur.inputs['Iterations'].default_value = rings
    group.links.new(center, blur.inputs['Value'])
    return _math(group, 'MAXIMUM', center, blur.outputs[0])


def _relax_in_plane(group, geometry, region, iterations, pin_boundary):
    """Relax points without leaving the symmetry plane."""
    return pinned_smooth(
        group,
        geometry,
        iterations,
        region,
        weight=0.5,
        movement=(1, 0, 1),
        pin_boundary=pin_boundary,
    )


def _relax_seam(group, geometry, seam_field, smooth):
    """Relax the cut seam in plane through a two-ring falloff."""
    return _relax_in_plane(group, geometry, _falloff(group, seam_field), smooth, pin_boundary=True)


def _relax_fill(group, geometry, boundary_field, smooth):
    """Smooth the junction band the wall builds from, including the symmetry axis."""
    region = _falloff(group, boundary_field, FILL_SMOOTH_RINGS)
    return _relax_in_plane(group, geometry, region, smooth, pin_boundary=False)


def _extrude_boundary(group, geometry, selection, y_factor):
    """Extrude selected edges to a fraction of their current height."""
    nodes, links = group.nodes, group.links
    extrude = nodes.new('GeometryNodeExtrudeMesh')
    extrude.mode = 'EDGES'
    extrude.inputs['Offset Scale'].default_value = 0
    links.new(geometry, extrude.inputs['Mesh'])
    links.new(selection, extrude.inputs['Selection'])
    position = nodes.new('GeometryNodeInputPosition')
    scale = nodes.new('ShaderNodeVectorMath')
    scale.operation = 'MULTIPLY'
    scale.inputs[1].default_value = (1, y_factor, 1)
    links.new(position.outputs[0], scale.inputs[0])
    place = nodes.new('GeometryNodeSetPosition')
    links.new(extrude.outputs['Mesh'], place.inputs['Geometry'])
    links.new(extrude.outputs['Top'], place.inputs['Selection'])
    links.new(scale.outputs[0], place.inputs['Position'])
    store = nodes.new('GeometryNodeStoreNamedAttribute')
    store.data_type, store.domain = 'FLOAT', 'FACE'
    store.inputs['Name'].default_value = FACE_ATTRIBUTE
    store.inputs['Value'].default_value = 3
    links.new(place.outputs[0], store.inputs['Geometry'])
    links.new(extrude.outputs['Side'], store.inputs['Selection'])
    return store.outputs[0], extrude.outputs['Top']


def _build_wall(group, front, boundary):
    """Attach the curved half wall to the front and mark its faces."""
    geometry, top = _extrude_boundary(group, front, boundary, .5)
    geometry, _ = _extrude_boundary(group, geometry, top, 0)
    return geometry


def build_image_cutout_symmetry_group():
    existing = bpy.data.node_groups.get(NODE_GROUP_NAME)
    if existing is not None:
        return existing
    group = bpy.data.node_groups.new(NODE_GROUP_NAME, 'GeometryNodeTree')
    group.is_modifier, group.use_fake_user = True, True
    group.color_tag = 'GEOMETRY'
    interface_socket(group, 'Geometry', 'INPUT', 'NodeSocketGeometry')
    vector_input(group, 'Direction', (0, 0, 1), subtype='DIRECTION').description = 'Align this direction with the local symmetry axis.'
    float_input(group, 'Offset', 0, minimum=-10, maximum=10, subtype='DISTANCE').description = 'Offset geometry along the aligned symmetry axis before scaling.'
    float_input(group, 'Scale', 1, maximum=2).description = 'Scale the offset geometry along the aligned symmetry axis.'
    options = group.interface.new_panel(name='Options', default_closed=True)
    bool_input(group, 'Fill Sides', True, parent=options).description = 'Connect matching front and mirrored boundaries.'
    smooth = interface_socket(group, 'Fill Smooth', 'INPUT', 'NodeSocketInt', options)
    smooth.default_value, smooth.min_value, smooth.max_value = 2, 0, 16
    smooth.description = 'Smooth the band around the front boundary and the symmetry seam over four rings. Zero keeps the raw outline.'
    seam = interface_socket(group, 'Seam Smooth', 'INPUT', 'NodeSocketInt', options)
    seam.default_value, seam.min_value, seam.max_value = 2, 0, 16
    seam.description = 'Smooth the cut seam inside the symmetry plane. Zero keeps the raw seam.'
    interface_socket(group, 'Geometry', 'OUTPUT', 'NodeSocketGeometry')
    nodes, links = group.nodes, group.links
    inputs = group_input(nodes, set()).outputs
    aligned = _align_front(group, inputs['Geometry'], inputs)
    edge = nodes.new('GeometryNodeInputMeshEdgeVertices')
    edge_length = nodes.new('ShaderNodeVectorMath')
    edge_length.operation = 'DISTANCE'
    links.new(edge.outputs['Position 1'], edge_length.inputs[0])
    links.new(edge.outputs['Position 2'], edge_length.inputs[1])
    statistics = nodes.new('GeometryNodeAttributeStatistic')
    statistics.data_type, statistics.domain = 'FLOAT', 'EDGE'
    links.new(aligned, statistics.inputs['Geometry'])
    links.new(edge_length.outputs['Value'], statistics.inputs['Attribute'])
    snap_tolerance = _math(
        group, 'MULTIPLY', _math(group, 'MAXIMUM', statistics.outputs['Median'], 1e-6), 1e-6,
    )
    is_boundary = _boundary_edge_field(group)
    boundary_field, seam_field = _boundary_fields(group, snap_tolerance)
    front = _fold_to_plane(group, aligned, snap_tolerance)
    shade = nodes.new('GeometryNodeSetShadeSmooth')
    shade.domain = 'FACE'
    links.new(front, shade.inputs['Geometry'])
    front = store_float_attribute(group, shade.outputs[0], FACE_ATTRIBUTE, 1., domain='FACE')
    front = _relax_seam(group, front, seam_field, inputs['Seam Smooth'])
    boundary = nodes.new('FunctionNodeBooleanMath')
    boundary.operation = 'AND'
    links.new(is_boundary, boundary.inputs[0])
    # The seam sits on the plane and welds through the merge, so only off-plane edges need walls.
    links.new(_off_plane_edge_field(group, snap_tolerance), boundary.inputs[1])
    fill_smooth = nodes.new('GeometryNodeSwitch')
    fill_smooth.input_type = 'INT'
    fill_smooth.inputs['False'].default_value = 0
    links.new(inputs['Fill Sides'], fill_smooth.inputs['Switch'])
    links.new(inputs['Fill Smooth'], fill_smooth.inputs['True'])
    front = _relax_fill(group, front, boundary_field, fill_smooth.outputs[0])
    walled = _build_wall(group, front, boundary.outputs[0])
    wall_shade = nodes.new('GeometryNodeSetShadeSmooth')
    wall_shade.domain = 'FACE'
    links.new(compare_node(group, 'EQUAL', read_float_attribute(group, FACE_ATTRIBUTE), 3), wall_shade.inputs['Selection'])
    links.new(walled, wall_shade.inputs['Geometry'])
    body = nodes.new('GeometryNodeSwitch')
    body.input_type = 'GEOMETRY'
    links.new(inputs['Fill Sides'], body.inputs['Switch'])
    links.new(front, body.inputs['False'])
    links.new(wall_shade.outputs[0], body.inputs['True'])
    mark_back = nodes.new('GeometryNodeStoreNamedAttribute')
    mark_back.data_type, mark_back.domain = 'FLOAT', 'FACE'
    mark_back.inputs['Name'].default_value = FACE_ATTRIBUTE
    mark_back.inputs['Value'].default_value = 2
    links.new(_mirror(group, body.outputs[0]), mark_back.inputs['Geometry'])
    links.new(compare_node(group, 'EQUAL', read_float_attribute(group, FACE_ATTRIBUTE), 1), mark_back.inputs['Selection'])
    join = nodes.new('GeometryNodeJoinGeometry')
    for geometry in (body.outputs[0], mark_back.outputs['Geometry']):
        links.new(geometry, join.inputs[0])
    position = nodes.new('GeometryNodeInputPosition')
    axis = nodes.new('ShaderNodeSeparateXYZ')
    links.new(position.outputs[0], axis.inputs[0])
    on_axis = compare_node(group, 'LESS_THAN', _math(group, 'ABSOLUTE', axis.outputs['Y']), snap_tolerance)
    merge = nodes.new('GeometryNodeMergeByDistance')
    merge.inputs['Distance'].default_value = 0.01
    links.new(join.outputs[0], merge.inputs['Geometry'])
    links.new(on_axis, merge.inputs['Selection'])
    geometry = store_float_attribute(group, merge.outputs[0], AXIS_ATTRIBUTE, 1., domain='FACE')
    orient = nodes.new('GeometryNodeTransform')
    orient.inputs['Rotation'].default_value = Matrix.Rotation(1.5707963267948966, 4, 'Z').to_euler()
    links.new(geometry, orient.inputs['Geometry'])
    # Fade the material normal map near the symmetry plane so the weld shades smoother.
    position = nodes.new('GeometryNodeInputPosition')
    components = nodes.new('ShaderNodeSeparateXYZ')
    links.new(position.outputs[0], components.inputs[0])
    ramp = nodes.new('ShaderNodeMapRange')
    ramp.interpolation_type = 'SMOOTHSTEP'
    ramp.clamp = True
    links.new(_math(group, 'ABSOLUTE', components.outputs['X']), ramp.inputs['Value'])
    ramp.inputs['From Min'].default_value = 0.0
    links.new(_math(group, 'MULTIPLY', statistics.outputs['Median'], SEAM_NORMAL_BAND), ramp.inputs['From Max'])
    ramp.inputs['To Min'].default_value = 1.0
    ramp.inputs['To Max'].default_value = 0.0
    reduction = _math(
        group, 'MAXIMUM',
        read_float_attribute(group, NORMAL_REDUCTION_ATTRIBUTE_NAME), ramp.outputs['Result'],
    )
    geometry = store_float_attribute(
        group, orient.outputs[0], NORMAL_REDUCTION_ATTRIBUTE_NAME, reduction,
    )
    links.new(geometry, nodes.new('NodeGroupOutput').inputs['Geometry'])
    prepare_node_group(group)
    return group