CBT_project/core/objective.py

from dataclasses import dataclass
from typing import Optional

import random

from scipy.ndimage import map_coordinates

import numpy as np
import torch

from core.cylinder import generate_cylinder_n_torch, generate_cylinder_o_torch, snap_to_discrete_values, generate_cylinder_tip_torch, snap_to_discrete_values_xfr
from core.intersection import center_line_intersections_torch
from core.scoring import cl_score_torch, cl_score_torch_xfr


# =====================================================================
# OptimizationContext: the single explicit container for all shared state
# needed to evaluate a candidate cylinder. New code should build one of
# these and bind it to an objective via make_objective_function[_xfr];
# no module globals are read during optimization.
# =====================================================================

@dataclass
class OptimizationContext:
    """Shared state for evaluating candidate cylinder placements.

    Required:
        cortical_tensor: cortical bone mask (uint8, 0/1)
        spine_tensor:    bone (binary) mask (uint8, 0/1)
        image2_shape:    (Z, Y, X) volume shape
        spacing:         voxel spacing, [sx, sy, sz] (mm)
        device:          torch device for tensor ops

    Optional:
        grid:            precomputed coordinate grid (z_t, y_t, x_t)
        use_tip_penalty: add the tip-cylinder penalty to the loss
        vbody_tensor:    VBODY (vertebral body) mask 0/1 — adds the VBODY
                         voxel rewards in cl_score_torch_xfr
        The *_array / spine_roi_tensor / diameter / length fields are kept
        for compatibility and debugging; the loss itself does not use them.
    """
    cortical_tensor: torch.Tensor
    spine_tensor: torch.Tensor
    image2_shape: tuple
    spacing: list
    device: torch.device
    grid: Optional[tuple] = None
    use_tip_penalty: bool = False
    vbody_tensor: Optional[torch.Tensor] = None
    spine_roi_tensor: Optional[torch.Tensor] = None
    image1_array: Optional[np.ndarray] = None
    image2_array: Optional[np.ndarray] = None
    image3_array: Optional[np.ndarray] = None
    diameter: Optional[float] = None
    length: Optional[float] = None


def cylinder_circle_line_intersection_loss_deductions_torch(
    ctx: OptimizationContext,
    diameter: float,
    length: float,
    params: list[float],
) -> float:
    """
    Computes the loss for a given set of cylinder params in PyTorch,
    returning a Python float for PSO consumption.

    All shared state (tensors, shape, spacing, device, grid, tip penalty)
    is taken from `ctx`; no module globals are involved.
    """
    position_z, position_y, position_x, azimuth, altitude = params

    cyl_fwd = generate_cylinder_n_torch(
        diameter,
        length,
        position_z,
        position_y,
        position_x,
        float(azimuth),
        float(altitude),
        ctx.image2_shape,
        ctx.spacing,
        ctx.device,
        ctx.grid
    )

    cyl_opp = generate_cylinder_o_torch(
        diameter,
        length,
        position_z,
        position_y,
        position_x,
        float(azimuth),
        float(altitude),
        ctx.image2_shape,
        ctx.spacing,
        ctx.device,
        ctx.grid
    )

    # We call the center_line_intersections in Torch mode
    intersections, _ = center_line_intersections_torch(
        position_z,
        position_y,
        position_x,
        azimuth,
        altitude,
        length,
        ctx.spine_tensor,
        ctx.spacing,
        ctx.device
    )

    cyl_tip = None
    if ctx.use_tip_penalty:
        cyl_tip = generate_cylinder_tip_torch(
            diameter, length,
            position_z, position_y, position_x,
            float(azimuth), float(altitude),
            ctx.image2_shape, ctx.spacing, ctx.device, ctx.grid
        )

    # loss_value = cl_score_torch(
    loss_value = cl_score_torch_xfr(
        ctx.cortical_tensor, ctx.spine_tensor,
        cyl_fwd, cyl_opp, intersections,
        cylinder_tip_torch=cyl_tip,
        vbody_tensor=ctx.vbody_tensor
    )

    return loss_value


# =====================================================================
# Context-bound objective builders (preferred API)
# =====================================================================

def _evaluate(params: list[float], ctx: OptimizationContext) -> float:
    """
    Core objective: params = [z, y, x, azimuth, altitude, diameter_raw, length_raw]
    """
    position_params = params[:5]  # [z, y, x, azimuth, altitude]
    diameter_raw = params[5]
    length_raw = params[6]

    # 將連續值轉換為離散值
    diameter_discrete, length_discrete = snap_to_discrete_values(diameter_raw, length_raw)

    return cylinder_circle_line_intersection_loss_deductions_torch(
        ctx,
        diameter_discrete,
        length_discrete,
        position_params
    )


