//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Tests/SimFactory/MakeSimulations.cpp
//! @brief     Implements simulation builders for standard tests.
//!
//! @homepage  http://www.bornagainproject.org
//! @license   GNU General Public License v3 or higher (see COPYING)
//! @copyright Forschungszentrum Jülich GmbH 2018
//! @authors   Scientific Computing Group at MLZ (see CITATION, AUTHORS)
//
//  ************************************************************************************************

#include "Tests/SimFactory/MakeSimulations.h"
#include "Base/Axis/MakeScale.h"
#include "Base/Axis/Scale.h"
#include "Base/Const/Units.h"
#include "Device/Beam/Beam.h"
#include "Device/Beam/FootprintGauss.h"
#include "Device/Beam/FootprintSquare.h"
#include "Device/Detector/OffspecDetector.h"
#include "Device/Detector/SphericalDetector.h"
#include "Device/Mask/Ellipse.h"
#include "Device/Mask/Line.h"
#include "Device/Mask/Polygon.h"
#include "Device/Mask/Rectangle.h"
#include "Device/Resolution/ResolutionFunction2DGaussian.h"
#include "Param/Distrib/Distributions.h"
#include "Resample/Option/SimulationOptions.h"
#include "Sim/Background/ConstantBackground.h"
#include "Sim/Scan/AlphaScan.h"
#include "Sim/Scan/QzScan.h"
#include "Sim/Simulation/includeSimulations.h"
#include <algorithm>
#include <map>
#include <numbers>

using std::numbers::pi;
using Units::angstrom;
using Units::deg;

namespace {

const R3 zplus(0.0, 0.0, 1.0);
const R3 yplus(0.0, 1.0, 0.0);

const std::map<std::string, std::pair<R3, R3>> ZPolarizationCases = {
    {"PP", {+zplus, +zplus}},
    {"PM", {+zplus, -zplus}},
    {"MP", {-zplus, +zplus}},
    {"MM", {-zplus, -zplus}},
};

const std::map<std::string, std::pair<R3, R3>> YPolarizationCases = {
    {"PP", {+yplus, +yplus}},
    {"PM", {+yplus, -yplus}},
    {"MP", {-yplus, +yplus}},
    {"MM", {-yplus, -yplus}},
};

const Beam stdBeam(1, 1 * angstrom, 0.2 * deg);
const Beam powerBeam(1e7, 1 * angstrom, 0.2 * deg);

} // namespace

//  ************************************************************************************************
//  GISAS simulations
//  ************************************************************************************************

//! GISAS simulation with small detector and phi[-2,2], theta[0,2].

std::unique_ptr<ScatteringSimulation> test::makeSimulation::MiniGISAS(const Sample& sample)
{
    SphericalDetector detector(25, -2 * deg, 2 * deg, 25, 0, 2 * deg);
    return std::make_unique<ScatteringSimulation>(stdBeam, sample, detector);
}


//! GISAS simulation with small detector and phi[-1,1], theta[0,1].

std::unique_ptr<ScatteringSimulation> test::makeSimulation::MiniGISAS_v2(const Sample& sample)
{
    SphericalDetector detector(25, -2 * deg, 2 * deg, 25, 0, 1 * deg);
    return std::make_unique<ScatteringSimulation>(stdBeam, sample, detector);
}

std::unique_ptr<ScatteringSimulation>
test::makeSimulation::MiniZPolarizedGISAS(const Sample& sample, const std::string& polCase)
{
    const auto zCase = ZPolarizationCases.at(polCase);
    std::unique_ptr<Beam> beam(stdBeam.clone());
    beam->setPolarization(zCase.first);
    SphericalDetector detector(25, -2 * deg, 2 * deg, 25, 0, 2 * deg);
    detector.setAnalyzer(zCase.second);
    auto sim = std::make_unique<ScatteringSimulation>(*beam, sample, detector);
    return sim;
}


//! Basic GISAS simulation with the detector phi[0,2], theta[0,2].

std::unique_ptr<ScatteringSimulation> test::makeSimulation::BasicGISAS(const Sample& sample)
{
    SphericalDetector detector(100, 0 * deg, 2 * deg, 100, 0 * deg, 2 * deg);
    return std::make_unique<ScatteringSimulation>(stdBeam, sample, detector);
}

//! Basic GISAS simulation for spin flip channel.

std::unique_ptr<ScatteringSimulation> test::makeSimulation::SpinflipGISAS(const Sample& sample)
{
    std::unique_ptr<Beam> beam(stdBeam.clone());
    beam->setPolarization(zplus);
    SphericalDetector detector(100, 0 * deg, 2 * deg, 100, 0 * deg, 2 * deg);
    detector.setAnalyzer(-1. * zplus);
    return std::make_unique<ScatteringSimulation>(*beam, sample, detector);
}

