//  ************************************************************************************************
//
//  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/PhysicalConstants.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/Simulation/includeSimulations.h"
#include <QUuid>

using PhysConsts::pi;

namespace {
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, ISimulation* 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, 0.01,
                                  RealLimits::limited(-1, 1), "polarizerBlochVector");
    m_analyzer_bloch_vector.init("Analyzer Bloch vector",
                                 "Polarization analyzer direction times efficiency", R3(), 3, 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::createScan(const Scale& axis) const
{
    auto result = std::make_unique<AlphaScan>(axis);

    result->setIntensity(scanItem()->intensity().dVal());
    result->setWavelength(scanItem()->wavelengthItem()->centralValue());

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

    {
        const BeamDistributionItem* it = scanItem()->wavelengthItem();
        ASSERT(it);
        const DistributionItem* distr_item = it->distributionItem();
        ASSERT(distr_item);
        const double scale = it->scaleFactor();
        if (std::unique_ptr<IDistribution1D> distr = distr_item->createDistribution(scale))
            result->setWavelengthDistribution(*distr);
    }

    {
        const BeamDistributionItem* it = scanItem()->grazingScanItem();
        ASSERT(it);
        const DistributionItem* distr_item = it->distributionItem();
        ASSERT(distr_item);
        const double scale = it->scaleFactor();
        if (std::unique_ptr<IDistribution1D> distr = distr_item->createDistribution(scale))
            result->setAngleDistribution(*distr);
    }

    return result;
}

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::BeamInclinationAngle,
                        beamItem()->inclinationAngleItem(), 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()->grazingScanItem());
        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
{
    BasicAxisItem* const axis_item = scanItem()->inclinationAxisItem();
    Scale xAxis = axis_item->makeScale("alpha_i (rad)");
    return std::make_unique<Frame>(xAxis.clone(),
                                   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 = createScan(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()->inclinationAxisItem()->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()->uniformAlphaAxisSelected() && axdim(0) == dfi->axdim(0);

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

    const auto* axisItem = dynamic_cast<const PointwiseAxisItem*>(scanItem()->alphaAxisItem());
    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
{
    BasicAxisItem* axis_item = scanItem()->inclinationAxisItem();
    ASSERT(axis_item);

    if (auto* pointwise_axis = dynamic_cast<PointwiseAxisItem*>(axis_item)) {
        const Scale* pAxis = pointwise_axis->axis();
        if (!pAxis) // workaround for loading project
            return {};

        if (pAxis->unit() == "1/nm") {
            double lambda = scanItem()->wavelengthItem()->centralValue();
            Scale ax =
                pAxis->transformedScale(Coordinate("alpha_i", "rad").label(), [lambda](double qz) {
                    return std::asin(qz * lambda / 4 / pi);
                });
            return std::make_unique<Frame>(ax.clone());
        }
        return std::make_unique<Frame>(pAxis->clone());
    }

    Scale axis = axis_item->makeScale("alpha_i (rad)");
    return std::make_unique<Frame>(axis.clone());
}

ISimulation* SpecularInstrumentItem::createSimulation(const Sample& sample) const
{
    std::unique_ptr<const Frame> frame = createFrame();
    std::unique_ptr<BeamScan> scan = createScan(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()->inclinationAxisItem();
    axisItem->setMin(0.0);
    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
{
    BasicAxisItem* const axis_item = scanItem()->inclinationAxisItem();
    Scale xAxis = axis_item->makeScale("alpha_i (rad)");
    Scale zAxis =
        EquiDivision("z (nm)", m_z_axis.nbins(), m_z_axis.min().dVal(), m_z_axis.max().dVal());
    std::vector<const Scale*> axes{xAxis.clone(), zAxis.clone()};
    return std::make_unique<Frame>(axes);
}

ISimulation* DepthprobeInstrumentItem::createSimulation(const Sample& sample) const
{
    const Frame frame = makeFrame();
    std::unique_ptr<BeamScan> scan = createScan(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();
    }
}
