/*//////////////////////////////////////////////////////////////////////////////
// CMT Cosmic Muon Tomography project //////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////

  Copyright (c) 2014, Universita' degli Studi di Padova, INFN sez. di Padova
  All rights reserved

  Authors: Andrea Rigoni Garola < andrea.rigoni@pd.infn.it >

  ------------------------------------------------------------------
  This library is free software;  you  can  redistribute  it  and/or
  modify it  under the  terms  of  the  GNU  Lesser  General  Public
  License as published  by  the  Free  Software  Foundation;  either
  version 3.0 of the License, or (at your option) any later version.

  This library is  distributed in  the hope that it will  be useful,
  but  WITHOUT ANY WARRANTY;  without  even  the implied warranty of
  MERCHANTABILITY  or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.

  You should have received a copy of  the GNU Lesser General  Public
  License along with this library.

//////////////////////////////////////////////////////////////////////////////*/

#include "vtkHandlerWidget.h"
#include <iostream>
#include <vtkActor.h>
#include <vtkArcSource.h>
#include <vtkArrowSource.h>
#include <vtkCallbackCommand.h>
#include <vtkCamera.h>
#include <vtkCellPicker.h>
#include <vtkConeSource.h>
#include <vtkCubeSource.h>
#include <vtkInteractorObserver.h>
#include <vtkMath.h>
#include <vtkMatrix4x4.h>
#include <vtkObjectFactory.h>
#include <vtkPlane.h>
#include <vtkPolyDataMapper.h>
#include <vtkProp3D.h>
#include <vtkPropPicker.h>
#include <vtkProperty.h>
#include <vtkRegularPolygonSource.h>
#include <vtkRenderWindow.h>
#include <vtkRenderWindowInteractor.h>
#include <vtkRenderer.h>
#include <vtkSmartPointer.h>
#include <vtkTransform.h>
#include "Math/Transform.h"
#include "Vtk/Math/vtkDense.h"

namespace uLib {
namespace Vtk {

struct HandlerWidgetData {
  vtkSmartPointer<::vtkRenderer> m_OverlayRenderer;
  ::vtkProp *m_HighlightedProp;

  // Visual components //
  vtkSmartPointer<::vtkActor> m_AxesX, m_AxesY, m_AxesZ;    // Arrows
  vtkSmartPointer<::vtkActor> m_RotX, m_RotY, m_RotZ;       // Rings
  vtkSmartPointer<::vtkActor> m_RotCam;                     // Camera ring
  vtkSmartPointer<::vtkActor> m_ScaleX, m_ScaleY, m_ScaleZ; // Cubes

  // cut plane to see only half of rotation handles
  vtkSmartPointer<::vtkPlane> m_ClipPlane;

  // picker to select the gizmo
  vtkSmartPointer<::vtkCellPicker> m_Picker;

  // initial transform of the object
  vtkSmartPointer<::vtkTransform> m_InitialTransform;
  
  // undo stack
  std::vector<uLib::TRS> m_UndoStack;

