//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      GUI/Model/Sim/InstrumentItems.cpp
//! @brief     Implement class InstrumentItem and all its children
//!
//! @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/Sim/InstrumentItems.h"
#include "Base/Axis/Frame.h"
#include "Base/Axis/MakeScale.h"
#include "Base/Axis/Scale.h"
#include "Base/Const/Units.h"
#include "Base/Util/Assert.h"
#include "Device/Beam/Beam.h"
#include "Device/Beam/IFootprint.h"
#include "Device/Data/Datafield.h"
#include "Device/Detector/IDetector.h"
#include "Device/Detector/OffspecDetector.h"
#include "GUI/Model/Axis/PointwiseAxisItem.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/Data/DataItem.h"
#include "GUI/Model/Detector/DetectorItem.h"
#include "GUI/Model/Detector/OffspecDetectorItem.h"
#include "GUI/Model/File/DatafileItem.h"
#include "GUI/Model/Sample/SampleItem.h"
#include "GUI/Model/Sim/InstrumentCatalog.h"
#include "GUI/Model/Sim/InstrumentItems.h"
#include "GUI/Model/Util/Backup.h"
#include "Param/Distrib/Distributions.h"
#include "Sample/Multilayer/Sample.h"
#include "Sim/Background/IBackground.h"
#include "Sim/Scan/AlphaScan.h"
#include "Sim/Scan/LambdaScan.h"
#include "Sim/Scan/QzScan.h"
#include "Sim/Simulation/includeSimulations.h"
#include <QUuid>

namespace {

std::unique_ptr<IDistribution1D> createDistrFromItem(const BeamDistributionItem* item)
{
    ASSERT(item);
    const DistributionItem* distr_item = item->distributionItem();
    ASSERT(distr_item);
    const double scale = item->scaleFactor();
    return distr_item->createDistribution(scale);
}

namespace Tag {

const QString AlphaAxis("AlphaAxis");
const QString AnalyzerBlochVector("AnalyzerBlochVector");
const QString Background("Background");
const QString BaseData("BaseData");
const QString Beam("Beam");
const QString Description("Description");
const QString Detector("Detector");
const QString ExpandBeamParametersGroupbox("ExpandBeamParametersGroupbox");
const QString ExpandDetectorGroupbox("ExpandDetectorGroupbox");
const QString ExpandEnvironmentGroupbox("ExpandEnvironmentGroupbox");
const QString ExpandInfoGroupbox("ExpandInfoGroupbox");
const QString ExpandPolarizerAnalyzerGroupbox("ExpandPolarizerAnalyzerGroupbox");
const QString Id("Id");
const QString Name("Name");
const QString PolarizerBlochVector("PolarizerBlochVector");
const QString Scan("Scan");
const QString WithPolarizer("WithPolarizer");
const QString WithAnalyzer("WithAnalyzer");
const QString ZAxis("ZAxis");

} // namespace Tag

void setBeamDistribution(ParameterDistribution::WhichParameter which,
                         const BeamDistributionItem* item, ScatteringSimulation* simulation)
{
    DistributionItem* di = item->distributionItem();
    ASSERT(di);
    if (std::unique_ptr<IDistribution1D> d = di->createDistribution(item->scaleFactor()))
        simulation->addParameterDistribution(which, *d);
}

} // namespace

//  ************************************************************************************************
//  class InstrumentItem
//  ************************************************************************************************

InstrumentItem::InstrumentItem()
    : m_with_polarizer(false)
    , m_with_analyzer(false)
{
    m_id = QUuid::createUuid().toString();
    m_polarizer_bloch_vector.init("Polarizer Bloch vector", "",
                                  "Beam polarizer direction times efficiency", R3(), 3, true, 0.01,
                                  RealLimits::limited(-1, 1), "polarizerBlochVector");
    m_analyzer_bloch_vector.init("Analyzer Bloch vector", "",
                                 "Polarization analyzer direction times efficiency", R3(), 3, true,
                                 0.01, RealLimits::limited(-1, 1), "analyzerBlochVector");
    m_background.simpleInit("Background", "", BackgroundCatalog::Type::Constant);
}

InstrumentItem::~InstrumentItem() = default;

InstrumentItem* InstrumentItem::clone() const
{
    const auto type = InstrumentCatalog::type(this);
    InstrumentItem* copy = InstrumentCatalog::create(type);
    GUI::Util::copyContents(this, copy);
    copy->setId(QUuid::createUuid().toString());
    return copy;
}

