#include "MuDecayChannel.hh"
#include "musrMuonium.hh"

#include "G4AntiNeutrinoMu.hh"
#include "G4DecayProducts.hh"
#include "G4NeutrinoE.hh"
#include "G4Positron.hh"
#include "Randomize.hh"

#include <array>
#include <cmath>
#include <iostream>
#include <memory>
#include <string>

namespace {

constexpr int numberOfDecays = 50000;
constexpr double expectedMeanScaledEnergy = 0.7;
constexpr double meanEnergyTolerance = 0.008;
constexpr double isotropyTolerance = 0.015;
constexpr double energyConservationTolerance = 1.e-9;

}  // namespace

int main()
{
  G4Positron::PositronDefinition();
  G4NeutrinoE::NeutrinoEDefinition();
  G4AntiNeutrinoMu::AntiNeutrinoMuDefinition();
  const auto* muonium = musrMuonium::Definition();

  MuDecayChannel decay("Mu",1.);
  G4Random::setTheSeed(314159265);

  const double endpoint = muonium->GetPDGMass()/2.;
  const std::array<std::string,3> expectedDaughters = {
    "e+", "nu_e", "anti_nu_mu"
  };

  double scaledEnergySum = 0.;
  G4ThreeVector directionSum;

  for (int event=0; event<numberOfDecays; ++event) {
    std::unique_ptr<G4DecayProducts> products(decay.DecayIt(muonium->GetPDGMass()));
    if (!products || products->entries()!=3) {
      std::cerr<<"decay "<<event<<": expected three daughters\n";
      return 1;
    }

    for (int daughter=0; daughter<products->entries(); ++daughter) {
      const auto* particle = (*products)[daughter];
      if (!particle) {
        std::cerr<<"decay "<<event<<": null daughter "<<daughter<<'\n';
        return 1;
      }

      const std::string actualName =
        particle->GetParticleDefinition()->GetParticleName();
      if (actualName!=expectedDaughters[daughter]) {
        std::cerr<<"decay "<<event<<", daughter "<<daughter
                 <<": expected "<<expectedDaughters[daughter]
                 <<", got "<<actualName<<'\n';
        return 1;
      }

      if (!std::isfinite(particle->GetKineticEnergy()) ||
          particle->GetKineticEnergy()<0.) {
        std::cerr<<"decay "<<event<<", daughter "<<daughter
                 <<": invalid kinetic energy "
                 <<particle->GetKineticEnergy()<<'\n';
        return 1;
      }
    }

    double daughterEnergy = 0.;
    for (int daughter=0; daughter<products->entries(); ++daughter) {
      daughterEnergy += (*products)[daughter]->GetTotalEnergy();
    }
    if (std::abs(daughterEnergy-muonium->GetPDGMass())>
        energyConservationTolerance) {
      std::cerr<<"decay "<<event
               <<": energy is not conserved; daughter sum is "
               <<daughterEnergy<<", parent is "<<muonium->GetPDGMass()<<'\n';
      return 1;
    }

    const auto* positron = (*products)[0];
    const double scaledEnergy = positron->GetTotalEnergy()/endpoint;
    if (scaledEnergy<0. || scaledEnergy>1.) {
      std::cerr<<"decay "<<event
               <<": positron energy outside the Michel endpoint: "
               <<scaledEnergy<<'\n';
      return 1;
    }

    scaledEnergySum += scaledEnergy;
    directionSum += positron->GetMomentumDirection();
  }

  const double meanScaledEnergy = scaledEnergySum/numberOfDecays;
  if (std::abs(meanScaledEnergy-expectedMeanScaledEnergy)>
      meanEnergyTolerance) {
    std::cerr<<"Michel spectrum mean: expected "
             <<expectedMeanScaledEnergy<<" +/- "<<meanEnergyTolerance
             <<", got "<<meanScaledEnergy<<'\n';
    return 1;
  }

  const G4ThreeVector meanDirection = directionSum/numberOfDecays;
  if (std::abs(meanDirection.x())>isotropyTolerance ||
      std::abs(meanDirection.y())>isotropyTolerance ||
      std::abs(meanDirection.z())>isotropyTolerance) {
    std::cerr<<"decay directions are not isotropic: mean direction "
             <<meanDirection<<'\n';
    return 1;
  }

  return 0;
}