/* SPDX-FileCopyrightText: Copyright (c) 2020 - 2026, the adamantine authors.
 * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 */

#include <ElectronBeamHeatSource.hh>
#include <instantiation.hh>
#include <types.hh>

#include <deal.II/base/utilities.h>

namespace adamantine
{

template <int dim>
ElectronBeamHeatSource<dim>::ElectronBeamHeatSource(
    boost::property_tree::ptree const &beam_database,
    boost::optional<boost::property_tree::ptree const &> const
        &units_optional_database)
    : HeatSource<dim>(beam_database, units_optional_database),
      _five_axis(this->_scan_path.is_five_axis())
{
}

template <int dim>
void ElectronBeamHeatSource<dim>::update_time(double time)
{
  double const segment_power_modifier =
      this->_scan_path.get_power_modifier(time);
  this->_source_on = (segment_power_modifier > 0.0);
  dealii::Point<3> const &path = this->_scan_path.value(time);
  _quaternion = this->_scan_path.get_quaternion(time);
  // Copy the scalar value of path to the vectorized data type
  for (unsigned int d = 0; d < 3; ++d)
  {
    _beam_center(d) = path(d);
  }
  _alpha =
      -this->_beam.absorption_efficiency * this->_beam.max_power *
      segment_power_modifier * _log_01 /
      (dealii::numbers::PI * this->_beam.radius_squared * this->_beam.depth);
  _depth = this->_beam.depth;
  _inv_depth = 1.0 / _depth;
  _radius_squared = this->_beam.radius_squared;
  _inv_radius_squared = 1.0 / _radius_squared;
}

template <int dim>
double ElectronBeamHeatSource<dim>::value(dealii::Point<dim> const &point) const
{
  dealii::Point<dim> rotated_point;
  if constexpr (dim == 2)
  {
    rotated_point = point;
  }
  else
  {
    rotated_point = _five_axis ? _quaternion.rotate(point) : point;
  }

  double const z = rotated_point[axis<dim>::z] - _beam_center[axis<dim>::z][0];
  if ((z + _depth[0]) < 0.)
  {
    return 0.;
  }
  else
  {
    double const z_rel = z * _inv_depth[0];
    double const distribution_z =
        -3. * dealii::Utilities::fixed_power<2>(z_rel) - 2. * z_rel + 1.;

    double xpy_squared = dealii::Utilities::fixed_power<2>(
        rotated_point[axis<dim>::x] - _beam_center[axis<dim>::x][0]);
    if constexpr (dim == 3)
    {
      xpy_squared += dealii::Utilities::fixed_power<2>(
          rotated_point[axis<dim>::y] - _beam_center[axis<dim>::y][0]);
    }

    // Evaluating the exponential is very expensive. Return early if we know
    // that the heat source will be small.
    if (xpy_squared > 5. * _radius_squared[0])
    {
      return 0.;
    }

    // Electron beam heat source equation
    double heat_source =
        _alpha[0] * std::exp(_log_01 * xpy_squared * _inv_radius_squared[0]) *
        distribution_z;

    return heat_source;
  }
}

template <int dim>
dealii::VectorizedArray<double> ElectronBeamHeatSource<dim>::value(
    dealii::Point<dim, dealii::VectorizedArray<double>> const &points) const
{
  dealii::Point<dim, dealii::VectorizedArray<double>> rotated_points;
  if constexpr (dim == 2)
  {
    rotated_points = points;
  }
  else
  {
    rotated_points = _five_axis ? _quaternion.rotate(points) : points;
  }

  auto const z = rotated_points[axis<dim>::z] - _beam_center[axis<dim>::z];
  auto const z_depth = z + _depth;
  dealii::VectorizedArray<double> depth_mask;
  for (unsigned int i = 0; i < depth_mask.size(); ++i)
  {
    depth_mask[i] = z_depth[i] < 0. ? 0. : 1.;
  }
  // If all the lanes of depth_mask are zero, we can return early
  if (depth_mask.sum() == 0)
  {
    return depth_mask;
  }

  auto xpy_squared = dealii::Utilities::fixed_power<2>(
      rotated_points[axis<dim>::x] - _beam_center[axis<dim>::x]);
  if constexpr (dim == 3)
  {
    auto const y_squared = dealii::Utilities::fixed_power<2>(
        rotated_points[axis<dim>::y] - _beam_center[axis<dim>::y]);
    xpy_squared += y_squared;
  }

  // Evaluating the exponential is very expensive. Return early if we know
  // that the heat source will be small.
  auto const xpy_area = 5. * _radius_squared - xpy_squared;
  dealii::VectorizedArray<double> xpy_mask;
  for (unsigned int i = 0; i < xpy_mask.size(); ++i)
  {
    xpy_mask[i] = xpy_area[i] < 0. ? 0. : 1.;
  }
  // If all the lanes of xpy_mask are zero, we can return early
  if (xpy_mask.sum() == 0)
  {
    return xpy_mask;
  }

  dealii::VectorizedArray<double> minus_three = -3.;
  dealii::VectorizedArray<double> one = 1.;
  dealii::VectorizedArray<double> z_rel = z * _inv_depth;
  dealii::VectorizedArray<double> distribution_z =
      dealii::Utilities::fixed_power<2>(z_rel);
  distribution_z *= minus_three;
  distribution_z -= z_rel;
  distribution_z -= z_rel;
  distribution_z += one;

  dealii::VectorizedArray<double> exponent = _log_01 * xpy_squared;
  exponent *= _inv_radius_squared;
  return depth_mask * _alpha * std::exp(exponent) * distribution_z;
}

template <int dim>
dealii::BoundingBox<dim>
ElectronBeamHeatSource<dim>::get_bounding_box(double const time,
                                              double const scaling_factor) const
{
  dealii::Point<3> const &beam_center = this->_scan_path.value(time, false);
  if constexpr (dim == 2)
  {
    return {{{beam_center[axis<dim>::x] - scaling_factor * this->_beam.radius,
              beam_center[axis<dim>::z] - scaling_factor * this->_beam.depth},
             {beam_center[axis<dim>::x] + scaling_factor * this->_beam.radius,
              beam_center[axis<dim>::z]}}};
  }
  else
  {
    dealii::Point<3> max_corner(
        beam_center[axis<dim>::x] + scaling_factor * this->_beam.radius,
        beam_center[axis<dim>::y] + scaling_factor * this->_beam.radius,
        beam_center[axis<dim>::z]);
    dealii::Point<3> min_corner(
        beam_center[axis<dim>::x] - scaling_factor * this->_beam.radius,
        beam_center[axis<dim>::y] - scaling_factor * this->_beam.radius,
        beam_center[axis<dim>::z] - scaling_factor * this->_beam.depth);
    if (this->_scan_path.is_five_axis())
    {
      // We need to rotate the box to match the scan path
      Quaternion quaternion = this->_scan_path.get_quaternion(time);
      dealii::Point<3> rotated_max_corner = quaternion.rotate(max_corner);
      dealii::Point<3> rotated_min_corner = quaternion.rotate(min_corner);

      // We need to recompute the min and max corners after rotation.
      dealii::Point<3> new_max_corner;
      dealii::Point<3> new_min_corner;
      for (int d = 0; d < 3; ++d)
      {
        new_max_corner[d] =
            std::max(rotated_max_corner[d], rotated_min_corner[d]);
        new_min_corner[d] =
            std::min(rotated_max_corner[d], rotated_min_corner[d]);
      }

      return {{new_min_corner, new_max_corner}};
    }
    else
    {
      return {{min_corner, max_corner}};
    }
  }
}
} // namespace adamantine

INSTANTIATE_DIM(ElectronBeamHeatSource)