//  ************************************************************************************************
//
//  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"
#include "Sim/Scan/LambdaScan.h"
#include "Sim/Scan/QzScan.h"
#include <numbers>

using std::numbers::pi;

namespace {

namespace Tag {

const QString AzimuthalAngle("AzimuthalAngle");
const QString GrazingAngle("GrazingAngle");
const QString BaseData("BaseData");
const QString BeamDistribution("BeamDistribution");
const QString ScanDistribution("ScanDistribution");
const QString ExpandBeamParametersGroupbox("ExpandBeamParametersGroupbox");
const QString ExpandFootprintGroupbox("ExpandFootprintGroupbox");
const QString Footprint("Footprint");
const QString Intensity("Intensity");
const QString IsUniformAxis("IsUniformAxis");
const QString ScanType("ScanType");
const QString ListScan("ListScan");
const QString PointwiseAxis("PointwiseAxis"); // used in pre-21
const QString UniformAxis("UniformAxis");
const QString Wavelength("Wavelength");

} // namespace Tag
} // namespace

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

SourceItem::SourceItem()
    : m_wavelength_item(std::make_unique<BeamDistributionItem>(1.))
    , m_grazing_angle_item(std::make_unique<BeamDistributionItem>(Units::deg))
    , m_azimuthal_angle_item(std::make_unique<BeamDistributionItem>(Units::deg))
{
    m_intensity.init("Intensity", "", "Beam intensity in neutrons/photons per sec.", 1e8, 3, false,
                     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);
    ASSERT(m_wavelength_item->distributionItem());
    m_wavelength_item->distributionItem()->setUnits("nm");
    m_wavelength_item->distributionItem()->center().setLimits(RealLimits::nonnegative());

    m_grazing_angle_item->resetToValue(0.2);
    ASSERT(m_grazing_angle_item->distributionItem());
    m_grazing_angle_item->distributionItem()->center().setLimits(RealLimits::limited(0., 90.));
    m_grazing_angle_item->distributionItem()->setUnits("deg");

    ASSERT(m_azimuthal_angle_item->distributionItem());
    m_azimuthal_angle_item->distributionItem()->center().setLimits(RealLimits::limited(-90., 90.));
    m_azimuthal_angle_item->distributionItem()->setUnits("deg");
}

