/*
 * (c) Copyright 2018 CORSIKA Project, corsika-project@lists.kit.edu
 *
 * See file AUTHORS for a list of contributors.
 *
 * This software is distributed under the terms of the GNU General Public
 * Licence version 3 (GPL Version 3). See file LICENSE for a full version of
 * the license.
 */

#include <corsika/process/very_simple_model/VerySimpleModel.h>
#include <corsika/setup/SetupStack.h>
#include <corsika/setup/SetupTrajectory.h>

#include <tuple>

using namespace corsika;
using namespace corsika::units::si;
using namespace corsika::process::very_simple_model;

using SetupParticle = corsika::setup::Stack::StackIterator;
using SetupProjectile = corsika::setup::StackView::StackIterator;

  void VerySimpleModel::Init() {}

  VerySimpleModel::VerySimpleModel() {}

  template <>
  process::EProcessReturn VerySimpleModel::DoInteraction(SetupProjectile& vP) {
    const auto projectileID = vP.GetPID();

    auto const posOrig = vP.GetPosition();
    auto const tOrig = vP.GetTime();
    auto const eProjectileLab = vP.GetEnergy();
    auto const pProjectileLab = vP.GetMomentum();

    const auto* currentNode = vP.GetNode();
    const auto& mediumComposition =
        currentNode->GetModelProperties().GetNuclearComposition();

    // get cross sections for target materials
    auto const& compVec = mediumComposition.GetComponents();
    std::vector<units::si::CrossSectionType> cross_section_of_components(compVec.size());

    for (size_t i = 0; i < compVec.size(); ++i) {
      cross_section_of_components[i] = GetCrossSection(projectileID, compVec[i]);
    }

    const auto corsikaTargetId =
        mediumComposition.SampleTarget(cross_section_of_components, fRNG);

    // add secondaries to corsika stack
    auto const sec1Code = particles::Code::MuPlus;
    auto const sec1E = eProjectileLab;
    auto const sec1P = pProjectileLab;

    auto secondary1 = vP.AddSecondary(
        std::tuple<particles::Code, units::si::HEPEnergyType, stack::MomentumVector,
              geometry::Point, units::si::TimeType>(sec1Code, sec1E, sec1P, posOrig,
                                                    tOrig));

    auto const sec2Code = particles::Code::MuMinus;
    auto const sec2E = eProjectileLab;
    auto const sec2P = pProjectileLab;
    
    auto secondary2 = vP.AddSecondary(
        std::tuple<particles::Code, units::si::HEPEnergyType, stack::MomentumVector,
              geometry::Point, units::si::TimeType>(sec2Code, sec2E, sec2E, posOrig,
                                                    tOrig));

    return process::EProcessReturn::eOk;
  }

  template <>
  corsika::units::si::CrossSectionType GetCrossSection(SetupParticle const& vProjectile,
                                                       particles::Code vTargetCode) {
    if (vProjectile.GetPID() == particles::MuPlus) {
      return 20_mb;
    } else if (vProjectile.GetPID() == particles::MuMinus) {
      return 10_mb;
    } else {
      return 0_mb;
    }
  }

  template <>
  corsika::units::si::GrammageType GetInteractionLength(Projectile const& vP) {
    auto const* currentNode = vP.GetNode();
    const auto& mediumComposition =
        currentNode->GetModelProperties().GetNuclearComposition();

    si::CrossSectionType weightedProdCrossSection = mediumComposition.WeightedSum(
        [=](particles::Code targetID) -> si::CrossSectionType {
          return this->GetCrossSection(corsikaBeamId, targetID);
        });

    GrammageType const int_length = mediumComposition.GetAverageMassNumber() *
                                    units::constants::u / weightedProdCrossSection;

    return int_length;
  }