/*===================================================================

The Medical Imaging Interaction Toolkit (MITK)

Copyright (c) German Cancer Research Center,
Division of Medical and Biological Informatics.
All rights reserved.

This software is distributed WITHOUT ANY WARRANTY; without
even the implied warranty of MERCHANTABILITY or FITNESS FOR
A PARTICULAR PURPOSE.

See LICENSE.txt or http://www.mitk.org for details.

===================================================================*/

#include "itkTractsToDWIImageFilter.h"
#include <boost/progress.hpp>
#include <vtkSmartPointer.h>
#include <vtkPolyData.h>
#include <vtkCellArray.h>
#include <vtkPoints.h>
#include <vtkPolyLine.h>
#include <itkImageRegionIteratorWithIndex.h>
#include <itkResampleImageFilter.h>
#include <itkNearestNeighborInterpolateImageFunction.h>
#include <itkBSplineInterpolateImageFunction.h>
#include <itkCastImageFilter.h>
#include <itkImageFileWriter.h>
#include <itkRescaleIntensityImageFilter.h>
#include <itkWindowedSincInterpolateImageFunction.h>
#include <itkResampleDwiImageFilter.h>
#include <itkKspaceImageFilter.h>
#include <itkDftImageFilter.h>
#include <itkAddImageFilter.h>
#include <itkConstantPadImageFilter.h>
#include <itkCropImageFilter.h>
#include <mitkAstroStickModel.h>
#include <vtkTransform.h>
#include <iostream>
#include <fstream>
#include <itkImageDuplicator.h>
#include <boost/lexical_cast.hpp>