void SourceItem::writeTo(QXmlStreamWriter* w) const
{
    m_intensity.writeTo2(w, Tag::Intensity);
    XML::writeTaggedElement(w, Tag::Wavelength, *m_wavelength_item);
    XML::writeTaggedElement(w, Tag::GrazingAngle, *m_grazing_angle_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::GrazingAngle)
            XML::readTaggedElement(r, tag, *m_grazing_angle_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::grazingAngleItem() const
{
    ASSERT(m_grazing_angle_item);
    return m_grazing_angle_item.get();
}

void SourceItem::setGrazingAngle(double value)
{
    ASSERT(m_grazing_angle_item);
    m_grazing_angle_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() {}

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

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

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

ScanItem::ScanItem()
    : m_current_axis_is_uniform_axis(true)
    , m_uniform_axis(std::make_unique<BasicAxisItem>())
    , m_scan_distribution_item(std::make_unique<BeamDistributionItem>())
{
    m_scan_type.simpleInit("Scan type", "", ScanTypeCatalog::Type::Alpha);
    setAxisPresentationDefaults(m_uniform_axis.get());
}

void ScanItem::setScan(const BeamScan* scan)
{
    setIntensity(scan->commonIntensity());

    BasicAxisItem* axis_item = currentAxisItem();
    const Scale* axis = scan->coordinateAxis();
    ASSERT(axis);
    ASSERT(axis->isEquiScan());

    axis_item->resize(static_cast<int>(axis->size()));

    if (const auto* s = dynamic_cast<const AlphaScan*>(scan)) {
        setWavelength(s->commonWavelength());
        setAzimuthalAngle(s->commonAzimuthalAngle());
        setFootprintItem(s->commonFootprint());

        axis_item->setMin(axis->min() / Units::deg);
        axis_item->setMax(axis->max() / Units::deg);

    } else if (const auto* s = dynamic_cast<const LambdaScan*>(scan)) {
        setGrazingAngle(s->commonGrazingAngle());
        setAzimuthalAngle(s->commonAzimuthalAngle());
        setFootprintItem(s->commonFootprint());

        axis_item->setMin(axis->min() / Units::nm);
        axis_item->setMax(axis->max() / Units::nm);

    } else if (dynamic_cast<const QzScan*>(scan)) {
        axis_item->setMin(axis->min() * Units::nm);
        axis_item->setMax(axis->max() * Units::nm);

    } else
        ASSERT_NEVER
}

void ScanItem::writeTo(QXmlStreamWriter* w) const
{
    XML::writeBaseElement<SourceItem>(w, XML::Tag::BaseData, this);
    XML::writeTaggedElement(w, Tag::ScanType, m_scan_type);
    XML::writeTaggedElement(w, Tag::ScanDistribution, *m_scan_distribution_item);
    XML::writeTaggedValue(w, Tag::IsUniformAxis, m_current_axis_is_uniform_axis);
    XML::writeTaggedElement(w, Tag::UniformAxis, *m_uniform_axis);
    if (m_pointwise_axis)
        XML::writeTaggedElement(w, Tag::ListScan, *m_pointwise_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::ScanDistribution)
            XML::readTaggedElement(r, tag, *m_scan_distribution_item);
        else if (tag == Tag::ScanType)
            XML::readTaggedElement(r, tag, m_scan_type);
        else if (tag == Tag::IsUniformAxis)
            m_current_axis_is_uniform_axis = XML::readTaggedBool(r, tag);
        else if (tag == Tag::UniformAxis) {
            m_uniform_axis = std::make_unique<BasicAxisItem>();
            setAxisPresentationDefaults(m_uniform_axis.get());
            XML::readTaggedElement(r, tag, *m_uniform_axis);
        } else if (tag == Tag::ListScan || tag == Tag::PointwiseAxis) { // compatibility with pre-21
            m_pointwise_axis = std::make_unique<PointwiseAxisItem>();
            XML::readTaggedElement(r, tag, *m_pointwise_axis);
        } else
            r->skipCurrentElement();
    }
}

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

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

BasicAxisItem* ScanItem::currentAxisItem() const
{
    return m_current_axis_is_uniform_axis ? m_uniform_axis.get() : m_pointwise_axis.get();
}

//! Returns cloned object
Scale* ScanItem::newUniformScale() const
{
    ScanTypeItem* scan_type = scanTypeSelection().certainItem();

    Scale* xAxis = nullptr;
    if (dynamic_cast<const AlphaScanTypeItem*>(scan_type))
        xAxis = m_uniform_axis->makeAlphaScale().clone();
    else if (dynamic_cast<const LambdaScanTypeItem*>(scan_type))
        xAxis = m_uniform_axis->makeLambdaScale().clone();
    else if (dynamic_cast<const QzScanTypeItem*>(scan_type))
        xAxis = m_uniform_axis->makeQzScale().clone();
    else
        ASSERT_NEVER;

    ASSERT(xAxis);
    return xAxis;
}

//! Returns cloned object
Scale* ScanItem::newPointwiseScale() const
{
    ASSERT(m_pointwise_axis);
    const Scale* pAxis = m_pointwise_axis->scale();
    if (!pAxis) // workaround for loading project
        return {};

    ScanTypeItem* scan_type = scanTypeSelection().certainItem();

    // transform q-coords to angular or spectral coords
    if (pAxis->unit() == "1/nm") {
        if (dynamic_cast<AlphaScanTypeItem*>(scan_type)) {
            double lambda = wavelengthItem()->centralValue();
            Scale ax =
                pAxis->transformedScale(Coordinate("alpha_i", "rad").label(), [lambda](double qz) {
                    if (lambda <= 0)
                        throw std::runtime_error("Nonpositive wavelength");
                    double s = qz * lambda / 4 / pi;
                    if (s > 1)
                        throw std::runtime_error(
                            "Q_z or wavelength are too big: the grazing angle is more than 90 deg");
                    return std::asin(s);
                });
            return ax.clone();

        } else if (dynamic_cast<LambdaScanTypeItem*>(scan_type)) {
            double alpha = Units::deg2rad(grazingAngleItem()->centralValue());

            // also reverse axis
            Scale ax = pAxis->reversedScale().transformedScale(
                Coordinate("lambda", "nm").label(), [alpha](double qz) {
                    if (alpha <= 0)
                        throw std::runtime_error("Nonpositive grazing angle");
                    if (qz <= 0)
                        throw std::runtime_error("Nonpositive q_z = " + std::to_string(qz));
                    return 4 * pi * std::sin(alpha) / qz;
                });
            return ax.clone();

        } else if (dynamic_cast<const QzScanTypeItem*>(scan_type)) {
            Scale ax("q_z (1/nm)", pAxis->bins());
            return ax.clone();

        } else
            ASSERT_NEVER;
    }

    return pAxis->clone();
}

bool ScanItem::pointwiseAxisDefined() const
{
    return (bool)m_pointwise_axis;
}

void ScanItem::selectUniformAxis()
{
    m_current_axis_is_uniform_axis = true;
}

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

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

void ScanItem::initListScan(const Scale& axis)
{
    if (!m_pointwise_axis)
        m_pointwise_axis = std::make_unique<PointwiseAxisItem>();

    m_pointwise_axis->setScale(axis);
}

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

void ScanItem::setAxisPresentationDefaults(BasicAxisItem* axisItem) const
{
    ASSERT(axisItem);
    if (dynamic_cast<PointwiseAxisItem*>(axisItem))
        return;

    // default for both alpha and lambda scans

    axisItem->resize(500);
    axisItem->setMin(0.01); // positive, because wavelength>0
    axisItem->setMax(3.0);
}