//  ************************************************************************************************
//
//  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/Axis/BasicAxisItem.h"
#include "GUI/Model/Beam/BeamAngleItems.h"
#include "GUI/Model/Beam/BeamWavelengthItem.h"
#include "GUI/Model/Beam/GrazingScanItem.h"
#include "GUI/Model/Beam/SourceItems.h"
#include "GUI/Model/Descriptor/DistributionItems.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/BackgroundItems.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/QzScan.h" // yet unused -> TODO support QzScans!
#include "Sim/Simulation/includeSimulations.h"

namespace {

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

        if (const auto* ellipse = dynamic_cast<const Ellipse*>(shape)) {
            auto* ellipseItem = new EllipseItem();
            ellipseItem->setXCenter(ellipse->getCenterX());
            ellipseItem->setYCenter(ellipse->getCenterY());
            ellipseItem->setXRadius(ellipse->radiusX());
            ellipseItem->setYRadius(ellipse->radiusY());
            ellipseItem->setAngle(ellipse->getTheta());
            ellipseItem->setMaskValue(mask_value);
            // TODO: why prepend the mask, instead of appending it?
            result->add_item(ellipseItem);
        }

        else if (const auto* rectangle = dynamic_cast<const Rectangle*>(shape)) {
            auto* rectangleItem = new RectangleItem();
            rectangleItem->setXLow(rectangle->getXlow());
            rectangleItem->setYLow(rectangle->getYlow());
            rectangleItem->setXUp(rectangle->getXup());
            rectangleItem->setYUp(rectangle->getYup());
            rectangleItem->setMaskValue(mask_value);
            // TODO: why prepend the mask, instead of appending it?
            result->add_item(rectangleItem);
        }

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

