//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      GUI/Model/Model/ParameterTreeUtil.cpp
//! @brief     Implements ParameterTreeUtil namespace
//!
//! @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/Model/ParameterTreeUtil.h"
#include "GUI/Model/Beam/BeamAngleItems.h"
#include "GUI/Model/Beam/BeamWavelengthItem.h"
#include "GUI/Model/Beam/SourceItems.h"
#include "GUI/Model/Descriptor/DistributionItems.h"
#include "GUI/Model/Detector/OffspecDetectorItem.h"
#include "GUI/Model/Detector/RectangularDetectorItem.h"
#include "GUI/Model/Detector/ResolutionFunctionItems.h"
#include "GUI/Model/Detector/SphericalDetectorItem.h"
#include "GUI/Model/Device/BackgroundItems.h"
#include "GUI/Model/Device/InstrumentItems.h"
#include "GUI/Model/Job/JobItem.h"
#include "GUI/Model/Job/ParameterTreeItems.h"
#include "GUI/Model/Sample/CompoundItem.h"
#include "GUI/Model/Sample/CoreAndShellItem.h"
#include "GUI/Model/Sample/LayerItem.h"
#include "GUI/Model/Sample/MaterialItem.h"
#include "GUI/Model/Sample/MesocrystalItem.h"
#include "GUI/Model/Sample/ParticleItem.h"
#include "GUI/Model/Sample/ParticleLayoutItem.h"
#include "GUI/Model/Sample/SampleItem.h"

using std::variant;

namespace {

QString labelWithUnit(const QString& label, variant<QString, Unit> unit)
{
    const QString s = std::holds_alternative<QString>(unit) ? std::get<QString>(unit)
                                                            : unitAsString(std::get<Unit>(unit));

    if (!s.isEmpty())
        return label + " [" + s + "]";

    return label;
}

template <typename Catalog>
ParameterLabelItem* addLabel(ParameterLabelItem* parent, const QString& category,
                             const typename Catalog::CatalogedType* p)
{
    const auto title = category + " (" + Catalog::uiInfo(Catalog::type(p)).menuEntry + ")";
    return new ParameterLabelItem(title, parent);
}

template <typename Catalog>
ParameterLabelItem* addLabel(ParameterLabelItem* parent, const typename Catalog::CatalogedType* p)
{
    const auto title = Catalog::uiInfo(Catalog::type(p)).menuEntry;
    return new ParameterLabelItem(title, parent);
}

} // namespace

//  ************************************************************************************************

ParameterTreeBuilder::ParameterTreeBuilder(JobItem* jobItem, bool recreateBackupValues)
    : m_jobItem(jobItem)
    , m_recreateBackupValues(recreateBackupValues)
{
}

void ParameterTreeBuilder::build()
{
    addMaterials();
    addSample();
    addInstrument();
}

void ParameterTreeBuilder::addMaterials()
{
    auto* materialTopLabel =
        new ParameterLabelItem("Materials", parameterContainerItem()->parameterTreeRoot());
    for (auto* item : m_jobItem->sampleItem()->materialModel().materialItems()) {
        auto* label = new ParameterLabelItem(item->matItemName(), materialTopLabel);
        if (item->hasRefractiveIndex()) {
            addParameterItem(label, item->delta());
            addParameterItem(label, item->beta());
        } else {
            addParameterItem(label, item->sldRe());
            addParameterItem(label, item->sldIm());
        }

        if (allowMagneticFields()) {
            // Processing z-magnetization is not implemented yet (see issue #654)
            // addParameterItem(label, item->magnetization());
            addMagnetizationNoZ(label, item->magnetization());
        }
    }
}