namespace itk
{

template< class PixelType >
TractsToDWIImageFilter< PixelType >::TractsToDWIImageFilter()
    : m_CircleDummy(false)
    , m_VolumeAccuracy(10)
    , m_AddGibbsRinging(false)
    , m_NumberOfRepetitions(1)
    , m_EnforcePureFiberVoxels(false)
    , m_InterpolationShrink(1000)
    , m_FiberRadius(0)
    , m_SignalScale(25)
    , m_kOffset(0)
    , m_tLine(1)
    , m_UseInterpolation(false)
    , m_SimulateRelaxation(true)
    , m_tInhom(50)
    , m_TE(100)
    , m_FrequencyMap(NULL)
    , m_EddyGradientStrength(0.001)
    , m_SimulateEddyCurrents(false)
    , m_Spikes(0)
    , m_Wrap(1.0)
    , m_NoiseModel(NULL)
    , m_SpikeAmplitude(1)
    , m_AddMotionArtifact(false)
{
    m_Spacing.Fill(2.5); m_Origin.Fill(0.0);
    m_DirectionMatrix.SetIdentity();
    m_ImageRegion.SetSize(0, 10);
    m_ImageRegion.SetSize(1, 10);
    m_ImageRegion.SetSize(2, 10);

    m_MaxTranslation.Fill(0.0);
    m_MaxRotation.Fill(0.0);

    m_RandGen = itk::Statistics::MersenneTwisterRandomVariateGenerator::New();
    m_RandGen->SetSeed();
}

template< class PixelType >
TractsToDWIImageFilter< PixelType >::~TractsToDWIImageFilter()
{

}

template< class PixelType >
TractsToDWIImageFilter< PixelType >::DoubleDwiType::Pointer TractsToDWIImageFilter< PixelType >::DoKspaceStuff( std::vector< DoubleDwiType::Pointer >& images )
{
    // create slice object
    ImageRegion<2> sliceRegion;
    sliceRegion.SetSize(0, m_UpsampledImageRegion.GetSize()[0]);
    sliceRegion.SetSize(1, m_UpsampledImageRegion.GetSize()[1]);
    Vector< double, 2 > sliceSpacing;
    sliceSpacing[0] = m_UpsampledSpacing[0];
    sliceSpacing[1] = m_UpsampledSpacing[1];

    // frequency map slice
    SliceType::Pointer fMapSlice = NULL;
    if (m_FrequencyMap.IsNotNull())
    {
        fMapSlice = SliceType::New();
        ImageRegion<2> region;
        region.SetSize(0, m_UpsampledImageRegion.GetSize()[0]);
        region.SetSize(1, m_UpsampledImageRegion.GetSize()[1]);
        fMapSlice->SetLargestPossibleRegion( region );
        fMapSlice->SetBufferedRegion( region );
        fMapSlice->SetRequestedRegion( region );
        fMapSlice->Allocate();
    }

    DoubleDwiType::Pointer newImage = DoubleDwiType::New();
    newImage->SetSpacing( m_Spacing );
    newImage->SetOrigin( m_Origin );
    newImage->SetDirection( m_DirectionMatrix );
    newImage->SetLargestPossibleRegion( m_ImageRegion );
    newImage->SetBufferedRegion( m_ImageRegion );
    newImage->SetRequestedRegion( m_ImageRegion );
    newImage->SetVectorLength( images.at(0)->GetVectorLength() );
    newImage->Allocate();

    MatrixType transform = m_DirectionMatrix;
    for (int i=0; i<3; i++)
        for (int j=0; j<3; j++)
        {
            if (j<2)
                transform[i][j] *= m_UpsampledSpacing[j];
            else
                transform[i][j] *= m_Spacing[j];
        }

    std::vector< unsigned int > spikeVolume;
    for (int i=0; i<m_Spikes; i++)
        spikeVolume.push_back(rand()%images.at(0)->GetVectorLength());
    std::sort (spikeVolume.begin(), spikeVolume.end());
    std::reverse (spikeVolume.begin(), spikeVolume.end());

    m_StatusText += "0%   10   20   30   40   50   60   70   80   90   100%\n";
    m_StatusText += "|----|----|----|----|----|----|----|----|----|----|\n*";
    unsigned long lastTick = 0;

    boost::progress_display disp(2*images.at(0)->GetVectorLength()*images.at(0)->GetLargestPossibleRegion().GetSize(2));
    for (unsigned int g=0; g<images.at(0)->GetVectorLength(); g++)
    {
        std::vector< int > spikeSlice;
        while (!spikeVolume.empty() && spikeVolume.back()==g)
        {
            spikeSlice.push_back(rand()%images.at(0)->GetLargestPossibleRegion().GetSize(2));
            spikeVolume.pop_back();
        }
        std::sort (spikeSlice.begin(), spikeSlice.end());
        std::reverse (spikeSlice.begin(), spikeSlice.end());

        for (unsigned int z=0; z<images.at(0)->GetLargestPossibleRegion().GetSize(2); z++)
        {
            std::vector< SliceType::Pointer > compartmentSlices;
            std::vector< double > t2Vector;

            for (unsigned int i=0; i<images.size(); i++)
            {
                DiffusionSignalModel<double>* signalModel;
                if (i<m_FiberModels.size())
                    signalModel = m_FiberModels.at(i);
                else
                    signalModel = m_NonFiberModels.at(i-m_FiberModels.size());

                SliceType::Pointer slice = SliceType::New();
                slice->SetLargestPossibleRegion( sliceRegion );
                slice->SetBufferedRegion( sliceRegion );
                slice->SetRequestedRegion( sliceRegion );
                slice->SetSpacing(sliceSpacing);
                slice->Allocate();
                slice->FillBuffer(0.0);

                // extract slice from channel g
                for (unsigned int y=0; y<images.at(0)->GetLargestPossibleRegion().GetSize(1); y++)
                    for (unsigned int x=0; x<images.at(0)->GetLargestPossibleRegion().GetSize(0); x++)
                    {
                        SliceType::IndexType index2D; index2D[0]=x; index2D[1]=y;
                        DoubleDwiType::IndexType index3D; index3D[0]=x; index3D[1]=y; index3D[2]=z;

                        slice->SetPixel(index2D, images.at(i)->GetPixel(index3D)[g]);

                        if (fMapSlice.IsNotNull() && i==0)
                            fMapSlice->SetPixel(index2D, m_FrequencyMap->GetPixel(index3D));
                    }

                compartmentSlices.push_back(slice);
                t2Vector.push_back(signalModel->GetT2());
            }

            if (this->GetAbortGenerateData())
                return NULL;

            // create k-sapce (inverse fourier transform slices)
            itk::Size<2> outSize; outSize.SetElement(0, m_ImageRegion.GetSize(0)); outSize.SetElement(1, m_ImageRegion.GetSize(1));
            itk::KspaceImageFilter< SliceType::PixelType >::Pointer idft = itk::KspaceImageFilter< SliceType::PixelType >::New();
            idft->SetCompartmentImages(compartmentSlices);
            idft->SetT2(t2Vector);
            idft->SetkOffset(m_kOffset);
            idft->SettLine(m_tLine);
            idft->SetTE(m_TE);
            idft->SetTinhom(m_tInhom);
            idft->SetSimulateRelaxation(m_SimulateRelaxation);
            idft->SetSimulateEddyCurrents(m_SimulateEddyCurrents);
            idft->SetEddyGradientMagnitude(m_EddyGradientStrength);
            idft->SetZ((double)z-(double)images.at(0)->GetLargestPossibleRegion().GetSize(2)/2.0);
            idft->SetDirectionMatrix(transform);
            idft->SetDiffusionGradientDirection(m_FiberModels.at(0)->GetGradientDirection(g));
            idft->SetFrequencyMap(fMapSlice);
            idft->SetSignalScale(m_SignalScale);
            idft->SetOutSize(outSize);
            int numSpikes = 0;
            while (!spikeSlice.empty() && spikeSlice.back()==z)
            {
                numSpikes++;
                spikeSlice.pop_back();
            }
            idft->SetSpikes(numSpikes);
            idft->SetSpikeAmplitude(m_SpikeAmplitude);
            idft->Update();

            ComplexSliceType::Pointer fSlice;
            fSlice = idft->GetOutput();

            ++disp;
            unsigned long newTick = 50*disp.count()/disp.expected_count();
            for (int tick = 0; tick<(newTick-lastTick); tick++)
                m_StatusText += "*";
            lastTick = newTick;

            // fourier transform slice
            SliceType::Pointer newSlice;
            itk::DftImageFilter< SliceType::PixelType >::Pointer dft = itk::DftImageFilter< SliceType::PixelType >::New();
            dft->SetInput(fSlice);
            dft->Update();
            newSlice = dft->GetOutput();

            // put slice back into channel g
            for (unsigned int y=0; y<fSlice->GetLargestPossibleRegion().GetSize(1); y++)
                for (unsigned int x=0; x<fSlice->GetLargestPossibleRegion().GetSize(0); x++)
                {
                    DoubleDwiType::IndexType index3D; index3D[0]=x; index3D[1]=y; index3D[2]=z;
                    SliceType::IndexType index2D; index2D[0]=x; index2D[1]=y;

                    DoubleDwiType::PixelType pix3D = newImage->GetPixel(index3D);
                    pix3D[g] = newSlice->GetPixel(index2D);
                    newImage->SetPixel(index3D, pix3D);
                }
            ++disp;
            newTick = 50*disp.count()/disp.expected_count();
            for (int tick = 0; tick<(newTick-lastTick); tick++)
                m_StatusText += "*";
            lastTick = newTick;
        }
    }
    m_StatusText += "\n\n";
    return newImage;
}

template< class PixelType >
void TractsToDWIImageFilter< PixelType >::GenerateData()
{
    m_StartTime = clock();
    m_StatusText = "Starting simulation\n";

    // check input data
    if (m_FiberBundle.IsNull())
        itkExceptionMacro("Input fiber bundle is NULL!");

    int numFibers = m_FiberBundle->GetNumFibers();
    if (numFibers<=0)
        itkExceptionMacro("Input fiber bundle contains no fibers!");

    if (m_FiberModels.empty())
        itkExceptionMacro("No diffusion model for fiber compartments defined!");

    if (m_EnforcePureFiberVoxels)
        while (m_FiberModels.size()>1)
            m_FiberModels.pop_back();

    if (m_NonFiberModels.empty())
        itkExceptionMacro("No diffusion model for non-fiber compartments defined!");

    int baselineIndex = m_FiberModels[0]->GetFirstBaselineIndex();
    if (baselineIndex<0)
        itkExceptionMacro("No baseline index found!");

    // initialize output dwi image
    ImageRegion<3> croppedRegion = m_ImageRegion; croppedRegion.SetSize(1, croppedRegion.GetSize(1)*m_Wrap);
    itk::Point<double,3> shiftedOrigin = m_Origin; shiftedOrigin[1] += (m_ImageRegion.GetSize(1)-croppedRegion.GetSize(1))*m_Spacing[1]/2;

    typename OutputImageType::Pointer outImage = OutputImageType::New();
    outImage->SetSpacing( m_Spacing );
    outImage->SetOrigin( shiftedOrigin );
    outImage->SetDirection( m_DirectionMatrix );
    outImage->SetLargestPossibleRegion( croppedRegion );
    outImage->SetBufferedRegion( croppedRegion );
    outImage->SetRequestedRegion( croppedRegion );
    outImage->SetVectorLength( m_FiberModels[0]->GetNumGradients() );
    outImage->Allocate();
    typename OutputImageType::PixelType temp;
    temp.SetSize(m_FiberModels[0]->GetNumGradients());
    temp.Fill(0.0);
    outImage->FillBuffer(temp);

    // ADJUST GEOMETRY FOR FURTHER PROCESSING
    // is input slize size a power of two?
    unsigned int x=m_ImageRegion.GetSize(0); unsigned int y=m_ImageRegion.GetSize(1);
    if ( x%2 == 1 )
        m_ImageRegion.SetSize(0, x+1);
    if ( y%2 == 1 )
        m_ImageRegion.SetSize(1, y+1);

    // apply in-plane upsampling
    double upsampling = 1;
    if (m_AddGibbsRinging)
    {
        m_StatusText += "Gibbs ringing enabled\n";
        MITK_INFO << "Adding ringing artifacts.";
        upsampling = 2;
    }
    m_UpsampledSpacing = m_Spacing;
    m_UpsampledSpacing[0] /= upsampling;
    m_UpsampledSpacing[1] /= upsampling;
    m_UpsampledImageRegion = m_ImageRegion;
    m_UpsampledImageRegion.SetSize(0, m_ImageRegion.GetSize()[0]*upsampling);
    m_UpsampledImageRegion.SetSize(1, m_ImageRegion.GetSize()[1]*upsampling);
    m_UpsampledOrigin = m_Origin;
    m_UpsampledOrigin[0] -= m_Spacing[0]/2; m_UpsampledOrigin[0] += m_UpsampledSpacing[0]/2;
    m_UpsampledOrigin[1] -= m_Spacing[1]/2; m_UpsampledOrigin[1] += m_UpsampledSpacing[1]/2;
    m_UpsampledOrigin[2] -= m_Spacing[2]/2; m_UpsampledOrigin[2] += m_UpsampledSpacing[2]/2;

    // generate double images to store the individual compartment signals
    std::vector< DoubleDwiType::Pointer > compartments;
    for (unsigned int i=0; i<m_FiberModels.size()+m_NonFiberModels.size(); i++)
    {
        DoubleDwiType::Pointer doubleDwi = DoubleDwiType::New();
        doubleDwi->SetSpacing( m_UpsampledSpacing );
        doubleDwi->SetOrigin( m_UpsampledOrigin );
        doubleDwi->SetDirection( m_DirectionMatrix );
        doubleDwi->SetLargestPossibleRegion( m_UpsampledImageRegion );
        doubleDwi->SetBufferedRegion( m_UpsampledImageRegion );
        doubleDwi->SetRequestedRegion( m_UpsampledImageRegion );
        doubleDwi->SetVectorLength( m_FiberModels[0]->GetNumGradients() );
        doubleDwi->Allocate();
        DoubleDwiType::PixelType pix;
        pix.SetSize(m_FiberModels[0]->GetNumGradients());
        pix.Fill(0.0);
        doubleDwi->FillBuffer(pix);
        compartments.push_back(doubleDwi);
    }

    // initialize volume fraction images
    m_VolumeFractions.clear();
    for (unsigned int i=0; i<m_FiberModels.size()+m_NonFiberModels.size(); i++)
    {
        ItkDoubleImgType::Pointer doubleImg = ItkDoubleImgType::New();
        doubleImg->SetSpacing( m_UpsampledSpacing );
        doubleImg->SetOrigin( m_UpsampledOrigin );
        doubleImg->SetDirection( m_DirectionMatrix );
        doubleImg->SetLargestPossibleRegion( m_UpsampledImageRegion );
        doubleImg->SetBufferedRegion( m_UpsampledImageRegion );
        doubleImg->SetRequestedRegion( m_UpsampledImageRegion );
        doubleImg->Allocate();
        doubleImg->FillBuffer(0);
        m_VolumeFractions.push_back(doubleImg);
    }

    // resample mask image and frequency map to fit upsampled geometry
    if (m_AddGibbsRinging)
    {
        if (m_TissueMask.IsNotNull())
        {
            // rescale mask image (otherwise there are problems with the resampling)
            itk::RescaleIntensityImageFilter<ItkUcharImgType,ItkUcharImgType>::Pointer rescaler = itk::RescaleIntensityImageFilter<ItkUcharImgType,ItkUcharImgType>::New();
            rescaler->SetInput(0,m_TissueMask);
            rescaler->SetOutputMaximum(100);
            rescaler->SetOutputMinimum(0);
            rescaler->Update();

            // resample mask image
            itk::ResampleImageFilter<ItkUcharImgType, ItkUcharImgType>::Pointer resampler = itk::ResampleImageFilter<ItkUcharImgType, ItkUcharImgType>::New();
            resampler->SetInput(rescaler->GetOutput());
            resampler->SetOutputParametersFromImage(m_TissueMask);
            resampler->SetSize(m_UpsampledImageRegion.GetSize());
            resampler->SetOutputSpacing(m_UpsampledSpacing);
            resampler->SetOutputOrigin(m_UpsampledOrigin);
            resampler->Update();
            m_TissueMask = resampler->GetOutput();
        }

        // resample frequency map
        if (m_FrequencyMap.IsNotNull())
        {
            itk::ResampleImageFilter<ItkDoubleImgType, ItkDoubleImgType>::Pointer resampler = itk::ResampleImageFilter<ItkDoubleImgType, ItkDoubleImgType>::New();
            resampler->SetInput(m_FrequencyMap);
            resampler->SetOutputParametersFromImage(m_FrequencyMap);
            resampler->SetSize(m_UpsampledImageRegion.GetSize());
            resampler->SetOutputSpacing(m_UpsampledSpacing);
            resampler->SetOutputOrigin(m_UpsampledOrigin);
            resampler->Update();
            m_FrequencyMap = resampler->GetOutput();
        }
    }

    // no input tissue mask is set -> create default
    bool maskImageSet = true;
    if (m_TissueMask.IsNull())
    {
        m_StatusText += "No tissue mask set\n";
        MITK_INFO << "No tissue mask set";
        m_TissueMask = ItkUcharImgType::New();
        m_TissueMask->SetSpacing( m_UpsampledSpacing );
        m_TissueMask->SetOrigin( m_UpsampledOrigin );
        m_TissueMask->SetDirection( m_DirectionMatrix );
        m_TissueMask->SetLargestPossibleRegion( m_UpsampledImageRegion );
        m_TissueMask->SetBufferedRegion( m_UpsampledImageRegion );
        m_TissueMask->SetRequestedRegion( m_UpsampledImageRegion );
        m_TissueMask->Allocate();
        m_TissueMask->FillBuffer(1);
        maskImageSet = false;
    }
    else
    {
        m_StatusText += "Using tissue mask\n";
        MITK_INFO << "Using tissue mask";
    }

    m_ImageRegion = croppedRegion;
    x=m_ImageRegion.GetSize(0); y=m_ImageRegion.GetSize(1);
    if ( x%2 == 1 )
        m_ImageRegion.SetSize(0, x+1);
    if ( y%2 == 1 )
        m_ImageRegion.SetSize(1, y+1);

    // resample fiber bundle for sufficient voxel coverage
    m_StatusText += "\n"+this->GetTime()+" > Resampling fibers ...\n";
    double segmentVolume = 0.0001;
    float minSpacing = 1;
    if(m_UpsampledSpacing[0]<m_UpsampledSpacing[1] && m_UpsampledSpacing[0]<m_UpsampledSpacing[2])
        minSpacing = m_UpsampledSpacing[0];
    else if (m_UpsampledSpacing[1] < m_UpsampledSpacing[2])
        minSpacing = m_UpsampledSpacing[1];
    else
        minSpacing = m_UpsampledSpacing[2];
    FiberBundleType fiberBundle = m_FiberBundle->GetDeepCopy();
    fiberBundle->ResampleFibers(minSpacing/m_VolumeAccuracy);
    double mmRadius = m_FiberRadius/1000;
    if (mmRadius>0)
        segmentVolume = M_PI*mmRadius*mmRadius*minSpacing/m_VolumeAccuracy;

    double interpFact = 2*atan(-0.5*m_InterpolationShrink);
    double maxVolume = 0;
    double voxelVolume = m_UpsampledSpacing[0]*m_UpsampledSpacing[1]*m_UpsampledSpacing[2];

    if (m_AddMotionArtifact)
    {
        if (m_RandomMotion)
        {
            m_StatusText += "Adding random motion artifacts:\n";
            m_StatusText += "Maximum rotation: +/-" + boost::lexical_cast<std::string>(m_MaxRotation) + "°\n";
            m_StatusText += "Maximum translation: +/-" + boost::lexical_cast<std::string>(m_MaxTranslation) + "mm\n";
        }
        else
        {
            m_StatusText += "Adding linear motion artifacts:\n";
            m_StatusText += "Maximum rotation: " + boost::lexical_cast<std::string>(m_MaxRotation) + "°\n";
            m_StatusText += "Maximum translation: " + boost::lexical_cast<std::string>(m_MaxTranslation) + "mm\n";
        }
        MITK_INFO << "Adding motion artifacts";
        MITK_INFO << "Maximum rotation: " << m_MaxRotation;
        MITK_INFO << "Maxmimum translation: " << m_MaxTranslation;
    }
    maxVolume = 0;

    m_StatusText += "\n"+this->GetTime()+" > Generating signal of " + boost::lexical_cast<std::string>(m_FiberModels.size()) + " fiber compartments\n";
    MITK_INFO << "Generating signal of " << m_FiberModels.size() << " fiber compartments";
    boost::progress_display disp(numFibers*m_FiberModels.at(0)->GetNumGradients());

    ofstream logFile;
    logFile.open("fiberfox_motion.log");
    logFile << "0 rotation: 0,0,0; translation: 0,0,0\n";

    // get transform for motion artifacts
    FiberBundleType fiberBundleTransformed = fiberBundle;
    VectorType rotation = m_MaxRotation/m_FiberModels.at(0)->GetNumGradients();
    VectorType translation = m_MaxTranslation/m_FiberModels.at(0)->GetNumGradients();

    // creat image to hold transformed mask (motion artifact)
    ItkUcharImgType::Pointer tempTissueMask = ItkUcharImgType::New();
    itk::ImageDuplicator<ItkUcharImgType>::Pointer duplicator = itk::ImageDuplicator<ItkUcharImgType>::New();
    duplicator->SetInputImage(m_TissueMask);
    duplicator->Update();
    tempTissueMask = duplicator->GetOutput();

    // second upsampling needed for motion artifacts
    ImageRegion<3>                      upsampledImageRegion = m_UpsampledImageRegion;
    itk::Vector<double,3>               upsampledSpacing = m_UpsampledSpacing;
    upsampledSpacing[0] /= 4;
    upsampledSpacing[1] /= 4;
    upsampledSpacing[2] /= 4;
    upsampledImageRegion.SetSize(0, m_UpsampledImageRegion.GetSize()[0]*4);
    upsampledImageRegion.SetSize(1, m_UpsampledImageRegion.GetSize()[1]*4);
    upsampledImageRegion.SetSize(2, m_UpsampledImageRegion.GetSize()[2]*4);
    itk::Point<double,3> upsampledOrigin = m_UpsampledOrigin;
    upsampledOrigin[0] -= m_UpsampledSpacing[0]/2; upsampledOrigin[0] += upsampledSpacing[0]/2;
    upsampledOrigin[1] -= m_UpsampledSpacing[1]/2; upsampledOrigin[1] += upsampledSpacing[1]/2;
    upsampledOrigin[2] -= m_UpsampledSpacing[2]/2; upsampledOrigin[2] += upsampledSpacing[2]/2;
    ItkUcharImgType::Pointer upsampledTissueMask = ItkUcharImgType::New();
    itk::ResampleImageFilter<ItkUcharImgType, ItkUcharImgType>::Pointer upsampler = itk::ResampleImageFilter<ItkUcharImgType, ItkUcharImgType>::New();
    upsampler->SetInput(m_TissueMask);
    upsampler->SetOutputParametersFromImage(m_TissueMask);
    upsampler->SetSize(upsampledImageRegion.GetSize());
    upsampler->SetOutputSpacing(upsampledSpacing);
    upsampler->SetOutputOrigin(upsampledOrigin);
    itk::NearestNeighborInterpolateImageFunction<ItkUcharImgType>::Pointer nn_interpolator
            = itk::NearestNeighborInterpolateImageFunction<ItkUcharImgType>::New();
    upsampler->SetInterpolator(nn_interpolator);
    upsampler->Update();
    upsampledTissueMask = upsampler->GetOutput();

    m_StatusText += "0%   10   20   30   40   50   60   70   80   90   100%\n";
    m_StatusText += "|----|----|----|----|----|----|----|----|----|----|\n*";
    unsigned int lastTick = 0;

    for (int g=0; g<m_FiberModels.at(0)->GetNumGradients(); g++)
    {
        vtkPolyData* fiberPolyData = fiberBundleTransformed->GetFiberPolyData();

        ItkDoubleImgType::Pointer intraAxonalVolume = ItkDoubleImgType::New();
        intraAxonalVolume->SetSpacing( m_UpsampledSpacing );
        intraAxonalVolume->SetOrigin( m_UpsampledOrigin );
        intraAxonalVolume->SetDirection( m_DirectionMatrix );
        intraAxonalVolume->SetLargestPossibleRegion( m_UpsampledImageRegion );
        intraAxonalVolume->SetBufferedRegion( m_UpsampledImageRegion );
        intraAxonalVolume->SetRequestedRegion( m_UpsampledImageRegion );
        intraAxonalVolume->Allocate();
        intraAxonalVolume->FillBuffer(0);

        // generate fiber signal
        for( int i=0; i<numFibers; i++ )
        {
            vtkCell* cell = fiberPolyData->GetCell(i);
            int numPoints = cell->GetNumberOfPoints();
            vtkPoints* points = cell->GetPoints();

            if (numPoints<2)
                continue;

            for( int j=0; j<numPoints; j++)
            {
                if (this->GetAbortGenerateData())
                {
                    m_StatusText += "\n"+this->GetTime()+" > Simulation aborted\n";
                    return;
                }

                double* temp = points->GetPoint(j);
                itk::Point<float, 3> vertex = GetItkPoint(temp);
                itk::Vector<double> v = GetItkVector(temp);

                itk::Vector<double, 3> dir(3);
                if (j<numPoints-1)
                    dir = GetItkVector(points->GetPoint(j+1))-v;
                else
                    dir = v-GetItkVector(points->GetPoint(j-1));

                if (dir.GetSquaredNorm()<0.0001 || dir[0]!=dir[0] || dir[1]!=dir[1] || dir[2]!=dir[2])
                    continue;

                itk::Index<3> idx;
                itk::ContinuousIndex<float, 3> contIndex;
                tempTissueMask->TransformPhysicalPointToIndex(vertex, idx);
                tempTissueMask->TransformPhysicalPointToContinuousIndex(vertex, contIndex);

                if (!m_UseInterpolation)    // use nearest neighbour interpolation
                {
                    if (!tempTissueMask->GetLargestPossibleRegion().IsInside(idx) || tempTissueMask->GetPixel(idx)<=0)
                        continue;

                    // generate signal for each fiber compartment
                    for (unsigned int k=0; k<m_FiberModels.size(); k++)
                    {
                        DoubleDwiType::Pointer doubleDwi = compartments.at(k);
                        m_FiberModels[k]->SetFiberDirection(dir);
                        DoubleDwiType::PixelType pix = doubleDwi->GetPixel(idx);
                        pix[g] += segmentVolume*m_FiberModels[k]->SimulateMeasurement(g);

                        if (pix[g]!=pix[g])
                        {
                            std::cout << "pix[g] " << pix[g] << std::endl;
                            std::cout << "dir " << dir << std::endl;
                            std::cout << "segmentVolume " << segmentVolume << std::endl;
                            std::cout << "m_FiberModels[k]->SimulateMeasurement(g) " << m_FiberModels[k]->SimulateMeasurement(g) << std::endl;
                        }

                        doubleDwi->SetPixel(idx, pix );

                        double vol = intraAxonalVolume->GetPixel(idx) + segmentVolume;
                        intraAxonalVolume->SetPixel(idx, vol );

                        if (g==0 && vol>maxVolume)
                            maxVolume = vol;
                    }
                    continue;
                }

                double frac_x = contIndex[0] - idx[0]; double frac_y = contIndex[1] - idx[1]; double frac_z = contIndex[2] - idx[2];
                if (frac_x<0)
                {
                    idx[0] -= 1;
                    frac_x += 1;
                }
                if (frac_y<0)
                {
                    idx[1] -= 1;
                    frac_y += 1;
                }
                if (frac_z<0)
                {
                    idx[2] -= 1;
                    frac_z += 1;
                }

                frac_x = atan((0.5-frac_x)*m_InterpolationShrink)/interpFact + 0.5;
                frac_y = atan((0.5-frac_y)*m_InterpolationShrink)/interpFact + 0.5;
                frac_z = atan((0.5-frac_z)*m_InterpolationShrink)/interpFact + 0.5;

                // use trilinear interpolation
                itk::Index<3> newIdx;
                for (int x=0; x<2; x++)
                {
                    frac_x = 1-frac_x;
                    for (int y=0; y<2; y++)
                    {
                        frac_y = 1-frac_y;
                        for (int z=0; z<2; z++)
                        {
                            frac_z = 1-frac_z;

                            newIdx[0] = idx[0]+x;
                            newIdx[1] = idx[1]+y;
                            newIdx[2] = idx[2]+z;

                            double frac = frac_x*frac_y*frac_z;

                            // is position valid?
                            if (!tempTissueMask->GetLargestPossibleRegion().IsInside(newIdx) || tempTissueMask->GetPixel(newIdx)<=0)
                                continue;

                            // generate signal for each fiber compartment
                            for (unsigned int k=0; k<m_FiberModels.size(); k++)
                            {
                                DoubleDwiType::Pointer doubleDwi = compartments.at(k);
                                m_FiberModels[k]->SetFiberDirection(dir);
                                DoubleDwiType::PixelType pix = doubleDwi->GetPixel(newIdx);
                                pix[g] += segmentVolume*frac*m_FiberModels[k]->SimulateMeasurement(g);
                                doubleDwi->SetPixel(newIdx, pix );

                                double vol = intraAxonalVolume->GetPixel(idx) + segmentVolume;
                                intraAxonalVolume->SetPixel(idx, vol );

                                if (g==0 && vol>maxVolume)
                                    maxVolume = vol;
                            }
                        }
                    }
                }
            }
            ++disp;
            unsigned long newTick = 50*disp.count()/disp.expected_count();
            for (int tick = 0; tick<(newTick-lastTick); tick++)
                m_StatusText += "*";
            lastTick = newTick;
        }

        // generate non-fiber signal
        ImageRegionIterator<ItkUcharImgType> it3(tempTissueMask, tempTissueMask->GetLargestPossibleRegion());
        double fact = 1;
        if (m_FiberRadius<0.0001)
            fact = voxelVolume/maxVolume;
        while(!it3.IsAtEnd())
        {
            if (it3.Get()>0)
            {
                DoubleDwiType::IndexType index = it3.GetIndex();

                // get fiber volume fraction
                DoubleDwiType::Pointer fiberDwi = compartments.at(0);
                DoubleDwiType::PixelType fiberPix = fiberDwi->GetPixel(index); // intra axonal compartment
                if (fact>1) // auto scale intra-axonal if no fiber radius is specified
                {
                    fiberPix[g] *= fact;
                    fiberDwi->SetPixel(index, fiberPix);
                }
                double f = intraAxonalVolume->GetPixel(index)*fact;

                if (f>voxelVolume || (f>0.0 && m_EnforcePureFiberVoxels) )  // more fiber than space in voxel?
                {
                    fiberPix[g] *= voxelVolume/f;
                    fiberDwi->SetPixel(index, fiberPix);
                    m_VolumeFractions.at(0)->SetPixel(index, 1);
                }
                else
                {
                    m_VolumeFractions.at(0)->SetPixel(index, f/voxelVolume);

                    double nonf = voxelVolume-f;    // non-fiber volume
                    double inter = 0;
                    if (m_FiberModels.size()>1)
                        inter = nonf * f/voxelVolume;   // inter-axonal fraction of non fiber compartment scales linearly with f
                    double other = nonf - inter;        // rest of compartment
                    double singleinter = inter/(m_FiberModels.size()-1);

                    // adjust non-fiber and intra-axonal signal
                    for (unsigned int i=1; i<m_FiberModels.size(); i++)
                    {
                        DoubleDwiType::Pointer doubleDwi = compartments.at(i);
                        DoubleDwiType::PixelType pix = doubleDwi->GetPixel(index);
                        if (f>0)
                            pix[g] /= f;
                        pix[g] *= singleinter;
                        doubleDwi->SetPixel(index, pix);
                        m_VolumeFractions.at(i)->SetPixel(index, singleinter/voxelVolume);
                    }
                    for (unsigned int i=0; i<m_NonFiberModels.size(); i++)
                    {
                        DoubleDwiType::Pointer doubleDwi = compartments.at(i+m_FiberModels.size());
                        DoubleDwiType::PixelType pix = doubleDwi->GetPixel(index);
                        //                        if (dynamic_cast< mitk::AstroStickModel<double>* >(m_NonFiberModels.at(i)))
                        //                        {
                        //                            mitk::AstroStickModel<double>* model = dynamic_cast< mitk::AstroStickModel<double>* >(m_NonFiberModels.at(i));
                        //                            model->SetSeed(8111984);
                        //                        }
                        pix[g] += m_NonFiberModels[i]->SimulateMeasurement(g)*other*m_NonFiberModels[i]->GetWeight();
                        doubleDwi->SetPixel(index, pix);
                        m_VolumeFractions.at(i+m_FiberModels.size())->SetPixel(index, other/voxelVolume*m_NonFiberModels[i]->GetWeight());
                    }
                }
            }
            ++it3;
        }

        // move fibers
        if (m_AddMotionArtifact)
        {
            if (m_RandomMotion)
            {
                fiberBundleTransformed = fiberBundle->GetDeepCopy();
                rotation[0] = m_RandGen->GetVariateWithClosedRange(m_MaxRotation[0]*2)-m_MaxRotation[0];
                rotation[1] = m_RandGen->GetVariateWithClosedRange(m_MaxRotation[1]*2)-m_MaxRotation[1];
                rotation[2] = m_RandGen->GetVariateWithClosedRange(m_MaxRotation[2]*2)-m_MaxRotation[2];
                translation[0] = m_RandGen->GetVariateWithClosedRange(m_MaxTranslation[0]*2)-m_MaxTranslation[0];
                translation[1] = m_RandGen->GetVariateWithClosedRange(m_MaxTranslation[1]*2)-m_MaxTranslation[1];
                translation[2] = m_RandGen->GetVariateWithClosedRange(m_MaxTranslation[2]*2)-m_MaxTranslation[2];
            }

            // rotate mask image
            if (maskImageSet)
            {
                ImageRegionIterator<ItkUcharImgType> maskIt(upsampledTissueMask, upsampledTissueMask->GetLargestPossibleRegion());
                tempTissueMask->FillBuffer(0);

                while(!maskIt.IsAtEnd())
                {
                    if (maskIt.Get()<=0)
                    {
                        ++maskIt;
                        continue;
                    }

                    DoubleDwiType::IndexType index = maskIt.GetIndex();
                    itk::Point<double, 3> point;
                    upsampledTissueMask->TransformIndexToPhysicalPoint(index, point);
                    if (m_RandomMotion)
                        point = fiberBundle->TransformPoint(point.GetVnlVector(), rotation[0],rotation[1],rotation[2],translation[0],translation[1],translation[2]);
                    else
                        point = fiberBundle->TransformPoint(point.GetVnlVector(), rotation[0]*(g+1),rotation[1]*(g+1),rotation[2]*(g+1),translation[0]*(g+1),translation[1]*(g+1),translation[2]*(g+1));

                    tempTissueMask->TransformPhysicalPointToIndex(point, index);
                    if (tempTissueMask->GetLargestPossibleRegion().IsInside(index))
                        tempTissueMask->SetPixel(index,100);
                    ++maskIt;
                }
            }

            // rotate fibers
            logFile << g+1 << " rotation:" << rotation[0] << "," << rotation[1] << "," << rotation[2] << ";";
            logFile << " translation:" << translation[0] << "," << translation[1] << "," << translation[2] << "\n";
            fiberBundleTransformed->TransformFibers(rotation[0],rotation[1],rotation[2],translation[0],translation[1],translation[2]);
        }
    }
    logFile.close();
    m_StatusText += "\n\n";
    if (this->GetAbortGenerateData())
    {
        m_StatusText += "\n"+this->GetTime()+" > Simulation aborted\n";
        return;
    }

    // do k-space stuff
    DoubleDwiType::Pointer doubleOutImage;
    if (m_Spikes>0 || m_FrequencyMap.IsNotNull() || m_kOffset>0 || m_SimulateRelaxation || m_SimulateEddyCurrents || m_AddGibbsRinging || m_Wrap<1.0)
    {
        m_StatusText += this->GetTime()+" > Adjusting complex signal\n";
        MITK_INFO << "Adjusting complex signal";
        doubleOutImage = DoKspaceStuff(compartments);
        m_SignalScale = 1;
    }
    else
    {
        m_StatusText += this->GetTime()+" > Summing compartments\n";
        MITK_INFO << "Summing compartments";
        doubleOutImage = compartments.at(0);

        for (unsigned int i=1; i<compartments.size(); i++)
        {
            itk::AddImageFilter< DoubleDwiType, DoubleDwiType, DoubleDwiType>::Pointer adder = itk::AddImageFilter< DoubleDwiType, DoubleDwiType, DoubleDwiType>::New();
            adder->SetInput1(doubleOutImage);
            adder->SetInput2(compartments.at(i));
            adder->Update();
            doubleOutImage = adder->GetOutput();
        }
    }
    if (this->GetAbortGenerateData())
    {
        m_StatusText += "\n"+this->GetTime()+" > Simulation aborted\n";
        return;
    }

    m_StatusText += this->GetTime()+" > Finalizing image\n";
    MITK_INFO << "Finalizing image";
    unsigned int window = 0;
    unsigned int min = itk::NumericTraits<unsigned int>::max();
    ImageRegionIterator<OutputImageType> it4 (outImage, outImage->GetLargestPossibleRegion());
    DoubleDwiType::PixelType signal; signal.SetSize(m_FiberModels[0]->GetNumGradients());
    boost::progress_display disp2(outImage->GetLargestPossibleRegion().GetNumberOfPixels());

    m_StatusText += "0%   10   20   30   40   50   60   70   80   90   100%\n";
    m_StatusText += "|----|----|----|----|----|----|----|----|----|----|\n*";
    lastTick = 0;

    while(!it4.IsAtEnd())
    {
        if (this->GetAbortGenerateData())
        {
            m_StatusText += "\n"+this->GetTime()+" > Simulation aborted\n";
            return;
        }

        ++disp2;
        unsigned long newTick = 50*disp2.count()/disp2.expected_count();
        for (int tick = 0; tick<(newTick-lastTick); tick++)
            m_StatusText += "*";
        lastTick = newTick;

        typename OutputImageType::IndexType index = it4.GetIndex();
        signal = doubleOutImage->GetPixel(index)*m_SignalScale;

        if (m_NoiseModel!=NULL)
        {
            DoubleDwiType::PixelType accu = signal; accu.Fill(0.0);
            for (unsigned int i=0; i<m_NumberOfRepetitions; i++)
            {
                DoubleDwiType::PixelType temp = signal;
                m_NoiseModel->AddNoise(temp);
                accu += temp;
            }
            signal = accu/m_NumberOfRepetitions;
        }

        for (unsigned int i=0; i<signal.Size(); i++)
        {
            if (signal[i]>0)
                signal[i] = floor(signal[i]+0.5);
            else
                signal[i] = ceil(signal[i]-0.5);

            if (!m_FiberModels.at(0)->IsBaselineIndex(i) && signal[i]>window)
                window = signal[i];
            if (!m_FiberModels.at(0)->IsBaselineIndex(i) && signal[i]<min)
                min = signal[i];
        }
        it4.Set(signal);
        ++it4;
    }
    window -= min;
    unsigned int level = window/2 + min;
    m_LevelWindow.SetLevelWindow(level, window);
    this->SetNthOutput(0, outImage);

    m_StatusText += "\n\n";
    m_StatusText += "Finished simulation\n";
    m_StatusText += "Simulation time: "+GetTime();
}

template< class PixelType >
itk::Point<float, 3> TractsToDWIImageFilter< PixelType >::GetItkPoint(double point[3])
{
    itk::Point<float, 3> itkPoint;
    itkPoint[0] = point[0];
    itkPoint[1] = point[1];
    itkPoint[2] = point[2];
    return itkPoint;
}

template< class PixelType >
itk::Vector<double, 3> TractsToDWIImageFilter< PixelType >::GetItkVector(double point[3])
{
    itk::Vector<double, 3> itkVector;
    itkVector[0] = point[0];
    itkVector[1] = point[1];
    itkVector[2] = point[2];
    return itkVector;
}

template< class PixelType >
vnl_vector_fixed<double, 3> TractsToDWIImageFilter< PixelType >::GetVnlVector(double point[3])
{
    vnl_vector_fixed<double, 3> vnlVector;
    vnlVector[0] = point[0];
    vnlVector[1] = point[1];
    vnlVector[2] = point[2];
    return vnlVector;
}

template< class PixelType >
vnl_vector_fixed<double, 3> TractsToDWIImageFilter< PixelType >::GetVnlVector(Vector<float,3>& vector)
{
    vnl_vector_fixed<double, 3> vnlVector;
    vnlVector[0] = vector[0];
    vnlVector[1] = vector[1];
    vnlVector[2] = vector[2];
    return vnlVector;
}

template< class PixelType >
std::string TractsToDWIImageFilter< PixelType >::GetTime()
{
    unsigned long total = (double)(clock() - m_StartTime)/CLOCKS_PER_SEC;
    unsigned long hours = total/3600;
    unsigned long minutes = (total%3600)/60;
    unsigned long seconds = total%60;
    std::string out = "";
    out.append(boost::lexical_cast<std::string>(hours));
    out.append(":");
    out.append(boost::lexical_cast<std::string>(minutes));
    out.append(":");
    out.append(boost::lexical_cast<std::string>(seconds));
    return out;
}

}


