uLib/src/Vtk/vtkObjectsContext.cpp

#include "vtkObjectsContext.h"
#include "Vtk/Math/vtkAssembly.h"
#include "Vtk/Math/vtkContainerBox.h"
#include "Vtk/Math/vtkCylinder.h"
#include "Vtk/Math/vtkVoxImage.h"

#include "Vtk/HEP/Detectors/vtkDetectorChamber.h"
#include "Vtk/HEP/Geant/vtkBoxSolid.h"
#include "Vtk/HEP/Geant/vtkTessellatedSolid.h"

#include <cstring>
#include <iostream>
#include <vtkAssembly.h>
#include <vtkPropCollection.h>

#include "HEP/Detectors/DetectorChamber.h"
#include "HEP/Geant/Solid.h"
#include "Math/Assembly.h"
#include "Math/ContainerBox.h"
#include "Math/Cylinder.h"
#include "Math/VoxImage.h"

namespace uLib {
namespace Vtk {

ObjectsContext::ObjectsContext(uLib::ObjectsContext *context)
    : m_Context(context), m_Assembly(::vtkAssembly::New()) {
  this->SetProp(m_Assembly);
  if (m_Context) {
    Object::connect(m_Context, &uLib::ObjectsContext::ObjectAdded, this,
                    &ObjectsContext::OnObjectAdded);
    Object::connect(m_Context, &uLib::ObjectsContext::ObjectRemoved, this,
                    &ObjectsContext::OnObjectRemoved);
    this->Synchronize();
  }
}

ObjectsContext::~ObjectsContext() {
  if (m_Context) {
    Object::disconnect(m_Context, &uLib::ObjectsContext::ObjectAdded, this,
                      &ObjectsContext::OnObjectAdded);
    Object::disconnect(m_Context, &uLib::ObjectsContext::ObjectRemoved, this,
                      &ObjectsContext::OnObjectRemoved);
  }
  for (auto const &[obj, prop3d] : m_Prop3Ds) {
    delete prop3d;
  }
  m_Assembly->Delete();
}

void ObjectsContext::Synchronize() {
  if (!m_Context)
    return;

  // 1. Identify objects to add and remove
  const auto &objects = m_Context->GetObjects();
  std::map<uLib::Object *, bool> currentObjects;
  for (const auto& obj : objects)
    currentObjects[obj.get()] = true;

  // Remove Prop3Ds for objects no longer in context
  for (auto it = m_Prop3Ds.begin(); it != m_Prop3Ds.end();) {
    if (currentObjects.find(it->first) == currentObjects.end()) {
      it->second->DisconnectRenderer(
          nullptr); // If we have a ref to a renderer we should disconnect but
                    // Prop3D doesn't store it easily
      // Actually Prop3D::DisconnectRenderer(vtkRenderer*) needs the renderer.
      // For now we just remove from assembly
      if (auto *p3d = vtkProp3D::SafeDownCast(it->second->GetProp()))
        m_Assembly->RemovePart(p3d);
      this->Prop3DRemoved(it->second);
      delete it->second;
      it = m_Prop3Ds.erase(it);
    } else {
      ++it;
    }
  }

  // Add Prop3Ds for new objects
  for (const auto& obj : objects) {
    if (m_Prop3Ds.find(obj.get()) == m_Prop3Ds.end()) {
      Prop3D *prop3d = this->CreateProp3D(obj.get());
      if (prop3d) {
        m_Prop3Ds[obj.get()] = prop3d;
        if (auto *p3d = vtkProp3D::SafeDownCast(prop3d->GetProp()))
          m_Assembly->AddPart(p3d);
        this->Prop3DAdded(prop3d);
      }
    }
  }
}

void ObjectsContext::OnObjectAdded(uLib::Object *obj) {
  if (!obj)
    return;
  if (m_Prop3Ds.find(obj) == m_Prop3Ds.end()) {
    Prop3D *prop3d = this->CreateProp3D(obj);
    if (prop3d) {
      m_Prop3Ds[obj] = prop3d;
      if (auto *p3d = vtkProp3D::SafeDownCast(prop3d->GetProp()))
        m_Assembly->AddPart(p3d);
      this->Prop3DAdded(prop3d);
    }
  }
}

void ObjectsContext::OnObjectRemoved(uLib::Object *obj) {
  if (!obj)
    return;
  auto it = m_Prop3Ds.find(obj);
  if (it != m_Prop3Ds.end()) {
    // For now we just remove from assembly.
    // Prop3D::DisconnectRenderer(vtkRenderer*) needs the renderer, but we don't
    // have it here easily.
    if (auto *p3d = vtkProp3D::SafeDownCast(it->second->GetProp()))
      m_Assembly->RemovePart(p3d);
    this->Prop3DRemoved(it->second);
    delete it->second;
    m_Prop3Ds.erase(it);
  }
}

Prop3D *ObjectsContext::GetProp3D(uLib::Object *obj) {
  auto it = m_Prop3Ds.find(obj);
  if (it != m_Prop3Ds.end())
    return it->second;
  return nullptr;
}

void ObjectsContext::Update() {
  for (auto const &[obj, prop3d] : m_Prop3Ds) {
    prop3d->Update();
  }
}

void ObjectsContext::SyncFromVtk() {
  for (auto const &[obj, prop3d] : m_Prop3Ds) {
    prop3d->SyncFromVtk();
  }
}

Prop3D *ObjectsContext::CreateProp3D(uLib::Object *obj) {
  if (!obj)
    return nullptr;

  if (auto* p3d = dynamic_cast<Prop3D*>(obj)) {
      return p3d;
  }

  if (auto *vox = dynamic_cast<uLib::Abstract::VoxImage *>(obj)) {
    return new VoxImage(vox);
  } else if (auto *box = dynamic_cast<uLib::ContainerBox *>(obj)) {
    return new ContainerBox(box);
  } else if (auto *chamber = dynamic_cast<uLib::DetectorChamber *>(obj)) {
    return new DetectorChamber(chamber);
  } else if (auto *cylinder = dynamic_cast<uLib::Cylinder *>(obj)) {
    return new Cylinder(cylinder);
  } else if (auto *assembly = dynamic_cast<uLib::Assembly *>(obj)) {
    return new Assembly(assembly);
  } else if (auto *pv = dynamic_cast<uLib::Geant::PhysicalVolume *>(obj)) {
    uLib::Geant::Solid *solid = pv->GetLogical()->GetSolid();
    if (auto *box = dynamic_cast<uLib::Geant::BoxSolid *>(solid)) {
      return new BoxSolid(pv);
    } else if (auto *tess = dynamic_cast<uLib::Geant::TessellatedSolid *>(solid)) {
      return new TessellatedSolid(pv);
    } else {
      return new GeantSolid(pv);
    }
  } else if (auto *box = dynamic_cast<uLib::Geant::BoxSolid *>(obj)) {
    return new BoxSolid(box);
  } else if (auto *solid = dynamic_cast<uLib::Geant::Solid *>(obj)) {
    // Unplaced solid visualization
    return new GeantSolid(solid);
  }

  // Fallback if we don't know the exact class but it might be a context itself
  if (auto subCtx = dynamic_cast<uLib::ObjectsContext *>(obj)) {
    return new ObjectsContext(subCtx);
  }

  return nullptr;
}

void ObjectsContext::Prop3DAdded(Prop3D *prop3d) {
  ULIB_SIGNAL_EMIT(ObjectsContext::Prop3DAdded, prop3d);
}

void ObjectsContext::Prop3DRemoved(Prop3D *prop3d) {
  ULIB_SIGNAL_EMIT(ObjectsContext::Prop3DRemoved, prop3d);
}

} // namespace Vtk
} // namespace uLib