void ParameterTreeBuilder::addSample()
{
    auto* label = new ParameterLabelItem("Sample", parameterContainerItem()->parameterTreeRoot());
    addParameterItem(label, m_jobItem->sampleItem()->crossCorrLength());

    // Processing external field is not implemented yet, so temporary disable it (see issue #654)
    // if (allowMagneticFields())
    //    addParameterItem(label, m_jobItem->sampleItem()->externalField());

    int iLayer = 0;
    for (auto* layer : m_jobItem->sampleItem()->layerItems()) {
        auto* layerLabel = new ParameterLabelItem("Layer" + QString::number(iLayer++), label);
        if (!layer->isTopLayer() && !layer->isBottomLayer())
            addParameterItem(layerLabel, layer->thickness());
        if (!layer->isTopLayer())
            if (auto* roughnessItem = layer->roughnessSelection().currentItem()) {
                auto* roughnessLabel = new ParameterLabelItem("Top roughness", layerLabel);
                for (auto* property : roughnessItem->roughnessProperties())
                    addParameterItem(roughnessLabel, *property);
            }

        int iLayout = 0;
        for (auto* layout : layer->layoutItems()) {
            auto* label = new ParameterLabelItem("Layout" + QString::number(iLayout++), layerLabel);
            if (!layout->totalDensityIsDefinedByInterference())
                addParameterItem(label, layout->ownDensity());

            addInterference(label, layout);

            for (auto* p : layout->itemsWithParticles())
                addItemWithParticles(label, p, true);
        }
    }
}

void ParameterTreeBuilder::addParameterItem(ParameterLabelItem* parent, DoubleProperty& d,
                                            const QString& label)
{
    auto* parameterItem = new ParameterItem(parent);
    parameterItem->setTitle(labelWithUnit(label.isEmpty() ? d.label() : label, d.unit()));
    parameterItem->linkToProperty(d);
    if (m_recreateBackupValues)
        m_jobItem->parameterContainerItem()->setBackupValue(parameterItem->link(), d.value());
}

void ParameterTreeBuilder::addParameterItem(ParameterLabelItem* parent, VectorProperty& d)
{
    auto* label = new ParameterLabelItem(d.label(), parent);
    addParameterItem(label, d.x());
    addParameterItem(label, d.y());
    addParameterItem(label, d.z());
}

void ParameterTreeBuilder::addMagnetizationNoZ(ParameterLabelItem* parent, VectorProperty& d)
{
    // Setting z-component is temporary disabled (see issue #654)
    // When interaction with magnetic field in fronting medium is implemented,
    // delete this method and use 'addParameterItem' instead

    auto* label = new ParameterLabelItem(d.label(), parent);
    addParameterItem(label, d.x());
    addParameterItem(label, d.y());
}


ParameterContainerItem* ParameterTreeBuilder::parameterContainerItem()
{
    return m_jobItem->parameterContainerItem();
}

bool ParameterTreeBuilder::allowMagneticFields() const
{
    // Depthprobe works only with scalar fluxes
    return !m_jobItem->instrumentItem()->is<DepthprobeInstrumentItem>();
}

