uLib/src/Vtk/HEP/Geant/vtkEmitterPrimary.cpp

#include "vtkEmitterPrimary.h"
#include <vtkActor.h>
#include <vtkArrowSource.h>
#include <vtkMatrix4x4.h>
#include <vtkNew.h>
#include <vtkPolyDataMapper.h>
#include <vtkTransform.h>
#include <vtkTransformPolyDataFilter.h>
#include <vtkPlaneSource.h>
#include <vtkAppendPolyData.h>
#include <vtkPolyData.h>
#include <vtkProperty.h>
#include <vtkCylinderSource.h>
#include "Vtk/Math/vtkDense.h"

namespace uLib {
namespace Vtk {

EmitterPrimary::EmitterPrimary(Geant::EmitterPrimary &emitter)
    : m_emitter(emitter), m_Actor(vtkActor::New()) {
    
    vtkNew<vtkArrowSource> arrow;
    
    // Default arrow is along X+. 
    // Scale it to 1 meter (1000 mm) and rotate to point towards Z- (local)
    vtkNew<vtkTransform> transform;
    transform->Scale(1000.0, 1000.0, 1000.0);
    transform->RotateY(90.0);
    
    vtkNew<vtkTransformPolyDataFilter> transformFilter;
    transformFilter->SetTransform(transform);
    transformFilter->SetInputConnection(arrow->GetOutputPort());
    transformFilter->Update();
    
    vtkSmartPointer<vtkPolyDataMapper> mapper = vtkSmartPointer<vtkPolyDataMapper>::New();
    mapper->SetInputData(transformFilter->GetOutput());
    m_Actor->SetMapper(mapper);
    m_Actor->SetScale(1.0);

    vtkNew<vtkMatrix4x4> vmat;
    Matrix4fToVtk(m_emitter.GetWorldMatrix(), vmat);
    m_Actor->SetUserMatrix(vmat);

    this->SetProp(m_Actor);
    
    Object::connect(&m_emitter, &Object::Updated, this, &EmitterPrimary::contentUpdate);
    this->contentUpdate();
}

EmitterPrimary::~EmitterPrimary() {
    Object::disconnect(&m_emitter, &Object::Updated, this, &EmitterPrimary::contentUpdate);
    m_Actor->Delete();
}

void EmitterPrimary::contentUpdate() {
    vtkMatrix4x4 *vmat = m_Actor->GetUserMatrix();
    if (!vmat) {
        vtkNew<vtkMatrix4x4> mat;
        m_Actor->SetUserMatrix(mat);
        vmat = mat;
    }

    Matrix4f transform = m_emitter.GetWorldMatrix();
    Matrix4fToVtk(transform, vmat);

    Prop3D::Update();
}

void EmitterPrimary::Update() {
    vtkMatrix4x4 *vmat = m_Actor->GetUserMatrix();
    if (!vmat) return;

    Matrix4f transform = VtkToMatrix4f(vmat);
    m_emitter.SetMatrix(transform);
    m_emitter.Updated();
}

// -------------------------------------------------------------------------- //
// vtkSkyPlaneEmitterPrimary

vtkSkyPlaneEmitterPrimary::vtkSkyPlaneEmitterPrimary(Geant::SkyPlaneEmitterPrimary &emitter)
    : EmitterPrimary(emitter), m_skyEmitter(emitter), m_PlaneSource(vtkPlaneSource::New()) {
    
    vtkNew<vtkAppendPolyData> append;
    // Base class constructor already set an arrow. We keep it as a directional indicator.
    append->AddInputData(vtkPolyData::SafeDownCast(m_Actor->GetMapper()->GetInput()));
    append->AddInputConnection(m_PlaneSource->GetOutputPort());
    
    vtkNew<vtkPolyDataMapper> mapper;
    mapper->SetInputConnection(append->GetOutputPort());
    m_Actor->SetMapper(mapper);
    
    m_Actor->GetProperty()->SetOpacity(0.5);
    m_Actor->GetProperty()->SetColor(0.2, 0.6, 0.9);
    
    this->contentUpdate();
}

vtkSkyPlaneEmitterPrimary::~vtkSkyPlaneEmitterPrimary() {
    m_PlaneSource->Delete();
}

void vtkSkyPlaneEmitterPrimary::contentUpdate() {
    uLib::Vector2f size = m_skyEmitter.GetSkySize();
    m_PlaneSource->SetOrigin(-size(0)/2.0, -size(1)/2.0, 0.0);
    m_PlaneSource->SetPoint1( size(0)/2.0, -size(1)/2.0, 0.0);
    m_PlaneSource->SetPoint2(-size(0)/2.0,  size(1)/2.0, 0.0);
    m_PlaneSource->Update();

    EmitterPrimary::contentUpdate();
}

// -------------------------------------------------------------------------- //
// vtkQuadMeshEmitterPrimary

vtkQuadMeshEmitterPrimary::vtkQuadMeshEmitterPrimary(Geant::QuadMeshEmitterPrimary &emitter)
    : EmitterPrimary(emitter), m_meshEmitter(emitter) {
    this->contentUpdate();
}

void vtkQuadMeshEmitterPrimary::contentUpdate() {
    // For now stick with the arrow. In the future visualize the mesh if useful.
    EmitterPrimary::contentUpdate();
}
 
// -------------------------------------------------------------------------- //
// vtkCylinderEmitterPrimary

vtkCylinderEmitterPrimary::vtkCylinderEmitterPrimary(Geant::CylinderEmitterPrimary &emitter)
    : EmitterPrimary(emitter), m_cylinderEmitter(emitter), m_CylinderSource(vtkCylinderSource::New()) {
    
    // vtkCylinderSource is along Y by default.
    // We will update its actual dimensions in contentUpdate().
    m_CylinderSource->SetResolution(64);

    // Rotate it to be along local Z
    vtkNew<vtkTransform> tcyl;
    tcyl->RotateX(90.0);
    
    vtkNew<vtkTransformPolyDataFilter> cylFilter;
    cylFilter->SetTransform(tcyl);
    cylFilter->SetInputConnection(m_CylinderSource->GetOutputPort());

    vtkNew<vtkAppendPolyData> append;
    // Keep the directional arrow from base class
    append->AddInputData(vtkPolyData::SafeDownCast(m_Actor->GetMapper()->GetInput()));
    append->AddInputConnection(cylFilter->GetOutputPort());
    
    vtkNew<vtkPolyDataMapper> mapper;
    mapper->SetInputConnection(append->GetOutputPort());
    m_Actor->SetMapper(mapper);
    
    m_Actor->GetProperty()->SetOpacity(0.5);
    m_Actor->GetProperty()->SetColor(0.2, 0.6, 0.9);
    
    this->contentUpdate();
}

vtkCylinderEmitterPrimary::~vtkCylinderEmitterPrimary() {
    m_CylinderSource->Delete();
}

void vtkCylinderEmitterPrimary::contentUpdate() {
    float r = m_cylinderEmitter.GetRadius();
    float h = m_cylinderEmitter.GetHeight();
    
    m_CylinderSource->SetRadius(r);
    m_CylinderSource->SetHeight(h);
    // vtkCylinderSource is along Y. To have the base at Y=0, we center it at h/2.
    m_CylinderSource->SetCenter(0, h/2.0, 0);
    m_CylinderSource->Update();

    EmitterPrimary::contentUpdate();
}

} // namespace Vtk
} // namespace uLib