/*//////////////////////////////////////////////////////////////////////////////
// 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 >

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

#include "Vtk/Math/vtkVoxImage.h"
#include "Vtk/uLibVtkViewer.h"
#include <cmath>
#include <iostream>
#include <vtkCallbackCommand.h>
#include <vtkProperty.h>
#include <vtkRenderWindow.h>
#include <vtkRenderWindowInteractor.h>
#include <vtkRenderer.h>
#include <vtkSmartPointer.h>

using namespace uLib;

struct AppState {
  std::vector<Vtk::VoxImage *> images;
  Vtk::Viewer *viewer;
};

void KeyPressCallbackFunction(vtkObject *caller, long unsigned int eventId,
                              void *clientData, void *callData) {
  auto *interactor = static_cast<vtkRenderWindowInteractor *>(caller);
  auto *state = static_cast<AppState *>(clientData);

  std::string key = interactor->GetKeySym();
  if (key == "w") {
    std::cout << "--> Switching all images to Wireframe Box" << std::endl;
    for (auto *img : state->images)
      img->SetRepresentation(Vtk::Prop3D::Wireframe);
    state->viewer->GetRenderWindow()->Render();
  } else if (key == "s") {
    std::cout << "--> Switching all images to Surface (Volume Rendering)"
              << std::endl;
    for (auto *img : state->images)
      img->SetRepresentation(Vtk::Prop3D::Surface);
    state->viewer->GetRenderWindow()->Render();
  } else if (key >= "0" && key <= "5") {
    int preset = key[0] - '0';
    std::cout << "--> Switching all images to Rendering Preset " << preset
              << std::endl;
    for (auto *img : state->images)
      img->setShadingPreset(preset);
    state->viewer->GetRenderWindow()->Render();
  }
}

int main(int argc, char **argv) {
  float factor = 1.0e6f;

  // --- Image 1: Spherical Shell ---
  Vector3i dims1(64, 64, 64);
  VoxImage<Voxel> img1(dims1);
  img1.SetSpacing(Vector3f(1.0, 1.0, 1.0));
  img1.SetPosition(Vector3f(-40, -32, -32));

  for (int z = 0; z < dims1(2); ++z) {
    for (int y = 0; y < dims1(1); ++y) {
      for (int x = 0; x < dims1(0); ++x) {
        float dx = x - 32.0f;
        float dy = y - 32.0f;
        float dz = z - 32.0f;
        float dist = std::sqrt(dx * dx + dy * dy + dz * dz);
        Voxel v;
        if (dist < 25.0f && dist > 10.0f) {
          v.Value = (40.0f * (25.0f - dist) / 15.0f) / factor;
        } else {
          v.Value = 0.0f;
        }
        img1[Vector3i(x, y, z)] = v;
      }
    }
  }

  // --- Image 2: Axes Gradient ---
  Vector3i dims2(64, 64, 64);
  VoxImage<Voxel> img2(dims2);
  img2.SetSpacing(Vector3f(1.0, 1.0, 1.0));
  img2.SetPosition(Vector3f(40, -32, -32));

  for (int z = 0; z < dims2(2); ++z) {
    for (int y = 0; y < dims2(1); ++y) {
      for (int x = 0; x < dims2(0); ++x) {
        Voxel v;
        // Linear gradient along X, Y, Z
        float val =
            (float(x) / dims2(0) + float(y) / dims2(1) + float(z) / dims2(2)) /
            3.0f;
        v.Value = (40.0f * val) / factor;
        img2[Vector3i(x, y, z)] = v;
      }
    }
  }

  Vtk::VoxImage vtk_img1(&img1);
  vtk_img1.setShadingPreset(0);

  Vtk::VoxImage vtk_img2(&img2);
  vtk_img2.setShadingPreset(1); // Use Composite without MIP for variety

  Vtk::Viewer viewer;
  viewer.GetRenderer()->SetBackground(0.05, 0.05, 0.1);
  viewer.AddProp3D(vtk_img1);
  viewer.AddProp3D(vtk_img2);

  // Setup KeyPress Callback
  AppState state;
  state.images.push_back(&vtk_img1);
  state.images.push_back(&vtk_img2);
  state.viewer = &viewer;

  vtkSmartPointer<vtkCallbackCommand> keyCallback =
      vtkSmartPointer<vtkCallbackCommand>::New();
  keyCallback->SetCallback(KeyPressCallbackFunction);
  keyCallback->SetClientData(&state);
  viewer.GetInteractor()->AddObserver(vtkCommand::KeyPressEvent, keyCallback);

  std::cout << "=========================================" << std::endl;
  std::cout << " VoxImage Interactive Viewer Test" << std::endl;
  std::cout << " [LEFT] Spherical Shell (MIP)" << std::endl;
  std::cout << " [RIGHT] Axes Gradient (Composite)" << std::endl;
  std::cout << "-----------------------------------------" << std::endl;
  std::cout << " Press [w] to show Wireframe Bounding Boxes" << std::endl;
  std::cout << " Press [s] to show Volume Rendering" << std::endl;
  std::cout << " Press [0..5] to switch Rendering Presets:" << std::endl;
  std::cout << "   0: MIP (Grayscale)" << std::endl;
  std::cout << "   1: Composite (Grayscale)" << std::endl;
  std::cout << "   2: Composite (Shaded)" << std::endl;
  std::cout << "   3: CT Bone/Tissue (Bone colors)" << std::endl;
  std::cout << "   4: MIP (Rainbow)" << std::endl;
  std::cout << "   5: Additive (Total path sum)" << std::endl;
  std::cout << " Press [q] to quit" << std::endl;
  std::cout << "=========================================" << std::endl;

  viewer.ZoomAuto();
  viewer.Start();

  return 0;
}