//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      GUI/Model/Beam/SourceItems.cpp
//! @brief     Implements BeamItem hierarchy.
//!
//! @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/Beam/SourceItems.h"
#include "Base/Axis/Frame.h"
#include "Base/Axis/Scale.h"
#include "Base/Const/Units.h"
#include "Base/Util/Assert.h"
#include "Device/Beam/Beam.h"
#include "Device/Beam/FootprintGauss.h"
#include "Device/Beam/FootprintSquare.h"
#include "GUI/Model/Axis/BasicAxisItem.h"
#include "GUI/Model/Axis/PointwiseAxisItem.h"
#include "GUI/Model/Beam/BeamDistributionItem.h"
#include "GUI/Model/Beam/DistributionItems.h"
#include "Param/Distrib/Distributions.h"
#include "Sim/Scan/AlphaScan.h"

namespace {

namespace Tag {

const QString AzimuthalAngle("AzimuthalAngle");
const QString BaseData("BaseData");
const QString BeamInclinationDistribution("BeamInclinationDistribution");
const QString Distribution("Distribution");
const QString ExpandBeamParametersGroupbox("ExpandBeamParametersGroupbox");
const QString ExpandFootprintGroupbox("ExpandFootprintGroupbox");
const QString Footprint("Footprint");
const QString GrazingScan("GrazingScan");
const QString InclinationAngle("InclinationAngle");
const QString Intensity("Intensity");
const QString IsUniformAxis("IsUniformAxis");
const QString ListScan("ListScan");
const QString PointwiseAxis("PointwiseAxis"); // used in pre-21
const QString UniformAxis("UniformAxis");
const QString Wavelength("Wavelength");

} // namespace Tag

void setAxisPresentationDefaults(BasicAxisItem* axisItem)
{
    ASSERT(axisItem);
    if (dynamic_cast<PointwiseAxisItem*>(axisItem))
        return;
    axisItem->setMin(0.0);
    axisItem->setMax(3.0);
    axisItem->resize(500);
}

} // namespace

//  ************************************************************************************************
//  SourceItem
//  ************************************************************************************************

SourceItem::SourceItem()
    : m_wavelength_item(std::make_unique<BeamDistributionItem>(1.))
    , m_azimuthal_angle_item(std::make_unique<BeamDistributionItem>(Units::deg))
{
    m_intensity.init("Intensity", "Beam intensity in neutrons/photons per sec.", 1e8, 3,
                     RealLimits::limited(0.0, 1e32), "intensity");

    m_footprint.simpleInit(
        "Footprint type",
        "Model for surface area where scattering takes place (\"beam footprint\")",
        FootprintCatalog::Type::Gaussian);

    m_wavelength_item->resetToValue(0.1);
}

void SourceItem::writeTo(QXmlStreamWriter* w) const
{
    m_intensity.writeTo2(w, Tag::Intensity);
    XML::writeTaggedElement(w, Tag::Wavelength, *m_wavelength_item);
    XML::writeTaggedElement(w, Tag::AzimuthalAngle, *m_azimuthal_angle_item);
    XML::writeTaggedValue(w, Tag::ExpandBeamParametersGroupbox, expandBeamParameters);
    XML::writeTaggedElement(w, Tag::Footprint, m_footprint);
    XML::writeTaggedValue(w, Tag::ExpandFootprintGroupbox, expandFootprint);
}

void SourceItem::readFrom(QXmlStreamReader* r)
{
    while (r->readNextStartElement()) {
        QString tag = r->name().toString();
        if (tag == Tag::Intensity)
            m_intensity.readFrom2(r, tag);
        else if (tag == Tag::Wavelength)
            XML::readTaggedElement(r, tag, *m_wavelength_item);
        else if (tag == Tag::AzimuthalAngle)
            XML::readTaggedElement(r, tag, *m_azimuthal_angle_item);
        else if (tag == Tag::ExpandBeamParametersGroupbox)
            expandBeamParameters = XML::readTaggedBool(r, tag);
        else if (tag == Tag::Footprint)
            XML::readTaggedElement(r, tag, m_footprint);
        else if (tag == Tag::ExpandFootprintGroupbox)
            expandFootprint = XML::readTaggedBool(r, tag);
        else
            r->skipCurrentElement();
    }
}

