"""Build a quad panorama with periodic sampling and radial depth separation."""

import math

import bpy

from ..common.depth_surface import (
    _math, split_depth_surface, build_surface_camera,
)
from ..common.nodes import float_input, group_input, interface_socket, prepare_node_group, evaluate_field
from ..common.material import inherit_material_slots
from ..common.smoothing import expand_smooth


def _quad_sphere(group, segments):
    nodes, links = group.nodes, group.links
    cube = nodes.new("GeometryNodeMeshCube")
    cube.inputs["Size"].default_value = (2, 2, 2)
    for axis in "XYZ":
        links.new(segments, cube.inputs[f"Vertices {axis}"])
    xyz = nodes.new("ShaderNodeSeparateXYZ")
    links.new(nodes.new("GeometryNodeInputPosition").outputs[0], xyz.inputs[0])
    squares = {axis: _math(group, "MULTIPLY", xyz.outputs[axis], xyz.outputs[axis]) for axis in "XYZ"}
    combine = nodes.new("ShaderNodeCombineXYZ")
    for axis, a, b in (("X", "Y", "Z"), ("Y", "Z", "X"), ("Z", "X", "Y")):
        half = _math(group, "DIVIDE", _math(group, "ADD", squares[a], squares[b]), 2)
        third = _math(group, "DIVIDE", _math(group, "MULTIPLY", squares[a], squares[b]), 3)
        root = _math(group, "SQRT", _math(group, "ADD", _math(group, "SUBTRACT", 1, half), third))
        links.new(_math(group, "MULTIPLY", xyz.outputs[axis], root), combine.inputs[axis])
    position = nodes.new("GeometryNodeSetPosition")
    links.new(cube.outputs["Mesh"], position.inputs["Geometry"])
    links.new(combine.outputs[0], position.inputs["Position"])
    return position.outputs["Geometry"]