            polygonItem->setMaskValue(mask_value);
            polygonItem->setIsClosed(true);
            // TODO: why prepend the mask, instead of appending it?
            result->add_item(polygonItem);
        }

        else if (const auto* vline = dynamic_cast<const VerticalLine*>(shape)) {
            auto* lineItem = new VerticalLineItem();
            lineItem->setPos(vline->getXpos());
            lineItem->setMaskValue(mask_value);
            // TODO: why prepend the mask, instead of appending it?
            result->add_item(lineItem);
        }

        else if (const auto* hline = dynamic_cast<const HorizontalLine*>(shape)) {
            auto* lineItem = new HorizontalLineItem();
            lineItem->setPos(hline->getYpos());
            lineItem->setMaskValue(mask_value);
            // TODO: why prepend the mask, instead of appending it?
            result->add_item(lineItem);
        }

        else if (const auto* plane = dynamic_cast<const InfinitePlane*>(shape)) {
            Q_UNUSED(plane);
            auto* planeItem = new FullframeItem();
            planeItem->setMaskValue(mask_value);
            result->add_item(planeItem);
        }

        else
            ASSERT_NEVER;
    }

    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->setXUp(xBounds.second);
        roiItem->setYUp(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));
}

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

    if (const auto* dd = dynamic_cast<const DistributionGate*>(d)) {
        auto* item = pdi->setDistributionItemType<DistributionGateItem>();
        item->setRange(factor * dd->min(), factor * dd->max());
    } else if (const auto* dd = dynamic_cast<const DistributionLorentz*>(d)) {
        auto* item = pdi->setDistributionItemType<DistributionLorentzItem>();
        item->setMean(factor * dd->mean());
        item->setHwhm(factor * dd->hwhm());
    } else if (const auto* dd = dynamic_cast<const DistributionGaussian*>(d)) {
        auto* item = pdi->setDistributionItemType<DistributionGaussianItem>();
        item->setMean(factor * dd->mean());
        item->setStandardDeviation(factor * dd->getStdDev());
    } else if (const auto* dd = dynamic_cast<const DistributionLogNormal*>(d)) {
        auto* item = pdi->setDistributionItemType<DistributionLogNormalItem>();
        item->setMedian(factor * dd->getMedian());
        item->setScaleParameter(dd->getScalePar());
    } else if (const auto* dd = dynamic_cast<const DistributionCosine*>(d)) {
        auto* item = pdi->setDistributionItemType<DistributionCosineItem>();
        item->setMean(factor * dd->mean());
        item->setHwhm(factor * dd->hwhm());
    } else if (const auto* dd = dynamic_cast<const DistributionTrapezoid*>(d)) {
        auto* item = pdi->setDistributionItemType<DistributionTrapezoidItem>();
        item->setCenter(factor * dd->mean());
        item->setLeftWidth(factor * dd->getLeftWidth());
        item->setMiddleWidth(factor * dd->getMiddleWidth());
        item->setRightWidth(factor * dd->getRightWidth());
    } 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());

    if (distItem->hasRelSamplingWidth())
        distItem->setRelSamplingWidth(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->setRelSamplingWidth(distribution->relSamplingWidth());
}

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->setInclinationAngle(Units::rad2deg(beam.alpha_i()));
    beam_item->setAzimuthalAngle(Units::rad2deg(beam.phi_i()));
    beam_item->setFootprint(beam.footprint());

    for (const ParameterDistribution& pd : simulation.paramDistributions()) {
        if (pd.whichParameter() == ParameterDistribution::BeamWavelength)
            setDistribution(beam_item->wavelengthItem(), pd);
        else if (pd.whichParameter() == ParameterDistribution::BeamInclinationAngle)
            setDistribution(beam_item->beamDistributionItem(), 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 =
                detector_item->setResolutionFunctionType<ResolutionFunction2DGaussianItem>();
            item->setSigmaX(resfunc->sigmaX());
            item->setSigmaY(resfunc->sigmaY());
        } 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().setWidth(Units::rad2deg(detector.axis(0).span()));
    detector_item->phiAxis().setCenter(Units::rad2deg(detector.axis(0).center()));
    detector_item->alphaAxis().setNbins(detector.axis(1).size());
    detector_item->alphaAxis().setWidth(Units::rad2deg(detector.axis(1).span()));
    detector_item->alphaAxis().setCenter(Units::rad2deg(detector.axis(1).center()));

    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();
    if (const auto* constant_bg = dynamic_cast<const ConstantBackground*>(bg)) {
        auto* constant_bg_item = instrument_item->setBackgroundItemType<ConstantBackgroundItem>();
        double value = constant_bg->backgroundValue();
        constant_bg_item->setBackgroundValue(value);
    } else if (dynamic_cast<const PoissonBackground*>(bg))
        instrument_item->setBackgroundItemType<PoissonBackgroundItem>();
}

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);
    return result;
}

OffspecInstrumentItem* createOffspecInstrumentItem(const OffspecSimulation& simulation)
{
    auto* result = new OffspecInstrumentItem();

    const BeamScan* scan = simulation.scan();
    result->scanItem()->setScan(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& phiAxisItem = detectorItem->phiAxis();
    phiAxisItem.setNbins(phi_axis.size());
    phiAxisItem.setMin(Units::rad2deg(phi_axis.min()));
    phiAxisItem.setMax(Units::rad2deg(phi_axis.max()));

    auto& alphaAxisItem = detectorItem->alphaAxis();
    alphaAxisItem.setNbins(alpha_axis.size());
    alphaAxisItem.setMin(Units::rad2deg(alpha_axis.min()));
    alphaAxisItem.setMax(Units::rad2deg(alpha_axis.max()));

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

SpecularInstrumentItem* createSpecularInstrumentItem(const SpecularSimulation& simulation)
{
    auto* result = new SpecularInstrumentItem();

    const BeamScan* scan = simulation.scan();
    result->scanItem()->setScan(scan);

    if (const auto* s2 = dynamic_cast<const AlphaScan*>(scan)) {
        if (const IDistribution1D* distribution = s2->wavelengthDistribution())
            addDistributionToItem(result->scanItem()->wavelengthItem(), distribution);
        if (const IDistribution1D* distribution = s2->alphaDistribution())
            addDistributionToItem(result->scanItem()->grazingScanItem(), distribution);
    }
    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
        ASSERT_NEVER;

    setBackground(result, simulation);

    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;
}