void SourceItem::setWavelength(double value)
{
    ASSERT(m_wavelength_item);
    m_wavelength_item->resetToValue(value);
}

BeamDistributionItem* SourceItem::wavelengthItem() const
{
    ASSERT(m_wavelength_item);
    return m_wavelength_item.get();
}

void SourceItem::setAzimuthalAngle(double value)
{
    ASSERT(m_azimuthal_angle_item);
    m_azimuthal_angle_item->resetToValue(value);
}

BeamDistributionItem* SourceItem::azimuthalAngleItem() const
{
    ASSERT(m_azimuthal_angle_item);
    return m_azimuthal_angle_item.get();
}

void SourceItem::setGaussianFootprint(double value)
{
    m_footprint.setCertainItem(new FootprintGaussianItem(value));
}

void SourceItem::setSquareFootprint(double value)
{
    m_footprint.setCertainItem(new FootprintSquareItem(value));
}

void SourceItem::setFootprintItem(const IFootprint* footprint)
{
    if (!footprint)
        return;

    if (const auto* const fp = dynamic_cast<const FootprintGauss*>(footprint))
        setGaussianFootprint(fp->widthRatio());
    else if (const auto* const fp = dynamic_cast<const FootprintSquare*>(footprint))
        setSquareFootprint(fp->widthRatio());
}

//  ************************************************************************************************
//  BeamItem
//  ************************************************************************************************

BeamItem::BeamItem()
    : m_inclination_angle_item(std::make_unique<BeamDistributionItem>(Units::deg))
{
    m_inclination_angle_item->resetToValue(0.2);
}

void BeamItem::writeTo(QXmlStreamWriter* w) const
{
    XML::writeBaseElement<SourceItem>(w, XML::Tag::BaseData, this);
    XML::writeTaggedElement(w, Tag::InclinationAngle, *m_inclination_angle_item);
}

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

void BeamItem::setInclinationAngle(double value)
{
    ASSERT(m_inclination_angle_item);
    m_inclination_angle_item->resetToValue(value);
}

BeamDistributionItem* BeamItem::inclinationAngleItem() const
{
    ASSERT(m_inclination_angle_item);
    return m_inclination_angle_item.get();
}

std::unique_ptr<Beam> BeamItem::createBeam() const
{
    double lambda = wavelengthItem()->centralValue();
    double inclination_angle = Units::deg2rad(inclinationAngleItem()->centralValue());
    double azimuthal_angle = Units::deg2rad(azimuthalAngleItem()->centralValue());

    auto result =
        std::make_unique<Beam>(intensity().dVal(), lambda, inclination_angle, azimuthal_angle);
    result->setFootprint(m_footprint.certainItem()->createFootprint().get());
    return result;
}

//  ************************************************************************************************
//  ScanItem
//  ************************************************************************************************

ScanItem::ScanItem()
    : m_current_axis_is_uniform_axis(true)
    , m_uniform_alpha_axis(std::make_unique<BasicAxisItem>())
    , m_grazing_scan_item(std::make_unique<BeamDistributionItem>(Units::deg))
{
    setAxisPresentationDefaults(m_uniform_alpha_axis.get());
}

void ScanItem::setScan(const BeamScan* scan)
{
    setIntensity(scan->commonIntensity());
    if (const auto* s = dynamic_cast<const AlphaScan*>(scan))
        setWavelength(s->commonWavelength());
    setAzimuthalAngle(0.0);

    BasicAxisItem* axis_item = inclinationAxisItem();
    const Scale* axis = scan->coordinateAxis();
    ASSERT(axis);
    // TODO do we need to check that axis is equidistant?
    axis_item->resize(static_cast<int>(axis->size()));
    axis_item->setMin(axis->min() / Units::deg);
    axis_item->setMax(axis->max() / Units::deg);

    setFootprintItem(scan->commonFootprint());
}