def build_image_depth_panorama_group():
    group = bpy.data.node_groups.new("O Image Depth Panorama", "GeometryNodeTree")
    group.is_modifier = group.use_fake_user = True
    interface_socket(group, "Geometry", "INPUT", "NodeSocketGeometry")
    subdivide = interface_socket(group, "Subdivide", "INPUT", "NodeSocketInt")
    subdivide.default_value, subdivide.min_value, subdivide.max_value = 7, 0, 8
    float_input(group, "Depth Scale", 1, maximum=1, subtype="FACTOR")
    float_input(group, "Split Threshold", 0, maximum=1, subtype="FACTOR")
    smooth = interface_socket(group, "Smooth", "INPUT", "NodeSocketInt")
    smooth.default_value, smooth.min_value, smooth.max_value = 0, 0, 20
    interface_socket(group, "Valid Only", "INPUT", "NodeSocketBool").default_value = True
    options = group.interface.new_panel("Options", default_closed=True)
    float_input(group, "Smooth Weight", 1, maximum=1, subtype="FACTOR", parent=options)
    float_input(group, "Validity Threshold", 0.9, maximum=1, subtype="FACTOR", parent=options)
    float_input(group, "Invalid Distance", 100, minimum=1e-5, subtype="DISTANCE", parent=options)
    data = group.interface.new_panel("Data", default_closed=True)
    interface_socket(group, "Depth Image", "INPUT", "NodeSocketImage", data).hide_in_modifier = True
    float_input(group, "Reference Distance", 1, minimum=1e-6, subtype="DISTANCE", parent=data).hide_in_modifier = True
    interface_socket(group, "Geometry", "OUTPUT", "NodeSocketGeometry")
    nodes, links = group.nodes, group.links

    def uv(direction):
        normalize = nodes.new("ShaderNodeVectorMath")
        normalize.operation = "NORMALIZE"
        links.new(direction, normalize.inputs[0])
        xyz = nodes.new("ShaderNodeSeparateXYZ")
        links.new(normalize.outputs[0], xyz.inputs[0])
        longitude = _math(group, "DIVIDE", _math(group, "ARCTAN2", xyz.outputs["Y"], xyz.outputs["X"]), 2 * math.pi)
        z = _math(group, "MAXIMUM", -1, _math(group, "MINIMUM", 1, xyz.outputs["Z"]))
        latitude = _math(group, "SUBTRACT", 1, _math(group, "DIVIDE", _math(group, "ARCCOSINE", z), math.pi))
        return longitude, latitude, xyz.outputs["Z"]

    def combine(u, v):
        node = nodes.new("ShaderNodeCombineXYZ")
        links.new(u, node.inputs["X"])
        links.new(v, node.inputs["Y"])
        return node.outputs[0]

    def texture(coordinates, interpolation):
        inputs = group_input(nodes, {"Depth Image"})
        info = nodes.new("GeometryNodeImageInfo")
        links.new(inputs.outputs["Depth Image"], info.inputs["Image"])
        axes = nodes.new("ShaderNodeSeparateXYZ")
        links.new(coordinates, axes.inputs[0])
        half_pixel = _math(group, "DIVIDE", 0.5, info.outputs["Height"])
        v = _math(group, "MAXIMUM", half_pixel,
                  _math(group, "MINIMUM", axes.outputs["Y"], _math(group, "SUBTRACT", 1, half_pixel)))
        node = nodes.new("GeometryNodeImageTexture")
        node.interpolation, node.extension = interpolation, "REPEAT"
        links.new(inputs.outputs["Depth Image"], node.inputs["Image"])
        links.new(combine(axes.outputs["X"], v), node.inputs["Vector"])
        return node.outputs["Color"]

    def depth(coordinates):
        linear, nearest = texture(coordinates, "Linear"), texture(coordinates, "Closest")
        alphas = []
        for color in (linear, nearest):
            separate = nodes.new("FunctionNodeSeparateColor")
            links.new(color, separate.inputs["Color"])
            alphas.append(separate.outputs["Alpha"])
        choose = nodes.new("GeometryNodeSwitch")
        choose.input_type = "RGBA"
        links.new(_math(group, "LESS_THAN", alphas[0],
                        _math(group, "SUBTRACT", alphas[1], 1e-6)), choose.inputs["Switch"])
        links.new(linear, choose.inputs["False"])
        links.new(nearest, choose.inputs["True"])
        return choose.outputs[0], alphas[1]

    inputs = group_input(nodes, {"Subdivide"})
    sphere = _quad_sphere(group, _math(group, "ADD", _math(group, "POWER", 2, inputs.outputs["Subdivide"]), 1))
    position = nodes.new("GeometryNodeInputPosition").outputs[0]
    longitude, latitude, z = uv(position)
    coordinates = evaluate_field(group, combine(longitude, latitude), "FLOAT_VECTOR", "POINT")
    face = nodes.new("GeometryNodeFieldOnDomain")
    face.data_type, face.domain = "FLOAT_VECTOR", "FACE"
    links.new(position, face.inputs["Value"])
    face_u, face_v, _ = uv(face.outputs["Value"])
    correction = _math(group, "ROUND", _math(group, "SUBTRACT", face_u, longitude))
    choose = nodes.new("GeometryNodeSwitch")
    choose.input_type = "FLOAT"
    links.new(_math(group, "GREATER_THAN", _math(group, "ABSOLUTE", z), 0.999999), choose.inputs["Switch"])
    links.new(_math(group, "ADD", longitude, correction), choose.inputs["False"])
    links.new(face_u, choose.inputs["True"])
    uv_map = nodes.new("GeometryNodeStoreNamedAttribute")
    uv_map.data_type, uv_map.domain = "FLOAT2", "CORNER"
    uv_map.inputs["Name"].default_value = "UVMap"
    links.new(sphere, uv_map.inputs["Geometry"])
    links.new(combine(choose.outputs[0], latitude), uv_map.inputs["Value"])
    pixel, _ = depth(coordinates)
    pixel = evaluate_field(group, pixel, "FLOAT_VECTOR", "POINT")
    nearest = nodes.new("FunctionNodeSeparateColor")
    links.new(texture(coordinates, "Closest"), nearest.inputs["Color"])
    controls = group_input(nodes, {"Validity Threshold"})
    invalid = _math(group, "MAXIMUM",
                    _math(group, "LESS_THAN", nearest.outputs["Alpha"], controls.outputs["Validity Threshold"]),
                    _math(group, "LESS_THAN", nearest.outputs["Red"], 1e-5))
    face_invalid = evaluate_field(group, invalid, "FLOAT", "POINT")
    any_invalid = _math(group, "GREATER_THAN", face_invalid, 0)
    mixed = _math(group, "MULTIPLY", any_invalid,
                  _math(group, "LESS_THAN", face_invalid, 1))
    selection = nodes.new("GeometryNodeSwitch")
    selection.input_type = "BOOLEAN"
    links.new(group_input(nodes, {"Valid Only"}).outputs["Valid Only"], selection.inputs["Switch"])
    links.new(mixed, selection.inputs["False"])
    links.new(any_invalid, selection.inputs["True"])
    delete = nodes.new("GeometryNodeDeleteGeometry")
    delete.domain, delete.mode = "FACE", "ALL"
    links.new(uv_map.outputs["Geometry"], delete.inputs["Geometry"])
    links.new(selection.outputs[0], delete.inputs["Selection"])
    controls = group_input(nodes, {"Split Threshold", "Reference Distance"})
    strength = controls.outputs["Split Threshold"]
    geometry, corner_camera, cut_vertex = split_depth_surface(group, delete.outputs[0], depth(combine(face_u, face_v))[0], strength,
                                   _math(group, "DIVIDE", 1, controls.outputs["Reference Distance"]),
                                   face_center=face.outputs["Value"])
    distance = build_surface_camera(group, pixel, corner_camera, cut_vertex, strength)
    axes = nodes.new("ShaderNodeSeparateXYZ")
    links.new(distance, axes.inputs[0])
    distance = nodes.new("GeometryNodeSwitch")
    distance.input_type = "FLOAT"
    links.new(face_invalid, distance.inputs["Switch"])
    links.new(axes.outputs["X"], distance.inputs["False"])
    links.new(group_input(nodes, {"Invalid Distance"}).outputs["Invalid Distance"], distance.inputs["True"])
    controls = group_input(nodes, {"Depth Scale"})
    radius = _math(group, "ADD", 1, _math(group, "MULTIPLY", controls.outputs["Depth Scale"], _math(group, "SUBTRACT", distance.outputs[0], 1)))
    radial = nodes.new("ShaderNodeVectorMath")
    radial.operation = "SCALE"
    links.new(position, radial.inputs[0])
    links.new(radius, radial.inputs[3])
    deform = nodes.new("GeometryNodeSetPosition")
    links.new(geometry, deform.inputs["Geometry"])
    links.new(radial.outputs[0], deform.inputs["Position"])
    geometry = expand_smooth(group, deform.outputs[0], iterations_name="Smooth", weight_name="Smooth Weight")
    flip = nodes.new("GeometryNodeFlipFaces")
    links.new(geometry, flip.inputs["Mesh"])
    smooth = nodes.new("GeometryNodeSetShadeSmooth")
    smooth.domain = "FACE"
    smooth.inputs["Shade Smooth"].default_value = True
    links.new(flip.outputs["Mesh"], smooth.inputs["Geometry"])
    geometry = inherit_material_slots(group, smooth.outputs[0], group_input(nodes, {"Geometry"}).outputs["Geometry"])
    output = nodes.new("NodeGroupOutput")
    links.new(geometry, output.inputs["Geometry"])
    prepare_node_group(group)
    return group