/**
 * @file
 * @copyright This code is licensed under the 3-clause BSD license.\n
 *            Copyright ETH Zurich, Laboratory for Physical Chemistry, Reiher Group.\n
 *            See LICENSE.txt for details.
 */

#include "NDDODipoleMatrixCalculator.h"
#include "Sparrow/Implementations/Nddo/Utils/DipoleUtils/AtomPairDipole.h"
#include "Sparrow/Implementations/Nddo/Utils/DipoleUtils/GTODipoleMatrixBlock.h"
#include <Sparrow/Implementations/Nddo/Am1/AM1Method.h>
#include <Sparrow/Implementations/Nddo/Mndo/MNDOMethod.h>
#include <Sparrow/Implementations/Nddo/Pm6/PM6Method.h>
#include <Sparrow/Implementations/Nddo/Utils/NDDOInitializer.h>
#include <Sparrow/Implementations/Nddo/Utils/ParameterUtils/ElementParameters.h>
#include <Utils/DataStructures/AtomsOrbitalsIndexes.h>
#include <Utils/Geometry.h>
#include <Utils/Typenames.h>
#include <Eigen/Core>

namespace Scine {
namespace Sparrow {

template<class NDDOMethod>
NDDODipoleMatrixCalculator<NDDOMethod>::~NDDODipoleMatrixCalculator() = default;

template<class NDDOMethod>
void NDDODipoleMatrixCalculator<NDDOMethod>::initialize() {
  nAtoms_ = aoIndexes_.getNAtoms();
  nAOs_ = aoIndexes_.getNAtomicOrbitals();
  dipoleMatrix_.reset(nAOs_);
  valid_ = false;
}
template<class NDDOMethod>
void NDDODipoleMatrixCalculator<NDDOMethod>::setIntegralMethod(const IntegralMethod& method) {
  integralMethod_ = method;
  valid_ = false;
}

template<class NDDOMethod>
void NDDODipoleMatrixCalculator<NDDOMethod>::fillDipoleMatrix(const Eigen::RowVector3d& dipoleEvaluationCoordinate) {
  initialize();
  valid_ = false;
  for (int atom_i = 0; atom_i < nAtoms_; ++atom_i) {
    auto const firstAOIndex_i = aoIndexes_.getFirstOrbitalIndex(atom_i);
    auto const& parA = elementParameters_.get(elementTypes_[atom_i]);
    auto const Ri = positions_.row(atom_i);

    for (int atom_j = atom_i; atom_j < nAtoms_; ++atom_j) {
      auto const firstAOIndex_j = aoIndexes_.getFirstOrbitalIndex(atom_j);
      auto const& parB = elementParameters_.get(elementTypes_[atom_j]);
      auto const Rj = positions_.row(atom_j);

      Eigen::RowVector3d Rij = Rj - Ri;

      AtomPairDipole::fillAtomPairDipoleBlock(dipoleMatrix_, firstAOIndex_i, firstAOIndex_j, integralMethod_,
                                              parA.GTOs(), parB.GTOs(), Ri, Rj, Rij, dipoleEvaluationCoordinate);
    }
  }
  valid_ = true;
}

template<class NDDOMethod>
const Utils::DipoleMatrix& NDDODipoleMatrixCalculator<NDDOMethod>::getAODipoleMatrix() const {
  return dipoleMatrix_;
}

template<class NDDOMethod>
void NDDODipoleMatrixCalculator<NDDOMethod>::setAODipoleMatrix(Utils::DipoleMatrix dipoleMatrix) {
  dipoleMatrix_ = std::move(dipoleMatrix);
  valid_ = true;
}

template<class NDDOMethod>
bool NDDODipoleMatrixCalculator<NDDOMethod>::isValid() const {
  return valid_;
}

template<class NDDOMethod>
NDDODipoleMatrixCalculator<NDDOMethod>::NDDODipoleMatrixCalculator(NDDOMethod& method)
  : aoIndexes_(method.getAtomsOrbitalsIndexesHolder()),
    elementTypes_(method.getElementTypes()),
    positions_(method.getPositions()),
    elementParameters_(method.getInitializer().getElementParameters()),
    overlapMatrix_(method.getOverlapMatrix()),
    molecularOrbitals_(method.getMolecularOrbitals()) {
  initialize();
}

template<class NDDOMethod>
std::unique_ptr<NDDODipoleMatrixCalculator<NDDOMethod>> NDDODipoleMatrixCalculator<NDDOMethod>::create(NDDOMethod& method) {
  NDDODipoleMatrixCalculator<NDDOMethod> instance(method);
  return std::make_unique<NDDODipoleMatrixCalculator<NDDOMethod>>(std::move(instance));
}

template<class NDDOMethod>
Utils::DipoleMatrix NDDODipoleMatrixCalculator<NDDOMethod>::getMODipoleMatrix() const {
  if (molecularOrbitals_.isRestricted())
    return calculateMODipoleMatrixRestricted();
  else
    return calculateMODipoleMatrixUnrestricted();
}

template<class NDDOMethod>
void NDDODipoleMatrixCalculator<NDDOMethod>::invalidate() {
  valid_ = false;
}

template<class NDDOMethod>
Utils::DipoleMatrix NDDODipoleMatrixCalculator<NDDOMethod>::calculateMODipoleMatrixRestricted() const {
  Utils::DipoleMatrix moDipoleMatrix;
  auto& MOMatrix = molecularOrbitals_.restrictedMatrix();
  moDipoleMatrix.reset(MOMatrix.cols());
  // Get the MO matrix as D_{MO} = C^T * D_{AO} * C
  for (int dimension = 0; dimension < 3; ++dimension) {
    moDipoleMatrix[dimension] = MOMatrix.transpose() * dipoleMatrix_[dimension] * MOMatrix;
  }
  return moDipoleMatrix;
}

template<class NDDOMethod>
Utils::DipoleMatrix NDDODipoleMatrixCalculator<NDDOMethod>::calculateMODipoleMatrixUnrestricted() const {
  Utils::DipoleMatrix moDipoleMatrix;
  auto& alphaMOMatrix = molecularOrbitals_.alphaMatrix();
  auto& betaMOMatrix = molecularOrbitals_.betaMatrix();
  moDipoleMatrix.reset(alphaMOMatrix.cols());
  // Get the MO matrix as D_{MO} = C_\alpha^T * D_{AO} * C_\alpha
  //                             + C_\beta^T  * D_{AO} * C_\beta
  for (int dimension = 0; dimension < 3; ++dimension) {
    auto alphaContribution = alphaMOMatrix.transpose() * dipoleMatrix_[dimension] * alphaMOMatrix;
    auto betaContribution = betaMOMatrix.transpose() * dipoleMatrix_[dimension] * betaMOMatrix;
    moDipoleMatrix[dimension] = alphaContribution + betaContribution;
  }
  return moDipoleMatrix;
}

template class NDDODipoleMatrixCalculator<nddo::PM6Method>;
template class NDDODipoleMatrixCalculator<nddo::AM1Method>;
template class NDDODipoleMatrixCalculator<nddo::MNDOMethod>;

} // namespace Sparrow
} // namespace Scine