bool InstrumentItem::alignedWith(const DatafileItem* dfi) const
{
    ASSERT(dfi);
    for (size_t i = 0; i < detectorRank(); ++i)
        if (axdim(i) != dfi->axdim(i))
            return false;
    return true;
}

void InstrumentItem::writeTo(QXmlStreamWriter* w) const
{
    XML::writeTaggedValue(w, Tag::Id, m_id);
    XML::writeTaggedValue(w, Tag::Name, name());
    XML::writeTaggedValue(w, Tag::Description, description());
    XML::writeTaggedValue(w, Tag::WithPolarizer, m_with_polarizer);
    XML::writeTaggedElement(w, Tag::PolarizerBlochVector, m_polarizer_bloch_vector);
    XML::writeTaggedValue(w, Tag::WithAnalyzer, m_with_analyzer);
    XML::writeTaggedElement(w, Tag::AnalyzerBlochVector, m_analyzer_bloch_vector);
    XML::writeTaggedElement(w, Tag::Background, m_background);
    XML::writeTaggedValue(w, Tag::ExpandInfoGroupbox, expandInfo);
    XML::writeTaggedValue(w, Tag::ExpandPolarizerAnalyzerGroupbox, expandPolarizerAnalyzer);
    XML::writeTaggedValue(w, Tag::ExpandEnvironmentGroupbox, expandEnvironment);
    XML::writeTaggedValue(w, Tag::ExpandDetectorGroupbox, expandDetector);
}

void InstrumentItem::readFrom(QXmlStreamReader* r)
{
    while (r->readNextStartElement()) {
        QString tag = r->name().toString();
        if (tag == Tag::Id)
            m_id = XML::readTaggedString(r, tag);
        else if (tag == Tag::Name)
            setName(XML::readTaggedString(r, tag));
        else if (tag == Tag::Description)
            setDescription(XML::readTaggedString(r, tag));
        else if (tag == Tag::WithPolarizer)
            m_with_polarizer = XML::readTaggedBool(r, tag);
        else if (tag == Tag::WithAnalyzer)
            m_with_analyzer = XML::readTaggedBool(r, tag);
        else if (tag == Tag::AnalyzerBlochVector)
            XML::readTaggedElement(r, tag, m_analyzer_bloch_vector);
        else if (tag == Tag::PolarizerBlochVector)
            XML::readTaggedElement(r, tag, m_polarizer_bloch_vector);
        else if (tag == Tag::Background)
            XML::readTaggedElement(r, tag, m_background);
        else if (tag == Tag::ExpandInfoGroupbox)
            expandInfo = XML::readTaggedBool(r, tag);
        else if (tag == Tag::ExpandPolarizerAnalyzerGroupbox)
            expandPolarizerAnalyzer = XML::readTaggedBool(r, tag);
        else if (tag == Tag::ExpandEnvironmentGroupbox)
            expandEnvironment = XML::readTaggedBool(r, tag);
        else if (tag == Tag::ExpandDetectorGroupbox)
            expandDetector = XML::readTaggedBool(r, tag);
        else
            r->skipCurrentElement();
    }
}

Frame InstrumentItem::makeFrame() const
{
    return *createFrame().release();
}

//  ************************************************************************************************
//  class ScanningInstrumentItem
//  ************************************************************************************************

ScanningInstrumentItem::ScanningInstrumentItem(double intensity)
    : m_scan(new ScanItem())
{
    m_scan->intensity().setDVal(intensity); // overwrite default value set by BeamItem c'tor
}

//! Takes ownership of argument 'axis'.
std::unique_ptr<PhysicalScan> ScanningInstrumentItem::createPhysicalScan(const Scale& axis) const
{
    ScanTypeItem* type = m_scan->scanTypeSelection().certainItem();

    if (dynamic_cast<const AlphaScanTypeItem*>(type)) {
        auto result = std::make_unique<AlphaScan>(axis);

        // distribution around scanned grazing position
        if (auto distr = ::createDistrFromItem(scanItem()->scanDistributionItem()))
            result->setGrazingAngleDistribution(*distr);

        result->setWavelength(scanItem()->wavelengthItem()->centralValue());
        // distribution around fixed wavelength position
        if (auto distr = ::createDistrFromItem(scanItem()->wavelengthItem()))
            result->setWavelengthDistribution(*distr);

        setupPhysicalScan(result.get());
        return result;

    } else if (dynamic_cast<const LambdaScanTypeItem*>(type)) {
        auto result = std::make_unique<LambdaScan>(axis);

        // distribution around scanned wavelength position
        if (auto distr = ::createDistrFromItem(scanItem()->scanDistributionItem()))
            result->setWavelengthDistribution(*distr);

        result->setGrazingAngle(Units::deg2rad(scanItem()->grazingAngleItem()->centralValue()));
        // distribution around fixed grazing position
        if (auto distr = ::createDistrFromItem(scanItem()->grazingAngleItem()))
            result->setGrazingAngleDistribution(*distr);

        setupPhysicalScan(result.get());
        return result;

    } else
        ASSERT_NEVER;
}

