// SPDX-License-Identifier: LGPL-3.0-or-later
// Copyright (C) 2022, 2023 Wenqing Fan, Barak Schmookler, Whitney Armstrong, Sylvester Joosten, Dmitry Romanov, Christopher Dilks, Wouter Deconinck

#include <Acts/Definitions/Algebra.hpp>
#include <Acts/Definitions/Direction.hpp>
#include <Acts/Definitions/TrackParametrization.hpp>
#include <Acts/Definitions/Units.hpp>
#if Acts_VERSION_MAJOR >= 46
#include <Acts/EventData/BoundTrackParameters.hpp>
#else
#include <Acts/EventData/GenericBoundTrackParameters.hpp>
#endif
#include <Acts/EventData/MultiTrajectoryHelpers.hpp>
#include <Acts/EventData/VectorMultiTrajectory.hpp>
#include <Acts/Geometry/GeometryIdentifier.hpp>
#include <Acts/Geometry/TrackingGeometry.hpp>
#include <Acts/MagneticField/MagneticFieldProvider.hpp>
#include <Acts/Material/MaterialInteraction.hpp>
#include <Acts/Propagator/ActorList.hpp>
#include <Acts/Propagator/EigenStepper.hpp>
#include <Acts/Propagator/MaterialInteractor.hpp>
#include <Acts/Propagator/Navigator.hpp>
#include <Acts/Propagator/Propagator.hpp>
#include <Acts/Propagator/PropagatorResult.hpp>
#include <Acts/Surfaces/CylinderBounds.hpp>
#include <Acts/Surfaces/CylinderSurface.hpp>
#include <Acts/Surfaces/DiscSurface.hpp>
#include <Acts/Surfaces/RadialBounds.hpp>
#include <Acts/Utilities/Logger.hpp>
#include <DD4hep/Handle.h>
#include <Evaluator/DD4hepUnits.h>
#include <edm4eic/Cov2f.h>
#include <edm4eic/Cov3f.h>
#include <edm4hep/Vector3f.h>
#include <edm4hep/utils/vector_utils.h>
#include <spdlog/common.h>
#include <Eigen/Core>
#include <Eigen/Geometry>
#include <algorithm>
#include <any>
#include <cmath>
#include <cstdint>
#include <functional>
#include <iterator>
#include <map>
#include <optional>
#include <stdexcept>
#include <string>
#include <tuple>
#include <typeinfo>
#include <utility>
#include <variant>

#include "algorithms/tracking/ActsGeometryProvider.h"
#include "algorithms/tracking/TrackPropagation.h"
#include "algorithms/tracking/TrackPropagationConfig.h"
#include "extensions/spdlog/SpdlogToActs.h"