#include <cctbx/boost_python/flex_fwd.h>

#include <cctbx/miller/phase_transfer.h>
#include <boost/python/def.hpp>

namespace cctbx { namespace miller { namespace boost_python {

  void wrap_phase_transfer()
  {
    using namespace boost::python;

    def("phase_transfer",
      (af::shared<std::complex<double> >(*)
        (sgtbx::space_group const&,
         af::const_ref<index<> > const&,
         af::const_ref<double> const&,
         af::const_ref<std::complex<double> > const&,
         double const&))
           phase_transfer);

    def("phase_transfer",
      (af::shared<std::complex<double> >(*)
        (sgtbx::space_group const&,
         af::const_ref<index<> > const&,
         af::const_ref<double> const&,
         af::const_ref<double> const&,
         bool))
           phase_transfer);
  }

}}} // namespace cctbx::miller::boost_python


#include <booster/BoosterNetworkSimulator.hpp>
#include <booster/BoosterImpactOptions.hpp>

#include <merlion/Merlion.hpp>
#include <merlion/BlasWrapper.hpp>
#include <merlion/TriSolve.hpp>
#include <merlion/MerlionDefines.hpp>
#include <merlion/SparseMatrix.hpp>

#include <merlionUtils/TaskTimer.hpp>

#include <set>
#include <exception>
#include <algorithm>