std::unique_ptr<BeamScan> ScanningInstrumentItem::createBeamScan(const Scale& axis) const
{
    ScanTypeItem* type = m_scan->scanTypeSelection().certainItem();

    if (auto* typeItem = dynamic_cast<const QzScanTypeItem*>(type)) {
        auto result = std::make_unique<QzScan>(axis);

        // distribution around scanned qz position
        if (auto distr = ::createDistrFromItem(scanItem()->scanDistributionItem())) {
            if (typeItem->isUseRelativeResolution())
                result->setRelativeQResolution(*distr, 1);
            else
                result->setAbsoluteQResolution(*distr, 1);
        }

        result->setIntensity(scanItem()->intensity().dVal());
        return result;

    } else
        return createPhysicalScan(axis);
}

//! Common settings of any PhysicalScan.
void ScanningInstrumentItem::setupPhysicalScan(PhysicalScan* scan) const
{
    scan->setIntensity(scanItem()->intensity().dVal());
    scan->setAzimuthalAngle(scanItem()->azimuthalAngleItem()->centralValue());
    if (auto distr = ::createDistrFromItem(scanItem()->azimuthalAngleItem()))
        scan->setAzimuthalAngleDistribution(*distr);

    FootprintItem* const footprint_item = scanItem()->footprintSelection().certainItem();
    scan->setFootprint(footprint_item->createFootprint().get());
}

void ScanningInstrumentItem::writeTo(QXmlStreamWriter* w) const
{
    XML::writeBaseElement<InstrumentItem>(w, XML::Tag::BaseData, this);
    XML::writeTaggedElement(w, Tag::Scan, *m_scan);
}

void ScanningInstrumentItem::readFrom(QXmlStreamReader* r)
{
    while (r->readNextStartElement()) {
        QString tag = r->name().toString();
        if (tag == Tag::BaseData)
            XML::readBaseElement<InstrumentItem>(r, tag, this);
        else if (tag == Tag::Scan)
            XML::readTaggedElement(r, tag, *m_scan);
        else
            r->skipCurrentElement();
    }
}

//  ************************************************************************************************
//  class Scatter2DInstrumentItem
//  ************************************************************************************************

Scatter2DInstrumentItem::Scatter2DInstrumentItem()
    : m_beam(std::make_unique<BeamItem>())
    , m_detector(std::make_unique<DetectorItem>())
{
    setName("GISAS");
}

size_t Scatter2DInstrumentItem::axdim(int i) const
{
    ASSERT(detectorItem());
    if (i == 0)
        return detectorItem()->phiAxis().nbins();
    else if (i == 1)
        return detectorItem()->alphaAxis().nbins();
    else
        ASSERT_NEVER;
}

void Scatter2DInstrumentItem::updateToRealData(const DatafileItem* dfi)
{
    ASSERT(dfi);
    ASSERT(dfi->rank() == 2);

    detectorItem()->phiAxis().setNbins(dfi->axdim(0));
    detectorItem()->alphaAxis().setNbins(dfi->axdim(1));
}

std::unique_ptr<Frame> Scatter2DInstrumentItem::createFrame() const
{
    return std::make_unique<Frame>(detectorItem()->createDetector()->clippedFrame());
}

ISimulation* Scatter2DInstrumentItem::createSimulation(const Sample& sample) const
{
    std::unique_ptr<Beam> beam = beamItem()->createBeam();
    if (withPolarizer())
        beam->setPolarization(m_polarizer_bloch_vector);
    std::unique_ptr<IDetector> detector = detectorItem()->createDetector();
    if (withAnalyzer())
        detector->setAnalyzer(m_analyzer_bloch_vector);
    auto* result = new ScatteringSimulation(*beam, sample, *detector);

    setBeamDistribution(ParameterDistribution::BeamWavelength, beamItem()->wavelengthItem(),
                        result);
    setBeamDistribution(ParameterDistribution::BeamGrazingAngle, beamItem()->grazingAngleItem(),
                        result);
    setBeamDistribution(ParameterDistribution::BeamAzimuthalAngle, beamItem()->azimuthalAngleItem(),
                        result);

    if (std::unique_ptr<const IBackground> background = backgroundItem()->createBackground())
        result->setBackground(*background);
    return result;
}