//! GISAS simulation with beam divergence applied.

std::unique_ptr<ScatteringSimulation>
test::makeSimulation::MiniGISASBeamDivergence(const Sample& sample)
{
    std::unique_ptr<ScatteringSimulation> result = MiniGISAS(sample);

    DistributionLogNormal wavelength_distr(1 * angstrom, 0.1, 5);
    DistributionGaussian alpha_distr(0.2 * deg, 0.02 * deg, 4);
    DistributionGate phi_distr(-0.1 * deg, 0.02 * deg, 3);

    result->addParameterDistribution(ParameterDistribution::BeamWavelength, wavelength_distr);
    result->addParameterDistribution(ParameterDistribution::BeamGrazingAngle, alpha_distr);
    result->addParameterDistribution(ParameterDistribution::BeamAzimuthalAngle, phi_distr);

    return result;
}

//! GISAS simulation with multiple masks on the detector plane.

std::unique_ptr<ScatteringSimulation> test::makeSimulation::GISASWithMasks(const Sample& sample)
{
    SphericalDetector detector(50, -1 * deg, 1 * deg, 50, 0, 2 * deg);

    detector.maskAll();
    // pacman
    detector.addMask(Ellipse(0.0 * deg, 1.0 * deg, 0.5 * deg, 0.5 * deg), false);
    detector.addMask(Ellipse(0.11 * deg, 1.25 * deg, 0.05 * deg, 0.05 * deg), true);

    std::vector<std::pair<double, double>> points = {{0.0 * deg, 1.0 * deg},
                                                     {0.5 * deg, 1.2 * deg},
                                                     {0.5 * deg, 0.8 * deg},
                                                     {0.0 * deg, 1.0 * deg}};
    detector.addMask(Polygon(points), true);

    detector.addMask(Rectangle(0.45 * deg, 0.95 * deg, 0.55 * deg, 1.05 * deg), false);
    detector.addMask(Rectangle(0.61 * deg, 0.95 * deg, 0.71 * deg, 1.05 * deg), false);
    detector.addMask(Rectangle(0.75 * deg, 0.95 * deg, 0.85 * deg, 1.05 * deg), false);

    // more masks
    detector.addMask(Ellipse(-0.5 * deg, 1.5 * deg, 0.3 * deg, 0.1 * deg, 45. * deg), false);
    detector.addMask(VerticalLine(-0.6 * deg), true);
    detector.addMask(HorizontalLine(0.3 * deg), false);

    return std::make_unique<ScatteringSimulation>(powerBeam, sample, detector);
}

//! GISAS simulation with detector resolution.

std::unique_ptr<ScatteringSimulation>
test::makeSimulation::MiniGISASDetectorResolution(const Sample& sample)
{
    std::unique_ptr<ScatteringSimulation> result = MiniGISAS(sample);
    ResolutionFunction2DGaussian resfunc(0.0025, 0.0025);
    result->detector().setResolutionFunction(resfunc);
    return result;
}

//! GISAS simulation with small detector and including specular peak.

std::unique_ptr<ScatteringSimulation>
test::makeSimulation::MiniGISASSpecularPeak(const Sample& sample)
{
    SphericalDetector detector(25, -2 * deg, 2 * deg, 25, 0, 2 * deg);
    std::unique_ptr<ScatteringSimulation> result =
        std::make_unique<ScatteringSimulation>(stdBeam, sample, detector);
    result->options().setIncludeSpecular(true);
    return result;
}

//! GISAS simulation with large detector to test performance.

std::unique_ptr<ScatteringSimulation> test::makeSimulation::MaxiGISAS(const Sample& sample)
{
    SphericalDetector detector(256, -2 * deg, 2 * deg, 256, 0, 2 * deg);
    return std::make_unique<ScatteringSimulation>(stdBeam, sample, detector);
}

//! Basic GISAS for polarization studies.

std::unique_ptr<ScatteringSimulation> test::makeSimulation::MaxiGISAS00(const Sample& sample)
{
    std::unique_ptr<Beam> beam(stdBeam.clone());
    beam->setPolarization(zplus);
    SphericalDetector detector(256, -2 * deg, 2 * deg, 256, 0, 2 * deg);
    detector.setAnalyzer(zplus);
    return std::make_unique<ScatteringSimulation>(*beam, sample, detector);
}

//! GISAS simulation with Monte-Carlo integration switched ON.