using namespace merlionUtils;

//forward declarations
void PerformBoosterSimulations(const BoosterNetworkSimulator& chl_net,
                               const NetworkSimulator& tox_net,
			       const BoosterImpactOptions& options,
                               bool pd_impacts=false);

int main(int argc, char** argv)
{
   // Handle the optional inputs to this executable
   BoosterImpactOptions options;
   options.ParseInputs(argc, argv);

   // Helpful interface for setting up the booster station simulations
   BoosterNetworkSimulator chl_net(options.disable_warnings);
   NetworkSimulator tox_net(true);

   if (options.logging) {
      chl_net.StartLogging(options.output_prefix+"boosterimpact.log");
      tox_net.StartLogging(options.output_prefix+"boosterimpact.log");
      options.PrintSummary(chl_net.Log());
   }

   std::cout << "\n@@@ PARSING INPUT FILES @@@" << std::endl;
   // Read Chlorine network info
   if ((options.chl_inp_filename != "") && (options.chl_wqm_filename != "")) {
      std::cerr << std::endl;
      std::cerr << "ERROR: Both a wqm file and an inp file cannot be specified for the boosters." << std::endl;
      std::cerr << std::endl;
      return 1;
   }
   if (options.chl_inp_filename != "") {
      // Read an epanet input file
      merlionUtils::TaskTimer Run_Hydraulics;
      chl_net.ReadINPFile(options.chl_inp_filename,
                          options.sim_duration_min,
                          options.qual_step_min,
                          options.chl_epanet_output_filename,
                          options.chl_merlion_save_filename,
                          -1,
                          -1,
                          options.ignore_merlion_warnings,
                          options.chl_decay_k);
      Run_Hydraulics.StopAndPrint(std::cout, "\n\t- Built booster Merlion Water Quality Model ");
   }
   else if (options.chl_wqm_filename != "") {
      if (options.chl_decay_k != -1.0f) {
         std::cerr << std::endl;
         std::cerr << "ERROR: A booster decay coefficient cannot be specified when using the --booster-wqm option." << std::endl;
         std::cerr << std::endl;
         return 1;
      }
      merlionUtils::TaskTimer Read_WQM;
      // read a merlion water quality model file
      chl_net.ReadWQMFile(options.chl_wqm_filename);
      Read_WQM.StopAndPrint(std::cout, "\n\t- Read booster Merlion Water Quality Model ");
   }
   else {
      std::cerr << std::endl;
      std::cerr << "ERROR: Missing network input file for the boosters." << std::endl;
      std::cerr << std::endl;
      return 1;
   }
   /////////////////

   // Read Toxin network info
   if ((options.tox_inp_filename != "") && (options.tox_wqm_filename != "")) {
      std::cerr << std::endl;
      std::cerr << "ERROR: Both a wqm file and an inp file cannot be specified for the toxin." << std::endl;
      std::cerr << std::endl;
      return 1;
   }
   if (options.tox_inp_filename != "") {
      // Read an epanet input file
      merlionUtils::TaskTimer Run_Hydraulics;
      tox_net.ReadINPFile(options.tox_inp_filename,
                          options.sim_duration_min,
                          options.qual_step_min,
                          options.tox_epanet_output_filename,
                          options.tox_merlion_save_filename,
                          -1,
                          -1,
                          options.ignore_merlion_warnings,
                          options.tox_decay_k);
      Run_Hydraulics.StopAndPrint(std::cout, "\n\t- Built toxin Merlion Water Quality Model ");
   }
   else if (options.tox_wqm_filename != "") {
      if (options.tox_decay_k != -1.0f) {
         std::cerr << std::endl;
         std::cerr << "ERROR: A toxin decay coefficient cannot be specified when using the --tox-wqm option." << std::endl;
         std::cerr << std::endl;
         return 1;
      }
      merlionUtils::TaskTimer Read_WQM;
      // read a merlion water quality model file
      tox_net.ReadWQMFile(options.tox_wqm_filename);
      Read_WQM.StopAndPrint(std::cout, "\n\t- Read toxin Merlion Water Quality Model ");
   }
   else {
      std::cerr << std::endl;
      std::cerr << "ERROR: Missing network input file for the toxin." << std::endl;
      std::cerr << std::endl;
      return 1;
   }
   ///////////////

   // **A weak way of checking the the models match (except for, of course, their decay coefficients)
   if ((tox_net.Model().N != chl_net.Model().N) ||
       (tox_net.Model().n_nodes != chl_net.Model().n_nodes) ||
       (tox_net.Model().n_links != chl_net.Model().n_links) ||
       (tox_net.Model().n_steps != chl_net.Model().n_steps) ||
       (tox_net.Model().qual_steps_per_hour != chl_net.Model().qual_steps_per_hour)) {
      std::cerr << std::endl;
      std::cerr << "ERROR: Toxin and Booster models don't match in network layout or simulation length." << std::endl;
      std::cerr << std::endl;
      exit(1);
   }

   // Read in the booster station specs
   if (options.booster_filename != "") {
      merlionUtils::TaskTimer Read_BoosterFile;
      // read in the file defining he booster station parameters
      // this function also defines the booster candidates
      // depending on what is in the booster .spec file. Default is all
      // NZD junctions
      chl_net.ReadBoosterFile(options.booster_filename);
      Read_BoosterFile.StopAndPrint(std::cout, "\n\t- Read water booster spec file ");
   }
   else {
      std::cerr << std::endl;
      std::cerr << "ERROR: Missing booster station specs file." << std::endl;
      std::cerr << std::endl;
      return 1;
   }

   bool detection_times_provided(false);
   // Read in the Scenarios
   if (options.tsg_filename != "") {
      merlionUtils::TaskTimer Read_Scen;
      // Read in the injection scenarios define by the *.tsg file
      tox_net.ReadTSGFile(options.tsg_filename);
      Read_Scen.StopAndPrint(std::cout, "\n\t- Read scenario data in .tsg file");
   }
   else if (options.tsi_filename != "") {
      merlionUtils::TaskTimer Read_tsi;
      // Read in the injection scenarios define by the *.tsi file
      if (options.isDefault("tsi-species-id")) {
         tox_net.ReadTSIFile(options.tsi_filename);
      }
      else {
         if (options.tsi_species_id < 0) {
            std::cerr << std::endl;
            std::cerr << "ERROR: A TSI species id must be positive" << std::endl;
            std::cerr << std::endl;
            return 1;
         }
         tox_net.ReadTSIFile(options.tsi_filename, false, options.tsi_species_id);
      }
      Read_tsi.StopAndPrint(std::cout, "\n\t- Read scenario data in .tsi file");
   }
   else if (options.scn_filename != "") {
      merlionUtils::TaskTimer Read_Scen;
      // Read in the injection scenarios define by the *.scn file
      tox_net.ReadSCNFile(options.scn_filename);
      Read_Scen.StopAndPrint(std::cout, "\n\t- Read scenario data in .scn file");
   }
   else if (options.dscn_filename != "") {
      merlionUtils::TaskTimer Read_Scen;
      // Read in the injection scenarios define by the *.dscn file
      tox_net.ReadSCNFile(options.dscn_filename);
      Read_Scen.StopAndPrint(std::cout, "\n\t- Read scenario data in .dscn file");
      detection_times_provided = true;
   }
   else {
      std::cerr << std::endl;
      std::cerr << "ERROR: Missing scenario file." << std::endl;
      std::cerr << std::endl;
      return 1;
   }
   int scen_file_cnt = (options.tsg_filename != "")?(1):(0);
   scen_file_cnt += (options.tsi_filename != "")?(1):(0);
   scen_file_cnt += (options.scn_filename != "")?(1):(0);
   scen_file_cnt += (options.dscn_filename != "")?(1):(0);
   if (scen_file_cnt > 1) {
      std::cerr << std::endl;
      std::cerr << "ERROR: Too many scenario options specified. Choose one: TSG, TSI, SCN, DSCN" << std::endl;
      std::cerr << std::endl;
      return 1;
   }

   // Read in the sensor node list or inverse
   if (options.wqm_inverse_filename != "") {
      merlionUtils::TaskTimer Read_Inverse;
      tox_net.ReadInverse(options.wqm_inverse_filename);
      Read_Inverse.StopAndPrint(std::cout, "\n\t- Read Merlion Inverse File ");
   }
   else if (options.sensor_filename != "") {
      merlionUtils::TaskTimer Read_Sensors;
      // Read the file defining the set of sensors
      tox_net.ReadSensorFile(options.sensor_filename);
      Read_Sensors.StopAndPrint(std::cout,"\n\t- Read sensor file ");
   }
   else if (!detection_times_provided) {
      std::cerr << std::endl;
      std::cerr << "ERROR: Missing sensor or inverse file which is required when detection times are not provided." << std::endl;
      std::cerr << std::endl;
      return 1;
   }
   if ((options.sensor_filename != "") && (options.wqm_inverse_filename != "")) {
      std::cerr << std::endl;
      std::cerr << "ERROR: Both a sensor file and an inverse file cannot be specified." << std::endl;
      std::cerr << std::endl;
      return 1;
   }

   std::cout << "\n@@@ DETERMINING SCENARIO DETECT TIMES @@@" << std::endl;
   if ((options.wqm_inverse_filename == "") && (!detection_times_provided)) {
      merlionUtils::TaskTimer Build_Inverse;

      // Build reduced inverse so sensor node detections can be rapidly
      // determined for each injection
      int min_inj_timestep = tox_net.Model().n_steps;
      for STL_CONST_ITERATE(InjScenList, p_scen, p_stop, tox_net.InjectionScenarios()) {
         int scen_start = (*p_scen)->EarliestInjectionTimestep(tox_net.Model());
         if (scen_start < min_inj_timestep) {
            min_inj_timestep = scen_start;
         }
      }

      // Build a list of node/time tuple where rows of the inverse are needed.
      if (tox_net.SensorNodeIDS().empty() && tox_net.GrabSampleIDS().empty()) {
         std::cerr << std::endl;
         std::cerr << "ERROR: A non-empty list of sensor nodes or grab samples is required when "
	           << "using a tsg or scn file." << std::endl;
         std::cerr << std::endl;
         return 1;
      }

      tox_net.GenerateReducedInverse(min_inj_timestep,tox_net.Model().n_steps-1,0,1,true);
      Build_Inverse.StopAndPrint(std::cout, "\n\t- Inverse Built ");
   }

   merlionUtils::TaskTimer Classify;
   // Apply the sensor layout and generate the new scenario data
   // based on detection times
   bool scenario_aggregation = true;
   tox_net.ClassifyScenarios(detection_times_provided,
                             tox_net.Model().n_steps-1,
                             options.detection_tol_gpm3,
                             options.detection_interval_min);
   chl_net.CopyScenarioList(tox_net);
   chl_net.ClassifyBoosterScenarios(true,
                                    scenario_aggregation,
                                    options.detection_tol_gpm3,
                                    options.detection_interval_min);
   Classify.StopAndPrint(std::cout, "\n\t- Classified Scenarios ");

   bool pd_impacts = false;
   if (options.demand_percapita_m3pmin ||
       options.ingestion_rate_m3pmin) {
      pd_impacts = true;
      merlionUtils::TaskTimer GeneratePopUsageData;
      tox_net.GeneratePopulationUsageData(options.demand_percapita_m3pmin,
                                          options.ingestion_rate_m3pmin);
      GeneratePopUsageData.StopAndPrint(std::cout, "\n\t- Generated Population Usage Data");
   }

   // Free the memory for the reduced system inverse
   tox_net.ResetInverse();
   // Free the memory for the list of sensors
   tox_net.ResetSensorData();

   double booster_seconds = 0.0;
   std::cout << "\n@@@ SIMULATING BOOSTER STATIONS @@@" << std::endl;
   merlionUtils::TaskTimer SimAndSave;
   PerformBoosterSimulations(chl_net,tox_net,options,pd_impacts);
   SimAndSave.StopAndSave(booster_seconds);

   double sims_per_second = (chl_net.BoosterCandidates().size()*chl_net.BoosterScenarios().size() + chl_net.DetectedScenarios().size())/booster_seconds;
   std::cout << "\n\t- Finished Booster Station Simulations (time in seconds): " << booster_seconds << std::endl;
   std::cout << "\n\t- Approximate Simulations Per Second: " << sims_per_second << std::endl;

   chl_net.StopLogging();
   tox_net.StopLogging();

   // Free all other memory being used
   chl_net.clear();
   tox_net.clear();

   return 0;
}