void ParameterTreeBuilder::addInterference(ParameterLabelItem* layoutLabel,
                                           const ParticleLayoutItem* layout)
{
    auto* interference = layout->interferenceSelection().currentItem();
    if (!interference)
        return;

    const auto itfType = InterferenceItemCatalog::type(interference);
    const QString title = InterferenceItemCatalog::uiInfo(itfType).menuEntry;

    auto* label = new ParameterLabelItem("Interference (" + title + ")", layoutLabel);

    if (auto* itf = dynamic_cast<InterferenceRadialParacrystalItem*>(interference)) {
        addParameterItem(label, itf->positionVariance());
        addParameterItem(label, itf->peakDistance());
        addParameterItem(label, itf->dampingLength());
        addParameterItem(label, itf->domainSize());
        addParameterItem(label, itf->kappa());

        auto* pdf = itf->probabilityDistributionSelection().currentItem();
        auto* pdfLabel = addLabel<Profile1DItemCatalog>(label, "PDF", pdf);
        for (auto* d : pdf->profileProperties())
            addParameterItem(pdfLabel, *d);
    } else if (auto* itf = dynamic_cast<Interference2DParacrystalItem*>(interference)) {
        addParameterItem(label, itf->positionVariance());
        addParameterItem(label, itf->dampingLength());
        addParameterItem(label, itf->domainSize1());
        addParameterItem(label, itf->domainSize2());
        addLattice(label, itf);

        auto* pdf1 = itf->probabilityDistributionSelection1().currentItem();
        auto* pdf2 = itf->probabilityDistributionSelection2().currentItem();
        const bool samePdfTypes =
            Profile2DItemCatalog::type(pdf1) == Profile2DItemCatalog::type(pdf2);
        auto* pdf1Label =
            addLabel<Profile2DItemCatalog>(label, samePdfTypes ? "PDF1" : "PDF", pdf1);
        for (auto* d : pdf1->profileProperties())
            addParameterItem(pdf1Label, *d);
        auto* pdf2Label =
            addLabel<Profile2DItemCatalog>(label, samePdfTypes ? "PDF2" : "PDF", pdf2);
        for (auto* d : pdf2->profileProperties())
            addParameterItem(pdf2Label, *d);
    } else if (auto* itf = dynamic_cast<Interference1DLatticeItem*>(interference)) {
        addParameterItem(label, itf->positionVariance());
        addParameterItem(label, itf->length());
        addParameterItem(label, itf->rotationAngle());

        auto* df = itf->decayFunctionSelection().currentItem();
        auto* dfLabel = addLabel<Profile1DItemCatalog>(label, "Decay function", df);
        for (auto* d : df->profileProperties())
            addParameterItem(dfLabel, *d);
    } else if (auto* itf = dynamic_cast<Interference2DLatticeItem*>(interference)) {
        addParameterItem(label, itf->positionVariance());
        addLattice(label, itf);

        auto* df = itf->decayFunctionSelection().currentItem();
        auto* dfLabel = addLabel<Profile2DItemCatalog>(label, "Decay function", df);
        for (auto* d : df->profileProperties())
            addParameterItem(dfLabel, *d);
    } else if (auto* itf = dynamic_cast<InterferenceFinite2DLatticeItem*>(interference)) {
        // domainSize1 and domainSize2 are of type UInt (not matching the double approach for tuning
        // and fitting). In BornAgain 1.18 these values have not been added to the tuning tree, and
        // also not to the fitting parameters. Maybe this should be necessary, but for now this
        // stays the same and the two sizes are not added
        addParameterItem(label, itf->positionVariance());
        addLattice(label, itf);
    } else if (auto* itf = dynamic_cast<InterferenceHardDiskItem*>(interference)) {
        addParameterItem(label, itf->positionVariance());
        addParameterItem(label, itf->radius());
        addParameterItem(label, itf->density());
    }
}

ParameterLabelItem* ParameterTreeBuilder::addItemWithParticles(ParameterLabelItem* parentLabel,
                                                               ItemWithParticles* p,
                                                               bool enableAbundance,
                                                               bool enablePosition)
{
    auto* label = addLabel<ItemWithParticlesCatalog>(parentLabel, p);

    if (enableAbundance)
        addParameterItem(label, p->abundance());
    if (enablePosition)
        addParameterItem(label, p->position());
    addRotation(label, p);

    if (const auto* particle = dynamic_cast<const ParticleItem*>(p)) {
        auto* formFactor = particle->formFactorItem();
        auto* ffLabel = addLabel<FormFactorItemCatalog>(label, "Formfactor", formFactor);
        for (auto* d : formFactor->geometryProperties())
            addParameterItem(ffLabel, *d);
    } else if (const auto* particleComposition = dynamic_cast<const CompoundItem*>(p)) {
        for (auto* p : particleComposition->itemsWithParticles())
            addItemWithParticles(label, p, false);
    } else if (const auto* coreShell = dynamic_cast<const CoreAndShellItem*>(p)) {
        auto* l = addItemWithParticles(label, coreShell->coreItem(), false);
        l->setTitle(l->title() + " (Core)");
        l = addItemWithParticles(label, coreShell->shellItem(), false, false);
        l->setTitle(l->title() + " (Shell)");
    } else if (auto* meso = dynamic_cast<MesocrystalItem*>(p)) {
        addParameterItem(label, meso->vectorA());
        addParameterItem(label, meso->vectorB());
        addParameterItem(label, meso->vectorC());

        auto* outerShape = meso->outerShapeSelection().currentItem();
        auto* ffLabel = addLabel<FormFactorItemCatalog>(label, "Outer shape", outerShape);
        for (auto* d : outerShape->geometryProperties())
            addParameterItem(ffLabel, *d);

        auto* l = addItemWithParticles(label, meso->basisItem(), false);
        l->setTitle(l->title() + " (Basis particle)");
    }

    return label;
}