  HandlerWidgetData() {
    m_Picker = vtkSmartPointer<::vtkCellPicker>::New();
    m_InitialTransform = vtkSmartPointer<::vtkTransform>::New();
    m_ClipPlane = vtkSmartPointer<::vtkPlane>::New();
    m_OverlayRenderer = vtkSmartPointer<::vtkRenderer>::New();
    m_HighlightedProp = nullptr;
  }
};

vtkStandardNewMacro(vtkHandlerWidget);

vtkHandlerWidget::vtkHandlerWidget() : d(new HandlerWidgetData()) {
  this->Interaction = IDLE;
  d->m_Picker->SetTolerance(0.01); // Increased tolerance for thin gizmos
  this->EventCallbackCommand->SetCallback(vtkHandlerWidget::ProcessEvents);
  this->EventCallbackCommand->SetClientData(this);
  this->m_Frame = LOCAL;
  d->m_OverlayRenderer->SetLayer(1);
  d->m_OverlayRenderer->EraseOff();
  d->m_OverlayRenderer->InteractiveOff();
  this->Priority = 50.0; // Higher priority to beat camera style
  this->m_TranslationEnabled = true;
  this->m_RotationEnabled = true;
  this->m_ScalingEnabled = true;
  this->CreateGizmos();
}

vtkHandlerWidget::~vtkHandlerWidget() {
  this->SetEnabled(0);
  delete d;
}

::vtkRenderer *vtkHandlerWidget::GetOverlayRenderer() {
  return d->m_OverlayRenderer;
}


void vtkHandlerWidget::SetProp3D(::vtkProp3D *prop) {
  if (this->Prop3D == prop) {
    return;
  }
  this->Prop3D = prop;
  if (this->Prop3D) {
    this->d->m_UndoStack.clear(); // Clear history when selecting new object
    this->UpdateGizmoPosition();
  }
  this->Modified();
}

void vtkHandlerWidget::SetEnabled(int enabling) {
  if (enabling) {
    if (this->Enabled)
      return;
    if (!this->Interactor) {
      vtkErrorMacro(
          << "The interactor must be set prior to enabling the widget");
      return;
    }

    if (!this->CurrentRenderer) {
      this->CurrentRenderer = this->Interactor->FindPokedRenderer(
          this->Interactor->GetLastEventPosition()[0],
          this->Interactor->GetLastEventPosition()[1]);
    }

    if (!this->CurrentRenderer) {
      return;
    }

    this->Enabled = 1;

    // Add observers
    ::vtkRenderWindowInteractor *i = this->Interactor;
    this->SetPriority(this->Priority);
    i->AddObserver(::vtkCommand::LeftButtonPressEvent,
                   this->EventCallbackCommand, this->Priority);
    i->AddObserver(::vtkCommand::LeftButtonReleaseEvent,
                   this->EventCallbackCommand, this->Priority);
    i->AddObserver(::vtkCommand::MouseMoveEvent, this->EventCallbackCommand,
                   this->Priority);
    i->AddObserver(::vtkCommand::RenderEvent, this->EventCallbackCommand,
                   this->Priority);
    i->AddObserver(::vtkCommand::KeyPressEvent, this->EventCallbackCommand,
                   this->Priority);

    this->UpdateGizmoPosition();
    
    // Manage Layers
    auto win = this->Interactor->GetRenderWindow();
    if (win->GetNumberOfLayers() < 2) {
      win->SetNumberOfLayers(2);
    }
    
    // Sync Viewport and Camera
    this->d->m_OverlayRenderer->SetViewport(this->CurrentRenderer->GetViewport());
    this->d->m_OverlayRenderer->SetActiveCamera(this->CurrentRenderer->GetActiveCamera());
    win->AddRenderer(this->d->m_OverlayRenderer);

    this->d->m_OverlayRenderer->AddActor(d->m_AxesX);
    this->d->m_OverlayRenderer->AddActor(d->m_AxesY);
    this->d->m_OverlayRenderer->AddActor(d->m_AxesZ);
    this->d->m_OverlayRenderer->AddActor(d->m_RotX);
    this->d->m_OverlayRenderer->AddActor(d->m_RotY);
    this->d->m_OverlayRenderer->AddActor(d->m_RotZ);
    this->d->m_OverlayRenderer->AddActor(d->m_RotCam);
    this->d->m_OverlayRenderer->AddActor(d->m_ScaleX);
    this->d->m_OverlayRenderer->AddActor(d->m_ScaleY);
    this->d->m_OverlayRenderer->AddActor(d->m_ScaleZ);

    this->UpdateVisibility();

    this->InvokeEvent(::vtkCommand::EnableEvent, nullptr);
  } else {
    if (!this->Enabled)
      return;

    this->Enabled = 0;
    this->Highlight(nullptr);
    if (this->Interactor) {
      this->Interactor->RemoveObserver(this->EventCallbackCommand);
      if (this->Interactor->GetRenderWindow()) {
          this->Interactor->GetRenderWindow()->MakeCurrent();
          this->Interactor->GetRenderWindow()->RemoveRenderer(this->d->m_OverlayRenderer);
      }
    }
    this->d->m_OverlayRenderer->RemoveAllViewProps();
    this->InvokeEvent(::vtkCommand::DisableEvent, nullptr);
  }

  if (this->Interactor) {
    this->Interactor->Render();
  }
}

void vtkHandlerWidget::ProcessEvents(::vtkObject *caller, unsigned long event,
                                     void *clientdata, void *calldata) {
  (void)caller;
  (void)calldata;
  vtkHandlerWidget *self = reinterpret_cast<vtkHandlerWidget *>(clientdata);
  switch (event) {
  case ::vtkCommand::LeftButtonPressEvent:
    self->OnLeftButtonDown();
    if (self->Interaction != ::uLib::Vtk::vtkHandlerWidget::IDLE)
      self->EventCallbackCommand->SetAbortFlag(1);
    break;
  case ::vtkCommand::LeftButtonReleaseEvent:
    self->OnLeftButtonUp();
    if (self->EventCallbackCommand->GetAbortFlag()) {
      // Don't let release bleed if press was captured
    }
    break;
  case ::vtkCommand::MouseMoveEvent:
    self->OnMouseMove();
    if (self->Interaction != ::uLib::Vtk::vtkHandlerWidget::IDLE)
      self->EventCallbackCommand->SetAbortFlag(1);
    break;
  case ::vtkCommand::RenderEvent:
    self->UpdateGizmoPosition();
    break;
  case ::vtkCommand::KeyPressEvent:
    self->OnKeyPress();
    break;
  }
}

void vtkHandlerWidget::OnKeyPress() {
  std::string key = this->Interactor->GetKeySym();
  bool ctrl = (this->Interactor->GetControlKey() != 0);
  
  if (ctrl && key == "z") {
    if (!this->d->m_UndoStack.empty()) {
      std::cout << "Undoing last transform action..." << std::endl;
      uLib::TRS target = this->d->m_UndoStack.back();
      this->d->m_UndoStack.pop_back();
      
      if (this->Prop3D) {
        this->Prop3D->SetPosition(target.position.x(), target.position.y(), target.position.z());
        // Convert Model Radians to VTK Degrees
        this->Prop3D->SetOrientation(target.rotation.x() / CLHEP::degree, 
                                   target.rotation.y() / CLHEP::degree, 
                                   target.rotation.z() / CLHEP::degree);
        this->Prop3D->SetScale(target.scaling.x(), target.scaling.y(), target.scaling.z());
        this->Prop3D->SetUserMatrix(nullptr);
        this->Prop3D->Modified();
        this->UpdateGizmoPosition();
        this->InvokeEvent(::vtkCommand::InteractionEvent, nullptr);
        this->Interactor->Render();
      }
    }
  }
}

void vtkHandlerWidget::OnLeftButtonDown() {
  int X = this->Interactor->GetEventPosition()[0];
  int Y = this->Interactor->GetEventPosition()[1];

  if (!this->CurrentRenderer) {
    this->CurrentRenderer = this->Interactor->FindPokedRenderer(X, Y);
  }

  this->d->m_Picker->Pick(X, Y, 0.0, this->d->m_OverlayRenderer);
  ::vtkProp *prop = this->d->m_Picker->GetViewProp();
  this->d->m_Picker->GetPickPosition(this->m_StartPickPosition);

  if (!prop)
    return;

  this->Interaction = IDLE;
  if (prop == d->m_AxesX)
    this->Interaction = TRANS_X;
  else if (prop == d->m_AxesY)
    this->Interaction = TRANS_Y;
  else if (prop == d->m_AxesZ)
    this->Interaction = TRANS_Z;
  else if (prop == d->m_RotX)
    this->Interaction = ROT_X;
  else if (prop == d->m_RotY)
    this->Interaction = ROT_Y;
  else if (prop == d->m_RotZ)
    this->Interaction = ROT_Z;
  else if (prop == d->m_ScaleX)
    this->Interaction = SCALE_X;
  else if (prop == d->m_ScaleY)
    this->Interaction = SCALE_Y;
  else if (prop == d->m_ScaleZ)
    this->Interaction = SCALE_Z;
  else if (prop == d->m_RotCam)
    this->Interaction = ROT_CAM;

  if (this->Interaction != IDLE) {
    this->StartEventPosition[0] = X;
    this->StartEventPosition[1] = Y;
    if (this->Prop3D) {
      // Capture current state for Undo
      this->d->m_UndoStack.push_back(uLib::TRS(uLib::Vtk::VtkToMatrix4f(this->Prop3D->GetMatrix())));
      if (this->d->m_UndoStack.size() > 50) this->d->m_UndoStack.erase(this->d->m_UndoStack.begin());

      // Use the prop's total matrix for calculation baseline
      this->d->m_InitialTransform->SetMatrix(this->Prop3D->GetMatrix());
    }
    this->EventCallbackCommand->SetAbortFlag(1);
    this->InvokeEvent(::vtkCommand::StartInteractionEvent, nullptr);
    this->Interactor->Render();
  }
}

void vtkHandlerWidget::OnLeftButtonUp() {
  if (this->Interaction == IDLE)
    return;

  this->Interaction = IDLE;
  this->EventCallbackCommand->SetAbortFlag(1);
  this->InvokeEvent(::vtkCommand::EndInteractionEvent, nullptr);
  this->Interactor->Render();
}

void vtkHandlerWidget::OnMouseMove() {
  if (!this->Prop3D || !this->CurrentRenderer)
    return;

  int X = this->Interactor->GetEventPosition()[0];
  int Y = this->Interactor->GetEventPosition()[1];

  if (this->Interaction == IDLE) {
    this->d->m_Picker->Pick(X, Y, 0.0, this->d->m_OverlayRenderer);
    ::vtkProp *prop = this->d->m_Picker->GetViewProp();
    this->Highlight(prop);
    this->UpdateGizmoPosition(); // Ensure camera adjustments happen
    return;
  }

  double dx = X - this->StartEventPosition[0];
  double dy = Y - this->StartEventPosition[1];

  if (dx == 0 && dy == 0) return;
  // std::cout << "Interaction " << this->Interaction << " dx=" << dx << " dy=" << dy << std::endl;

  // Get current gizmo properties from its actors
  vtkMatrix4x4 *gizmo_mat = d->m_AxesX->GetUserMatrix();
  if (!gizmo_mat)
    return;

  double gpos[3] = {gizmo_mat->GetElement(0, 3), gizmo_mat->GetElement(1, 3),
                    gizmo_mat->GetElement(2, 3)};

  // Normalized gizmo axes
  double gx[3] = {gizmo_mat->GetElement(0, 0), gizmo_mat->GetElement(1, 0),
                  gizmo_mat->GetElement(2, 0)};
  double gy[3] = {gizmo_mat->GetElement(0, 1), gizmo_mat->GetElement(1, 1),
                  gizmo_mat->GetElement(2, 1)};
  double gz[3] = {gizmo_mat->GetElement(0, 2), gizmo_mat->GetElement(1, 2),
                  gizmo_mat->GetElement(2, 2)};

  auto get_motion_magnitude = [&](double axis[3], double origin[3]) {
    double p1[3] = {origin[0], origin[1], origin[2]};
    double p2[3] = {origin[0] + axis[0], origin[1] + axis[1],
                    origin[2] + axis[2]};

    double d1[3], d2[3];
    this->ComputeWorldToDisplay(this->CurrentRenderer, p1[0], p1[1], p1[2], d1);
    this->ComputeWorldToDisplay(this->CurrentRenderer, p2[0], p2[1], p2[2], d2);

    double v[2] = {d2[0] - d1[0], d2[1] - d1[1]};
    double v_mag_sq = v[0] * v[0] + v[1] * v[1];
    if (v_mag_sq < 1.0)
      return 0.0;
    return (dx * v[0] + dy * v[1]) / v_mag_sq;
  };

  auto get_rotation_magnitude = [&](double axis[3]) -> double {
      // 1. Get Mouse Ray in World Space
      double worldNear[4], worldFar[4];
      this->ComputeDisplayToWorld(static_cast<double>(X), static_cast<double>(Y), 0.0, worldNear);
      this->ComputeDisplayToWorld(static_cast<double>(X), static_cast<double>(Y), 1.0, worldFar);
      
      double camPos[3];
      this->CurrentRenderer->GetActiveCamera()->GetPosition(camPos);
      
      double rayDir[3];
      for(int i=0; i<3; ++i) rayDir[i] = worldFar[i] - worldNear[i];
      vtkMath::Normalize(rayDir);

      // 2. Intersect Ray with Rotation Plane (Center = gpos, Normal = axis)
      double planeNormal[3] = {axis[0], axis[1], axis[2]};
      vtkMath::Normalize(planeNormal);
      
      double denom = vtkMath::Dot(rayDir, planeNormal);
      if (std::abs(denom) < 1e-6) return 0.0;
      
      double t = (vtkMath::Dot(gpos, planeNormal) - vtkMath::Dot(camPos, planeNormal)) / denom;
      double p_current[3];
      for(int i=0; i<3; ++i) p_current[i] = camPos[i] + t * rayDir[i];

      // 3. Calculate Angular Displacement in Plane
      double v_start[3] = {this->m_StartPickPosition[0] - gpos[0],
                          this->m_StartPickPosition[1] - gpos[1],
                          this->m_StartPickPosition[2] - gpos[2]};
      double v_end[3] = {p_current[0] - gpos[0],
                        p_current[1] - gpos[1],
                        p_current[2] - gpos[2]};

      double r_start = vtkMath::Norm(v_start);
      double r_end = vtkMath::Norm(v_end);
      if (r_start < 1e-9 || r_end < 1e-9) return 0.0;
      
      for(int i=0; i<3; ++i) { v_start[i] /= r_start; v_end[i] /= r_end; }

      double cross[3];
      vtkMath::Cross(v_start, v_end, cross);
      double sinAng = vtkMath::Dot(cross, planeNormal);
      double cosAng = vtkMath::Dot(v_start, v_end);
      
      return vtkMath::DegreesFromRadians(atan2(sinAng, cosAng));
  };

  // Create a transform that represents the operation in world space around gpos
  vtkNew<vtkTransform> op;
  op->PostMultiply();
  op->Translate(-gpos[0], -gpos[1], -gpos[2]);

  double mag = 0;
  switch (this->Interaction) {
  case TRANS_X:
    mag = get_motion_magnitude(gx, gpos);
    op->Translate(mag * gx[0], mag * gx[1], mag * gx[2]);
    break;
  case TRANS_Y:
    mag = get_motion_magnitude(gy, gpos);
    op->Translate(mag * gy[0], mag * gy[1], mag * gy[2]);
    break;
  case TRANS_Z:
    mag = get_motion_magnitude(gz, gpos);
    op->Translate(mag * gz[0], mag * gz[1], mag * gz[2]);
    break;
  case ROT_X:
    mag = get_rotation_magnitude(gx);
    {
      double ax[3] = {gx[0], gx[1], gx[2]};
      vtkMath::Normalize(ax);
      op->RotateWXYZ(mag, ax[0], ax[1], ax[2]);
    }
    break;
  case ROT_Y:
    mag = get_rotation_magnitude(gy);
    {
      double ax[3] = {gy[0], gy[1], gy[2]};
      vtkMath::Normalize(ax);
      op->RotateWXYZ(mag, ax[0], ax[1], ax[2]);
    }
    break;
  case ROT_Z:
    mag = get_rotation_magnitude(gz);
    {
      double ax[3] = {gz[0], gz[1], gz[2]};
      vtkMath::Normalize(ax);
      op->RotateWXYZ(mag, ax[0], ax[1], ax[2]);
    }
    break;
  case SCALE_X:
  case SCALE_Y:
  case SCALE_Z:
    {
        // 1. Calculate magnitude
        if (this->Interaction == SCALE_X) mag = get_motion_magnitude(gx, gpos);
        else if (this->Interaction == SCALE_Y) mag = get_motion_magnitude(gy, gpos);
        else mag = get_motion_magnitude(gz, gpos);
        
        double s = std::max(0.1, 1.0 + mag);
        
        // 2. Build a strictly orthonormal basis from the gizmo axes
        double X[3] = {gx[0], gx[1], gx[2]};
        double Y[3] = {gy[0], gy[1], gy[2]};
        double Z[3];
        
        vtkMath::Normalize(X);
        double dot = vtkMath::Dot(X, Y);
        for(int i=0; i<3; ++i) Y[i] -= dot * X[i];
        vtkMath::Normalize(Y);
        vtkMath::Cross(X, Y, Z); // Z is now orthogonal to X and Y
        
        vtkNew<vtkMatrix4x4> basis;
        basis->Identity();
        for(int i=0; i<3; ++i) {
            basis->SetElement(i, 0, X[i]);
            basis->SetElement(i, 1, Y[i]);
            basis->SetElement(i, 2, Z[i]);
        }
        
        vtkNew<vtkMatrix4x4> basis_inv;
        vtkMatrix4x4::Invert(basis, basis_inv);
        
        // 3. Assemble oriented scale: T(gpos) * basis * S * basis_inv * T(-gpos)
        // With PostMultiply: Result = T(gpos) * (basis * (S * (basis_inv * (T(-gpos) * Ident))))
        op->Concatenate(basis_inv);
        if (this->Interaction == SCALE_X) op->Scale(s, 1.0, 1.0);
        else if (this->Interaction == SCALE_Y) op->Scale(1.0, s, 1.0);
        else op->Scale(1.0, 1.0, s);
        op->Concatenate(basis);
    }
    break;
  case ROT_CAM: {
    double camPos[3];
    this->CurrentRenderer->GetActiveCamera()->GetPosition(camPos);
    double dir[3] = {camPos[0] - gpos[0], camPos[1] - gpos[1], camPos[2] - gpos[2]};
    vtkMath::Normalize(dir);
    
    // For camera ring, use screen-space angular delta
    double c_disp[3];
    this->ComputeWorldToDisplay(gpos[0], gpos[1], gpos[2], c_disp);
    double v1[2] = {this->StartEventPosition[0] - c_disp[0], this->StartEventPosition[1] - c_disp[1]};
    double v2[2] = {static_cast<double>(X) - c_disp[0], static_cast<double>(Y) - c_disp[1]};
    if (vtkMath::Norm2D(v1) > 1.0 && vtkMath::Norm2D(v2) > 1.0) {
        double d_ang = vtkMath::DegreesFromRadians(atan2(v2[1], v2[0]) - atan2(v1[1], v1[0]));
        while (d_ang > 180) d_ang -= 360;
        while (d_ang < -180) d_ang += 360;
        mag = d_ang;
    }
    op->RotateWXYZ(mag, dir[0], dir[1], dir[2]);
    break;
  }
}

  op->Translate(gpos[0], gpos[1], gpos[2]);

  // Total transform = Base * Chain * Interaction
  vtkNew<vtkTransform> total;
  total->PostMultiply();
  total->SetMatrix(this->d->m_InitialTransform->GetMatrix()); // d->m_InitialTransform is already Base*Chain
  total->Concatenate(op);

  if (this->Prop3D) {
      vtkNew<vtkMatrix4x4> result;
      total->GetMatrix(result); 
      this->Prop3D->SetUserMatrix(result);
      
      // Reset individual TRS components so UserMatrix is the single source of truth
      this->Prop3D->SetPosition(0, 0, 0);
      this->Prop3D->SetOrientation(0, 0, 0);
      this->Prop3D->SetScale(1, 1, 1);
  }
  
  this->Prop3D->Modified();
  this->UpdateGizmoPosition();

  this->InvokeEvent(::vtkCommand::InteractionEvent, nullptr);
  this->Interactor->Render();
}

void vtkHandlerWidget::SetReferenceFrame(ReferenceFrame frame) {
  this->m_Frame = frame;
  this->UpdateGizmoPosition();
  if (this->Interactor)
    this->Interactor->Render();
}

void vtkHandlerWidget::PlaceWidget(double bounds[6]) {
  (void)bounds;
  this->UpdateGizmoPosition();
}

void vtkHandlerWidget::SetTranslationEnabled(bool enabled) {
  this->m_TranslationEnabled = enabled;
  this->UpdateVisibility();
  if (this->Interactor) this->Interactor->Render();
}

void vtkHandlerWidget::SetRotationEnabled(bool enabled) {
  this->m_RotationEnabled = enabled;
  this->UpdateVisibility();
  if (this->Interactor) this->Interactor->Render();
}

void vtkHandlerWidget::SetScalingEnabled(bool enabled) {
  this->m_ScalingEnabled = enabled;
  this->UpdateVisibility();
  if (this->Interactor) this->Interactor->Render();
}

void vtkHandlerWidget::UpdateVisibility() {
    if (!d->m_AxesX) return;
    
    d->m_AxesX->SetVisibility(m_TranslationEnabled);
    d->m_AxesY->SetVisibility(m_TranslationEnabled);
    d->m_AxesZ->SetVisibility(m_TranslationEnabled);
    
    d->m_RotX->SetVisibility(m_RotationEnabled);
    d->m_RotY->SetVisibility(m_RotationEnabled);
    d->m_RotZ->SetVisibility(m_RotationEnabled);
    d->m_RotCam->SetVisibility(m_RotationEnabled);
    
    d->m_ScaleX->SetVisibility(m_ScalingEnabled);
    d->m_ScaleY->SetVisibility(m_ScalingEnabled);
    d->m_ScaleZ->SetVisibility(m_ScalingEnabled);
    
    // Update picker list
    if (d->m_Picker) {
        d->m_Picker->InitializePickList();
        if (m_TranslationEnabled) {
            d->m_Picker->AddPickList(d->m_AxesX);
            d->m_Picker->AddPickList(d->m_AxesY);
            d->m_Picker->AddPickList(d->m_AxesZ);
        }
        if (m_RotationEnabled) {
            d->m_Picker->AddPickList(d->m_RotX);
            d->m_Picker->AddPickList(d->m_RotY);
            d->m_Picker->AddPickList(d->m_RotZ);
            d->m_Picker->AddPickList(d->m_RotCam);
        }
        if (m_ScalingEnabled) {
            d->m_Picker->AddPickList(d->m_ScaleX);
            d->m_Picker->AddPickList(d->m_ScaleY);
            d->m_Picker->AddPickList(d->m_ScaleZ);
        }
        d->m_Picker->PickFromListOn();
    }
}

void vtkHandlerWidget::PlaceWidget() { this->UpdateGizmoPosition(); }

void vtkHandlerWidget::SetTransform(::vtkTransform *t) {
  if (!t || !this->Prop3D)
    return;
  this->Prop3D->SetUserMatrix(t->GetMatrix());
  this->UpdateGizmoPosition();
}

void vtkHandlerWidget::GetTransform(::vtkTransform *t) {
  if (!t || !this->Prop3D)
    return;
  t->SetMatrix(this->Prop3D->GetMatrix());
}

void vtkHandlerWidget::CreateGizmos() {
  auto create_arrow = [](double dir[3], double color[3]) {
    auto arrow = vtkSmartPointer<::vtkArrowSource>::New();
    arrow->SetTipLength(0.2);
    arrow->SetTipRadius(0.06);
    arrow->SetShaftRadius(0.015);

    auto mapper = vtkSmartPointer<::vtkPolyDataMapper>::New();
    mapper->SetInputConnection(arrow->GetOutputPort());
    auto actor = vtkSmartPointer<::vtkActor>::New();
    actor->SetMapper(mapper);
    actor->GetProperty()->SetColor(color);
    actor->GetProperty()->SetLighting(0); // Saturated colors, no shadows

    if (dir[1] > 0)
      actor->RotateZ(90);
    else if (dir[2] > 0)
      actor->RotateY(-90);
    return actor;
  };

  auto create_ring = [&](int axis, double color[3]) {
    auto circle = vtkSmartPointer<::vtkRegularPolygonSource>::New();
    circle->SetNumberOfSides(64);
    circle->SetRadius(1.0);
    circle->SetCenter(0, 0, 0);
    circle->GeneratePolygonOff();
    circle->GeneratePolylineOn();

    if (axis == 0) circle->SetNormal(1, 0, 0);
    else if (axis == 1) circle->SetNormal(0, 1, 0);
    else if (axis == 2) circle->SetNormal(0, 0, 1);

    auto mapper = vtkSmartPointer<::vtkPolyDataMapper>::New();
    mapper->SetInputConnection(circle->GetOutputPort());
    mapper->AddClippingPlane(this->d->m_ClipPlane);

    auto actor = vtkSmartPointer<::vtkActor>::New();
    actor->SetMapper(mapper);
    actor->GetProperty()->SetColor(color);
    actor->GetProperty()->SetLineWidth(3);
    actor->GetProperty()->SetLighting(0);
    return actor;
  };

  double red[] = {1.0, 0.0, 0.0}, green[] = {0.0, 1.0, 0.0},
         blue[] = {0.0, 0.0, 1.0}, white[] = {1.0, 1.0, 1.0};

  double x[] = {1, 0, 0}, y[] = {0, 1, 0}, z[] = {0, 0, 1};
  d->m_AxesX = create_arrow(x, red);
  d->m_AxesY = create_arrow(y, green);
  d->m_AxesZ = create_arrow(z, blue);

  d->m_RotX = create_ring(0, red);
  d->m_RotY = create_ring(1, green);
  d->m_RotZ = create_ring(2, blue);

  d->m_RotCam = vtkSmartPointer<::vtkActor>::New();
  {
    auto circle = vtkSmartPointer<::vtkRegularPolygonSource>::New();
    circle->SetNumberOfSides(64);
    circle->SetRadius(1.3); // Slightly larger
    circle->SetCenter(0, 0, 0);
    circle->GeneratePolygonOff();
    circle->GeneratePolylineOn();
    auto mapper = vtkSmartPointer<::vtkPolyDataMapper>::New();
    mapper->SetInputConnection(circle->GetOutputPort());
    d->m_RotCam->SetMapper(mapper);
    d->m_RotCam->GetProperty()->SetColor(white);
    d->m_RotCam->GetProperty()->SetLineWidth(2);
    d->m_RotCam->GetProperty()->SetLighting(0);
  }

  auto create_cube = [](double pos[3], double color[3]) {
    auto cube = vtkSmartPointer<::vtkCubeSource>::New();
    cube->SetCenter(pos[0], pos[1], pos[2]);
    cube->SetXLength(0.08);
    cube->SetYLength(0.08);
    cube->SetZLength(0.08);
    auto mapper = vtkSmartPointer<::vtkPolyDataMapper>::New();
    mapper->SetInputConnection(cube->GetOutputPort());
    auto actor = vtkSmartPointer<::vtkActor>::New();
    actor->SetMapper(mapper);
    actor->GetProperty()->SetColor(color);
    actor->GetProperty()->SetLighting(0);
    return actor;
  };

  double px[] = {1.2, 0, 0}, py[] = {0, 1.2, 0}, pz[] = {0, 0, 1.2};
  d->m_ScaleX = create_cube(px, red);
  d->m_ScaleY = create_cube(py, green);
  d->m_ScaleZ = create_cube(pz, blue);

  // Configure picker to only see gizmo actors (Pick-Through)
  d->m_Picker->InitializePickList();
  d->m_Picker->AddPickList(d->m_AxesX);
  d->m_Picker->AddPickList(d->m_AxesY);
  d->m_Picker->AddPickList(d->m_AxesZ);
  d->m_Picker->AddPickList(d->m_RotX);
  d->m_Picker->AddPickList(d->m_RotY);
  d->m_Picker->AddPickList(d->m_RotZ);
  d->m_Picker->AddPickList(d->m_RotCam);
  d->m_Picker->AddPickList(d->m_ScaleX);
  d->m_Picker->AddPickList(d->m_ScaleY);
  d->m_Picker->AddPickList(d->m_ScaleZ);
  d->m_Picker->PickFromListOn();
}

void vtkHandlerWidget::UpdateGizmoPosition() {
  if (!this->Prop3D)
    return;

  // Calculate scaling factor: min(object_bbox, 1/5 of viewport)
  double bboxSize = 1.0;
  double bounds[6];
  this->Prop3D->GetBounds(bounds);
  if (vtkMath::AreBoundsInitialized(bounds)) {
    bboxSize = std::max({bounds[1] - bounds[0], bounds[3] - bounds[2], bounds[5] - bounds[4]});
  }
  if (bboxSize < 1e-6) bboxSize = 1.0;

  double screenLimit = bboxSize * 2.0; // Default if no renderer
    if (this->CurrentRenderer && this->CurrentRenderer->GetRenderWindow() && this->CurrentRenderer->GetRenderWindow()->GetInteractor()) {
      int *sz = this->CurrentRenderer->GetSize();
      if (sz[1] > 0) {
        double pixelSize = std::min(sz[0], sz[1]) / 5.0;
        vtkCamera *cam = this->CurrentRenderer->GetActiveCamera();
        if (cam) {
          if (cam->GetParallelProjection()) {
            screenLimit = (pixelSize / (double)sz[1]) * 2.0 * cam->GetParallelScale();
          } else {
            double dist = cam->GetDistance();
            double angleRad = vtkMath::Pi() * cam->GetViewAngle() / 180.0;
            double viewHeightAtDist = 2.0 * dist * tan(angleRad / 2.0);
            screenLimit = (pixelSize / (double)sz[1]) * viewHeightAtDist;
          }
        }
      }
    }
  double scaleFactor = std::min(bboxSize, screenLimit);

  vtkNew<vtkMatrix4x4> mat_gizmo;
  mat_gizmo->Identity();

  double center[3];
  center[0] = (bounds[0] + bounds[1]) / 2.0;
  center[1] = (bounds[2] + bounds[3]) / 2.0;
  center[2] = (bounds[4] + bounds[5]) / 2.0;

  double pos[3];
  this->Prop3D->GetPosition(pos);

  if (m_Frame == LOCAL) {
    ::vtkMatrix4x4 *mat = this->Prop3D->GetMatrix();
    mat_gizmo->DeepCopy(mat);
    // Remove scaling
    for (int j = 0; j < 3; ++j) {
      double v[3] = {mat->GetElement(0, j), mat->GetElement(1, j),
                     mat->GetElement(2, j)};
      double len = vtkMath::Norm(v);
      if (len > 1e-6) {
        mat_gizmo->SetElement(0, j, v[0] / len);
        mat_gizmo->SetElement(1, j, v[1] / len);
        mat_gizmo->SetElement(2, j, v[2] / len);
      }
    }
  } else if (m_Frame == CENTER_LOCAL) {
    ::vtkMatrix4x4 *mat = this->Prop3D->GetMatrix();
    mat_gizmo->DeepCopy(mat);
    // Remove scaling
    for (int j = 0; j < 3; ++j) {
      double v[3] = {mat->GetElement(0, j), mat->GetElement(1, j),
                     mat->GetElement(2, j)};
      double len = vtkMath::Norm(v);
      if (len > 1e-6) {
        mat_gizmo->SetElement(0, j, v[0] / len);
        mat_gizmo->SetElement(1, j, v[1] / len);
        mat_gizmo->SetElement(2, j, v[2] / len);
      }
    }
    // Set position to center
    mat_gizmo->SetElement(0, 3, center[0]);
    mat_gizmo->SetElement(1, 3, center[1]);
    mat_gizmo->SetElement(2, 3, center[2]);
  } else if (m_Frame == GLOBAL) {
    ::vtkMatrix4x4 *mat = this->Prop3D->GetMatrix();
    mat_gizmo->Identity();
    mat_gizmo->SetElement(0, 3, mat->GetElement(0, 3));
    mat_gizmo->SetElement(1, 3, mat->GetElement(1, 3));
    mat_gizmo->SetElement(2, 3, mat->GetElement(2, 3));
  } else if (m_Frame == CENTER) {
    mat_gizmo->Identity();
    mat_gizmo->SetElement(0, 3, center[0]);
    mat_gizmo->SetElement(1, 3, center[1]);
    mat_gizmo->SetElement(2, 3, center[2]);
  }

  // Apply scaleFactor to the gizmo matrix (only top-left 3x3)
  for (int j = 0; j < 3; ++j) {
    for (int i = 0; i < 3; ++i) {
      mat_gizmo->SetElement(i, j, mat_gizmo->GetElement(i, j) * scaleFactor);
    }
  }

  d->m_AxesX->SetUserMatrix(mat_gizmo);
  d->m_AxesY->SetUserMatrix(mat_gizmo);
  d->m_AxesZ->SetUserMatrix(mat_gizmo);
  d->m_RotX->SetUserMatrix(mat_gizmo);
  d->m_RotY->SetUserMatrix(mat_gizmo);
  d->m_RotZ->SetUserMatrix(mat_gizmo);
  d->m_ScaleX->SetUserMatrix(mat_gizmo);
  d->m_ScaleY->SetUserMatrix(mat_gizmo);
  d->m_ScaleZ->SetUserMatrix(mat_gizmo);

  // Sync Overlay Renderer with Main Renderer
  if (this->CurrentRenderer && this->d->m_OverlayRenderer) {
    this->d->m_OverlayRenderer->SetViewport(this->CurrentRenderer->GetViewport());
    this->d->m_OverlayRenderer->SetAspect(this->CurrentRenderer->GetAspect());
    this->d->m_OverlayRenderer->ComputeAspect();

    if (this->d->m_OverlayRenderer->GetActiveCamera() != this->CurrentRenderer->GetActiveCamera()) {
        this->d->m_OverlayRenderer->SetActiveCamera(this->CurrentRenderer->GetActiveCamera());
    }
  }

  // Camera ring always faces camera
  if (this->CurrentRenderer) {
    double camPos[3];
    this->CurrentRenderer->GetActiveCamera()->GetPosition(camPos);
    double dir[3] = {camPos[0] - mat_gizmo->GetElement(0, 3),
                     camPos[1] - mat_gizmo->GetElement(1, 3),
                     camPos[2] - mat_gizmo->GetElement(2, 3)};
    vtkMath::Normalize(dir);

    // Orient RotCam actor to face 'dir'
    vtkNew<vtkTransform> tcam;
    tcam->PostMultiply();
    tcam->Scale(scaleFactor, scaleFactor, scaleFactor);
    // Default circle is in XY plane (Normal Z: 0,0,1)
    double z[3] = {0, 0, 1};
    double cross[3];
    vtkMath::Cross(z, dir, cross);
    double cross_mag = vtkMath::Norm(cross);
    if (cross_mag > 1e-6) {
      double angle = vtkMath::DegreesFromRadians(acos(vtkMath::Dot(z, dir)));
      tcam->RotateWXYZ(angle, cross[0], cross[1], cross[2]);
    } else if (vtkMath::Dot(z, dir) < 0) {
      tcam->RotateX(180);
    }
    tcam->Translate(mat_gizmo->GetElement(0, 3), mat_gizmo->GetElement(1, 3),
                    mat_gizmo->GetElement(2, 3));
    d->m_RotCam->SetUserMatrix(tcam->GetMatrix());

    // Update clipping plane for axes rings
    this->d->m_ClipPlane->SetOrigin(mat_gizmo->GetElement(0, 3),
                                 mat_gizmo->GetElement(1, 3),
                                 mat_gizmo->GetElement(2, 3));
    this->d->m_ClipPlane->SetNormal(dir);
  }
}

void vtkHandlerWidget::Highlight(::vtkProp *prop) {
  if (this->d->m_HighlightedProp == prop)
    return;

  // Restore previous
  if (this->d->m_HighlightedProp) {
    ::vtkActor *actor = ::vtkActor::SafeDownCast(this->d->m_HighlightedProp);
    if (actor) {
      actor->GetProperty()->SetColor(m_OriginalColor);
      actor->GetProperty()->SetLineWidth(3);
    }
  }

  this->d->m_HighlightedProp = nullptr;

  // Highlight new if it belongs to us
  if (prop == d->m_AxesX || prop == d->m_AxesY || prop == d->m_AxesZ || prop == d->m_RotX ||
      prop == d->m_RotY || prop == d->m_RotZ || prop == d->m_RotCam || prop == d->m_ScaleX ||
      prop == d->m_ScaleY || prop == d->m_ScaleZ) {
    this->d->m_HighlightedProp = prop;
    ::vtkActor *actor = ::vtkActor::SafeDownCast(prop);
    if (actor) {
      actor->GetProperty()->GetColor(m_OriginalColor);
      double h[3] = {m_OriginalColor[0] + 0.2, m_OriginalColor[1] + 0.2,
                     m_OriginalColor[2] + 0.2};
      for (int i = 0; i < 3; ++i)
        h[i] = std::min(1.0, h[i]);
      actor->GetProperty()->SetColor(h);
      actor->GetProperty()->SetLineWidth(5);
    }
  }

  if (this->Interactor)
    this->Interactor->Render();
}

} // namespace Vtk
} // namespace uLib