//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      GUI/Model/Job/ParameterTreeBuilder.cpp
//! @brief     Implements class ParameterTreeBuilder.
//!
//! @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/Job/ParameterTreeBuilder.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/Detector/DetectorItem.h"
#include "GUI/Model/Detector/OffspecDetectorItem.h"
#include "GUI/Model/Detector/ResolutionFunctionItems.h"
#include "GUI/Model/Job/JobItem.h"
#include "GUI/Model/Material/MaterialItem.h"
#include "GUI/Model/Par/ParameterTreeItems.h"
#include "GUI/Model/Sample/CompoundItem.h"
#include "GUI/Model/Sample/CoreAndShellItem.h"
#include "GUI/Model/Sample/LayerItem.h"
#include "GUI/Model/Sample/MesocrystalItem.h"
#include "GUI/Model/Sample/ParticleItem.h"
#include "GUI/Model/Sample/ParticleLayoutItem.h"
#include "GUI/Model/Sample/RoughnessItems.h"
#include "GUI/Model/Sample/SampleItem.h"
#include "GUI/Model/Sim/InstrumentItems.h"

using std::variant;

namespace {

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

} // namespace


ParameterTreeBuilder::ParameterTreeBuilder(JobItem* jobItem)
    : m_job_item(jobItem)
{
    ASSERT(m_job_item);
}

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