def _evaluate_xfr(params: list[float], ctx: OptimizationContext, y_indices) -> float:
    """
    Core xfr objective: params = [z, x, azimuth, altitude, diameter_raw, length_raw];
    y is derived from the per-column `y_indices` surface.
    """
    z, x, azimuth, altitude, diameter_raw, length_raw = params
    y = y_indices[round(z), round(x)] #+ random.uniform(-0.5, 0.5)

    position_params = [z, y, x, azimuth, altitude]

    # 將連續值轉換為離散值
    diameter_discrete, length_discrete = snap_to_discrete_values_xfr(diameter_raw, length_raw)

    diameter_loss = .9*diameter_discrete + .1*diameter_raw
    length_loss   = .9*  length_discrete + .1*  length_raw

    return cylinder_circle_line_intersection_loss_deductions_torch(
        ctx,
        diameter_loss,
        length_loss,
        position_params
    )


def make_objective_function(ctx: OptimizationContext):
    """
    Return an objective function bound to `ctx`.
    Pass the returned callable directly to pso / differential_evolution / minimize.
    """
    def objective(params: list[float]) -> float:
        return _evaluate(params, ctx)
    return objective


def make_objective_function_xfr(ctx: OptimizationContext, y_indices):
    """
    Return an xfr objective function bound to `ctx` and the `y_indices` surface.
    Pass the returned callable directly to pso / differential_evolution / minimize.
    """
    def objective(params: list[float]) -> float:
        return _evaluate_xfr(params, ctx, y_indices)
    return objective


# =====================================================================
# Legacy API (backward compatible)
#
# The old flow mutated module attributes on this file from the optimizers
# ("跨檔案注入變數"), which was fragile: any forgotten attribute showed up
# deep inside an optimizer callback as a NoneType/NameError. It is kept so
# existing callers (set_global_context + objective_function) keep working,
# but new code should use OptimizationContext + the make_* factories.
# =====================================================================

# Module-level state for the legacy path only.
cortical_tensor = None
spine_tensor = None
image1_array = None      # cortical_nii.gz
image2_array = None      # binarynii.gz
image2_shape = None
image3_array = None      # roi2.nii.gz
diameter = None
length = None
spacing = [0.5, 0.5, 0.5]
device = None
grid = None
USE_TIP_PENALTY = None

_current_context: Optional[OptimizationContext] = None


def set_global_context(
    cortical,
    spine,
    shape,
    spacing_,
    device_,
    grid_,
    use_tip_penalty=False
):
    """
    Legacy: set the process-wide context used by the legacy
    objective_function / objective_function_xfr wrappers.

    Returns the built OptimizationContext for convenience.
    """
    global _current_context
    global cortical_tensor, spine_tensor, image2_shape, spacing, device, grid, USE_TIP_PENALTY

    _current_context = OptimizationContext(
        cortical_tensor=cortical,
        spine_tensor=spine,
        image2_shape=shape,
        spacing=spacing_,
        device=device_,
        grid=grid_,
        use_tip_penalty=use_tip_penalty,
    )

    # mirror into legacy module attributes for any code that reads them
    cortical_tensor = cortical
    spine_tensor = spine
    image2_shape = shape
    spacing = spacing_
    device = device_
    grid = grid_
    USE_TIP_PENALTY = use_tip_penalty

    return _current_context


def _active_context() -> OptimizationContext:
    """Resolve the context for the legacy wrappers, with a clear error."""
    if _current_context is not None:
        return _current_context

    # Fallback: a caller (old-style optimizers) may have set the legacy
    # module attributes directly — rebuild a context from them.
    if cortical_tensor is not None and spine_tensor is not None:
        return OptimizationContext(
            cortical_tensor=cortical_tensor,
            spine_tensor=spine_tensor,
            image2_shape=image2_shape,
            spacing=spacing,
            device=device,
            grid=grid,
            use_tip_penalty=bool(USE_TIP_PENALTY),
        )

    raise RuntimeError(
        "No optimization context available. Either call set_global_context(...) "
        "or (preferred) build an OptimizationContext and use "
        "make_objective_function / make_objective_function_xfr."
    )


def objective_function(params: list[float]) -> float:
    """
    Legacy wrapper: evaluates against the context set by set_global_context.
    params = [z, y, x, azimuth, altitude, diameter_raw, length_raw]
    """
    return _evaluate(params, _active_context())


def objective_function_xfr(params: list[float], y_indices) -> float:
    """
    Legacy wrapper: evaluates against the context set by set_global_context.
    params = [z, x, azimuth, altitude, diameter_raw, length_raw]
    """
    return _evaluate_xfr(params, _active_context(), y_indices)