//forward declarations
void PerformBoosterSimulations(const BoosterNetworkSimulator& chl_net,
                               const NetworkSimulator& tox_net,
			       const BoosterImpactOptions& options,
                               bool pd_impacts/*=false*/)
{
   // Create a timer but don't start it
   TaskTimer timer(false);

   const BoosterNetworkSimulator& net = chl_net;
   const MerlionModelContainer& model = chl_net.Model();
   const MerlionModelContainer& tox_model = tox_net.Model();

   const float zero_conc_gpm3 = std::max(options.zero_conc_tol_gpm3,ZERO_CONC_GPM3);

   int num_booster_stations = net.BoosterCandidates().size();
   std::cout << "\nNetwork Stats:\n";
   std::cout << "\tNumber of Booster Station Candidates - " << num_booster_stations << "\n";
   std::cout << "\tNumber of Junctions                  - " << model.junctions.size() << "\n";
   std::cout << "\tNumber of Nonzero Demand Junctions   - " << model.nzd_junctions.size() << "\n";
   std::cout << "\tNumber of Tanks                      - " << model.tanks.size() << "\n";
   std::cout << "\tNumber of Reservoirs                 - " << model.reservoirs.size() << "\n";
   std::cout << "\tWater Quality Timestep (minutes)     - " << model.qual_step_minutes << "\n";
   std::cout << std::endl;

   // A shortcut for specifying all possible
   // right hand sides in one call of the multiple
   // right hand side triangular solver
   int nrhs = options.max_rhs;
   if (nrhs <= 0) {
      nrhs = num_booster_stations;
   }

   // Allocate the right hand side and solution vector for toxin injection sims
   float* tox_mass_inj_gpmin = new float[model.N];
   BlasInitZero(model.N, tox_mass_inj_gpmin);

   std::set<int> positive_demand_u_ids;
   for (int u_idx = 0; u_idx < model.N; ++u_idx) {
      if (model.demand_ut_m3[u_idx] > 0.0f) {
         positive_demand_u_ids.insert(u_idx);
      }
   }

   std::vector<double> volume_ingested_ut_m3;
   std::vector<double> node_total_dose_g;
   std::vector<double> node_total_pre_dose_g;
   if (pd_impacts) {
      volume_ingested_ut_m3.clear();
      volume_ingested_ut_m3.resize(model.N,0.0);
      node_total_dose_g.clear();
      node_total_dose_g.resize(model.n_nodes, 0.0);
      node_total_pre_dose_g.clear();
      node_total_pre_dose_g.resize(model.n_nodes, 0.0);
      for (int nt_idx = 0; nt_idx < model.N; ++nt_idx) {
         if (tox_net.VolumeIngested_m3()[nt_idx] > 0.0) {
            volume_ingested_ut_m3[model.perm_nt_to_upper[nt_idx]] = tox_net.VolumeIngested_m3()[nt_idx];
         }
      }
   }

   // Allocate memory for the rhs and solution vector for the booster sims
   // depending on the number right hand sides specified in options
   float* chl_mass_inj_gpmin(NULL);
   if (nrhs > 1) {
      int TMP_RHS = (nrhs<num_booster_stations)?nrhs:num_booster_stations;
      chl_mass_inj_gpmin = new float[model.N * TMP_RHS];
      BlasInitZero(model.N*TMP_RHS, chl_mass_inj_gpmin);
   }
   else {
      chl_mass_inj_gpmin = new float[model.N];
      BlasInitZero(model.N, chl_mass_inj_gpmin);
   }

   timer.Start();
   for STL_CONST_ITERATE(BoosterScenList, p_boost, p_boost_stop, net.BoosterScenarios()) {
      timer.Start();
      const BoosterScenario& booster_scen = **p_boost;

      const int sim_stop_timestep = booster_scen.sim_stop_timestep;
      const int min_row = model.N-model.n_nodes*(sim_stop_timestep+1);

      // the set of node,time ids in upper trianglular orientation where
      // chlorine is present (0 indicates no chlorine)
      std::vector<int> disinf_u_ids(model.N,0);

      if (nrhs > 1 && !net.BoosterCandidates().empty()) {
         // Code optimized for solving booster station sims using set of matrix rhs
         // FASTED METHOD WHERE MEMORY IS LEVERAGED WITH SPEED
         //////////////////////////////////////////////
         if (nrhs > num_booster_stations) {
            nrhs = num_booster_stations;
         }
         const int total_blocks = (num_booster_stations/nrhs) + (((num_booster_stations%nrhs)>0)?1:0);
	 const int odd_count = num_booster_stations%nrhs;
	 int current_block(0);

	 std::vector<int>::const_iterator p_bid = chl_net.BoosterCandidates().begin();
	 std::vector<int>::const_iterator p_bid_stop = chl_net.BoosterCandidates().end();
         std::vector<int>::const_iterator p_bid_block_start;
	 Injection booster_injection = *(booster_scen.Injections().front());

         std::vector<int> block_booster_ids;
         while (current_block < total_blocks) {
            block_booster_ids.resize(nrhs);
            p_bid_block_start = p_bid;
	    // set the rhs matrix
	    int max_row_chl = 0;
	    for (int r = 0; (r < nrhs) && (p_bid != p_bid_stop); ++r, ++p_bid) {
	       float mass_injected_g(0.0);
	       int max_row_index(model.N-1);
	       booster_injection.NodeName() = model.NodeName(*p_bid);
	       booster_injection.SetMultiArray(model,nrhs,r,chl_mass_inj_gpmin,mass_injected_g,max_row_index);
	       block_booster_ids[r] = *p_bid;
	       if (max_row_index > max_row_chl) {
		  max_row_chl = max_row_index;
	       }
            }

            // do the backsolve
            usolvem( model.N,
		     min_row,
		     max_row_chl,
		     nrhs,
		     model.G->Values(),
		     model.G->iRows(),
		     model.G->pCols(),
		     chl_mass_inj_gpmin );

            // Create a reference variable to make units more obvious. The linear solver
            // converts the rhs vector (units gpmin) into the solution vector (gpm3). We will
            // use this reference when it make sense for units (e.g below the linear solve).
            float*& chl_conc_gpm3 = chl_mass_inj_gpmin;

            // analyze the solution
	    // keep track of the nodes/times where chlorine is present
            for (int i = min_row*nrhs, i_stop = (max_row_chl+1)*nrhs - 1; i <= i_stop; ++i) {
	       if (chl_conc_gpm3[i] > zero_conc_gpm3) {
		  disinf_u_ids[i/nrhs] += 1;
	       }
            }

	    // reset the sparse set of locations where the injections modifed the array,
	    // these may have been outside the region of interest in the linear solver,
	    // in which case, the BlasResetZero code after this loop would not have
	    // zero that region of the array
	    p_bid = p_bid_block_start;
	    for (int r = 0; (r < nrhs) && (p_bid != p_bid_stop); ++r, ++p_bid) {
	       booster_injection.NodeName() = model.NodeName(*p_bid);
	       booster_injection.ClearMultiArray(model, nrhs, r, chl_conc_gpm3);
	    }
            // reset to zero only the values where chl_conc_gpm3
            // was possibly modified by linear solver, this saves time
            const int num_elements = nrhs*(max_row_chl-min_row+1);
            const int p_offset = min_row*nrhs;
            BlasResetZero(num_elements, chl_conc_gpm3+p_offset);

	    ++current_block;
	    // for the last multiple-rhs block, reduce the leading dimension
	    // of the rhs/sol matrix to match the number of rhs left over.
	    // The extra memory that was allocated will just be left unused
	    if (current_block == total_blocks-1) {
	       nrhs = (odd_count)?(odd_count):(nrhs);
	    }
         }
      }
      else if (!net.BoosterCandidates().empty()) {
         // Code optimized for solving booster station sims using all single rhs solves
         // SLOWEST METHOD BUT IN SOME CASES MAY BE OPTIMIZED FOR DATA ACCESS
         //////////////////////////////////////////////
         Injection booster_injection = *(booster_scen.Injections().front());
         for STL_CONST_ITERATE(std::vector<int>, p_bid, p_bid_end, chl_net.BoosterCandidates()) {

            int bnode_id = *p_bid;
            booster_injection.NodeName() = model.NodeName(bnode_id);
	    float mass_injected_g(0.0);
	    int max_row_chl(model.N-1);
            booster_injection.SetArray(model,chl_mass_inj_gpmin,mass_injected_g,max_row_chl);

            // do the backsolve
            usolve ( model.N,
		     min_row,
		     max_row_chl,
		     model.G->Values(),
		     model.G->iRows(),
		     model.G->pCols(),
		     chl_mass_inj_gpmin );

            // Create a reference variable to make units make more obvious. The linear solver
            // converts the rhs vector (units gpmin) into the solution vector (gpm3). We will
            // use this reference when it make sense for units (e.g below the linear solve).
            float*& chl_conc_gpm3 = chl_mass_inj_gpmin;

            // analyze the solution
            for (int u_idx = min_row; u_idx <= max_row_chl; ++u_idx) {
               if (chl_conc_gpm3[u_idx] > zero_conc_gpm3) {
                  disinf_u_ids[u_idx] += 1;
               }
            }

	    // reset the sparse set of locations where the injections modifed the array,
	    // these may have been outside the region of interest in the linear solver,
	    // in which case, the BlasResetZero code after this loop would not have
	    // zero that region of the array
            booster_injection.ClearArray(model, chl_conc_gpm3);
	    // reset to zero only the values where chl_conc_gpm3
            // was possibly modified by linear solver, this saves time
            const int num_elements = max_row_chl-min_row+1;
            const int p_offset = min_row;
            BlasResetZero(num_elements, chl_conc_gpm3+p_offset);
         }
      }

      timer.Stop();
      if (options.report_scenario_timing) {
         std::cout << "\tBOOSTER SCENARIO ID:    " << (*p_boost)->Name() << std::endl;
         std::cout << "\t- Booster Simulations:  " << timer.LastLap() << std::endl;
      }
      timer.Start();

      // this is the maximum row after which chlorine injections have not begun
      const int chl_start_timestep = booster_scen.Injections().front()->StartTimestep(model);
      const int max_row_chl = (model.N-1)-model.n_nodes*chl_start_timestep;
      const InjScenList& toxin_scen_list = booster_scen.ToxinScenarios();
      for STL_CONST_ITERATE(InjScenList, ptox_scen_pos, ptox_scen_pos_stop, toxin_scen_list) {
         const PInjScenario ptox_scen = *ptox_scen_pos;
         // set the rhs for the toxin injection
         float scenario_mass_injected_g(0.0);
         int max_row_tox(model.N-1);
         ptox_scen->SetArray(model, tox_mass_inj_gpmin, scenario_mass_injected_g, max_row_tox);
         ptox_scen->AddImpact(scenario_mass_injected_g, "Response-Window Mass Injected Grams");

         // Solver the linear system
         usolve ( model.N,
		  min_row,
		  max_row_tox,
		  tox_model.G->Values(),
		  tox_model.G->iRows(),
		  tox_model.G->pCols(),
		  tox_mass_inj_gpmin );

         // Create a reference variable to make units make more obvious. The linear solver
         // converts the rhs vector (units gpmin) into the solution vector (gpm3). We will
         // use this reference when it make sense for units (e.g below the linear solve).
         float*& tox_conc_gpm3 = tox_mass_inj_gpmin;

         const int u_idx_start = model.N-model.n_nodes*(sim_stop_timestep+1);
         std::set<int>::const_iterator p_dem_u_idx = positive_demand_u_ids.begin();
         while(*p_dem_u_idx < u_idx_start) {++p_dem_u_idx;}
         std::set<int>::const_iterator p_dem_u_idx_stop = --(positive_demand_u_ids.end());
         while(*p_dem_u_idx_stop > max_row_tox) {--p_dem_u_idx_stop;}
         ++p_dem_u_idx_stop;
         // determine full scenario impact
         float scenario_mass_consumed_g(0.0);
         // Check the mass balance
         float mass_check(0.0);
         if (!pd_impacts) {
            for (; p_dem_u_idx != p_dem_u_idx_stop; ++p_dem_u_idx) {
               int u_idx = *p_dem_u_idx;
               if (tox_conc_gpm3[u_idx] > zero_conc_gpm3) {
                  mass_check += tox_conc_gpm3[u_idx]*model.demand_ut_m3[u_idx];
                  if (!disinf_u_ids[u_idx]) {
                     // mc
                     scenario_mass_consumed_g += tox_conc_gpm3[u_idx]*model.demand_ut_m3[u_idx];
                  }
               }
            }
         }
         else { // Includes the code above, just minimizes checking if pd_impacts
            for (; p_dem_u_idx != p_dem_u_idx_stop; ++p_dem_u_idx) {
               int u_idx = *p_dem_u_idx;
               if (tox_conc_gpm3[u_idx] > zero_conc_gpm3) {
                  mass_check += tox_conc_gpm3[u_idx]*model.demand_ut_m3[u_idx];
                  if (!disinf_u_ids[u_idx]) {
                     // mc
                     scenario_mass_consumed_g += tox_conc_gpm3[u_idx]*model.demand_ut_m3[u_idx];
                     // pd
                     int n = model.perm_upper_to_nt[u_idx] / model.n_steps;
                     node_total_dose_g[n] += tox_conc_gpm3[u_idx] * volume_ingested_ut_m3[u_idx];
                  }
               }
            }
         }

         // determine impact before booster stations activate
         p_dem_u_idx = p_dem_u_idx_stop;
         --p_dem_u_idx;
         while(*p_dem_u_idx >= (max_row_chl+1)) {--p_dem_u_idx;}
         ++p_dem_u_idx;
         float scenario_pre_booster_mass_consumed_g(0.0);
         if (!pd_impacts) {
            for (; p_dem_u_idx != p_dem_u_idx_stop; ++p_dem_u_idx) {
               int u_idx = *p_dem_u_idx;
               if (tox_conc_gpm3[u_idx] > zero_conc_gpm3) {
                  scenario_pre_booster_mass_consumed_g += tox_conc_gpm3[u_idx]*model.demand_ut_m3[u_idx];
               }
            }
         }
         else { // Includes the code above, just minimizes checking if pd_impacts
            for (; p_dem_u_idx != p_dem_u_idx_stop; ++p_dem_u_idx) {
               int u_idx = *p_dem_u_idx;
               if (tox_conc_gpm3[u_idx] > zero_conc_gpm3) {
                  // pre mc
                  scenario_pre_booster_mass_consumed_g += tox_conc_gpm3[u_idx]*model.demand_ut_m3[u_idx];
                  // pre pd
                  int n = model.perm_upper_to_nt[u_idx] / model.n_steps;
                  node_total_pre_dose_g[n] += tox_conc_gpm3[u_idx] * volume_ingested_ut_m3[u_idx];
               }
            }
         }

         // Include mass remaining in tanks in mass balance check
         float tank_mass(0.0);
         for (std::vector<std::pair<int,int> >::const_iterator pos = model.tanks.begin(), stop = model.tanks.end(); pos != stop; ++pos) {
            int n = pos->first;
            int tank_n = pos->second;
            int t = sim_stop_timestep;
            int nt_tank_idx = tank_n*model.n_steps + t;
            int nt_idx = n*model.n_steps + t;
            int u_idx = model.perm_nt_to_upper[nt_idx];
            if (((tox_conc_gpm3[u_idx] > zero_conc_gpm3) && (!disinf_u_ids[u_idx]))) {
               tank_mass += tox_conc_gpm3[u_idx]*model.tank_volume_m3[nt_tank_idx];
            }
            mass_check += tox_conc_gpm3[u_idx]*model.tank_volume_m3[nt_tank_idx];
         }
         ptox_scen->AddImpact(scenario_mass_consumed_g, "Mass Consumed Grams");
         ptox_scen->AddImpact(scenario_pre_booster_mass_consumed_g, "Pre-Booster Mass Consumed Grams");
         ptox_scen->AddImpact(tank_mass, "Final Mass In Tanks Grams");
         ptox_scen->AddImpact(mass_check/scenario_mass_injected_g*100.0f, "Simulation Mass Balance");

         if (pd_impacts) {
            double population_dosed = 0.0;
            for (int n = 0; n < model.n_nodes; ++n) {
               if (node_total_dose_g[n] > options.population_dosed_threshold_g) {
                  population_dosed += tox_net.NodePopulation()[n];
               }
            }
            ptox_scen->AddImpact(population_dosed, "Population Dosed");
            double pre_population_dosed = 0.0;
            for (int n = 0; n < model.n_nodes; ++n) {
               if (node_total_pre_dose_g[n] > options.population_dosed_threshold_g) {
                  pre_population_dosed += tox_net.NodePopulation()[n];
               }
            }
            ptox_scen->AddImpact(pre_population_dosed, "Pre-Population Dosed");
            // reset the node dose vectors
            BlasResetZero(model.n_nodes, &(node_total_dose_g[0]));
            BlasResetZero(model.n_nodes, &(node_total_pre_dose_g[0]));
         }

	 // reset the sparse set of locations where the injections modifed the array,
	 // these may have been outside the region of interest in the linear solver,
	 // in which case, the BlasResetZero code after this loop would not have
	 // zero that region of the array
	 ptox_scen->ClearArray(model, tox_mass_inj_gpmin);
         // reset the rhs for the toxin backsolves
         const int num_elements = max_row_tox-min_row+1;
         const int p_offset = min_row;
         BlasResetZero(num_elements, tox_conc_gpm3+p_offset);
      }
      timer.Stop();
      if (options.report_scenario_timing) {
         std::cout << "\t- Toxin Simulations:  " << timer.LastLap() << std::endl;
         std::cout << std::endl;
      }
   }

   timer.Start();
   for STL_CONST_ITERATE(InjScenList, ptox_scen_pos, ptox_scen_pos_stop, chl_net.NonDetectedScenarios()) {
      const PInjScenario ptox_scen = *ptox_scen_pos;

      // We simulate until the end of the time horizon
      const int sim_stop_timestep = model.n_steps-1;
      const int min_row = 0;

      // Set the mass injected label for the non-detected scenarios
      ptox_scen->AddImpact(ptox_scen->MassInjected(chl_net.Model()),"Mass Injected Grams");

      // set the rhs for the toxin injection
      float scenario_mass_injected_g(0.0);
      int max_row_tox(model.N-1);
      ptox_scen->SetArray(model, tox_mass_inj_gpmin, scenario_mass_injected_g, max_row_tox);

      // Solver the linear system
      usolve (model.N,
              min_row,
              max_row_tox,
              tox_model.G->Values(),
              tox_model.G->iRows(),
              tox_model.G->pCols(),
              tox_mass_inj_gpmin );

      // Create a reference variable to make units make more obvious. The linear solver
      // converts the rhs vector (units gpmin) into the solution vector (gpm3). We will
      // use this reference when it make sense for units (e.g below the linear solve).
      float*& tox_conc_gpm3 = tox_mass_inj_gpmin;

      std::set<int>::const_iterator p_dem_u_idx = positive_demand_u_ids.begin();
      while(*p_dem_u_idx < min_row) {++p_dem_u_idx;}
      std::set<int>::const_iterator p_dem_u_idx_stop = --(positive_demand_u_ids.end());
      while(*p_dem_u_idx_stop > max_row_tox) {--p_dem_u_idx_stop;}
      ++p_dem_u_idx_stop;
      // determine full scenario impact
      float scenario_mass_consumed_g(0.0);
      // Check the mass balance
      float mass_check(0.0);
      for (; p_dem_u_idx != p_dem_u_idx_stop; ++p_dem_u_idx) {
         int u_idx = *p_dem_u_idx;
         if (tox_conc_gpm3[u_idx] > zero_conc_gpm3) {
            mass_check += tox_conc_gpm3[u_idx]*model.demand_ut_m3[u_idx];
            // mc
            scenario_mass_consumed_g += tox_conc_gpm3[u_idx]*model.demand_ut_m3[u_idx];
            // pd
            if (pd_impacts) {
               int n = model.perm_upper_to_nt[u_idx] / model.n_steps;
               node_total_dose_g[n] += tox_conc_gpm3[u_idx] * volume_ingested_ut_m3[u_idx];
            }
         }
      }

      // Include mass remaining in tanks in mass balance check
      float tank_mass(0.0);
      for (std::vector<std::pair<int,int> >::const_iterator pos = model.tanks.begin(), stop = model.tanks.end(); pos != stop; ++pos) {
         int n = pos->first;
         int tank_n = pos->second;
         int t = sim_stop_timestep;
         int nt_tank_idx = tank_n*model.n_steps + t;
         int nt_idx = n*model.n_steps + t;
         int u_idx = model.perm_nt_to_upper[nt_idx];
         if (tox_conc_gpm3[u_idx] > zero_conc_gpm3) {
            tank_mass += tox_conc_gpm3[u_idx]*model.tank_volume_m3[nt_tank_idx];
         }
         mass_check += tox_conc_gpm3[u_idx]*model.tank_volume_m3[nt_tank_idx];
      }
      ptox_scen->AddImpact(scenario_mass_consumed_g, "Mass Consumed Grams");
      ptox_scen->AddImpact(tank_mass, "Final Mass In Tanks Grams");
      ptox_scen->AddImpact(mass_check/scenario_mass_injected_g*100.0f, "Simulation Mass Balance");

      if (pd_impacts) {
         double population_dosed = 0.0;
         for (int n = 0; n < model.n_nodes; ++n) {
            if (node_total_dose_g[n] > options.population_dosed_threshold_g) {
               population_dosed += tox_net.NodePopulation()[n];
            }
         }
         ptox_scen->AddImpact(population_dosed, "Population Dosed");
      }

      // reset the node dose vector
      BlasResetZero(model.n_nodes, &(node_total_dose_g[0]));

      // reset the sparse set of locations where the injections modifed the array,
      // these may have been outside the region of interest in the linear solver,
      // in which case, the BlasResetZero code after this loop would not have
      // zero that region of the array
      ptox_scen->ClearArray(model, tox_mass_inj_gpmin);
      // reset the rhs for the toxin backsolves
      const int num_elements = max_row_tox-min_row+1;
      const int p_offset = min_row;
      BlasResetZero(num_elements, tox_conc_gpm3+p_offset);
   }
   timer.Stop();
   if (options.report_scenario_timing) {
      std::cout << "\t- Non-Detected Toxin Simulations:  " << timer.LastLap() << std::endl;
      std::cout << std::endl;
   }

   //print scenario impact
   std::ofstream out;
   out.precision(12);
   out.setf(std::ios::scientific,std::ios::floatfield);
   if (options.yaml) {
      std::string fname = options.output_prefix+"DetectedScenarioImpacts.yml";
      out.open(fname.c_str(), std::ios_base::out | std::ios_base::trunc);
      merlionUtils::PrintScenarioListYAML(net.DetectedScenarios(), out);
      out.close();
      fname = options.output_prefix+"NonDetectedScenarioImpacts.yml";
      out.open(fname.c_str(), std::ios_base::out | std::ios_base::trunc);
      merlionUtils::PrintScenarioListYAML(net.NonDetectedScenarios(), out);
      out.close();
      fname = options.output_prefix+"DiscardedScenarioImpacts.yml";
      out.open(fname.c_str(), std::ios_base::out | std::ios_base::trunc);
      merlionUtils::PrintScenarioListYAML(net.DiscardedScenarios(), out);
      out.close();
   }
   if (options.json) {
      std::string fname = options.output_prefix+"DetectedScenarioImpacts.json";
      out.open(fname.c_str(), std::ios_base::out | std::ios_base::trunc);
      merlionUtils::PrintScenarioListJSON(net.DetectedScenarios(), out);
      out.close();
      fname = options.output_prefix+"NonDetectedScenarioImpacts.json";
      out.open(fname.c_str(), std::ios_base::out | std::ios_base::trunc);
      merlionUtils::PrintScenarioListJSON(net.NonDetectedScenarios(), out);
      out.close();
      fname = options.output_prefix+"DiscardedScenarioImpacts.json";
      out.open(fname.c_str(), std::ios_base::out | std::ios_base::trunc);
      merlionUtils::PrintScenarioListJSON(net.DiscardedScenarios(), out);
      out.close();
   }
   if (!options.yaml && !options.json) {

      std::string fname = options.output_prefix+"ScenarioImpacts.txt";
      out.open(fname.c_str(), std::ios_base::out | std::ios_base::trunc);
      out << "Detected " << net.DetectedScenarios().size() << "\n";
      out << "label   weight   normalized-weight   orig-mass-injected   window-mass-injected   mass-in-tanks   percent-balance\n";
      for STL_CONST_ITERATE(InjScenList, p_scen, p_stop, net.DetectedScenarios()) {
         out << (*p_scen)->Name() << " "
             << (*p_scen)->Impact("Weight") << " "
             << (*p_scen)->Impact("Normalized Weight") << " "
             << (*p_scen)->Impact("Original Mass Injected Grams") << " "
             << (*p_scen)->Impact("Response-Window Mass Injected Grams") << " "
             << (*p_scen)->Impact("Final Mass In Tanks Grams") << " "
             << (*p_scen)->Impact("Simulation Mass Balance") << "\n";
      }
      out << "\n";
      out << "UnDetected " << net.NonDetectedScenarios().size() << std::endl;
      out << "label   weight   normalized-weight   mass-injected   mass-in-tanks   percent-balance\n";
      for STL_CONST_ITERATE(InjScenList, p_scen, p_stop, net.NonDetectedScenarios()) {
         out << (*p_scen)->Name() << " "
             << (*p_scen)->Impact("Weight") << " "
             << (*p_scen)->Impact("Normalized Weight") << " "
             << (*p_scen)->Impact("Mass Injected Grams") << " "
             << (*p_scen)->Impact("Final Mass In Tanks Grams") << " "
             << (*p_scen)->Impact("Simulation Mass Balance") << "\n";
      }
      out << "\n";
      out << "Discarded " << net.DiscardedScenarios().size() << std::endl;
      out << "label   weight   normalized-weight   percent-balance\n";
      for STL_CONST_ITERATE(InjScenList, p_scen, p_stop, net.DiscardedScenarios()) {
         out << (*p_scen)->Name() << " "
             << (*p_scen)->Impact("Weight") << " "
             << (*p_scen)->Impact("Normalized Weight") << " "
             << 0.0 << "\n";
      }
      out.close();

      fname = options.output_prefix+"ScenarioImpacts_MC.txt";
      out.open(fname.c_str(), std::ios_base::out | std::ios_base::trunc);
      out << "Detected " << net.DetectedScenarios().size() << std::endl;
      out << "label   mass-consumed  pre-booster-mass-consumed\n";
      for STL_CONST_ITERATE(InjScenList, p_scen, p_stop, net.DetectedScenarios()) {
         out << (*p_scen)->Name() << " "
             << (*p_scen)->Impact("Mass Consumed Grams") << " "
             << (*p_scen)->Impact("Pre-Booster Mass Consumed Grams") << "\n";
      }
      out << "\n";
      out << "UnDetected " << net.NonDetectedScenarios().size() << std::endl;
      out << "label   mass-consumed\n";
      for STL_CONST_ITERATE(InjScenList, p_scen, p_stop, net.NonDetectedScenarios()) {
         out << (*p_scen)->Name() << " "
             << (*p_scen)->Impact("Mass Consumed Grams") << "\n";
      }
      out << "\n";
      out << "Discarded " << net.DiscardedScenarios().size() << std::endl;
      out << "label\n";
      for STL_CONST_ITERATE(InjScenList, p_scen, p_stop, net.DiscardedScenarios()) {
         out << (*p_scen)->Name() << "\n";
      }
      out.close();

      if (pd_impacts) {
         fname = options.output_prefix+"ScenarioImpacts_PD.txt";
         out.open(fname.c_str(), std::ios_base::out | std::ios_base::trunc);
         out << "Detected " << net.DetectedScenarios().size() << std::endl;
         out << "label   population-dosed  pre-booster-population-dosed\n";
         for STL_CONST_ITERATE(InjScenList, p_scen, p_stop, net.DetectedScenarios()) {
            out << (*p_scen)->Name() << " "
                << (*p_scen)->Impact("Population Dosed") << " "
                << (*p_scen)->Impact("Pre-Population Dosed") << "\n";
         }
         out << "\n";
         out << "UnDetected " << net.NonDetectedScenarios().size() << std::endl;
         out << "label   population-dosed\n";
         for STL_CONST_ITERATE(InjScenList, p_scen, p_stop, net.NonDetectedScenarios()) {
            out << (*p_scen)->Name() << " "
                << (*p_scen)->Impact("Population Dosed") << "\n";
         }
         out << "\n";
         out << "Discarded " << net.DiscardedScenarios().size() << std::endl;
         out << "label\n";
         for STL_CONST_ITERATE(InjScenList, p_scen, p_stop, net.DiscardedScenarios()) {
            out << (*p_scen)->Name() << "\n";
         }
      }
   }
}