void ParameterTreeBuilder::addMaterials()
{
    auto* materialTopLabel =
        new ParameterLabelItem("Materials", parameterContainerItem()->parameterTreeRoot());

    QVector<MaterialItem*> materials_in_use;
    for (const auto& item : m_job_item->sampleItem()->itemsWithMaterial())
        if (!materials_in_use.contains(item->materialItem()))
            materials_in_use.append(item->materialItem());

    for (auto* item : materials_in_use) {
        auto* label = new ParameterLabelItem(item->matItemName(), materialTopLabel);
        QColor color = item->color();
        color.setAlpha(100); // make more transparent
        label->setColor(color);

        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());

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

    int iLayer = 0;
    for (auto* layer : m_job_item->sampleItem()->uniqueLayerItems()) {
        auto* layerLabel = new ParameterLabelItem("Layer" + QString::number(iLayer++), label);
        layerLabel->setColor(layer->color());
        if (!layer->isAmbient() && !layer->isSubstrate())
            addParameterItem(layerLabel, layer->thickness());
        if (!layer->isAmbient())
            if (auto* roughnessItem = layer->roughnessSelection().certainItem()) {
                auto* roughnessLabel = new ParameterLabelItem("Top roughness", layerLabel);
                for (auto* property : roughnessItem->roughnessProperties())
                    addParameterItem(roughnessLabel, *property);
            }

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

            addInterference(label2, layout);

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

void ParameterTreeBuilder::addParameterItem(ParameterLabelItem* parent, DoubleProperty& d,
                                            const QString& label)
{
    auto* parameterItem = new ParameterItem(parent);
    parameterItem->setTitle(label.isEmpty() ? d.label() : label);
    parameterItem->linkToProperty(d);
}

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_job_item->parameterContainerItem();
}

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

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

    const auto itfType = InterferenceCatalog::type(interference);
    const QString title = InterferenceCatalog::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().certainItem();
        auto* pdfLabel = addLabel2<Profile1DItem, Profile1DCatalog>(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().certainItem();
        auto* pdf2 = itf->probabilityDistributionSelection2().certainItem();
        const bool samePdfTypes = Profile2DCatalog::type(pdf1) == Profile2DCatalog::type(pdf2);
        auto* pdf1Label =
            addLabel2<Profile2DItem, Profile2DCatalog>(label, samePdfTypes ? "PDF1" : "PDF", pdf1);
        for (auto* d : pdf1->profileProperties())
            addParameterItem(pdf1Label, *d);
        auto* pdf2Label =
            addLabel2<Profile2DItem, Profile2DCatalog>(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().certainItem();
        auto* dfLabel = addLabel2<Profile1DItem, Profile1DCatalog>(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().certainItem();
        auto* dfLabel = addLabel2<Profile2DItem, Profile2DCatalog>(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)
{
    auto* label = addLabel2<ItemWithParticles, ParticleCatalog>(parentLabel, "", p);

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

    addParameterItem(label, p->position());
    addRotation(label, p);

    if (const auto* particle = dynamic_cast<const ParticleItem*>(p)) {
        auto* formFactor = particle->formFactorItem();
        auto* ffLabel =
            addLabel2<FormfactorItem, FormfactorCatalog>(label, "Formfactor", formFactor);
        for (auto* d : formFactor->geometryProperties())
            addParameterItem(ffLabel, *d);
    } else if (const auto* particleComposition = dynamic_cast<const CompoundItem*>(p)) {
        for (auto* pp : particleComposition->itemsWithParticles())
            addItemWithParticles(label, pp, 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);
        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().certainItem();
        auto* ffLabel =
            addLabel2<FormfactorItem, FormfactorCatalog>(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 = addLabel2<Lattice2DItem, Lattice2DCatalog>(parentLabel, "Lattice", lattice);
    for (auto* d : lattice->geometryValues(!itf->xiIntegration()))
        addParameterItem(label, *d);
}

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

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

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

    if (auto* iiI = dynamic_cast<Scatter2DInstrumentItem*>(iI)) {
        auto* beamItem = iiI->beamItem();
        auto* beamLabel = new ParameterLabelItem("Beam", label);
        addParameterItem(beamLabel, beamItem->intensity());
        addBeamDistribution(beamLabel, beamItem->wavelengthItem(), "Wavelength");
        addBeamDistribution(beamLabel, beamItem->grazingAngleItem(), "Grazing angle");
        addBeamDistribution(beamLabel, beamItem->azimuthalAngleItem(), "Azimuthal angle");
        addFootprint(beamLabel, beamItem->footprintSelection().certainItem());
        addBackground(label, iiI->backgroundItem());
        addDetector(label, iiI->detectorItem());
        addPolarization(label, iiI);
    } else if (auto* iiI = dynamic_cast<SpecularInstrumentItem*>(iI)) {
        auto* scanItem = iiI->scanItem();
        auto* beamLabel = new ParameterLabelItem("Beam", label);
        addParameterItem(beamLabel, scanItem->intensity());
        addScanDistributions(beamLabel, scanItem);
        addScanFootprint(beamLabel, scanItem);
        addBackground(label, iiI->backgroundItem());
        addPolarization(label, iiI);
    } else if (auto* iiI = dynamic_cast<OffspecInstrumentItem*>(iI)) {
        auto* scanItem = iiI->scanItem();
        auto* beamLabel = new ParameterLabelItem("Beam", label);
        addParameterItem(beamLabel, scanItem->intensity());
        addScanDistributions(beamLabel, scanItem);
        addBeamDistribution(beamLabel, scanItem->azimuthalAngleItem(), "Azimuthal angle");
        addScanFootprint(beamLabel, scanItem);
        addBackground(label, iiI->backgroundItem());
        addOffspecDetector(label, iiI->detectorItem());
        addPolarization(label, iiI);
    } else if (auto* iiI = dynamic_cast<DepthprobeInstrumentItem*>(iI)) {
        auto* scanItem = iiI->scanItem();
        auto* beamLabel = new ParameterLabelItem("Beam", label);
        addParameterItem(beamLabel, scanItem->intensity());
        addScanDistributions(beamLabel, scanItem);
        addPolarization(label, iiI);
    } else
        ASSERT_NEVER;
}

void ParameterTreeBuilder::addBeamDistribution(ParameterLabelItem* parentLabel,
                                               BeamDistributionItem* distributionItem,
                                               const QString& label, bool withMean)
{
    auto* distribution = distributionItem->distributionItem();
    if (auto* dn = dynamic_cast<DistributionDeltaItem*>(distribution)) {
        if (withMean)
            addParameterItem(parentLabel, dn->mean(), label);
    } else {
        const auto type = DistributionCatalog::type(distribution);
        const auto name = DistributionCatalog::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::addScanDistributions(ParameterLabelItem* parentLabel, ScanItem* scanItem)
{
    ScanTypeItem* scan_type = scanItem->scanTypeSelection().certainItem();

    if (dynamic_cast<AlphaScanTypeItem*>(scan_type)) {
        addBeamDistribution(parentLabel, scanItem->wavelengthItem(), "Wavelength");
        addBeamDistribution(parentLabel, scanItem->scanDistributionItem(), "Grazing angle", false);
    } else if (dynamic_cast<LambdaScanTypeItem*>(scan_type)) {
        addBeamDistribution(parentLabel, scanItem->scanDistributionItem(), "Wavelength", false);
        addBeamDistribution(parentLabel, scanItem->grazingAngleItem(), "Grazing angle");
    } else if (dynamic_cast<QzScanTypeItem*>(scan_type)) {
        addBeamDistribution(parentLabel, scanItem->scanDistributionItem(), "Qz", false);
    } else
        ASSERT_NEVER;
}

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

    auto* label = new ParameterLabelItem("Detector", parentLabel);
    auto* phiLabel = new ParameterLabelItem("Phi span", label);
    addParameterItem(phiLabel, detector->phiAxis().min());
    addParameterItem(phiLabel, detector->phiAxis().max());
    auto* alphaLabel = new ParameterLabelItem("Alpha span", label);
    addParameterItem(alphaLabel, detector->alphaAxis().min());
    addParameterItem(alphaLabel, detector->alphaAxis().max());
    addResolutionFunction(label);
}

void ParameterTreeBuilder::addOffspecDetector(ParameterLabelItem* parentLabel,
                                              OffspecDetectorItem* detector)
{
    auto* label = new ParameterLabelItem("Detector", parentLabel);
    auto* phiLabel = new ParameterLabelItem("Phi axis", label);
    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());
}

void ParameterTreeBuilder::addFootprint(ParameterLabelItem* instrumentLabel,
                                        FootprintItem* footprintItem)
{
    if (auto* fg = dynamic_cast<FootprintGaussianItem*>(footprintItem))
        addParameterItem(instrumentLabel, fg->gaussianFootprintValue());
    else if (auto* fs = dynamic_cast<FootprintSquareItem*>(footprintItem))
        addParameterItem(instrumentLabel, fs->squareFootprintValue());
}

void ParameterTreeBuilder::addScanFootprint(ParameterLabelItem* instrumentLabel, ScanItem* scanItem)
{
    ScanTypeItem* scan_type = scanItem->scanTypeSelection().certainItem();
    if (dynamic_cast<QzScanTypeItem*>(scan_type))
        return;

    addFootprint(instrumentLabel, scanItem->footprintSelection().certainItem());
}

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());
}