void ScanItem::writeTo(QXmlStreamWriter* w) const
{
    XML::writeBaseElement<SourceItem>(w, XML::Tag::BaseData, this);
    XML::writeTaggedElement(w, Tag::GrazingScan, *m_grazing_scan_item);
    XML::writeTaggedValue(w, Tag::IsUniformAxis, m_current_axis_is_uniform_axis);
    if (m_uniform_alpha_axis)
        XML::writeTaggedElement(w, Tag::UniformAxis, *m_uniform_alpha_axis);
    if (m_pointwise_alpha_axis)
        XML::writeTaggedElement(w, Tag::ListScan, *m_pointwise_alpha_axis);
}

void ScanItem::readFrom(QXmlStreamReader* r)
{
    while (r->readNextStartElement()) {
        QString tag = r->name().toString();
        if (tag == Tag::BaseData)
            XML::readBaseElement<SourceItem>(r, tag, this);
        else if (tag == Tag::GrazingScan)
            XML::readTaggedElement(r, tag, *m_grazing_scan_item);
        else if (tag == Tag::IsUniformAxis)
            m_current_axis_is_uniform_axis = XML::readTaggedBool(r, tag);
        else if (tag == Tag::UniformAxis) {
            m_uniform_alpha_axis = std::make_unique<BasicAxisItem>();
            setAxisPresentationDefaults(m_uniform_alpha_axis.get());
            XML::readTaggedElement(r, tag, *m_uniform_alpha_axis);
        } else if (tag == Tag::ListScan || tag == Tag::PointwiseAxis) { // compatibility with pre-21
            m_pointwise_alpha_axis = std::make_unique<PointwiseAxisItem>();
            setAxisPresentationDefaults(m_pointwise_alpha_axis.get());
            XML::readTaggedElement(r, tag, *m_pointwise_alpha_axis);
        } else
            r->skipCurrentElement();
    }
}

BasicAxisItem* ScanItem::inclinationAxisItem() const
{
    return alphaAxisItem();
}

void ScanItem::updateToData(const Scale& axis)
{
    if (axis.unit() == "bin") {
        initUniformAxis(axis);
        selectUniformAxis();
    } else {
        initListScan(axis);
        selectListScan();
    }
}

int ScanItem::nBins() const
{
    return alphaAxisItem()->size();
}

BasicAxisItem* ScanItem::alphaAxisItem() const
{
    return m_current_axis_is_uniform_axis ? m_uniform_alpha_axis.get()
                                          : m_pointwise_alpha_axis.get();
}

bool ScanItem::pointwiseAlphaAxisDefined() const
{
    return (bool)m_pointwise_alpha_axis;
}

bool ScanItem::pointwiseAlphaAxisSelected() const
{
    return !m_current_axis_is_uniform_axis;
}
void ScanItem::selectUniformAxis()
{
    m_current_axis_is_uniform_axis = true;
}

void ScanItem::selectListScan()
{
    ASSERT(pointwiseAlphaAxisDefined());
    m_current_axis_is_uniform_axis = false;
}

void ScanItem::initUniformAxis(const Scale& axis)
{
    m_uniform_alpha_axis->resize(static_cast<int>(axis.size()));
}

void ScanItem::initListScan(const Scale& axis)
{
    if (!m_pointwise_alpha_axis) {
        m_pointwise_alpha_axis = std::make_unique<PointwiseAxisItem>();
        setAxisPresentationDefaults(m_pointwise_alpha_axis.get());
    }
    m_pointwise_alpha_axis->setAxis(axis);
}

void ScanItem::updateAxIndicators(const Frame& frame)
{
    if (!m_pointwise_alpha_axis)
        return;
    m_pointwise_alpha_axis->updateAxIndicators(frame);
}