std::unique_ptr<ScatteringSimulation>
test::makeSimulation::MiniGISASMonteCarlo(const Sample& sample)
{
    std::unique_ptr<ScatteringSimulation> result = MiniGISAS(sample);
    result->options().setMonteCarloIntegration(true, 100);
    return result;
}

//! GISAS simulation with spherical detector, region of interest and mask.

std::unique_ptr<ScatteringSimulation>
test::makeSimulation::SphericalDetWithRoi(const Sample& sample)
{
    SphericalDetector detector(40, -2 * deg, 2 * deg, 30, -1 * deg, 2 * deg);
    detector.addMask(Rectangle(-0.5 * deg, 0.3 * deg, -0.2 * deg, 0.6 * deg));
    detector.setRegionOfInterest(-1.499 * deg, 0.251 * deg, 1.499 * deg, 1.749 * deg);

    return std::make_unique<ScatteringSimulation>(stdBeam, sample, detector);
}

//! GISAS simulation with rectangular detector, region of interest and mask.

std::unique_ptr<ScatteringSimulation> test::makeSimulation::RectDetWithRoi(const Sample& sample)
{
    SphericalDetector detector(40, -2 * deg, 2 * deg, 30, -1 * deg, 2 * deg);
    detector.addMask(Rectangle(0.35 * deg, -1.15 * deg, 0.75 * deg, 1.95 * deg));
    detector.setRegionOfInterest(-2.55 * deg, 0.25 * deg, 1.75 * deg, 1.85 * deg);
    return std::make_unique<ScatteringSimulation>(stdBeam, sample, detector);
}

//! GISAS simulation with an extra long wavelength

std::unique_ptr<ScatteringSimulation>
test::makeSimulation::ExtraLongWavelengthGISAS(const Sample& sample)
{
    std::unique_ptr<Beam> beam(Beam(1e8, 12 * Units::nm, 0.2 * deg).clone());
    SphericalDetector detector(100, -2 * deg, 2 * deg, 100, 0, 2 * deg);
    auto simulation = std::make_unique<ScatteringSimulation>(*beam, sample, detector);
    simulation->options().setIncludeSpecular(true);
    return simulation;
}

//! ISimulation with fitting.
//! Beam intensity set to provide reasonably large values in detector channels.
std::unique_ptr<ScatteringSimulation> test::makeSimulation::MiniGISASFit(const Sample& sample)
{
    SphericalDetector detector(25, -2 * deg, 2 * deg, 25, 0, 2 * deg);
    return std::make_unique<ScatteringSimulation>(powerBeam, sample, detector);
}

//  ************************************************************************************************
//  off-specular simulations
//  ************************************************************************************************

std::unique_ptr<OffspecSimulation> test::makeSimulation::MiniOffspec(const Sample& sample)
{
    OffspecDetector detector(9, -0.1 * deg, 0.1 * deg, 25, 0.1 * deg, 5.1 * deg);
    AlphaScan scan(20, 0.2 * deg, 4 * deg);
    scan.setWavelength(5 * angstrom);
    scan.setIntensity(1e9);
    auto result = std::make_unique<OffspecSimulation>(scan, sample, detector);
    result->options().setIncludeSpecular(true);
    return result;
}

//  ************************************************************************************************
//  specular simulations
//  ************************************************************************************************

BeamScan* test::makeSimulation::BasicSpecularScan(bool vsQ)
{
    const double wavelength = 1.54 * angstrom;
    const int n = 2000;
    const double max_angle = 5 * deg;

    if (vsQ) {
        Scale angle_axis = EquiScan("q_z (1/nm)", n, max_angle / n, max_angle);
        auto angles = angle_axis.binCenters();
        std::vector<double> qs(angle_axis.size(), 0.0);
        for (size_t i = 0, size = qs.size(); i < size; ++i)
            qs[i] = 4 * pi * std::sin(angles[i]) / wavelength;
        return new QzScan(qs);
    }
    auto* result = new AlphaScan(EquiScan("alpha_i (rad)", n, max_angle / n, max_angle));
    result->setWavelength(wavelength);
    return result;
}

std::unique_ptr<SpecularSimulation> test::makeSimulation::BasicSpecular(const Sample& sample,
                                                                        bool vsQ, double intensity)
{
    std::unique_ptr<BeamScan> scan(BasicSpecularScan(vsQ));
    scan->setIntensity(intensity);

    auto result = std::make_unique<SpecularSimulation>(*scan, sample);
    result->options().setUseAvgMaterials(true);
    return result;
}

