import torch
import numpy as np

from core.cylinder import generate_cylinder_n_torch, snap_to_discrete_values
from core.scoring import compute_overlap_ratio_from_cylinder_mask


def set_axes_equal_3d(ax):
    """
    Make axes of 3D plot have equal scale so that spheres appear as spheres,
    cubes are cubes, etc.
    """
    x_limits = ax.get_xlim3d()
    y_limits = ax.get_ylim3d()
    z_limits = ax.get_zlim3d()

    x_range = abs(x_limits[1] - x_limits[0])
    x_middle = np.mean(x_limits)
    y_range = abs(y_limits[1] - y_limits[0])
    y_middle = np.mean(y_limits)
    z_range = abs(z_limits[1] - z_limits[0])
    z_middle = np.mean(z_limits)

    plot_radius = 0.5*max([x_range, y_range, z_range])

    ax.set_xlim3d([x_middle - plot_radius, x_middle + plot_radius])
    ax.set_ylim3d([y_middle - plot_radius, y_middle + plot_radius])
    ax.set_zlim3d([z_middle - plot_radius, z_middle + plot_radius])

    try:
        ax.set_box_aspect([1, 1, 1])
    except AttributeError:
        pass


def res_plt_2_torch(spine_tensor, cortical_tensor, image_shape, image2_path,
                    base_folder, label_str, diameter_l, length_l,
                    diameter_r, length_r, best_position_l, best_position_r,
                    swarm_size, max_iter, total_time, spacing, CBT, device,
                    grid):
    """相容入口（舊簽名）：實作已合併進 res_bone_figure.render_bone_figure。
    label_str / image_shape 不再使用（volume / level 由 image2_path 反推）。"""
    from visualization.res_bone_figure import render_bone_figure
    return render_bone_figure(
        None, None, spine_tensor, cortical_tensor, base_folder,
        spacing=spacing, way='CBT' if CBT else 'TPS',
        best_position_l=best_position_l, best_position_r=best_position_r,
        diameter_l=diameter_l, length_l=length_l,
        diameter_r=diameter_r, length_r=length_r,
        image2_path=image2_path, device=device, grid=grid,
        swarm_size=swarm_size, max_iter=max_iter, total_time=total_time,
    )


def eval_overlap_from_position(
    pos,
    optimize_size: bool,
    spine_tensor: torch.Tensor,
    image_shape,
    spacing,
    device: torch.device,
    grid=None,
    fixed_diameter: float | None = None,
    fixed_length: float | None = None,
):
    """
    根據 position 生成 cylinder mask，再算 overlap ratio
    """

    if optimize_size:
        d, L = snap_to_discrete_values(pos[5], pos[6])
        params_5 = pos[:5]
    else:
        if fixed_diameter is None or fixed_length is None:
            raise ValueError("fixed_diameter and fixed_length must be provided when optimize_size=False")
        d, L = fixed_diameter, fixed_length
        params_5 = pos

    z, y, x, az, alt = params_5

    cyl_mask = generate_cylinder_n_torch(
        d, L,
        z, y, x,
        az, alt,
        image_shape, spacing,
        device=device,
        grid=grid
    )

    overlap = compute_overlap_ratio_from_cylinder_mask(cyl_mask, spine_tensor)
    return overlap, d, L