void ParameterTreeBuilder::addLattice(ParameterLabelItem* parentLabel,
                                      const Interference2DAbstractLatticeItem* itf)
{
    auto* lattice = itf->latticeTypeItem();
    auto* label = addLabel<Lattice2DItemCatalog>(parentLabel, "Lattice", lattice);
    for (auto* d : lattice->geometryValues(!itf->xiIntegration()))
        addParameterItem(label, *d);
}

void ParameterTreeBuilder::addRotation(ParameterLabelItem* parentLabel, ItemWithParticles* p)
{
    auto* r = p->rotationSelection().currentItem();
    if (!r)
        return;

    auto* label = addLabel<RotationItemCatalog>(parentLabel, "Rotation", r);
    for (auto* d : r->rotationProperties())
        addParameterItem(label, *d);
}

void ParameterTreeBuilder::addInstrument()
{
    auto* iI = m_jobItem->instrumentItem();
    auto* label = new ParameterLabelItem(iI->instrumentType() + " instrument",
                                         parameterContainerItem()->parameterTreeRoot());

    if (auto* iiI = dynamic_cast<GISASInstrumentItem*>(iI)) {
        auto* beamItem = iiI->beamItem();
        auto* beamLabel = new ParameterLabelItem("Beam", label);
        addParameterItem(beamLabel, beamItem->intensity());
        addBeamDistribution(beamLabel, beamItem->wavelengthItem(), "Wavelength");
        addBeamDistribution(beamLabel, beamItem->beamDistributionItem(), "Inclination angle");
        addBeamDistribution(beamLabel, beamItem->azimuthalAngleItem(), "Azimuthal angle");
        addDetector(label, iiI->detectorItem());
        addPolarization(label, iiI);
        addBackground(label, iiI->backgroundItem());
    } else if (auto* iiI = dynamic_cast<SpecularInstrumentItem*>(iI)) {
        auto* beamLabel = new ParameterLabelItem("Beam", label);
        addParameterItem(beamLabel, iiI->scanItem()->intensity());
        addBeamDistribution(beamLabel, iiI->scanItem()->wavelengthItem(), "Wavelength");
        // TODO implement correctly "Inclination angle" which is scanned
        // https://jugit.fz-juelich.de/mlz/bornagain/-/issues/301
        // addBeamDistribution(beamLabel, beamItem->grazingScanItem(), "Inclination angle");
        addPolarization(label, iiI);
        addBackground(label, iiI->backgroundItem());
    } else if (auto* iiI = dynamic_cast<OffspecInstrumentItem*>(iI)) {
        auto* beamLabel = new ParameterLabelItem("Beam", label);
        addParameterItem(beamLabel, iiI->scanItem()->intensity());
        addBeamDistribution(beamLabel, iiI->scanItem()->wavelengthItem(), "Wavelength");
        addBeamDistribution(beamLabel, iiI->scanItem()->azimuthalAngleItem(), "Azimuthal angle");
        addOffspecDetector(label, iiI->detectorItem());
        addPolarization(label, iiI);
    } else if (auto* iiI = dynamic_cast<DepthprobeInstrumentItem*>(iI)) {
        auto* beamLabel = new ParameterLabelItem("Parameters", label);
        addBeamDistribution(beamLabel, iiI->scanItem()->wavelengthItem(), "Wavelength");
        // TODO implement correctly "Inclination angle" which is scanned
        // https://jugit.fz-juelich.de/mlz/bornagain/-/issues/301
        // addBeamDistribution(beamLabel, scanItem->grazingScanItem(), "Inclination angle");
        addPolarization(label, iiI);
    } else
        ASSERT(false);
}