void Scatter2DInstrumentItem::writeTo(QXmlStreamWriter* w) const
{
    XML::writeBaseElement<InstrumentItem>(w, XML::Tag::BaseData, this);
    XML::writeTaggedElement(w, Tag::Beam, *m_beam);
    XML::writeTaggedElement(w, Tag::Detector, *m_detector);
}

void Scatter2DInstrumentItem::readFrom(QXmlStreamReader* r)
{
    while (r->readNextStartElement()) {
        QString tag = r->name().toString();
        if (tag == Tag::BaseData)
            XML::readBaseElement<InstrumentItem>(r, tag, this);
        else if (tag == Tag::Beam)
            XML::readTaggedElement(r, tag, *m_beam);
        else if (tag == Tag::Detector)
            XML::readTaggedElement(r, tag, *m_detector);
        else
            r->skipCurrentElement();
    }
}

//  ************************************************************************************************
//  class OffspecInstrumentItem
//  ************************************************************************************************

OffspecInstrumentItem::OffspecInstrumentItem()
    : ScanningInstrumentItem(1e8)
    , m_detector(new OffspecDetectorItem)
{
    setName("Offspec");
}

size_t OffspecInstrumentItem::axdim(int i) const
{
    if (i == 0) {
        ASSERT(scanItem() && scanItem()->scanDistributionItem());
        return scanItem()->nBins();
    } else if (i == 1) {
        ASSERT(detectorItem());
        return detectorItem()->alphaAxis().nbins();
    } else
        ASSERT_NEVER;
}

void OffspecInstrumentItem::updateToRealData(const DatafileItem* dfi)
{
    ASSERT(dfi);
    ASSERT(dfi->rank() == 2);

    detectorItem()->alphaAxis().setNbins(dfi->axdim(1));

    // Here we suppose that OffspecInstrument can only have _uniform_ grazing angle scan, so we
    // adjust only number of points. Othewise see implementation of SpecularInstrument.
    scanItem()->initUniformAxis(dfi->dataItem()->c_field()->axis(0));
}

std::unique_ptr<Frame> OffspecInstrumentItem::createFrame() const
{
    const Scale* xAxis = scanItem()->newUniformScale();
    return std::make_unique<Frame>(xAxis, detectorItem()->createOffspecDetector()->axis(1).clone());
}

ISimulation* OffspecInstrumentItem::createSimulation(const Sample& sample) const
{
    const Frame frame = makeFrame();
    std::unique_ptr<OffspecDetector> detector = detectorItem()->createOffspecDetector();
    std::unique_ptr<PhysicalScan> scan = createPhysicalScan(frame.axis(0));
    if (withPolarizer())
        scan->setPolarization(m_polarizer_bloch_vector);
    if (withAnalyzer())
        detector->setAnalyzer(m_analyzer_bloch_vector); // offspec uses detector analyzer, not scan
    auto* result = new OffspecSimulation(*scan, sample, *detector);

    if (std::unique_ptr<IBackground> background = backgroundItem()->createBackground())
        result->setBackground(*background);
    return result;
}

void OffspecInstrumentItem::writeTo(QXmlStreamWriter* w) const
{
    XML::writeBaseElement<ScanningInstrumentItem>(w, XML::Tag::BaseData, this);
    XML::writeTaggedElement(w, Tag::Detector, *m_detector);
}

void OffspecInstrumentItem::readFrom(QXmlStreamReader* r)
{
    while (r->readNextStartElement()) {
        QString tag = r->name().toString();
        if (tag == Tag::BaseData)
            XML::readBaseElement<ScanningInstrumentItem>(r, tag, this);
        else if (tag == Tag::Detector)
            XML::readTaggedElement(r, tag, *m_detector);
        else
            r->skipCurrentElement();
    }
}

//  ************************************************************************************************
//  class SpecularInstrumentItem
//  ************************************************************************************************

SpecularInstrumentItem::SpecularInstrumentItem()
    : ScanningInstrumentItem(1e6)
{
    setName("Specular");
}

size_t SpecularInstrumentItem::axdim(int i) const
{
    ASSERT(i == 0);
    return scanItem()->currentAxisItem()->size();
}

void SpecularInstrumentItem::updateToRealData(const DatafileItem* dfi)
{
    ASSERT(dfi);
    ASSERT(dfi->rank() == 1);

    const Scale& dataAxis = dfi->dataItem()->c_field()->axis(0);
    scanItem()->updateToData(dataAxis);
    scanItem()->updateAxIndicators(makeFrame());
}

