//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      GUI/Model/FromCore/ItemizeSimulation.cpp
//! @brief     Implements namespace GUI::FromCore.
//!
//! @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 "GUI/Model/FromCore/ItemizeSimulation.h"
#include "Base/Axis/Scale.h"
#include "Base/Const/Units.h"
#include "Device/Beam/Beam.h"
#include "Device/Detector/IDetector.h"
#include "Device/Detector/OffspecDetector.h"
#include "Device/Mask/Ellipse.h"
#include "Device/Mask/InfinitePlane.h"
#include "Device/Mask/Line.h"
#include "Device/Mask/MaskStack.h"
#include "Device/Mask/Polygon.h"
#include "Device/Mask/Rectangle.h"
#include "Device/Resolution/ConvolutionDetectorResolution.h"
#include "Device/Resolution/ResolutionFunction2DGaussian.h"
#include "GUI/Model/Beam/BeamDistributionItem.h"
#include "GUI/Model/Beam/DistributionItems.h"
#include "GUI/Model/Beam/FootprintItems.h"
#include "GUI/Model/Beam/SourceItems.h"
#include "GUI/Model/Detector/DetectorItem.h"
#include "GUI/Model/Detector/OffspecDetectorItem.h"
#include "GUI/Model/Detector/ResolutionFunctionItems.h"
#include "GUI/Model/Mask/MasksSet.h"
#include "GUI/Model/Sim/InstrumentsSet.h"
#include "GUI/Model/Sim/SimulationOptionsItem.h"
#include "Param/Distrib/Distributions.h"
#include "Resample/Option/SimulationOptions.h"
#include "Sim/Background/ConstantBackground.h"
#include "Sim/Background/PoissonBackground.h"
#include "Sim/Scan/AlphaScan.h"
#include "Sim/Scan/LambdaScan.h"
#include "Sim/Scan/QzScan.h"
#include "Sim/Simulation/includeSimulations.h"