namespace eicrecon {

template <typename... L> struct multilambda : L... {
  using L::operator()...;
  constexpr multilambda(L... lambda) : L(std::move(lambda))... {}
};

void TrackPropagation::init() {
  const auto* detector = m_detector;

  std::map<uint32_t, std::size_t> system_id_layers;

  multilambda _toDouble = {
      [](const std::string& v) { return dd4hep::_toDouble(v); },
      [](const double& v) { return v; },
  };

  auto _toActsSurface =
      [&_toDouble, &detector, &system_id_layers](
          const std::variant<CylinderSurfaceConfig, DiscSurfaceConfig> surface_variant)
      -> std::shared_ptr<Acts::Surface> {
    if (std::holds_alternative<CylinderSurfaceConfig>(surface_variant)) {
      CylinderSurfaceConfig surface = std::get<CylinderSurfaceConfig>(surface_variant);
      const double rmin =
          std::visit(_toDouble, surface.rmin) / dd4hep::mm * Acts::UnitConstants::mm;
      const double zmin =
          std::visit(_toDouble, surface.zmin) / dd4hep::mm * Acts::UnitConstants::mm;
      const double zmax =
          std::visit(_toDouble, surface.zmax) / dd4hep::mm * Acts::UnitConstants::mm;
      const uint32_t system_id = detector->constant<uint32_t>(surface.id);
      auto bounds              = std::make_shared<Acts::CylinderBounds>(rmin, (zmax - zmin) / 2);
      auto t                   = Acts::Translation3(Acts::Vector3(0, 0, (zmax + zmin) / 2));
      auto tf                  = Acts::Transform3(t);
      auto acts_surface        = Acts::Surface::makeShared<Acts::CylinderSurface>(tf, bounds);
      acts_surface->assignGeometryId(
          Acts::GeometryIdentifier().withExtra(system_id).withLayer(++system_id_layers[system_id]));
      return acts_surface;
    }
    if (std::holds_alternative<DiscSurfaceConfig>(surface_variant)) {
      DiscSurfaceConfig surface = std::get<DiscSurfaceConfig>(surface_variant);
      const double zmin =
          std::visit(_toDouble, surface.zmin) / dd4hep::mm * Acts::UnitConstants::mm;
      const double rmin =
          std::visit(_toDouble, surface.rmin) / dd4hep::mm * Acts::UnitConstants::mm;
      const double rmax =
          std::visit(_toDouble, surface.rmax) / dd4hep::mm * Acts::UnitConstants::mm;
      const uint32_t system_id = detector->constant<uint32_t>(surface.id);
      auto bounds              = std::make_shared<Acts::RadialBounds>(rmin, rmax);
      auto t                   = Acts::Translation3(Acts::Vector3(0, 0, zmin));
      auto tf                  = Acts::Transform3(t);
      auto acts_surface        = Acts::Surface::makeShared<Acts::DiscSurface>(tf, bounds);
      acts_surface->assignGeometryId(
          Acts::GeometryIdentifier().withExtra(system_id).withLayer(++system_id_layers[system_id]));
      return acts_surface;
    }
    throw std::domain_error("Unknown surface type");
  };
  m_target_surfaces.resize(m_cfg.target_surfaces.size());
  std::ranges::transform(m_cfg.target_surfaces, m_target_surfaces.begin(), _toActsSurface);
  m_filter_surfaces.resize(m_cfg.filter_surfaces.size());
  std::ranges::transform(m_cfg.filter_surfaces, m_filter_surfaces.begin(), _toActsSurface);

  trace("Initialized");
}

void TrackPropagation::propagateToSurfaceList(const Input& input, const Output& output) const {
  const auto [tracks, track_states, tracks_acts] = input;
  auto [track_segments]                          = output;

  // logging
  trace("Propagate tracks: --------------------");
  trace("number of tracks: {}", tracks->size());

  // Construct ConstTrackContainer from underlying containers
  auto trackStateContainer = std::make_shared<Acts::ConstVectorMultiTrajectory>(*track_states);
  auto trackContainer      = std::make_shared<Acts::ConstVectorTrackContainer>(*tracks_acts);
  ActsExamples::ConstTrackContainer constTracks(trackContainer, trackStateContainer);

  // loop over input tracks
  std::size_t i = 0;
  for (const auto& track : constTracks) {

    // check if this track can be propagated to any filter surface
    bool track_reaches_filter_surface{false};
    for (const auto& filter_surface : m_filter_surfaces) {
      auto point = propagate(edm4eic::Track{}, track, constTracks, filter_surface);
      if (point) {
        track_reaches_filter_surface = true;
        break;
      }
    }
    if (!track_reaches_filter_surface) {
      ++i;
      continue;
    }

    // start a mutable TrackSegment
    auto track_segment = track_segments->create();

    // corresponding track
    if (tracks->size() == constTracks.size()) {
      trace("track segment connected to track {}", i);
      track_segment.setTrack((*tracks)[i]);
      ++i;
    }

    // zero measurements of segment length
    decltype(edm4eic::TrackSegmentData::length) length            = 0;
    decltype(edm4eic::TrackSegmentData::lengthError) length_error = 0;

    // loop over projection-target surfaces
    for (const auto& target_surface : m_target_surfaces) {

      // project the track to this surface
      auto point = propagate(edm4eic::Track{}, track, constTracks, target_surface);
      if (!point) {
        trace("<> Failed to propagate track to this plane");
        continue;
      }

      // logging
      trace("<> track: x=( {:>10.2f} {:>10.2f} {:>10.2f} )", point->position.x, point->position.y,
            point->position.z);
      trace("               p=( {:>10.2f} {:>10.2f} {:>10.2f} )", point->momentum.x,
            point->momentum.y, point->momentum.z);

      // track point cut
      if (!m_cfg.track_point_cut(*point)) {
        trace("                 => REJECTED by trackPointCut");
        if (m_cfg.skip_track_on_track_point_cut_failure) {
          break;
        }
        continue;
      }

      // update the `TrackSegment` length
      // FIXME: `length` and `length_error` are currently not used by any callers, and may not be correctly calculated here
      if (track_segment.points_size() > 0) {
        auto pos0 = point->position;
        auto pos1 = std::prev(track_segment.points_end())->position;
        auto dist = edm4hep::utils::magnitude(pos0 - pos1);
        length += dist;
        trace("               dist to previous point: {}", dist);
      }

      // add the `TrackPoint` to the `TrackSegment`
      track_segment.addToPoints(*point);

    } // end `targetSurfaces` loop

    // set final length and length error
    track_segment.setLength(length);
    track_segment.setLengthError(length_error);

  } // end loop over input tracks
}

std::unique_ptr<edm4eic::TrackPoint>
TrackPropagation::propagate(const edm4eic::Track& /* track */,
                            const ActsExamples::ConstTrackProxy& acts_track,
                            const ActsExamples::ConstTrackContainer& trackContainer,
                            const std::shared_ptr<const Acts::Surface>& targetSurf) const {

  auto tipIndex = acts_track.tipIndex();

  trace("  Propagating track with tip index {}", tipIndex);

  // Collect the trajectory summary info
  auto trajState =
      Acts::MultiTrajectoryHelpers::trajectoryState(trackContainer.trackStateContainer(), tipIndex);
  int m_nMeasurements = trajState.nMeasurements;
  int m_nStates       = trajState.nStates;

  trace("  Num measurement in trajectory: {}", m_nMeasurements);
  trace("  Num states in trajectory     : {}", m_nStates);

  // Get track state at last measurement surface
  // For last measurement surface, filtered and smoothed results are equivalent
  auto trackState        = trackContainer.trackStateContainer().getTrackState(tipIndex);
  auto initSurface       = trackState.referenceSurface().getSharedPtr();
  const auto& initParams = trackState.filtered();
  const auto& initCov    = trackState.filteredCovariance();

  Acts::BoundTrackParameters initBoundParams(initSurface, initParams, initCov,
                                             acts_track.particleHypothesis());

  // Get pathlength of last track state with respect to perigee surface
  const auto initPathLength = trackState.pathLength();

  trace("    TrackPropagation. Propagating to surface # {}", typeid(targetSurf->type()).name());

  std::shared_ptr<const Acts::TrackingGeometry> trackingGeometry   = m_geoSvc->trackingGeometry();
  std::shared_ptr<const Acts::MagneticFieldProvider> magneticField = m_geoSvc->getFieldProvider();

  // Convert algorithm log level to Acts log level
  const auto spdlog_level = static_cast<spdlog::level::level_enum>(this->level());
  const auto acts_level   = eicrecon::SpdlogToActsLevel(spdlog_level);
  ACTS_LOCAL_LOGGER(Acts::getDefaultLogger("PROP", acts_level));

  using Propagator        = Acts::Propagator<Acts::EigenStepper<>, Acts::Navigator>;
  using PropagatorOptions = Propagator::template Options<Acts::ActorList<Acts::MaterialInteractor>>;
  Propagator propagator(Acts::EigenStepper<>(magneticField),
                        Acts::Navigator({.trackingGeometry = m_geoSvc->trackingGeometry()},
                                        logger().cloneWithSuffix("Navigator")),
                        logger().cloneWithSuffix("Propagator"));

  // Get run-scoped contexts from service
  const auto& gctx = m_geoSvc->getActsGeometryContext();
  const auto& mctx = m_geoSvc->getActsMagneticFieldContext();

  PropagatorOptions propagationOptions(gctx, mctx);

  // Some target surfaces (e.g. DIRC) may be inside the last measurement surface (e.g. BIC),
  // so we use a straight line intersection from the last measurement surface to the target
  // surface to determine if we have to propagate backwards.
  auto initPosition  = initBoundParams.position(gctx);
  auto initDirection = initBoundParams.direction();
  auto intersections = targetSurf->intersect(gctx, initPosition, initDirection);

  // Determine closest forward intersection (positive pathlength from perigee)
  auto intersection = intersections.closestForward();
  auto difference   = intersection.position() - initPosition;
  auto dot          = difference.dot(initBoundParams.direction());

  // Propagate forwards by default
  propagationOptions.direction = Acts::Direction::Forward();

  // but invert if the position difference is opposite to direction
  if (intersection.isValid() && dot < 0) {

    // The extra fields of the surface geometry ID contain the DD4hep system
    auto initSurfaceExtra   = initSurface->geometryId().extra();
    auto targetSurfaceExtra = targetSurf->geometryId().extra();
    debug("    inverting direction for propagator from surface {} to {}", initSurfaceExtra,
          targetSurfaceExtra);
    auto p1 = initBoundParams.position(gctx);
    debug("      initial position {} {} {}", p1.x(), p1.y(), p1.z());
    auto p2 = intersection.position();
    debug("      straight line intersection at {} {} {}", p2.x(), p2.y(), p2.z());

    // Propagate backwards
    propagationOptions.direction = Acts::Direction::Backward();
  }

  auto result = propagator.propagate(initBoundParams, *targetSurf, propagationOptions);

  // check propagation result
  if (!result.ok()) {
    trace("    propagation failed (!result.ok())");
    return nullptr;
  }
  trace("    propagation result is OK");

  // Pulling results to convenient variables
  auto trackStateParams  = *((*result).endParameters);
  const auto& parameter  = trackStateParams.parameters();
  const auto& covariance = *trackStateParams.covariance();

  // Path length
  const float pathLength      = initPathLength + (*result).pathLength;
  const float pathLengthError = 0;
  trace("    path len = {}", pathLength);

  // Position:
  auto projectionPos = trackStateParams.position(gctx);
  const decltype(edm4eic::TrackPoint::position) position{static_cast<float>(projectionPos(0)),
                                                         static_cast<float>(projectionPos(1)),
                                                         static_cast<float>(projectionPos(2))};
  const decltype(edm4eic::TrackPoint::positionError) positionError{0, 0, 0};
  trace("    pos x = {}", position.x);
  trace("    pos y = {}", position.y);
  trace("    pos z = {}", position.z);

  // Momentum
  const decltype(edm4eic::TrackPoint::momentum) momentum = edm4hep::utils::sphericalToVector(
      static_cast<float>(1.0 / std::abs(parameter[Acts::eBoundQOverP])),
      static_cast<float>(parameter[Acts::eBoundTheta]),
      static_cast<float>(parameter[Acts::eBoundPhi]));
  const decltype(edm4eic::TrackPoint::momentumError) momentumError{
      static_cast<float>(covariance(Acts::eBoundTheta, Acts::eBoundTheta)),
      static_cast<float>(covariance(Acts::eBoundPhi, Acts::eBoundPhi)),
      static_cast<float>(covariance(Acts::eBoundQOverP, Acts::eBoundQOverP)),
      static_cast<float>(covariance(Acts::eBoundTheta, Acts::eBoundPhi)),
      static_cast<float>(covariance(Acts::eBoundTheta, Acts::eBoundQOverP)),
      static_cast<float>(covariance(Acts::eBoundPhi, Acts::eBoundQOverP))};

  // time
  const float time{static_cast<float>(parameter(Acts::eBoundTime))};
  const float timeError{static_cast<float>(sqrt(covariance(Acts::eBoundTime, Acts::eBoundTime)))};

  // Direction
  const float theta(parameter[Acts::eBoundTheta]);
  const float phi(parameter[Acts::eBoundPhi]);
  const decltype(edm4eic::TrackPoint::directionError) directionError{
      static_cast<float>(covariance(Acts::eBoundTheta, Acts::eBoundTheta)),
      static_cast<float>(covariance(Acts::eBoundPhi, Acts::eBoundPhi)),
      static_cast<float>(covariance(Acts::eBoundTheta, Acts::eBoundPhi))};

  // >oO debug print
  trace("    loc 0   = {:.4f}", parameter[Acts::eBoundLoc0]);
  trace("    loc 1   = {:.4f}", parameter[Acts::eBoundLoc1]);
  trace("    phi     = {:.4f}", parameter[Acts::eBoundPhi]);
  trace("    theta   = {:.4f}", parameter[Acts::eBoundTheta]);
  trace("    q/p     = {:.4f}", parameter[Acts::eBoundQOverP]);
  trace("    p       = {:.4f}", 1.0 / parameter[Acts::eBoundQOverP]);
  trace("    err phi = {:.4f}", sqrt(covariance(Acts::eBoundPhi, Acts::eBoundPhi)));
  trace("    err th  = {:.4f}", sqrt(covariance(Acts::eBoundTheta, Acts::eBoundTheta)));
  trace("    err q/p = {:.4f}", sqrt(covariance(Acts::eBoundQOverP, Acts::eBoundQOverP)));
  trace("    chi2    = {:.4f}", trajState.chi2Sum);
  trace("    loc err = {:.4f}", static_cast<float>(covariance(Acts::eBoundLoc0, Acts::eBoundLoc0)));
  trace("    loc err = {:.4f}", static_cast<float>(covariance(Acts::eBoundLoc1, Acts::eBoundLoc1)));
  trace("    loc err = {:.4f}", static_cast<float>(covariance(Acts::eBoundLoc0, Acts::eBoundLoc1)));

  uint64_t surface = targetSurf->geometryId().value();
  uint32_t system  = 0; // default value...will be set in TrackPropagation factory

  return std::make_unique<edm4eic::TrackPoint>(
      edm4eic::TrackPoint{.surface         = surface,
                          .system          = system,
                          .position        = position,
                          .positionError   = positionError,
                          .momentum        = momentum,
                          .momentumError   = momentumError,
                          .time            = time,
                          .timeError       = timeError,
                          .theta           = theta,
                          .phi             = phi,
                          .directionError  = directionError,
                          .pathlength      = pathLength,
                          .pathlengthError = pathLengthError});
}

} // namespace eicrecon