bool SpecularInstrumentItem::alignedWith(const DatafileItem* dfi) const
{
    ASSERT(dfi);
    if (!dfi->holdsDimensionalData())
        return scanItem()->uniformAxisSelected() && (axdim(0) == dfi->axdim(0));

    if (!scanItem()->pointwiseAxisSelected())
        return false;

    const auto* axisItem = dynamic_cast<const PointwiseAxisItem*>(scanItem()->currentAxisItem());
    ASSERT(axisItem);

    const Scale* instrumentAxis = axisItem->axis();
    if (!instrumentAxis)
        return false;

    const Scale& real_data_axis = dfi->dataItem()->c_field()->axis(0);
    return *instrumentAxis == real_data_axis;
}

std::unique_ptr<Frame> SpecularInstrumentItem::createFrame() const
{
    if (dynamic_cast<const PointwiseAxisItem*>(scanItem()->currentAxisItem()))
        return std::make_unique<Frame>(scanItem()->newPointwiseScale());

    return std::make_unique<Frame>(scanItem()->newUniformScale());
}

ISimulation* SpecularInstrumentItem::createSimulation(const Sample& sample) const
{
    std::unique_ptr<const Frame> frame = createFrame();
    std::unique_ptr<BeamScan> scan = createBeamScan(frame->axis(0));
    if (withPolarizer())
        scan->setPolarization(m_polarizer_bloch_vector);
    if (withAnalyzer())
        scan->setAnalyzer(m_analyzer_bloch_vector);
    auto* result = new SpecularSimulation(*scan, sample);

    if (std::unique_ptr<IBackground> background = backgroundItem()->createBackground())
        result->setBackground(*background);
    return result;
}

void SpecularInstrumentItem::writeTo(QXmlStreamWriter* w) const
{
    XML::writeBaseElement<ScanningInstrumentItem>(w, XML::Tag::BaseData, this);
}

void SpecularInstrumentItem::readFrom(QXmlStreamReader* r)
{
    while (r->readNextStartElement()) {
        QString tag = r->name().toString();
        if (tag == Tag::BaseData)
            XML::readBaseElement<ScanningInstrumentItem>(r, tag, this);
        else
            r->skipCurrentElement();
    }
}

//  ************************************************************************************************
//  class DepthprobeInstrumentItem
//  ************************************************************************************************

DepthprobeInstrumentItem::DepthprobeInstrumentItem()
    : ScanningInstrumentItem(1e8)
    , m_z_axis("depth", "nm", -100, 100, RealLimits::limitless())
{
    setName("Depthprobe");

    BasicAxisItem* axisItem = scanItem()->currentAxisItem();
    axisItem->setMin(0.01);
    axisItem->setMax(1.0);
    axisItem->resize(500);
}

size_t DepthprobeInstrumentItem::axdim(int) const
{
    ASSERT_NEVER; // not yet implemented; not clear whether we want to support linking to real data
    return 0;
}

void DepthprobeInstrumentItem::updateToRealData(const DatafileItem* dfi)
{
    ASSERT(dfi);
    ASSERT_NEVER;
    // TODO ...
}

std::unique_ptr<Frame> DepthprobeInstrumentItem::createFrame() const
{
    const Scale* xAxis = scanItem()->newUniformScale();
    Scale zAxis =
        EquiDivision("z (nm)", m_z_axis.nbins(), m_z_axis.min().dVal(), m_z_axis.max().dVal());
    return std::make_unique<Frame>(xAxis, zAxis.clone());
}

ISimulation* DepthprobeInstrumentItem::createSimulation(const Sample& sample) const
{
    const Frame frame = makeFrame();
    std::unique_ptr<BeamScan> scan = createBeamScan(frame.axis(0));
    return new DepthprobeSimulation(*scan, sample, frame.axis(1));
}

void DepthprobeInstrumentItem::writeTo(QXmlStreamWriter* w) const
{
    XML::writeBaseElement<ScanningInstrumentItem>(w, XML::Tag::BaseData, this);
    XML::writeTaggedElement(w, Tag::ZAxis, m_z_axis);
}

void DepthprobeInstrumentItem::readFrom(QXmlStreamReader* r)
{
    while (r->readNextStartElement()) {
        QString tag = r->name().toString();
        if (tag == Tag::BaseData)
            XML::readBaseElement<ScanningInstrumentItem>(r, tag, this);
        else if (tag == Tag::ZAxis)
            XML::readTaggedElement(r, tag, m_z_axis);
        else
            r->skipCurrentElement();
    }
}