namespace {

//! Sets masks in MasksSet according to given IDetector core object.
std::unique_ptr<MasksSet> getMasksSet(const IDetector& detector)
{
    std::unique_ptr<MasksSet> result = std::make_unique<MasksSet>();
    const MaskStack* maskStack = detector.detectorMask();
    for (size_t i = 0; i < maskStack->numberOfMasks(); ++i) {
        const auto [shape, mask_value] = maskStack->patternAt(i);

        MaskItem* mask_item = nullptr;
        if (const auto* ellipse = dynamic_cast<const Ellipse*>(shape)) {
            auto* m = new EllipseItem;
            m->setXCenter(ellipse->getCenterX());
            m->setYCenter(ellipse->getCenterY());
            m->setXRadius(ellipse->radiusX());
            m->setYRadius(ellipse->radiusY());
            m->setAngle(ellipse->getTheta());
            mask_item = m;
        }

        else if (const auto* rectangle = dynamic_cast<const Rectangle*>(shape)) {
            auto* m = new RectangleItem;
            m->setXLow(rectangle->getXlow());
            m->setYLow(rectangle->getYlow());
            m->setXHig(rectangle->getXup());
            m->setYHig(rectangle->getYup());
            mask_item = m;
        }

        else if (const auto* polygon = dynamic_cast<const Polygon*>(shape)) {
            auto* m = new PolygonItem;
            std::vector<double> xpos, ypos;
            polygon->getPoints(xpos, ypos);
            for (size_t ii = 0; ii < xpos.size(); ++ii)
                m->addPoint(xpos[ii], ypos[ii]);
            m->setIsClosed(true);
            mask_item = m;
        }

        else if (const auto* vline = dynamic_cast<const VerticalLine*>(shape))
            mask_item = new VerticalLineItem(vline->getXpos());

        else if (const auto* hline = dynamic_cast<const HorizontalLine*>(shape))
            mask_item = new HorizontalLineItem(hline->getYpos());

        else if (dynamic_cast<const InfinitePlane*>(shape))
            mask_item = new FullframeItem;
        else
            ASSERT_NEVER;

        ASSERT(mask_item);
        mask_item->setMaskValue(mask_value);
        result->add_item(mask_item);
    }

    if (detector.hasExplicitRegionOfInterest()) {
        const auto xBounds = detector.regionOfInterestBounds(0);
        const auto yBounds = detector.regionOfInterestBounds(1);

        auto* roiItem = new RegionOfInterestItem;
        roiItem->setXLow(xBounds.first);
        roiItem->setYLow(yBounds.first);
        roiItem->setXHig(xBounds.second);
        roiItem->setYHig(yBounds.second);
        result->add_item(roiItem);
    }

    return result;
}

//! Sets masks in DetectorItem according to given IDetector core object.
void setMaskStacks(DetectorItem* detector_item, const IDetector& detector)
{
    if ((detector.detectorMask() && detector.detectorMask()->hasMasks())
        || detector.hasExplicitRegionOfInterest())
        detector_item->setMasks(getMasksSet(detector).get());
}

//! Sets BeamDistributionItem according to given IDistribution1D core object.
void setDistributionTypeAndPars(BeamDistributionItem* pdi, const IDistribution1D* d)
{
    const double factor = 1 / pdi->scaleFactor();
    PolyPtr<DistributionItem, DistributionCatalog>& di = pdi->distributionSelection();

    if (const auto* dd = dynamic_cast<const DistributionGate*>(d)) {
        auto* item = new DistributionGateItem();
        item->center().setDVal(factor * (dd->min() + dd->max()) / 2);
        item->halfwidth().setDVal(factor * (dd->max() - dd->min()) / 2);
        di.setCertainItem(item);
    } else if (const auto* dd = dynamic_cast<const DistributionLorentz*>(d)) {
        auto* item = new DistributionLorentzItem();
        item->mean().setDVal(factor * dd->mean());
        item->hwhm().setDVal(factor * dd->hwhm());
        di.setCertainItem(item);
    } else if (const auto* dd = dynamic_cast<const DistributionGaussian*>(d)) {
        auto* item = new DistributionGaussianItem();
        item->mean().setDVal(factor * dd->mean());
        item->standardDeviation().setDVal(factor * dd->getStdDev());
        di.setCertainItem(item);
    } else if (const auto* dd = dynamic_cast<const DistributionLogNormal*>(d)) {
        auto* item = new DistributionLogNormalItem();
        item->median().setDVal(factor * dd->getMedian());
        item->scaleParameter().setDVal(dd->getScalePar());
        di.setCertainItem(item);
    } else if (const auto* dd = dynamic_cast<const DistributionCosine*>(d)) {
        auto* item = new DistributionCosineItem();
        item->mean().setDVal(factor * dd->mean());
        item->hwhm().setDVal(factor * dd->hwhm());
        di.setCertainItem(item);
    } else
        ASSERT_NEVER;
}

void setDistribution(BeamDistributionItem* pdi, ParameterDistribution par_distr)
{
    setDistributionTypeAndPars(pdi, par_distr.getDistribution());

    DistributionItem* distItem = pdi->distributionItem();

    distItem->setNumberOfSamples((int)par_distr.nDraws());
    distItem->relSamplingWidth().setDVal(par_distr.relSamplingWidth());
}

void addDistributionToItem(BeamDistributionItem* pdi, const IDistribution1D* distribution)
{
    if (!pdi)
        return;
    setDistributionTypeAndPars(pdi, distribution);

    DistributionItem* distItem = pdi->distributionItem();

    distItem->setNumberOfSamples((int)distribution->nSamples());
    distItem->relSamplingWidth().setDVal(distribution->relSamplingWidth());
}

void setupScanItem(ScanItem* item, const BeamScan* scan)
{
    item->setScan(scan);

    if (const auto* s2 = dynamic_cast<const AlphaScan*>(scan)) {
        item->scanTypeSelection().setCertainItem(new AlphaScanTypeItem);

        if (const IDistribution1D* distribution = s2->wavelengthDistribution())
            addDistributionToItem(item->wavelengthItem(), distribution);
        if (const IDistribution1D* distribution = s2->azimuthalAngleDistribution())
            addDistributionToItem(item->azimuthalAngleItem(), distribution);

        // distribution of the scanned parameter
        if (const IDistribution1D* distribution = s2->grazingAngleDistribution())
            addDistributionToItem(item->scanDistributionItem(), distribution);

    } else if (const auto* s2 = dynamic_cast<const LambdaScan*>(scan)) {
        item->scanTypeSelection().setCertainItem(new LambdaScanTypeItem);

        if (const IDistribution1D* distribution = s2->grazingAngleDistribution())
            addDistributionToItem(item->grazingAngleItem(), distribution);
        if (const IDistribution1D* distribution = s2->azimuthalAngleDistribution())
            addDistributionToItem(item->azimuthalAngleItem(), distribution);

        // distribution of the scanned parameter
        if (const IDistribution1D* distribution = s2->wavelengthDistribution())
            addDistributionToItem(item->scanDistributionItem(), distribution);

    } else if (const auto* s2 = dynamic_cast<const QzScan*>(scan)) {
        auto* typeItem = new QzScanTypeItem;
        typeItem->setUseRelativeResolution(s2->resolution_is_relative());
        item->scanTypeSelection().setCertainItem(new QzScanTypeItem);

        // distribution of the scanned parameter
        if (const IDistribution1D* distribution = s2->qzDistribution())
            addDistributionToItem(item->scanDistributionItem(), distribution);

    } else
        ASSERT_NEVER
}

void setGISASBeamItem(BeamItem* beam_item, const ScatteringSimulation& simulation)
{
    ASSERT(beam_item);
    const Beam& beam = simulation.beam();

    beam_item->setIntensity(beam.intensity());
    beam_item->setWavelength(beam.wavelength());
    beam_item->setGrazingAngle(Units::rad2deg(beam.alpha_i()));
    beam_item->setAzimuthalAngle(Units::rad2deg(beam.phi_i()));
    beam_item->setFootprintItem(beam.footprint());

    for (const ParameterDistribution& pd : simulation.paramDistributions()) {
        if (pd.whichParameter() == ParameterDistribution::BeamWavelength)
            setDistribution(beam_item->wavelengthItem(), pd);
        else if (pd.whichParameter() == ParameterDistribution::BeamGrazingAngle)
            setDistribution(beam_item->grazingAngleItem(), pd);
        else if (pd.whichParameter() == ParameterDistribution::BeamAzimuthalAngle)
            setDistribution(beam_item->azimuthalAngleItem(), pd);
        else
            ASSERT_NEVER;
    }
}

void setDetectorResolution(DetectorItem* detector_item, const IDetector& detector)
{
    const IDetectorResolution* resfunc = detector.detectorResolution();

    if (!resfunc)
        return;

    if (const auto* convfunc = dynamic_cast<const ConvolutionDetectorResolution*>(resfunc)) {
        if (const auto* resfunc = dynamic_cast<const ResolutionFunction2DGaussian*>(
                convfunc->getResolutionFunction2D())) {
            auto* item = new ResolutionFunction2DGaussianItem();
            item->setSigmaX(resfunc->sigmaX());
            item->setSigmaY(resfunc->sigmaY());
            detector_item->resolutionFunctionSelection().setCertainItem(item);
        } else
            ASSERT_NEVER;
    } else
        ASSERT_NEVER;
}

void setPolarizer2(InstrumentItem* instrument_item, const PolFilter& analyzer)
{
    instrument_item->setAnalyzerBlochVector(analyzer.BlochVector());
}

void updateDetector(Scatter2DInstrumentItem* instrument_item, const IDetector& detector)
{
    auto* detector_item = instrument_item->detectorItem();

    detector_item->phiAxis().setNbins(detector.axis(0).size());
    detector_item->phiAxis().min().setDVal(Units::rad2deg(detector.axis(0).min()));
    detector_item->phiAxis().max().setDVal(Units::rad2deg(detector.axis(0).max()));
    detector_item->alphaAxis().setNbins(detector.axis(1).size());
    detector_item->alphaAxis().min().setDVal(Units::rad2deg(detector.axis(1).min()));
    detector_item->alphaAxis().max().setDVal(Units::rad2deg(detector.axis(1).max()));

    setDetectorResolution(detector_item, detector);
    setMaskStacks(detector_item, detector);
    setPolarizer2(instrument_item, detector.analyzer());
}

void setBackground(InstrumentItem* instrument_item, const ISimulation& simulation)
{
    const auto* bg = simulation.background();
    PolyPtr<BackgroundItem, BackgroundCatalog>& ib = instrument_item->backgroundSelection();
    if (const auto* constant_bg = dynamic_cast<const ConstantBackground*>(bg)) {
        auto* item = new ConstantBackgroundItem();
        item->setBackgroundValue(constant_bg->backgroundValue());
        ib.setCertainItem(item);
    } else if (dynamic_cast<const PoissonBackground*>(bg)) {
        auto item = new PoissonBackgroundItem();
        ib.setCertainItem(item);
    }
}

Scatter2DInstrumentItem* createScatter2DInstrumentItem(const ScatteringSimulation& simulation)
{
    auto* result = new Scatter2DInstrumentItem;
    setGISASBeamItem(result->beamItem(), simulation);
    result->setPolarizerBlochVector(simulation.beam().polVector());
    const auto* det = dynamic_cast<const IDetector*>(simulation.getDetector());
    ASSERT(det);
    updateDetector(result, *det);
    result->setWithPolarizer(true);
    result->setWithAnalyzer(true);
    setBackground(result, simulation);
    return result;
}

OffspecInstrumentItem* createOffspecInstrumentItem(const OffspecSimulation& simulation)
{
    auto* result = new OffspecInstrumentItem;
    setupScanItem(result->scanItem(), simulation.scan());

    const OffspecDetector& detector = simulation.detector();
    OffspecDetectorItem* detectorItem = result->detectorItem();

    const Scale& phi_axis = detector.axis(0);
    const Scale& alpha_axis = detector.axis(1);

    auto& phiAxisProperty = detectorItem->phiAxis();
    phiAxisProperty.setNbins(phi_axis.size());
    phiAxisProperty.min().setDVal(Units::rad2deg(phi_axis.min()));
    phiAxisProperty.max().setDVal(Units::rad2deg(phi_axis.max()));

    auto& alphaAxisProperty = detectorItem->alphaAxis();
    alphaAxisProperty.setNbins(alpha_axis.size());
    alphaAxisProperty.min().setDVal(Units::rad2deg(alpha_axis.min()));
    alphaAxisProperty.max().setDVal(Units::rad2deg(alpha_axis.max()));

    setPolarizer2(result, detector.analyzer());
    result->setWithPolarizer(true);
    result->setWithAnalyzer(true);
    setBackground(result, simulation);
    return result;
}

SpecularInstrumentItem* createSpecularInstrumentItem(const SpecularSimulation& simulation)
{
    auto* result = new SpecularInstrumentItem;
    setupScanItem(result->scanItem(), simulation.scan());

    // TODO set polarizer & analyzer&

    setBackground(result, simulation);
    return result;
}

DepthprobeInstrumentItem* createDepthprobeInstrumentItem(const DepthprobeSimulation& simulation)
{
    auto* result = new DepthprobeInstrumentItem;
    setupScanItem(result->scanItem(), simulation.scan());

    const Scale& z_scale = simulation.z_axis();
    result->zAxis().setNbins(z_scale.size());
    result->zAxis().min().setDVal(z_scale.min());
    result->zAxis().max().setDVal(z_scale.max());

    return result;
}

} // namespace