void ParameterTreeBuilder::addBeamDistribution(ParameterLabelItem* parentLabel,
                                               BeamDistributionItem* distributionItem,
                                               const QString& label, bool withMean)
{
    auto* distribution = distributionItem->distributionItem();
    if (auto* dn = dynamic_cast<DistributionNoneItem*>(distribution)) {
        if (withMean)
            addParameterItem(parentLabel, dn->mean(), label);
    } else {
        const auto type = DistributionItemCatalog::type(distribution);
        const auto name = DistributionItemCatalog::uiInfo(type).menuEntry;
        auto* item = new ParameterLabelItem(QString("%1 (%2 distribution)").arg(label).arg(name),
                                            parentLabel);
        for (auto* d : distribution->distributionValues(withMean))
            addParameterItem(item, *d);
    }
}

void ParameterTreeBuilder::addDetector(ParameterLabelItem* parentLabel, DetectorItem* detector)
{
    const auto addResolutionFunction = [this, detector](ParameterLabelItem* detLabel) {
        if (auto* r = dynamic_cast<ResolutionFunction2DGaussianItem*>(
                detector->resolutionFunctionSelection().currentItem())) {
            auto* label = new ParameterLabelItem("Resolution (Gaussian)", detLabel);
            addParameterItem(label, r->sigmaX());
            addParameterItem(label, r->sigmaY());
        }
    };

    if (auto* rectDetector = dynamic_cast<RectangularDetectorItem*>(detector)) {
        auto* label = new ParameterLabelItem("Detector (rectangular)", parentLabel);
        addParameterItem(label, rectDetector->width());
        addParameterItem(label, rectDetector->height());
        addResolutionFunction(label);
        if (rectDetector->detectorAlignment() == RectangularDetector::GENERIC) {
            addParameterItem(label, rectDetector->normalVector());
            addParameterItem(label, rectDetector->directionVector());
            addParameterItem(label, rectDetector->u0());
            addParameterItem(label, rectDetector->v0());
        } else {
            addParameterItem(label, rectDetector->u0());
            addParameterItem(label, rectDetector->v0());
            addParameterItem(label, rectDetector->distance());
        }
    } else if (auto* spherDetector = dynamic_cast<SphericalDetectorItem*>(detector)) {
        auto* label = new ParameterLabelItem("Detector (spherical)", parentLabel);
        auto* phiLabel = new ParameterLabelItem("Phi axis", label);
        const QString unit = unitAsString(Unit::degree);
        addParameterItem(phiLabel, spherDetector->phiAxis().min());
        addParameterItem(phiLabel, spherDetector->phiAxis().max());
        auto* alphaLabel = new ParameterLabelItem("Alpha axis", label);
        addParameterItem(alphaLabel, spherDetector->alphaAxis().min());
        addParameterItem(alphaLabel, spherDetector->alphaAxis().max());
        addResolutionFunction(label);
    } else
        ASSERT(false);
}

void ParameterTreeBuilder::addOffspecDetector(ParameterLabelItem* parentLabel,
                                              OffspecDetectorItem* detector)
{
    auto* label = new ParameterLabelItem("Detector", parentLabel);
    auto* phiLabel = new ParameterLabelItem("Phi axis", label);
    const QString unit = unitAsString(Unit::degree);
    addParameterItem(phiLabel, detector->phiAxis().min());
    addParameterItem(phiLabel, detector->phiAxis().max());
    auto* alphaLabel = new ParameterLabelItem("Alpha axis", label);
    addParameterItem(alphaLabel, detector->alphaAxis().min());
    addParameterItem(alphaLabel, detector->alphaAxis().max());
}

void ParameterTreeBuilder::addBackground(ParameterLabelItem* instrumentLabel,
                                         BackgroundItem* backgroundItem)
{
    if (auto* b = dynamic_cast<ConstantBackgroundItem*>(backgroundItem))
        addParameterItem(instrumentLabel, b->backgroundValue(),
                         labelWithUnit("Constant background", b->backgroundValue().unit()));
}

void ParameterTreeBuilder::addPolarization(ParameterLabelItem* instrumentLabel,
                                           InstrumentItem* instrument)
{
    if (!instrument->withPolarizer() && !instrument->withAnalyzer())
        return;

    auto* label = new ParameterLabelItem("Polarization analysis", instrumentLabel);

    if (instrument->withPolarizer())
        addParameterItem(label, instrument->polarizerBlochVector());

    if (instrument->withAnalyzer())
        addParameterItem(label, instrument->analyzerBlochVector());
}
