#!/usr/bin/env python3
"""Reconstruct the regular ring body by fitting smooth radial/Z profiles by angle."""

from __future__ import annotations

import json
import math
import sys
from pathlib import Path

import bpy
import numpy as np


def angular_distance(values: np.ndarray, center: float) -> np.ndarray:
    return np.abs((values - center + math.pi) % (2 * math.pi) - math.pi)


def smooth_periodic(values: np.ndarray, weights: np.ndarray, sigma_bins: float) -> np.ndarray:
    radius = max(1, int(math.ceil(sigma_bins * 3)))
    offsets = np.arange(-radius, radius + 1)
    kernel = np.exp(-0.5 * (offsets / sigma_bins) ** 2)
    kernel /= kernel.sum()
    numerator = np.zeros_like(values)
    denominator = np.zeros_like(values)
    for offset, factor in zip(offsets, kernel):
        numerator += np.roll(values * weights, offset) * factor
        denominator += np.roll(weights, offset) * factor
    return numerator / np.maximum(denominator, 1e-8)


def main() -> None:
    args = sys.argv[sys.argv.index("--") + 1 :]
    source = Path(args[0]).resolve()
    output = Path(args[1]).resolve()
    output.parent.mkdir(parents=True, exist_ok=True)

    bpy.ops.import_scene.gltf(filepath=str(source))
    obj = max((candidate for candidate in bpy.context.scene.objects if candidate.type == "MESH"), key=lambda item: len(item.data.vertices))
    coordinates = np.array([obj.matrix_world @ vertex.co for vertex in obj.data.vertices], dtype=np.float64)
    center = (coordinates[:, :2].min(axis=0) + coordinates[:, :2].max(axis=0)) * 0.5
    local_xy = coordinates[:, :2] - center
    radii = np.linalg.norm(local_xy, axis=1)
    angles = np.arctan2(local_xy[:, 1], local_xy[:, 0])
    z = coordinates[:, 2]

    bins = 720
    indices = np.floor(((angles + math.pi) / (2 * math.pi)) * bins).astype(int) % bins
    mean_r = np.bincount(indices, weights=radii, minlength=bins) / np.maximum(np.bincount(indices, minlength=bins), 1)
    mean_z = np.bincount(indices, weights=z, minlength=bins) / np.maximum(np.bincount(indices, minlength=bins), 1)
    counts = np.bincount(indices, minlength=bins).astype(np.float64)

    profile_r = smooth_periodic(mean_r, counts, sigma_bins=26.0)
    profile_z = smooth_periodic(mean_z, counts, sigma_bins=26.0)
    sampled_r = profile_r[indices]
    sampled_z = profile_z[indices]

    # Preserve the front ornament around -Y. Fair the regular body elsewhere.
    distance = angular_distance(angles, -math.pi / 2)
    protected = math.radians(38)
    fully_rebuilt = math.radians(74)
    t = np.clip((distance - protected) / (fully_rebuilt - protected), 0.0, 1.0)
    blend = t * t * (3.0 - 2.0 * t)

    residual_r = radii - sampled_r
    residual_z = z - sampled_z
    residual_r_center = np.zeros(bins)
    residual_z_center = np.zeros(bins)
    for index in range(bins):
        mask = indices == index
        residual_r_center[index] = np.median(residual_r[mask]) if mask.any() else 0.0
        residual_z_center[index] = np.median(residual_z[mask]) if mask.any() else 0.0
    residual_r -= residual_r_center[indices]
    residual_z -= residual_z_center[indices]

    target_r = sampled_r + residual_r
    target_z = sampled_z + residual_z
    final_r = radii * (1.0 - blend) + target_r * blend
    final_z = z * (1.0 - blend) + target_z * blend

    new_xy = local_xy / np.maximum(radii[:, None], 1e-9) * final_r[:, None] + center
    inverse = obj.matrix_world.inverted()
    for vertex, x, y, zz in zip(obj.data.vertices, new_xy[:, 0], new_xy[:, 1], final_z):
        vertex.co = inverse @ __import__("mathutils").Vector((float(x), float(y), float(zz)))
    obj.data.update()
    for polygon in obj.data.polygons:
        polygon.use_smooth = True

    bpy.ops.object.select_all(action="DESELECT")
    obj.select_set(True)
    bpy.context.view_layer.objects.active = obj
    bpy.ops.export_scene.gltf(
        filepath=str(output), export_format="GLB", use_selection=True, export_apply=True,
        export_yup=True, export_normals=True, export_tangents=False, export_texcoords=False,
        export_materials="EXPORT", export_animations=False, export_cameras=False,
        export_lights=False, export_draco_mesh_compression_enable=False,
    )
    result = {
        "source": str(source), "output": str(output), "bytes": output.stat().st_size,
        "vertices": len(obj.data.vertices), "protected_degrees": 38,
        "fully_rebuilt_degrees": 74, "profile_bins": bins,
    }
    output.with_suffix(".json").write_text(json.dumps(result, indent=2), encoding="utf-8")
    print("TN_RADIAL_RECONSTRUCTION", json.dumps(result))


if __name__ == "__main__":
    main()