InstrumentItem* GUI::FromCore::itemizeInstrument(const ISimulation& simulation)
{
    InstrumentItem* result;

    if (const auto* sim = dynamic_cast<const ScatteringSimulation*>(&simulation))
        result = createScatter2DInstrumentItem(*sim);
    else if (const auto* sim = dynamic_cast<const OffspecSimulation*>(&simulation))
        result = createOffspecInstrumentItem(*sim);
    else if (const auto* sim = dynamic_cast<const SpecularSimulation*>(&simulation))
        result = createSpecularInstrumentItem(*sim);
    else if (const auto* sim = dynamic_cast<const DepthprobeSimulation*>(&simulation))
        result = createDepthprobeInstrumentItem(*sim);
    else
        ASSERT_NEVER;

    return result;
}

SimulationOptionsItem* GUI::FromCore::itemizeOptions(const ISimulation& simulation)
{
    auto* result = new SimulationOptionsItem;

    if (simulation.options().isIntegrate())
        result->setUseMonteCarloIntegration(
            static_cast<unsigned>(simulation.options().getMcPoints()));
    else
        result->setUseAnalytical();

    result->setUseAverageMaterials(simulation.options().useAvgMaterials());
    result->setIncludeSpecularPeak(simulation.options().includeSpecular());

    return result;
}