std::unique_ptr<SpecularSimulation>
test::makeSimulation::BasicYPolarizedSpecular(const Sample& sample, const std::string& polCase,
                                              bool vsQ)
{
    const auto yCase = YPolarizationCases.at(polCase);
    std::unique_ptr<BeamScan> scan(BasicSpecularScan(vsQ));
    scan->setPolarization(yCase.first);
    scan->setAnalyzer(yCase.second);
    return std::make_unique<SpecularSimulation>(*scan, sample);
}

std::unique_ptr<SpecularSimulation>
test::makeSimulation::SpecularWithGaussianBeam(const Sample& sample)
{
    const double wavelength = 1.54 * angstrom;
    const int n = 2000;
    const double max_angle = 5 * deg;
    auto gaussian_ff = std::make_unique<FootprintGauss>(1.0);
    AlphaScan scan(EquiScan("alpha_i (rad)", n, max_angle / n, max_angle));
    scan.setWavelength(wavelength);
    scan.setFootprint(gaussian_ff.get());

    return std::make_unique<SpecularSimulation>(scan, sample);
}

std::unique_ptr<SpecularSimulation>
test::makeSimulation::SpecularWithSquareBeam(const Sample& sample)
{
    const double wavelength = 1.54 * angstrom;
    const int n = 2000;
    const double max_angle = 5 * deg;
    auto square_ff = std::make_unique<FootprintSquare>(1.0);
    AlphaScan scan(EquiScan("alpha_i (rad)", n, max_angle / n, max_angle));
    scan.setWavelength(wavelength);
    scan.setFootprint(square_ff.get());

    return std::make_unique<SpecularSimulation>(scan, sample);
}

std::unique_ptr<SpecularSimulation>
test::makeSimulation::SpecularDivergentBeam(const Sample& sample)
{
    const double wavelength = 1.54 * angstrom;
    const int n = 20;
    const size_t n_integration_points = 10;
    const double max_angle = 5 * deg;
    const double wavelength_stddev = 0.1 * angstrom;
    const double alpha_stddev = 0.1 * deg;
    AlphaScan scan(EquiScan("alpha_i (rad)", n, max_angle / n, max_angle));
    scan.setWavelength(wavelength);

    // Here central value for the wavelength distribution should be set equal to scan.walength(),
    // because during itemization GUI only relies on distribution parameters.
    DistributionGaussian wavelength_distr(wavelength, wavelength_stddev, n_integration_points, 2.);
    DistributionGaussian alpha_distr(0, alpha_stddev, n_integration_points, 2.);

    scan.setWavelengthDistribution(wavelength_distr);
    scan.setGrazingAngleDistribution(alpha_distr);

    return std::make_unique<SpecularSimulation>(scan, sample);
}

std::unique_ptr<SpecularSimulation>
test::makeSimulation::TOFRWithRelativeResolution(const Sample& sample)
{
    Scale qs = EquiScan("alpha_i (rad)", 500, 0.0, 1.0);
    QzScan scan(qs);
    scan.setRelativeQResolution(DistributionGaussian(0., 1., 20, 2.0), 0.03);

    auto result = std::make_unique<SpecularSimulation>(scan, sample);
    result->options().setUseAvgMaterials(true);
    return result;
}

std::unique_ptr<SpecularSimulation>
test::makeSimulation::TOFRWithPointwiseResolution(const Sample& sample)
{
    Scale qs = EquiScan("alpha_i (rad)", 500, 0.0, 1.0);
    QzScan scan(qs);

    std::vector<double> resolutions;
    resolutions.reserve(qs.size());
    auto qs_vector = qs.binCenters();
    std::for_each(qs_vector.begin(), qs_vector.end(),
                  [&resolutions](double q_val) { resolutions.push_back(0.03 * q_val); });
    scan.setVectorResolution(DistributionGaussian(0., 1., 20, 2.0), resolutions);

    auto result = std::make_unique<SpecularSimulation>(scan, sample);
    result->options().setUseAvgMaterials(true);
    return result;
}

//  ************************************************************************************************
//  depth probe simulations
//  ************************************************************************************************


std::unique_ptr<DepthprobeSimulation>
test::makeSimulation::BasicDepthprobeShiftedZ(const Sample& sample, double shift)
{
    AlphaScan scan(20, 0.025 * deg, 0.975 * deg);
    scan.setWavelength(1.);

    Scale zaxis = EquiDivision("z (nm)", 20, -100. + shift, +100. + shift);

    return std::make_unique<DepthprobeSimulation>(scan, sample, zaxis);
}

std::unique_ptr<DepthprobeSimulation> test::makeSimulation::BasicDepthprobe(const Sample& sample)
{
    return test::makeSimulation::BasicDepthprobeShiftedZ(sample, 0);
}