//==================================================================================================
/**
  Copyright 2016 Numscale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_ARCH_X86_SSE2_SIMD_FUNCTION_INTERLEAVE_EVEN_HPP_INCLUDED
#define BOOST_SIMD_ARCH_X86_SSE2_SIMD_FUNCTION_INTERLEAVE_EVEN_HPP_INCLUDED

#include <boost/simd/detail/overload.hpp>

namespace boost { namespace simd { namespace ext
{
  namespace bd =  boost::dispatch;
  namespace bs =  boost::simd;

  BOOST_DISPATCH_OVERLOAD ( interleave_even_
                          , (typename A0)
                          , bs::sse2_
                          , bs::pack_<bd::double_<A0>, bs::sse_>
                          , bs::pack_<bd::double_<A0>, bs::sse_>
                         )
  {
    BOOST_FORCEINLINE A0 operator()(const A0 & a0, const A0 & a1) const BOOST_NOEXCEPT
    {
      return _mm_unpacklo_pd(a0,a1);
    }
  };

  BOOST_DISPATCH_OVERLOAD ( interleave_even_
                          , (typename A0)
                          , bs::sse2_
                          , bs::pack_<bd::ints32_<A0>, bs::sse_>
                          , bs::pack_<bd::ints32_<A0>, bs::sse_>
                         )
  {
    BOOST_FORCEINLINE A0 operator()(const A0 & a0, const A0 & a1) const BOOST_NOEXCEPT
    {
      return _mm_unpacklo_epi32 ( _mm_castps_si128(_mm_shuffle_ps ( _mm_castsi128_ps(a0)
                                                                  , _mm_castsi128_ps(a0)
                                                                  , _MM_SHUFFLE(2,0,2,0)
                                                                  )
                                                  )
                                , _mm_castps_si128(_mm_shuffle_ps ( _mm_castsi128_ps(a1)
                                                                  , _mm_castsi128_ps(a1)
                                                                  , _MM_SHUFFLE(2,0,2,0)
                                                                  )
                                                  )
                                );
    }
  };

  BOOST_DISPATCH_OVERLOAD ( interleave_even_
                          , (typename A0)
                          , bs::sse2_
                          , bs::pack_<bd::ints64_<A0>, bs::sse_>
                          , bs::pack_<bd::ints64_<A0>, bs::sse_>
                         )
  {
    BOOST_FORCEINLINE A0 operator()(const A0 & a0, const A0 & a1) const BOOST_NOEXCEPT
    {
      return _mm_unpacklo_epi64(a0,a1);
    }
  };
} } }

#endif