/* Copyright (C) 2017 IBM Corp.
 *  Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *     http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied. See the License for the specific
 * language governing permissions and limitations under the License.
 */
/***********************************************************************
CRTmatrix: implements A matrix in CRT representation
************************************************************************/

#include <vector>
#include <cassert>
#include <cstdlib>
//#include <mpfr.h>
#include <NTL/mat_lzz_p.h>
#include <NTL/BasicThreadPool.h>
#include <NTL/ZZ.h>
#include <NTL/FFT.h>
#include <NTL/SmartPtr.h>

NTL_CLIENT
#include "utils/tools.h"
#include "DGaussSampler.h"
#include "CRTmatrix.h"
#include "TDMatrixParams.h"

//#define DEBUGPRINT

CRTmatrix& CRTmatrix::leftMultBy(const CRTmatrix& M) // A := M*A
{
    FHE_TIMER_START;
    assert(length()==params->kFactors);
    NTL::zz_pPush ppush; // backup NTL's current modulus

    // multi-threaded implementation, threads handle different moduli

    EXEC_RANGE(length(), first, last)
    //EXEC_RANGE(length(), first, last)
    for (long i=first; i<last; i++)
    {
        params->zzp_context[i].restore();
        mat_zz_p& matMod = (*this)[i];
        const mat_zz_p& Mmod = M[i];
        mul(matMod, Mmod, matMod);
    }
    //EXEC_RANGE_END
    EXEC_RANGE_END
    return *this;
}

CRTmatrix& CRTmatrix::leftMultBy(const mat_l& M) // A := M*A
{
    FHE_TIMER_START;
    assert(length()==params->kFactors);
    NTL::zz_pPush ppush; // backup NTL's current modulus

    // multi-threaded implementation, threads handle different moduli
    EXEC_RANGE(length(), first, last)
    for (long i=first; i<last; i++)
    {
        params->zzp_context[i].restore();
        mat_zz_p& matMod = (*this)[i];
        mat_zz_p Mmod = conv<mat_zz_p>(M);
        mul(matMod, Mmod, matMod);
    }
    EXEC_RANGE_END

    return *this;
}

CRTmatrix& CRTmatrix::leftMultBy(const vec_l& u) // A := u*A
{
    FHE_TIMER_START;
    assert(length()==params->kFactors);
    NTL::zz_pPush ppush; // backup NTL's current modulus

    // multi-threaded implementation, threads handle different moduli
    EXEC_RANGE(length(), first, last)

    for (long i=first; i<last; i++)
    {
        params->zzp_context[i].restore();
        mat_zz_p& matMod = (*this)[i]; // A
        const vec_zz_p umod = conv<vec_zz_p>(u);
        vec_zz_p vTmp;
        mul(vTmp, umod, matMod); // tmp = u*A

        // copy vTmp back to matMod
        matMod.SetDims(1, vTmp.length());
        matMod[0] = vTmp;
    }
    EXEC_RANGE_END

    return *this;
}

// Convert Mat<long> to CRT format, we assume that params are already set
CRTmatrix& CRTmatrix::operator=(const mat_l& M) // A := M
{
    FHE_TIMER_START;
    assert(params != NULL);

    NTL::zz_pPush ppush; // backup NTL's current modulus
    SetLength(params->kFactors);
    for (long i=0; i<length(); i++)
    {
        params->zzp_context[i].restore();
        conv((*this)[i], M);
    }

    return *this;
}

// Default assignment should work: NTL now guarantees
// that assignment and default constructors work, even
// "out of context"
CRTmatrix& CRTmatrix::operator=(const CRTmatrix& other) // A := M
{
    FHE_TIMER_START;
    params = other.params;
    NTL::zz_pPush ppush; // backup NTL's current modulus
    SetLength(other.length());
    for (long i=0; i<params->kFactors; i++)
    {
        params->zzp_context[i].restore();
        (*this)[i] = other[i];
    }
    return *this;
}

CRTmatrix& CRTmatrix::operator*=(const mat_l& M) // A := M*A
{
    FHE_TIMER_START;
    assert(length()==params->kFactors);
    NTL::zz_pPush ppush; // backup NTL's current modulus

    // multi-threaded implementation, threads handle different moduli
    EXEC_RANGE(length(), first, last)
    mat_zz_p Mmod;
    for (long i=first; i<last; i++)
    {
        params->zzp_context[i].restore();
        mat_zz_p& Amod = (*this)[i];
        Mmod = conv<mat_zz_p>(M);
        mul(Amod, Amod, Mmod);
    }
    EXEC_RANGE_END

    return *this;
}

CRTmatrix& CRTmatrix::operator*=(const CRTmatrix& M) // A := A*M
{
    FHE_TIMER_START;
    assert(length()==params->kFactors
           && M.length()==M.params->kFactors
           && params==M.params);
    NTL::zz_pPush ppush; // dibackup NTL's current modulus

    // multi-threaded implementation, threads handle different moduli
    EXEC_RANGE(length(), first, last)
    for (long i=first; i<last; i++)
    {
        params->zzp_context[i].restore();
        mat_zz_p& Amod = (*this)[i];
        const mat_zz_p& Mmod = M[i];
        mul(Amod, Amod, Mmod);
    }
    EXEC_RANGE_END

    return *this;
}

CRTmatrix& CRTmatrix::operator+=(const mat_l& M) // A := A+M
{
    FHE_TIMER_START;
    assert(length()==params->kFactors);
    NTL::zz_pPush ppush; // backup NTL's current modulus

    // multi-threaded implementation, threads handle different moduli
    EXEC_RANGE(length(), first, last)
    //for (long i=0; i<length(); i++)
    for (long i=first; i<last; i++)
    {
        mat_zz_p& matMod = (*this)[i];
        assert(matMod.NumRows()==M.NumRows() && matMod.NumCols()==M.NumCols());

        params->zzp_context[i].restore();
        for (long j=0; j<matMod.NumRows(); j++)
            for (long k=0; k<matMod.NumCols(); k++)
                matMod[j][k] += M[j][k];
    }
    EXEC_RANGE_END
    return *this;
}
CRTmatrix& CRTmatrix::operator-=(const mat_l& M) // A := A+M
{
    FHE_TIMER_START;
    assert(length()==params->kFactors);
    NTL::zz_pPush ppush; // backup NTL's current modulus

    EXEC_RANGE(length(), first, last)
    //for (long i=0; i<length(); i++)
    for (long i=first; i<last; i++)
    {
        mat_zz_p& matMod = (*this)[i];
        assert(matMod.NumRows()==M.NumRows() && matMod.NumCols()==M.NumCols());

        params->zzp_context[i].restore();
        for (long j=0; j<matMod.NumRows(); j++)
            for (long k=0; k<matMod.NumCols(); k++)
                matMod[j][k] -= M[j][k];
    }
    EXEC_RANGE_END
    return *this;
}

CRTmatrix& CRTmatrix::operator+=(const CRTmatrix& M) // A := A+M
{
    FHE_TIMER_START;
    assert(length()==params->kFactors
           && M.length()==M.params->kFactors
           && params==M.params);
    NTL::zz_pPush ppush; // backup NTL's current modulus

    EXEC_RANGE(length(), first, last)
    //for (long i=0; i<length(); i++)
    for (long i=first; i<last; i++)
    {
        mat_zz_p& Amod = (*this)[i];
        const mat_zz_p& Mmod = M[i];

        assert(Amod.NumRows()==Mmod.NumRows() && Amod.NumCols()==Mmod.NumCols());
        params->zzp_context[i].restore();
        Amod += Mmod;
    }
    EXEC_RANGE_END
    return *this;
}

CRTmatrix& CRTmatrix::operator-=(const CRTmatrix& M) // A := A+M
{
    FHE_TIMER_START;
    assert(length()==params->kFactors
           && M.length()==M.params->kFactors
           && params==M.params);
    NTL::zz_pPush ppush; // backup NTL's current modulus

    EXEC_RANGE(length(), first, last)
    //for (long i=0; i<length(); i++)
    for (long i=first; i<last; i++)
    {
        mat_zz_p& Amod = (*this)[i];
        const mat_zz_p& Mmod = M[i];

        assert(Amod.NumRows()==Mmod.NumRows() && Amod.NumCols()==Mmod.NumCols());
        params->zzp_context[i].restore();
        Amod -= Mmod;
    }
    EXEC_RANGE_END
    return *this;
}

// Get a CRT column: crtCol[i] is the j'th column of this[i]
void CRTmatrix::getColumn(Vec<vec_zz_p>& crtCol, long j)
{
    FHE_TIMER_START;
    crtCol.SetLength(length());
    for (long n=0; n<length(); n++)
    {
        vec_zz_p& colMod = crtCol[n];
        mat_zz_p& matMod = (*this)[n];

        colMod.SetLength(matMod.NumRows());
        // NOTE: this is OK, but beware that this would
        // not work for ZZ_p's without restoring context

        for (long i=0; i<matMod.NumRows(); i++)
            colMod[i] = matMod[i][j];
    }
}

//checks if all values of the matrix are smaller than the modolus
bool CRTmatrix::isSmall(long bitGap) const
{
    FHE_TIMER_START;
    mat_ZZ M; // reconstruct the matrix in ZZ format
    ZZ q;     // the modulus

    convert(M, q, *this); // cnovert to ZZ format

    q >>= bitGap;         // q / 2^{bitGap}

    for (long i=0; i<M.NumRows(); i++) for (long j=0; j<M.NumCols(); j++)
        {
            if (abs(M[i][j]) > q)
            {
                return false;
            }
        }
    return true;
}

// Choose a random n-by-m matrix in CRT representation
void CRTmatrix::randomFill(long n, long m)
{
    FHE_TIMER_START;

    const TDMatrixParams& prms = *params;
    SetLength(prms.kFactors);

    zz_pPush ppush; // backup NTL's current modulus

    EXEC_RANGE(length(), first, last)
    //for (long i=0; i< prms.kFactors; i++)
    for (long i=first; i< last; i++)
    {

        prms.zzp_context[i].restore();
        (*this)[i].SetDims(n, m);
        RandomFill((*this)[i]);
    }
    EXEC_RANGE_END
}

// Invert M modulo the product of the factors
bool CRTmatrix::invert(const mat_l& M)
{
    FHE_TIMER_START;
    SetLength(params->kFactors);

    zz_pPush ppush; // backup NTL's current modulus

    std::atomic<bool> result(true);

    EXEC_RANGE(params->kFactors, first, last)

    // Invert M wrt each factor separately
    //for (long i=0; i< params->kFactors; i++)
     for (long i=first; i< last; i++)
    {
        params->zzp_context[i].restore();
        // NOTE: Alternatively, can call relaxed_inv directly now
        if (!invMod((*this)[i], M))   // not invertible
        {
            SetLength(0);
            result = false;
            //return false;
        }
    }
    EXEC_RANGE_END

    return result;
    //return true;
}

void CRTmatrix::identity(long n)
{
    FHE_TIMER_START;
    zz_pPush ppush; // backup NTL's current modulus
    SetLength(params->kFactors);

    // Invert M wrt each factor separately
    for (long i=0; i< params->kFactors; i++)
    {
        params->zzp_context[i].restore();
        ident( (*this)[i], n ); // set as the identity mod f_i
    }
}

bool operator==(const CRTmatrix& A, const CRTmatrix& B)
{
    FHE_TIMER_START;
    if (A.params != B.params) return false;

    zz_pPush push; // backup the NTL current modulus
    for (long i=0; i < A.params->kFactors; i++)
    {
        (A.params->zzp_context)[i].restore();
#ifdef DEBUGPRINT
        cout << "A[i]=" << A[i] << ", B[i]=" << B[i] << endl;
#endif // DEBUGPRINT
        if (A[i] != B[i]) return false;
    }
    return true;
}

//invert a CRT matrix
bool CRTmatrix::invert(const CRTmatrix& A)
{
    FHE_TIMER_START;
    params = A.params;
    SetLength(params->kFactors);

    zz_pPush push; // backup the NTL current modulus

    for (long i=0; i < params->kFactors; i++)
    {
        const mat_zz_p& Amod = A[i];
        params->zzp_context[i].restore();

        zz_p d;
        //    if (isInvertible(Amod, params->factors[i]))
        //inv((*this)[i],Amod); // raises error if Amod is singular
        relaxed_inv(d,(*this)[i],Amod); //raises error if Amod is singular

        if (IsZero(d))
            return false; //non invertible
    }
    return true;
}

// reconstruct the matrix in ZZ fromat, returning in q the modulus
void convert(mat_ZZ& to, ZZ& q, const CRTmatrix& from)
{
    FHE_TIMER_START;
    q = to_ZZ(1L); // initalize to 1
    if (from.length() <=0)
    {
        to.kill(); // set as 0-by-0 matrix
        return;
    }

    to.SetDims(from[0].NumRows(), from[0].NumCols());
    clear(to);     // initialize to zero

    NTL::zz_pPush ppush; // backup NTL's current modulus
    for (long k=0; k<from.length(); k++)   // go over the factors one at a time
    {
        from.params->zzp_context[k].restore();
        NTL::CRT(to, q, from[k]); // incremental CRT
    }
}

// reconstruct the matrix in CRT fromat
// FIXME: can be optimized for many factors
void convert(CRTmatrix& to, const mat_ZZ& from, const TDMatrixParams& prms)
{
    FHE_TIMER_START;
    to.params = (TDMatrixParams*) &prms;
    to.SetLength(prms.kFactors);
    if (prms.kFactors==0) return; // nothing to do

    NTL::zz_pPush ppush; // backup NTL's current modulus

    EXEC_RANGE(to.length(), first, last)

    for (long k=first; k<last; k++)   // go over the factors one at a time
    {
        prms.zzp_context[k].restore();
        conv(to[k], from); // compute M mod current factor
    }

    EXEC_RANGE_END
}

void generateMultiPair(CRTmatrix& A, CRTmatrix& Ainv, long n)
{
    FHE_TIMER_START;

    NTL::Vec<CRTmatrix> Amats(INIT_SIZE, 12, A);
    NTL::Vec<CRTmatrix> Ainvmats(INIT_SIZE, 12, Ainv);

    EXEC_RANGE(12, first, last)
    for (long i=first; i<last; i++)
    {

        CRTmatrix &Apair = Amats[i];
        CRTmatrix &AinvPair = Ainvmats[i];
        generateMatrixPair(Apair, AinvPair, n);
    }

    EXEC_RANGE_END
}

void generateMatrixPair(CRTmatrix& A, CRTmatrix& Ainv, long n)
{
    FHE_TIMER_START;
    const TDMatrixParams& prms = *(A.params);

    A.SetLength(prms.kFactors);
    Ainv.SetLength(prms.kFactors);

    zz_pPush push; // backup the NTL current modulus

    // It's possible to do parallilizing at this level
    // need to call SetNumThreads(nt) in main.

    // NOTE: the zz_pPush in the "main" thread is sufficient,
    // as the "worker" threads do not (and should not) assume
    // any particular contextual state

#ifdef DEBUGPRINT
        cerr << "*** " << prms.kFactors << " " << AvailableThreads() << "\n";
#endif // DEBUGPRINT

    EXEC_RANGE(prms.kFactors, first, last)
    for (long i=first; i<last; i++)
    {

        mat_zz_p& Amod = A[i];
        mat_zz_p& AinvMod = Ainv[i];
        prms.zzp_context[i].restore(); // set i'th factor as NTL's current modulus
        GenerateMatrixPair(Amod, AinvMod, n);
    }
    EXEC_RANGE_END

    FHE_TIMER_STOP;
};

// binary I/O
#undef CRTMAT_CONVtoZZ
long CRTmatrix::writeToFile(FILE* handle) const
{
    FHE_TIMER_START;

    // write the parameters first
    long count = params->writeToFile(handle);
    assert (this->length() == params->kFactors);

#ifndef CRTMAT_CONVtoZZ
    long n,m;
    if (this->length()<=0)
        n = m = 0;
    else
    {
        n = (*this)[0].NumRows();
        m = (*this)[0].NumCols();
    }
    count += fwrite(&n, sizeof(n), 1, handle); // how many rows
    count += fwrite(&m, sizeof(m), 1, handle); // how many cols

    unsigned char buf[params->e];
    NTL::zz_pPush ppush; // backup NTL's current modulus
    for (long f=0; f<this->length(); f++)   // write one factor at a time
    {
        params->zzp_context[f].restore();
        const Mat<zz_p>& mat = (*this)[f];
        for (long i=0; i<n; i++) for (long j=0; j<m; j++)   // write every entry
            {
                for (long ie=0; ie < params->e; ie++) // write every byte
                    buf[ie] = (rep(mat[i][j]) >> (ie*8)) & 0xff; // extract one byte
                count += fwrite(buf, params->e, 1, handle); // do the actual writing
            }
#ifdef DEBUGPRINT
        //cout << "*this[f]="  << (*this)[f] << endl;
#endif // DEBUGPRINT
    }

#else
    ZZ q;
    mat_ZZ zM;
    convert(zM, q, *this); // convert to ZZ representation, then write it

    long qBytes = NumBytes(q); // # of bytes to represent q
    count += fwrite(&qBytes, sizeof(qBytes),1, handle); // how many bytes per int

    long n = zM.NumRows();
    long m = zM.NumCols();
    count += fwrite(&n, sizeof(n), 1, handle); // how many rows
    count += fwrite(&m, sizeof(m), 1, handle); // how many cols

    unsigned char buf[qBytes];
    for (long i=0; i<n; i++)
        for (long j=0; j<m; j++)
        {
            if (sign(zM[i][j])<0) zM[i][j] += q; // map to interval [0,q-1]
            BytesFromZZ(buf, zM[i][j], qBytes);      // get binary representation
            count += fwrite(buf, qBytes, 1, handle); // write it to file
        }
#endif
    return count;
}

long CRTmatrix::readFromFile(FILE* handle, TDMatrixParams* prmBuf)
{
    FHE_TIMER_START;

#ifdef DEBUGPRINT
    cout << "reading" << endl;
#endif

    assert(params != NULL || prmBuf != NULL); // some pointer must be provided
    long count;

    // If buffer is given, make params point to it and don't overwrite from iput
    if (prmBuf != NULL)
    {
        TDMatrixParams p;
        count = p.readFromFile(handle);
        assert(p == *prmBuf); // sanity check
        params = prmBuf;      // point to given params
    }
    else
        count = params->readFromFile(handle); // overwrite params from input

#ifndef CRTMAT_CONVtoZZ
    long n,m;
    count += fread(&n, sizeof(n), 1, handle); // how many rows
    count += fread(&m, sizeof(m), 1, handle); // how many column
    if (count == 0) return count; //nothing is read

    this->SetLength(params->kFactors);
    if (params->kFactors==0) return count; // nothing to do

    unsigned char buf[params->e];
    NTL::zz_pPush ppush; // backup NTL's current modulus
    for (long f=0; f<this->length(); f++)   // write one factor at a time
    {
        params->zzp_context[f].restore();
        (*this)[f].SetDims(n,m);
        for (long i=0; i<n; i++) for (long j=0; j<m; j++)   // write every entry
            {
                count += fread(buf, params->e, 1, handle); // read from file

                long zp = buf[(params->e)-1];  // assemble the integer, byte by byte
                //for (long ie=1; ie < params->e; ie++)
                for (long ie=((params->e)-2); ie >=0 ; ie--)
                {
                    zp <<= 8;
                    zp += buf[ie];
                }
                (*this)[f][i][j].LoopHole() = zp; // avoid calling rem()
            }
    }
#else
    long qBytes, n, m;
    count += fread(&qBytes, sizeof(qBytes), 1, handle); // how many bytes per int
    count += fread(&n, sizeof(n), 1, handle); // how many rows
    count += fread(&m, sizeof(m), 1, handle); // how many column

    if (count == 0) return count; //nothing is read

    unsigned char buf[qBytes];
    mat_ZZ zM(INIT_SIZE, n, m);
    for (long i=0; i<n; i++) for (long j=0; j<m; j++)
        {
            count += fread(buf, qBytes, 1, handle); // read from file
            ZZFromBytes(zM[i][j], buf, qBytes);     // make a ZZ object
        }

    convert(*this, zM, *params); // convert to CRT representation
#endif
    return count;
}


/******************************************************************************
**
** Copyright (C) 2009-2011 Kyle Lutz <kyle.r.lutz@gmail.com>
** All rights reserved.
**
** This file is a part of the chemkit project. For more information
** see <http://www.chemkit.org>.
**
** Redistribution and use in source and binary forms, with or without
** modification, are permitted provided that the following conditions
** are met:
**
**   * Redistributions of source code must retain the above copyright
**     notice, this list of conditions and the following disclaimer.
**   * Redistributions in binary form must reproduce the above copyright
**     notice, this list of conditions and the following disclaimer in the
**     documentation and/or other materials provided with the distribution.
**   * Neither the name of the chemkit project nor the names of its
**     contributors may be used to endorse or promote products derived
**     from this software without specific prior written permission.
**
** THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
** "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
** LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
** A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
** OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
** SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
** LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
** DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
** THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
** (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
** OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
**
******************************************************************************/

#include "uffcalculation.h"

#include <boost/algorithm/string.hpp>

#include "uffatomtyper.h"
#include "uffforcefield.h"
#include "uffparameters.h"

#include <chemkit/topology.h>
#include <chemkit/constants.h>
#include <chemkit/cartesiancoordinates.h>

// === UffCalculation ====================================================== //
UffCalculation::UffCalculation(int type, int atomCount, int parameterCount)
    : ForceFieldCalculation(type, atomCount, parameterCount)
{
}

// Returns the parameters for the given atom.
const UffAtomParameters* UffCalculation::parameters(const std::string &type) const
{
    const UffForceField *forceField = static_cast<const UffForceField *>(this->forceField());

    return forceField->parameters()->parameters(type);
}

// Returns the bond order of the bond between atom's a and b. If both
// atoms have a resonant type the bond order returned is 1.5.
// Otherwise the integer value of the bond order is returned.
chemkit::Real UffCalculation::bondOrder(size_t a, size_t b) const
{
    const boost::shared_ptr<chemkit::Topology> &topology = this->topology();

    int type = topology->bondedInteractionType(a, b);

    if(type == UffAtomTyper::Resonant){
        return 1.5; // resonant
    }
    else{
        return type;
    }
}

// Returns the length of the bond between two atoms.
chemkit::Real UffCalculation::bondLength(const UffAtomParameters *a, const UffAtomParameters *b, chemkit::Real bondOrder) const
{
    // r_ij = r_i + r_j + r_bo - r_en
    chemkit::Real r_bo = -0.1332 * (a->r + b->r) * log(bondOrder);
    chemkit::Real r_en = ((a->r * b->r) * pow((sqrt(a->X) - sqrt(b->X)), 2)) / (a->X*a->r + b->X*b->r);

    chemkit::Real r_ij = a->r + b->r + r_bo - r_en;

    return r_ij;
}

// === UffBondStrechCalculation ============================================ //
UffBondStrechCalculation::UffBondStrechCalculation(size_t a, size_t b)
    : UffCalculation(BondStrech, 2, 2)
{
    setAtom(0, a);
    setAtom(1, b);
}

bool UffBondStrechCalculation::setup()
{
    const UffAtomParameters *pa = parameters(atomType(0));
    const UffAtomParameters *pb = parameters(atomType(1));

    if(!pa || !pb){
        return false;
    }

    // n = bondorder (1.5 for aromatic, 1.366 for amide)
    chemkit::Real bondorder = bondOrder(atom(0), atom(1));

    chemkit::Real r0 = bondLength(pa, pb, bondorder);

    // parameter(1) = k_ij = 664.12 * (Z*_i * Z*_j) / r_ij^3
    chemkit::Real za = pa->Z;
    chemkit::Real zb = pb->Z;
    chemkit::Real kb = 664.12 * (za * zb) / pow(r0, 3);

    setParameter(0, kb);
    setParameter(1, r0);

    return true;
}

chemkit::Real UffBondStrechCalculation::energy(const chemkit::CartesianCoordinates *coordinates) const
{
    size_t a = atom(0);
    size_t b = atom(1);

    chemkit::Real kb = parameter(0);
    chemkit::Real r0 = parameter(1);
    chemkit::Real r = coordinates->distance(a, b);

    return 0.5 * kb * pow(r - r0, 2);
}

std::vector<chemkit::Vector3> UffBondStrechCalculation::gradient(const chemkit::CartesianCoordinates *coordinates) const
{
    size_t a = atom(0);
    size_t b = atom(1);

    chemkit::Real kb = parameter(0);
    chemkit::Real r0 = parameter(1);
    chemkit::Real r = coordinates->distance(a, b);

    // dE/dr
    chemkit::Real de_dr = kb * (r - r0);

    boost::array<chemkit::Vector3, 2> gradient = coordinates->distanceGradient(a, b);

    gradient[0] *= de_dr;
    gradient[1] *= de_dr;

    return std::vector<chemkit::Vector3>(gradient.begin(), gradient.end());
}

// === UffAngleBendCalculation ============================================= //
UffAngleBendCalculation::UffAngleBendCalculation(size_t a, size_t b, size_t c)
    : UffCalculation(AngleBend, 3, 4)
{
    setAtom(0, a);
    setAtom(1, b);
    setAtom(2, c);
}

bool UffAngleBendCalculation::setup()
{
    const UffAtomParameters *pa = parameters(atomType(0));
    const UffAtomParameters *pb = parameters(atomType(1));
    const UffAtomParameters *pc = parameters(atomType(2));

    if(!pa || !pb || !pc){
        return false;
    }

    chemkit::Real theta0 = pb->theta * chemkit::constants::DegreesToRadians;

    chemkit::Real bo_ij = bondOrder(atom(0), atom(1));
    chemkit::Real bo_jk = bondOrder(atom(1), atom(2));

    chemkit::Real r_ab = bondLength(pa, pb, bo_ij);
    chemkit::Real r_bc = bondLength(pb, pc, bo_jk);
    chemkit::Real r_ac = sqrt(pow(r_ab, 2)  + pow(r_bc, 2) - (2.0 * r_ab * r_bc * cos(theta0)));

    chemkit::Real beta = 664.12 / (r_ab * r_bc);

    chemkit::Real z_a = pa->Z;
    chemkit::Real z_c = pc->Z;

    // equation 13
    chemkit::Real ka = beta * ((z_a * z_c) / pow(r_ac, 5)) * r_ab * r_bc * (3.0 * r_ab * r_bc * (1.0 - pow(cos(theta0), 2.0)) - (pow(r_ac, 2.0) * cos(theta0)));

    setParameter(0, ka);

    chemkit::Real sinTheta0 = sin(theta0);

    // clamp sin(theta0) to 1e-3 because for some atoms theta0 == pi which
    // would lead to a division by zero error when calculating c2 below
    if(std::abs(sinTheta0) < 1e-3){
        sinTheta0 = 1e-3;
    }

    chemkit::Real c2 = 1 / (4 * pow(sinTheta0, 2));
    chemkit::Real c1 = -4 * c2 * cos(theta0);
    chemkit::Real c0 = c2 * (2 * pow(cos(theta0), 2) + 1);

    setParameter(1, c0);
    setParameter(2, c1);
    setParameter(3, c2);

    return true;
}

chemkit::Real UffAngleBendCalculation::energy(const chemkit::CartesianCoordinates *coordinates) const
{
    size_t a = atom(0);
    size_t b = atom(1);
    size_t c = atom(2);

    chemkit::Real ka = parameter(0);
    chemkit::Real c0 = parameter(1);
    chemkit::Real c1 = parameter(2);
    chemkit::Real c2 = parameter(3);

    chemkit::Real theta = coordinates->angleRadians(a, b, c);

    return ka * (c0 + (c1 * cos(theta)) + (c2 * cos(2*theta)));
}

std::vector<chemkit::Vector3> UffAngleBendCalculation::gradient(const chemkit::CartesianCoordinates *coordinates) const
{
    size_t a = atom(0);
    size_t b = atom(1);
    size_t c = atom(2);

    chemkit::Real ka = parameter(0);
    chemkit::Real c1 = parameter(2);
    chemkit::Real c2 = parameter(3);

    chemkit::Real theta = coordinates->angleRadians(a, b, c);

    // dE/dtheta
    chemkit::Real de_dtheta = -ka * (c1 * sin(theta) + 2 * c2 * sin(2 * theta));

    boost::array<chemkit::Vector3, 3> gradient = coordinates->angleGradientRadians(a, b, c);

    gradient[0] *= de_dtheta;
    gradient[1] *= de_dtheta;
    gradient[2] *= de_dtheta;

    return std::vector<chemkit::Vector3>(gradient.begin(), gradient.end());
}

// === UffTorsionCalculation =============================================== //
UffTorsionCalculation::UffTorsionCalculation(size_t a, size_t b, size_t c, size_t d)
    : UffCalculation(Torsion, 4, 3)
{
    setAtom(0, a);
    setAtom(1, b);
    setAtom(2, c);
    setAtom(3, d);
}

bool UffTorsionCalculation::setup()
{
    UffForceField *forceField = static_cast<UffForceField *>(this->forceField());
    const boost::shared_ptr<chemkit::Topology> &topology = this->topology();

    size_t b = atom(1);
    size_t c = atom(2);

    std::string typeB = topology->type(b);
    std::string typeC = topology->type(c);

    if(typeB.length() < 3 || typeC.length() < 3){
        return false;
    }

    const UffAtomParameters *pb = parameters(typeB);
    const UffAtomParameters *pc = parameters(typeC);

    chemkit::Real V = 0;
    chemkit::Real n = 0;
    chemkit::Real phi0 = 0;

    // sp3-sp3
    if(typeB[2] == '3' && typeC[2] == '3'){

        // exception for two group six atoms
        if(forceField->isGroupSix(atom(1)) && forceField->isGroupSix(atom(2))){
            if(boost::starts_with(typeB, "O_") && boost::starts_with(typeC, "O_")){
                V = 2; // sqrt(2*2)
            }
            else if(boost::starts_with(typeB, "O_") || boost::starts_with(typeC, "O_")){
                V = sqrt(2 * 6.8);
            }
            else{
                V = sqrt(6.8 * 6.8);
            }

            n = 2;
            phi0 = 90;
        }

        // general case
        else{
            // equation 16
            V = sqrt(pb->V * pc->V);

            n = 3;
            phi0 = 180 * chemkit::constants::DegreesToRadians;
        }
    }
    // sp2-sp2
    else if((typeB[2] == '2' || typeB[2] == 'R') && (typeC[2] == '2' || typeC[2] == 'R')){
        chemkit::Real bondorder = bondOrder(b, c);

        // equation 17
        V = 5 * sqrt(pb->U * pc->U) * (1 + 4.18 * log(bondorder));

        n = 2;
        phi0 = 180 * chemkit::constants::DegreesToRadians;
    }
    // group 6 sp3 - any sp2 or R
    else if((forceField->isGroupSix(b) && (typeC[2] == '2' || typeC[2] == 'R')) ||
            (forceField->isGroupSix(c) && (typeB[2] == '2' || typeB[2] == 'R'))){
        chemkit::Real bondorder = bondOrder(b, c);

        // equation 17
        V = 5 * sqrt(pb->U * pc->U) * (1 + 4.18 * log(bondorder));

        n = 2;
        phi0 = 90 * chemkit::constants::DegreesToRadians;
    }
    // sp3-sp2
    else if((typeB[2] == '3' && (typeC[2] == '2' || typeC[2] == 'R')) ||
            (typeC[2] == '3' && (typeB[2] == '2' || typeB[2] == 'R'))){
        V = 1;
        n = 6;
        phi0 = 0;
    }
    else{
        return false;
    }

    setParameter(0, V);
    setParameter(1, n);
    setParameter(2, phi0);

    return true;
}

chemkit::Real UffTorsionCalculation::energy(const chemkit::CartesianCoordinates *coordinates) const
{
    size_t a = atom(0);
    size_t b = atom(1);
    size_t c = atom(2);
    size_t d = atom(3);

    chemkit::Real V = parameter(0);
    chemkit::Real n = parameter(1);
    chemkit::Real phi0 = parameter(2);

    chemkit::Real phi = coordinates->torsionAngleRadians(a, b, c, d);

    return 0.5 * V * (1 - cos(n * phi0) * cos(n * phi));
}

std::vector<chemkit::Vector3> UffTorsionCalculation::gradient(const chemkit::CartesianCoordinates *coordinates) const
{
    size_t a = atom(0);
    size_t b = atom(1);
    size_t c = atom(2);
    size_t d = atom(3);

    chemkit::Real V = parameter(0);
    chemkit::Real n = parameter(1);
    chemkit::Real phi0 = parameter(2);

    chemkit::Real phi = coordinates->torsionAngleRadians(a, b, c, d);

    // dE/dphi
    chemkit::Real de_dphi = 0.5 * V * n * cos(n * phi0) * sin(n * phi);

    boost::array<chemkit::Vector3, 4> gradient = coordinates->torsionAngleGradientRadians(a, b, c, d);

    gradient[0] *= de_dphi;
    gradient[1] *= de_dphi;
    gradient[2] *= de_dphi;
    gradient[3] *= de_dphi;

    return std::vector<chemkit::Vector3>(gradient.begin(), gradient.end());
}

// === UffInversionCalculation ============================================= //
UffInversionCalculation::UffInversionCalculation(size_t a, size_t b, size_t c, size_t d)
    : UffCalculation(Inversion, 4, 4)
{
    setAtom(0, a);
    setAtom(1, b);
    setAtom(2, c);
    setAtom(3, d);
}

bool UffInversionCalculation::setup()
{
    const boost::shared_ptr<chemkit::Topology> &topology = this->topology();

    // b is the center atom
    size_t a = atom(0);
    size_t b = atom(1);
    size_t c = atom(2);
    size_t d = atom(3);

    std::string typeA = topology->type(a);
    std::string typeB = topology->type(b);
    std::string typeC = topology->type(c);
    std::string typeD = topology->type(d);

    chemkit::Real k = 0;
    chemkit::Real c0 = 0;
    chemkit::Real c1 = 0;
    chemkit::Real c2 = 0;

    // sp2 carbon
    if(typeB == "C_2" || typeB == "C_R"){
        if(typeA == "O_2" || typeC == "O_2" || typeD == "O_2"){
            k = 50;
        }
        else{
            k = 6;
        }

        c0 = 1;
        c1 = -1;
        c2 = 0;
    }

    // divide by 3
    k /= 3;

    setParameter(0, k);
    setParameter(1, c0);
    setParameter(2, c1);
    setParameter(3, c2);

    return true;
}

chemkit::Real UffInversionCalculation::energy(const chemkit::CartesianCoordinates *coordinates) const
{
    size_t a = atom(0);
    size_t b = atom(1);
    size_t c = atom(2);
    size_t d = atom(3);

    chemkit::Real k = parameter(0);
    chemkit::Real c0 = parameter(1);
    chemkit::Real c1 = parameter(2);
    chemkit::Real c2 = parameter(3);

    chemkit::Real w = coordinates->wilsonAngleRadians(a, b, c, d);
    chemkit::Real y = w + (chemkit::constants::Pi / 2.0);

    return k * (c0 + c1 * sin(y) + c2 * cos(2 * y));
}

std::vector<chemkit::Vector3> UffInversionCalculation::gradient(const chemkit::CartesianCoordinates *coordinates) const
{
    size_t a = atom(0);
    size_t b = atom(1);
    size_t c = atom(2);
    size_t d = atom(3);

    chemkit::Real k = parameter(0);
    chemkit::Real c1 = parameter(2);
    chemkit::Real c2 = parameter(3);

    chemkit::Real w = coordinates->wilsonAngleRadians(a, b, c, d);
    chemkit::Real y = w + (chemkit::constants::Pi / 2.0);

    // dE/dw
    chemkit::Real de_dw = k * (c1 * cos(y) - 2 * c2 * sin(2 * y));

    boost::array<chemkit::Vector3, 4> gradient = coordinates->wilsonAngleGradientRadians(a, b, c, d);

    gradient[0] *= de_dw;
    gradient[1] *= de_dw;
    gradient[2] *= de_dw;
    gradient[3] *= de_dw;

    return std::vector<chemkit::Vector3>(gradient.begin(), gradient.end());
}

// === UffVanDerWaalsCalculation =========================================== //
UffVanDerWaalsCalculation::UffVanDerWaalsCalculation(size_t a, size_t b)
    : UffCalculation(VanDerWaals, 2, 2)
{
    setAtom(0, a);
    setAtom(1, b);
}

bool UffVanDerWaalsCalculation::setup()
{
    const UffAtomParameters *pa = parameters(atomType(0));
    const UffAtomParameters *pb = parameters(atomType(1));
    if(!pa || !pb){
        return false;
    }

    // equation 22
    chemkit::Real d = sqrt(pa->D * pb->D);

    // equation 21b
    chemkit::Real x = sqrt(pa->x * pb->x);

    setParameter(0, d);
    setParameter(1, x);

    return true;
}

chemkit::Real UffVanDerWaalsCalculation::energy(const chemkit::CartesianCoordinates *coordinates) const
{
    size_t a = atom(0);
    size_t b = atom(1);

    chemkit::Real d = parameter(0);
    chemkit::Real x = parameter(1);
    chemkit::Real r = coordinates->distance(a, b);

    return d * (-2 * pow(x/r, 6) + pow(x/r, 12));
}

std::vector<chemkit::Vector3> UffVanDerWaalsCalculation::gradient(const chemkit::CartesianCoordinates *coordinates) const
{
    size_t a = atom(0);
    size_t b = atom(1);

    chemkit::Real d = parameter(0);
    chemkit::Real x = parameter(1);
    chemkit::Real r = coordinates->distance(a, b);

    // dE/dr
    chemkit::Real de_dr = -12 * d * x / pow(r, 2) * (pow(x/r, 11) - pow(x/r, 5));

    boost::array<chemkit::Vector3, 2> gradient = coordinates->distanceGradient(a, b);

    gradient[0] *= de_dr;
    gradient[1] *= de_dr;

    return std::vector<chemkit::Vector3>(gradient.begin(), gradient.end());
}

// === UffElectrostaticCalculation ========================================= //
UffElectrostaticCalculation::UffElectrostaticCalculation(size_t a, size_t b)
    : UffCalculation(Electrostatic, 2, 2)
{
    setAtom(0, a);
    setAtom(1, b);
}

bool UffElectrostaticCalculation::setup()
{
    return false;
}

chemkit::Real UffElectrostaticCalculation::energy(const chemkit::CartesianCoordinates *coordinates) const
{
    size_t a = atom(0);
    size_t b = atom(1);

    chemkit::Real qa = parameter(0);
    chemkit::Real qb = parameter(1);

    chemkit::Real e = 1;
    chemkit::Real r = coordinates->distance(a, b);

    return 332.037 * (qa * qb) / (e * r);
}


// Compare SoA, AoS, and Eigen (with expression parser)
// Ideal situation: everything known at compile time
//
// g++ ampBench.cpp -O3 -march=native -std=c++17 -ffast-math -I/usr/local/opt/eigen/include/eigen3 -I/usr/local/Cellar/boost/1.72.0/include/boost/ -Rpass-missed=loop-vectorize -lboost_timer -o ampBench
//

// SoA:
//  1.167841s wall, 1.100000s user + 0.060000s system = 1.160000s CPU (99.3%)
//
// SoA (array):
//  0.637546s wall, 0.570000s user + 0.060000s system = 0.630000s CPU (98.8%)
//
// SoA (stack):
//  0.926747s wall, 0.870000s user + 0.060000s system = 0.930000s CPU (100.4%)
//
// AoS:
//  14.228890s wall, 12.890000s user + 1.260000s system = 14.150000s CPU (99.4%)
//
// AoS (stack):
//  1.105483s wall, 0.910000s user + 0.180000s system = 1.090000s CPU (98.6%)
//
// Eigen:
//  1.072285s wall, 0.930000s user + 0.140000s system = 1.070000s CPU (99.8%)

#include <iostream>
#include <vector>
#include <array>
#include <memory>

#include <Eigen/Dense>

#include <boost/timer/timer.hpp>

// #include <blaze/math/StaticVector.h>
// #include <blaze/math/DynamicVector.h>

#include <blaze/Math.h>

#include "xtensor/xtensor.hpp"
#include "xtensor/xfixed.hpp"
#include "xtensor/xarray.hpp"
#include "xtensor/xio.hpp"

using blaze::StaticVector;
using blaze::DynamicVector;

static const int n_global = 1E6;

// Spin 5 - no need for template here
float legFunc(float cosHel)
{
    return -32.0*(63.0*cosHel*cosHel*cosHel*cosHel*cosHel - 70.0*cosHel*cosHel*cosHel + 15.0*cosHel)/63.0;
}

// AoS
struct ResParams
{
public:

    float mass;
    float qTerm;
    float ffRatioP;
    float ffRatioR;
    float spinTerms;

    float ampRe;
    float ampIm;

    ResParams()
    : mass(0),
      qTerm(0),
      ffRatioP(0),
      ffRatioR(0),
      spinTerms(0),
      ampRe(0),
      ampIm(0) { }

};

struct ResParamsSoA
{
public:

    float resMass;
    float resWidth;

    std::vector<float> * mass;
    std::vector<float> * qTerm;
    std::vector<float> * ffRatioP;
    std::vector<float> * ffRatioR;
    std::vector<float> * spinTerms;

    std::vector<float> * ampRe;
    std::vector<float> * ampIm;

    // Deal with these guys later...
    static const int spin = 4;
    static const int spinType = 1;

    ResParamsSoA(int s)
    {
        mass = new std::vector<float>(s);
        qTerm = new std::vector<float>(s);
        ffRatioP = new std::vector<float>(s);
        ffRatioR = new std::vector<float>(s);
        spinTerms = new std::vector<float>(s);

        ampRe = new std::vector<float>(s);
        ampIm = new std::vector<float>(s);
    }

    ~ResParamsSoA()
    {
        delete mass;
        delete qTerm;
        delete ffRatioP;
        delete ffRatioR;
        delete spinTerms;

        delete ampRe;
        delete ampIm;
    }

};

struct ResParamsSoAArray
{
public:

    float resMass;
    float resWidth;

    std::array<float, n_global> * mass;
    std::array<float, n_global> * qTerm;
    std::array<float, n_global> * ffRatioP;
    std::array<float, n_global> * ffRatioR;
    std::array<float, n_global> * spinTerms;

    std::array<float, n_global> * ampRe;
    std::array<float, n_global> * ampIm;

    // Deal with these guys later...
    static const int spin = 4;
    static const int spinType = 1;

    ResParamsSoAArray(int s)
    {
        mass = new std::array<float, n_global>();
        qTerm = new std::array<float, n_global>();
        ffRatioP = new std::array<float, n_global>();
        ffRatioR = new std::array<float, n_global>();
        spinTerms = new std::array<float, n_global>();

        ampRe = new std::array<float, n_global>();
        ampIm = new std::array<float, n_global>();
    }

    ~ResParamsSoAArray()
    {
        delete mass;
        delete qTerm;
        delete ffRatioP;
        delete ffRatioR;
        delete spinTerms;

        delete ampRe;
        delete ampIm;
    }

};

struct ResParamsSoAStack
{
public:

    float resMass;
    float resWidth;

    std::vector<float> mass;
    std::vector<float> qTerm;
    std::vector<float> ffRatioP;
    std::vector<float> ffRatioR;
    std::vector<float> spinTerms;

    std::vector<float> ampRe;
    std::vector<float> ampIm;

    // Deal with these guys later...
    static const int spin = 4;
    static const int spinType = 1;

    ResParamsSoAStack(int n) {

        mass.resize(n);
        qTerm.resize(n);
        ffRatioP.resize(n);
        ffRatioR.resize(n);
        spinTerms.resize(n);

        ampRe.resize(n);
        ampIm.resize(n);
    }

    ~ResParamsSoAStack() { }

};

struct ResParamsEigen
{
public:

    float resMass;
    float resWidth;

    // Too large to allocate on stack (?)
    // Possible to do more with static allocation?

    Eigen::Array<float,Eigen::Dynamic,1> mass;
    Eigen::Array<float,Eigen::Dynamic,1> qTerm;
    Eigen::Array<float,Eigen::Dynamic,1> ffRatioP;
    Eigen::Array<float,Eigen::Dynamic,1> ffRatioR;
    Eigen::Array<float,Eigen::Dynamic,1> spinTerms;

    Eigen::Array<float,Eigen::Dynamic,1> ampRe;
    Eigen::Array<float,Eigen::Dynamic,1> ampIm;

    // Deal with these guys later...
    static const int spin = 4;
    static const int spinType = 1;

    ResParamsEigen(int n)
    {

        resMass = 0;
        resWidth = 0;

        mass = Eigen::Array<float,Eigen::Dynamic,1>(n);
        qTerm = Eigen::Array<float,Eigen::Dynamic,1>(n);
        ffRatioP = Eigen::Array<float,Eigen::Dynamic,1>(n);
        ffRatioR = Eigen::Array<float,Eigen::Dynamic,1>(n);
        spinTerms = Eigen::Array<float,Eigen::Dynamic,1>(n);

        ampRe = Eigen::Array<float,Eigen::Dynamic,1>(n);
        ampIm = Eigen::Array<float,Eigen::Dynamic,1>(n);
    }

    ~ResParamsEigen()
    {

    }

};

struct ResParamsBlaze
{
public:

    float resMass;
    float resWidth;

    static const int n = n_global;

    // DynamicVector<float> mass;
    // DynamicVector<float> qTerm;
    // DynamicVector<float> ffRatioP;
    // DynamicVector<float> ffRatioR;
    // DynamicVector<float> spinTerms;
    //
    // DynamicVector<float> ampRe;
    // DynamicVector<float> ampIm;

    StaticVector<float, n> mass;
    StaticVector<float, n> qTerm;
    StaticVector<float, n> ffRatioP;
    StaticVector<float, n> ffRatioR;
    StaticVector<float, n> spinTerms;

    StaticVector<float, n> ampRe;
    StaticVector<float, n> ampIm;

    // Deal with these guys later...
    static const int spin = 4;
    static const int spinType = 1;

    ResParamsBlaze()
    {

        resMass = 0;
        resWidth = 0;

        // mass = DynamicVector<float>(this->n);
        // qTerm = DynamicVector<float>(this->n);
        // ffRatioP = DynamicVector<float>(this->n);
        // ffRatioR = DynamicVector<float>(this->n);
        // spinTerms = DynamicVector<float>(this->n);
        //
        // ampRe = DynamicVector<float>(this->n);
        // ampIm = DynamicVector<float>(this->n);

    }

    ~ResParamsBlaze()
    {

    }

};

struct ResParamsXTensor
{
public:

    float resMass;
    float resWidth;

    static const int n = n_global;

    // DynamicVector<float> mass;
    // DynamicVector<float> qTerm;
    // DynamicVector<float> ffRatioP;
    // DynamicVector<float> ffRatioR;
    // DynamicVector<float> spinTerms;
    //
    // DynamicVector<float> ampRe;
    // DynamicVector<float> ampIm;

    xt::xtensor_fixed<double, xt::xshape<n, 1>> mass;
    xt::xtensor_fixed<double, xt::xshape<n, 1>> qTerm;
    xt::xtensor_fixed<double, xt::xshape<n, 1>> ffRatioP;
    xt::xtensor_fixed<double, xt::xshape<n, 1>> ffRatioR;
    xt::xtensor_fixed<double, xt::xshape<n, 1>> spinTerms;

    xt::xtensor_fixed<double, xt::xshape<n, 1>> ampRe;
    xt::xtensor_fixed<double, xt::xshape<n, 1>> ampIm;

    // Deal with these guys later...
    static const int spin = 4;
    static const int spinType = 1;

    ResParamsXTensor()
    {

        resMass = 0;
        resWidth = 0;

        // mass = DynamicVector<float>(this->n);
        // qTerm = DynamicVector<float>(this->n);
        // ffRatioP = DynamicVector<float>(this->n);
        // ffRatioR = DynamicVector<float>(this->n);
        // spinTerms = DynamicVector<float>(this->n);
        //
        // ampRe = DynamicVector<float>(this->n);
        // ampIm = DynamicVector<float>(this->n);

    }

    ~ResParamsXTensor()
    {

    }

};

void calcAmpSoA(const ResParamsSoA & inParams)
{
    float resMass = inParams.resMass;
    float resWidth = inParams.resWidth;

    for (int i = 0; i < inParams.mass->size(); i++) {
        float totWidth = resWidth * inParams.qTerm->at(i);
        totWidth *= (resMass / inParams.mass->at(i));
        totWidth *= inParams.ffRatioP->at(i) * inParams.ffRatioR->at(i);

        float m2 = inParams.mass->at(i) * inParams.mass->at(i);
        float m2Term = resMass * resMass - m2;

        float scale = inParams.spinTerms->at(i);
        scale /= m2Term * m2Term + resMass * resMass * totWidth * totWidth;
        scale *= inParams.ffRatioP->at(i) * inParams.ffRatioR->at(i); // Optional -> template specialise?
        scale *= legFunc(inParams.spinTerms->at(i));

        inParams.ampRe->at(i) = m2Term * scale;
        inParams.ampIm->at(i) = resMass * totWidth * scale;
    }
}

void calcAmpSoA(const ResParamsSoAArray & inParams)
{
    float resMass = inParams.resMass;
    float resWidth = inParams.resWidth;

    #pragma clang loop vectorize(assume_safety)
    for (int i = 0; i < n_global; i++) {
        float totWidth = resWidth * inParams.qTerm->at(i);
        totWidth *= (resMass / inParams.mass->at(i));
        totWidth *= inParams.ffRatioP->at(i) * inParams.ffRatioR->at(i);

        float m2 = inParams.mass->at(i) * inParams.mass->at(i);
        float m2Term = resMass * resMass - m2;

        float scale = inParams.spinTerms->at(i);
        scale /= m2Term * m2Term + resMass * resMass * totWidth * totWidth;
        scale *= inParams.ffRatioP->at(i) * inParams.ffRatioR->at(i); // Optional -> template specialise?
        scale *= legFunc(inParams.spinTerms->at(i));

        inParams.ampRe->at(i) = m2Term * scale;
        inParams.ampIm->at(i) = resMass * totWidth * scale;
    }
}

void calcAmpSoAStack(ResParamsSoAStack & inParams)
{
    float resMass = inParams.resMass;
    float resWidth = inParams.resWidth;

    for (int i = 0; i < inParams.mass.size(); i++) {
        float totWidth = resWidth * inParams.qTerm.at(i);
        totWidth *= (resMass / inParams.mass.at(i));
        totWidth *= inParams.ffRatioP.at(i) * inParams.ffRatioR.at(i);

        float m2 = inParams.mass.at(i) * inParams.mass.at(i);
        float m2Term = resMass * resMass - m2;

        float scale = inParams.spinTerms.at(i);
        scale /= m2Term * m2Term + resMass * resMass * totWidth * totWidth;
        scale *= inParams.ffRatioP.at(i) * inParams.ffRatioR.at(i); // Optional -> template specialise?
        scale *= legFunc(inParams.spinTerms.at(i));

        inParams.ampRe.at(i) = m2Term * scale;
        inParams.ampIm.at(i) = resMass * totWidth * scale;
    }
}

void calcAmpAoS(const std::vector<std::unique_ptr<ResParams> > & inParams, float resMass, float resWidth)
{

    for (int i = 0; i < inParams.size(); i++) {
        float totWidth = resWidth * inParams.at(i)->qTerm;
        totWidth *= (resMass / inParams.at(i)->mass);
        totWidth *= inParams.at(i)->ffRatioP * inParams.at(i)->ffRatioR;

        float m2 = inParams.at(i)->mass * inParams.at(i)->mass;
        float m2Term = resMass * resMass - m2;

        float scale = inParams.at(i)->spinTerms;
        scale /= m2Term * m2Term + resMass * resMass * totWidth * totWidth;
        scale *= inParams.at(i)->ffRatioP * inParams.at(i)->ffRatioR; // Optional -> template specialise?
        scale *= legFunc(inParams.at(i)->spinTerms);

        inParams.at(i)->ampRe = m2Term * scale;
        inParams.at(i)->ampIm = resMass * totWidth * scale;
    }
}

void calcAmpAoSStack(std::vector<ResParams> & inParams, float resMass, float resWidth)
{

    for (int i = 0; i < inParams.size(); i++) {
        float totWidth = resWidth * inParams.at(i).qTerm;
        totWidth *= (resMass / inParams.at(i).mass);
        totWidth *= inParams.at(i).ffRatioP * inParams.at(i).ffRatioR;

        float m2 = inParams.at(i).mass * inParams.at(i).mass;
        float m2Term = resMass * resMass - m2;

        float scale = inParams.at(i).spinTerms;
        scale /= m2Term * m2Term + resMass * resMass * totWidth * totWidth;
        scale *= inParams.at(i).ffRatioP * inParams.at(i).ffRatioR; // Optional -> template specialise?
        scale *= legFunc(inParams.at(i).spinTerms);

        inParams.at(i).ampRe = m2Term * scale;
        inParams.at(i).ampIm = resMass * totWidth * scale;
    }
}

void calcAmpEigen(ResParamsEigen & inParams)
{
    float resMass = inParams.resMass;
    float resWidth = inParams.resWidth;

    Eigen::Array<float,Eigen::Dynamic,1> totWidth = resWidth * (inParams.qTerm);
    totWidth *= (resMass / (inParams.mass));
    totWidth *= totWidth * (inParams.ffRatioP) * (inParams.ffRatioR);

    Eigen::Array<float,Eigen::Dynamic,1> m2 = (inParams.mass) * (inParams.mass);
    Eigen::Array<float,Eigen::Dynamic,1> m2Term = resMass * resMass - m2;

    Eigen::Array<float,Eigen::Dynamic,1> scale = (inParams.spinTerms);
    scale /= m2Term * m2Term + resMass * resMass * totWidth * totWidth;
    scale *= (inParams.ffRatioP) * (inParams.ffRatioR);

    inParams.ampRe = m2Term * scale;
    inParams.ampIm = resMass * totWidth * scale;
}

void calcAmpBlaze(ResParamsBlaze & inParams)
{
    float resMass = inParams.resMass;
    float resWidth = inParams.resWidth;

    StaticVector<float, n_global> totWidth = resWidth * (inParams.qTerm);
    totWidth *= (resMass / (inParams.mass));
    totWidth *= totWidth * (inParams.ffRatioP) * (inParams.ffRatioR);

    StaticVector<float, n_global> m2 = (inParams.mass) * (inParams.mass);
    StaticVector<float, n_global> m2Term = resMass * resMass - m2;

    StaticVector<float, n_global> scale = (inParams.spinTerms);
    scale /= m2Term * m2Term + resMass * resMass * totWidth * totWidth;
    scale *= (inParams.ffRatioP) * (inParams.ffRatioR);

    inParams.ampRe = m2Term * scale;
    inParams.ampIm = resMass * totWidth * scale;
}

void calcAmpXTensor(ResParamsXTensor & inParams)
{
    float resMass = inParams.resMass;
    float resWidth = inParams.resWidth;

    xt::xtensor_fixed<double, xt::xshape<n_global, 1>> totWidth = resWidth * (inParams.qTerm);
    totWidth *= (resMass / (inParams.mass));
    totWidth *= totWidth * (inParams.ffRatioP) * (inParams.ffRatioR);

    xt::xtensor_fixed<double, xt::xshape<n_global, 1>> m2 = (inParams.mass) * (inParams.mass);
    xt::xtensor_fixed<double, xt::xshape<n_global, 1>> m2Term = resMass * resMass - m2;

    xt::xtensor_fixed<double, xt::xshape<n_global, 1>> scale = (inParams.spinTerms);
    scale /= m2Term * m2Term + resMass * resMass * totWidth * totWidth;
    scale *= (inParams.ffRatioP) * (inParams.ffRatioR);

    inParams.ampRe = m2Term * scale;
    inParams.ampIm = resMass * totWidth * scale;
}

void benchAoS()
{
    static const int n = n_global;

    std::vector<std::unique_ptr<ResParams>> resParamsAoS;
    resParamsAoS.resize(n);

    for (auto & p : resParamsAoS) {

        p = std::make_unique<ResParams>();

        p->qTerm = 0.05;
        p->mass = 0.3;
        p->ffRatioP = 3.3;
        p->ffRatioR = 1.0;
        p->spinTerms = 1.0;

        p->ampRe = 1.0;
        p->ampIm = 1.0;
    }

    float resMass = 1.0;
    float resWidth = 0.1;

    calcAmpAoS(resParamsAoS, resMass, resWidth);

    // std::cout << resParamsAoS.at(5)->ampRe << std::endl;

}

void benchAoSStack()
{
    int n = n_global;

    std::vector<ResParams> resParamsAoSStack;
    resParamsAoSStack.reserve(n);

    for (int i = 0; i < n; i++) {

        ResParams params;

        params.qTerm = 0.05;
        params.mass = 0.3;
        params.ffRatioP = 3.3;
        params.ffRatioR = 1.0;
        params.spinTerms = 1.0;

        params.ampRe = 1.0;
        params.ampIm = 1.0;

        resParamsAoSStack.push_back(params);

    }

    float resMass = 1.0;
    float resWidth = 0.1;

    calcAmpAoSStack(resParamsAoSStack, resMass, resWidth);

    // std::cout << resParamsAoSStack.at(5).ampRe << std::endl;
    // std::cout << resParamsAoSStack.at(5).ampIm << std::endl;
}

void benchSoA()
{
    ResParamsSoA parsR(n_global);

    std::random_device r;

    std::default_random_engine e1(r());
    std::uniform_real_distribution<float> uniform_dist(0, 1);

    std::fill(parsR.qTerm->begin(), parsR.qTerm->end(), uniform_dist(e1));
    std::fill(parsR.mass->begin(), parsR.mass->end(), uniform_dist(e1));
    std::fill(parsR.ffRatioP->begin(), parsR.ffRatioP->end(), uniform_dist(e1));
    std::fill(parsR.ffRatioR->begin(), parsR.ffRatioR->end(), uniform_dist(e1));
    std::fill(parsR.spinTerms->begin(), parsR.spinTerms->end(), uniform_dist(e1));

    std::fill(parsR.ampRe->begin(), parsR.ampRe->end(), uniform_dist(e1));
    std::fill(parsR.ampIm->begin(), parsR.ampIm->end(), uniform_dist(e1));

    parsR.resMass =uniform_dist(e1);
    parsR.resWidth =uniform_dist(e1);

    calcAmpSoA(parsR);

    // std::cout << (parsR.ampRe)->at(5) << std::endl;
    // std::cout << (parsR.ampIm)->at(5) << std::endl;
}

void benchSoAArray()
{
    ResParamsSoAArray parsR(n_global);

    std::random_device r;

    std::default_random_engine e1(r());
    std::uniform_real_distribution<float> uniform_dist(0, 1);

    std::fill(parsR.qTerm->begin(), parsR.qTerm->end(), uniform_dist(e1));
    std::fill(parsR.mass->begin(), parsR.mass->end(), uniform_dist(e1));
    std::fill(parsR.ffRatioP->begin(), parsR.ffRatioP->end(), uniform_dist(e1));
    std::fill(parsR.ffRatioR->begin(), parsR.ffRatioR->end(), uniform_dist(e1));
    std::fill(parsR.spinTerms->begin(), parsR.spinTerms->end(), uniform_dist(e1));

    std::fill(parsR.ampRe->begin(), parsR.ampRe->end(), uniform_dist(e1));
    std::fill(parsR.ampIm->begin(), parsR.ampIm->end(), uniform_dist(e1));

    parsR.resMass = uniform_dist(e1);
    parsR.resWidth = uniform_dist(e1);

    calcAmpSoA(parsR);

    // std::cout << (parsR.ampRe)->at(5) << std::endl;
    // std::cout << (parsR.ampIm)->at(5) << std::endl;
}

void benchSoAStack()
{
    ResParamsSoAStack parsR(n_global);

    std::random_device r;

    std::default_random_engine e1(r());
    std::uniform_real_distribution<float> uniform_dist(0, 1);

    std::fill(parsR.qTerm.begin(), parsR.qTerm.end(), uniform_dist(e1));
    std::fill(parsR.mass.begin(), parsR.mass.end(), uniform_dist(e1));
    std::fill(parsR.ffRatioP.begin(), parsR.ffRatioP.end(), uniform_dist(e1));
    std::fill(parsR.ffRatioR.begin(), parsR.ffRatioR.end(), uniform_dist(e1));
    std::fill(parsR.spinTerms.begin(), parsR.spinTerms.end(), uniform_dist(e1));

    std::fill(parsR.ampRe.begin(), parsR.ampRe.end(), uniform_dist(e1));
    std::fill(parsR.ampIm.begin(), parsR.ampIm.end(), uniform_dist(e1));

    parsR.resMass = uniform_dist(e1);
    parsR.resWidth = uniform_dist(e1);

    calcAmpSoAStack(parsR);

    // std::cout << parsR.ampRe.at(5) << std::endl;
    // std::cout << parsR.ampIm.at(5) << std::endl;
}

void benchEigen()
{
    int n = n_global;
    ResParamsEigen parsR(n);

    parsR.qTerm = Eigen::ArrayXf::Ones(n, 1);
    parsR.mass = Eigen::ArrayXf::Ones(n, 1);
    parsR.ffRatioP = Eigen::ArrayXf::Ones(n, 1);
    parsR.ffRatioR = Eigen::ArrayXf::Ones(n, 1);
    parsR.spinTerms = Eigen::ArrayXf::Ones(n, 1);

    parsR.ampRe = Eigen::ArrayXf::Ones(n, 1);
    parsR.ampIm = Eigen::ArrayXf::Ones(n, 1);

    parsR.resMass = 1.0;
    parsR.resWidth = 0.1;

    calcAmpEigen(parsR);

    // std::cout << (*parsR.ampRe)[5] << std::endl;
    // std::cout << (*parsR.ampIm)[5] << std::endl;
}

void benchBlaze()
{
    ResParamsBlaze parsR;

    std::random_device r;

    std::default_random_engine e1(r());
    std::uniform_real_distribution<float> uniform_dist(0, 1);

    parsR.qTerm = uniform_dist(e1);
    parsR.mass = uniform_dist(e1);
    parsR.ffRatioP = uniform_dist(e1);
    parsR.ffRatioR = uniform_dist(e1);
    parsR.spinTerms = uniform_dist(e1);

    parsR.ampRe = uniform_dist(e1);
    parsR.ampIm = uniform_dist(e1);

    parsR.resMass = uniform_dist(e1);
    parsR.resWidth = uniform_dist(e1);

    calcAmpBlaze(parsR);

    // std::cout << parsR.ampRe[5] << std::endl;
    // std::cout << parsR.ampIm[5] << std::endl;
}

void benchXTensor()
{
    ResParamsXTensor parsR;

    std::random_device r;

    std::default_random_engine e1(r());
    std::uniform_real_distribution<float> uniform_dist(0.0, 1.0);

    parsR.qTerm.fill(uniform_dist(e1));
    parsR.mass.fill(uniform_dist(e1));
    parsR.ffRatioP.fill(uniform_dist(e1));
    parsR.ffRatioR.fill(uniform_dist(e1));
    parsR.spinTerms.fill(uniform_dist(e1));

    parsR.ampRe.fill(uniform_dist(e1));
    parsR.ampIm.fill(uniform_dist(e1));

    parsR.resMass = uniform_dist(e1);
    parsR.resWidth = uniform_dist(e1);

    calcAmpXTensor(parsR);

    // std::cout << parsR.ampRe[5] << std::endl;
    // std::cout << parsR.ampIm[5] << std::endl;
}

int main(int argc, char const *argv[]) {

    int n_itr = 100;

    std::cout<< "SoA:" << std::endl;
    {
    boost::timer::auto_cpu_timer t;
    for (int i = 0; i < n_itr; i++) { benchSoA(); }
    }

    std::cout << std::endl;

    std::cout<< "SoA (array):" << std::endl;
    {
    boost::timer::auto_cpu_timer t;
    for (int i = 0; i < n_itr; i++) { benchSoAArray(); }
    }

    std::cout << std::endl;

    std::cout<< "SoA (stack):" << std::endl;
    {
    boost::timer::auto_cpu_timer t;
    for (int i = 0; i < n_itr; i++) { benchSoAStack(); }
    }

    std::cout << std::endl;

    std::cout<< "AoS:" << std::endl;
    {
    boost::timer::auto_cpu_timer t;
    for (int i = 0; i < n_itr; i++) { benchAoS(); }
    }

    std::cout << std::endl;

    std::cout<< "AoS (stack):" << std::endl;
    {
    boost::timer::auto_cpu_timer t;
    for (int i = 0; i < n_itr; i++) { benchAoSStack(); }
    }

    std::cout << std::endl;

    // Without Legendre term

    std::cout<< "Eigen:" << std::endl;
    {
    boost::timer::auto_cpu_timer t;
    for (int i = 0; i < n_itr; i++) { benchEigen(); }
    }

    // Segfaults for some reason?

    // std::cout << std::endl;
    //
    // std::cout<< "Blaze:" << std::endl;
    // {
    // boost::timer::auto_cpu_timer t;
    // for (int i = 0; i < n_itr; i++) { benchBlaze(); }
    // }
    //
    // std::cout << std::endl;
    //
    // std::cout<< "XTensor:" << std::endl;
    // {
    // boost::timer::auto_cpu_timer t;
    // for (int i = 0; i < n_itr; i++) { benchXTensor(); }
    // }

    return 0;
}


// Copyright(c) 2020-Present, Matthew R. Hennefarth
// Distributed under the MIT License (http://opensource.org/licenses/MIT)

#include "FieldLocations.h"

/* EXTERNAL LIBRARY HEADER FILES */
#include <Eigen/Dense>
#include <matplot/matplot.h>

/* CPET HEADER FILES */
#include "System.h"

namespace cpet {

FieldLocations FieldLocations::fromSimple(
    const std::vector<std::string>& options) {
  FieldLocations fl;
  constexpr auto create_location = [](const std::string& location) -> AtomID {
    return AtomID{location};
  };
  std::transform(options.begin(), options.end(),
                 std::back_inserter(fl.locations_), create_location);
  return fl;
}

FieldLocations FieldLocations::fromBlock(
    const std::vector<std::string>& options) {
  constexpr const char* PLOT_KEY = "plot";
  constexpr const char* LOCATIONS_KEY = "locations";
  constexpr const char* OUTPUT_KEY = "output";

  FieldLocations fl;
  for (const auto& line : options) {
    const auto tokens = util::split(line, ' ');
    if (tokens.size() < 2) {
      continue;
    }

    const auto key = util::tolower(tokens[0]);
    const std::vector<std::string> key_options{tokens.begin() + 1,
                                               tokens.end()};

    if (key == LOCATIONS_KEY) {
      constexpr auto create_location =
          [](const std::string& location) -> AtomID {
        return AtomID{location};
      };

      std::transform(key_options.begin(), key_options.end(),
                     std::back_inserter(fl.locations_), create_location);
    } else if (key == PLOT_KEY) {
      fl.plotStyle_ = decodePlotStyle_(key_options);
    } else if (key == OUTPUT_KEY) {
      fl.output_ = *key_options.begin();
    }
  }
  return fl;
}
void FieldLocations::computeEFieldsWith(
    const std::vector<System>& systems) const {
  std::vector<std::vector<Eigen::Vector3d>> results;
  for (const auto& point : locations_) {
    SPDLOG_INFO("=~=~=~=~[Field at {}]=~=~=~=~", point.ID());
    std::vector<Eigen::Vector3d> fieldTrajectoryAtPoint;

    for (const auto& system : systems) {
      Eigen::Vector3d location;
      if (point.position()) {
        location = *(point.position());
      } else {
        location = system.frame().find(point)->coordinate;
      }

      Eigen::Vector3d field = system.electricFieldAt(location);
      SPDLOG_INFO("{} [{}]", field.transpose(), field.norm());
      fieldTrajectoryAtPoint.emplace_back(field);
    }
    results.push_back(fieldTrajectoryAtPoint);
  }
  if (output_) {
    writeOutput_(results);
  }
  if (showPlots()) {
    plot_(results);
  }
}
void FieldLocations::writeOutput_(
    const std::vector<std::vector<Eigen::Vector3d>>& results) const {
  if (!output_) {
    return;
  }

  std::ofstream outFile(*output_, std::ios::out);
  if (outFile.is_open()) {
    for (size_t i = 0; i < results.size(); i++) {
      outFile << '#' << locations_[i].ID() << '\n';
      for (const Eigen::Vector3d& field : results[i]) {
        outFile << field.transpose() << '\n';
      }
    }
    outFile << std::flush;
  } else {
    SPDLOG_ERROR("Could not open file {}", *output_);
    throw cpet::io_error("Could not open file " + *output_);
  }
}
void FieldLocations::plot_(
    const std::vector<std::vector<Eigen::Vector3d>>& results) const {
  const auto numberOfPlots =
      util::countSetBits(static_cast<unsigned int>(plotStyle_));
  if (numberOfPlots <= 0 || numberOfPlots > 4) {
    SPDLOG_WARN(
        "Number of plots less than 1 or greater than 4! Error in logic");
    return;
  }

  auto figure = matplot::figure();
  if (numberOfPlots < 3) {
    figure->tiledlayout(numberOfPlots, 1);
  } else {
    figure->tiledlayout(2, 2);
  }

  /* can make this an option eventually
   figure->size(500,500); */

  auto current_ax = matplot::nexttile(0);

  for (const auto& data : results) {
    std::array<std::vector<double>, 4> rotatedElectricFields;
    for (size_t index = 0; index < 3; index++) {
      const auto extract_index =
          [&index](const Eigen::Vector3d& vector) -> double {
        return vector[static_cast<long>(index)];
      };

      std::transform(data.begin(), data.end(),
                     std::back_inserter(rotatedElectricFields.at(index)),
                     extract_index);
    }

    constexpr auto compute_magnitude =
        [](const Eigen::Vector3d& vector) -> double { return vector.norm(); };

    std::transform(data.begin(), data.end(),
                   std::back_inserter(*rotatedElectricFields.rbegin()),
                   compute_magnitude);

    size_t plot_index = 0;

    // constexpr std::array<const char*, 4> titles{"X", "Y", "Z", "Magnitude"};
    const auto plot = [&rotatedElectricFields, &plot_index, &current_ax,
                       &figure](size_t index) {
      constexpr std::array<const char*, 4> titles{"X", "Y", "Z", "Magnitude"};
      current_ax = figure->nexttile(plot_index);
      matplot::hold(current_ax, matplot::on);
      matplot::plot(current_ax, rotatedElectricFields.at(index));
      current_ax->xlabel("Frame");
      current_ax->ylabel("Magnitude (V/Ang)");
      current_ax->title(titles.at(index));
      ++plot_index;
    };

    if (plotX_()) {
      plot(0);
    }
    if (plotY_()) {
      plot(1);
    }
    if (plotZ_()) {
      plot(2);
    }
    if (plotM_()) {
      plot(3);
    }
  }

  std::vector<std::string> legend_list;
  constexpr auto get_string_representation = [](const AtomID& aid) {
    return aid.ID();
  };
  std::transform(locations_.begin(), locations_.end(),
                 std::back_inserter(legend_list), get_string_representation);

  current_ax->legend(legend_list);
  matplot::show();
}

}  // namespace cpet


/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */
/*                                                                           */
/*  This file is part of the library KASKADE 7                               */
/*    see http://www.zib.de/Numerik/numsoft/kaskade7/                        */
/*                                                                           */
/*  Copyright (C) 2002-2015 Zuse Institute Berlin                            */
/*                                                                           */
/*  KASKADE 7 is distributed under the terms of the ZIB Academic License.    */
/*    see $KASKADE/academic.txt                                              */
/*                                                                           */
/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */

#ifndef ERROR_DISTRIBUTION_HH
#define ERROR_DISTRIBUTION_HH

#include <string>

#include <boost/utility.hpp>
#include <boost/fusion/include/as_vector.hpp>

#include "fem/assemble.hh"
#include "fem/functionspace.hh"
#include "fem/lagrangespace.hh"
#include "fem/shapefunctioncache.hh"

// forward declarations
namespace Dune
{
  template <class,int> class FieldVector;
  template <class,int,int> class FieldMatrix;
  template <class, int> class QuadratureRule;
}

namespace Kaskade
{
  using namespace boost::fusion;
  using namespace AssemblyDetail;

  enum class ErrorNorm { Energy, L2, H1, H1_half };

  template <class OriginalEvaluators, class ExtensionEvaluators, class Functional,
            class ArgYH, class ArgYE, class ArgUH, class ArgUE, class ArgPH, class ArgPE>
  void evaluateData(OriginalEvaluators& originalEvaluators, ExtensionEvaluators& extendedEvaluators, Functional const& f,
                    ArgYH& yl, ArgYE& yh, ArgUH& ul, ArgUE& uh, ArgPH& pl, ArgPE& ph, double w)
  {
      if( f.considerControlVariable_ )
      {
        uh.value += w * at_c<Functional::uIdx>(f.errorEstimateH.data).value(at_c<Functional::uSHIdx>(extendedEvaluators));
        ul.value += w * at_c<Functional::uIdx>(f.errorEstimateL.data).value(at_c<Functional::uSLIdx>(originalEvaluators));
        uh.gradient += w * at_c<Functional::uIdx>(f.errorEstimateH.data).gradient(at_c<Functional::uSHIdx>(extendedEvaluators));
        ul.gradient += w * at_c<Functional::uIdx>(f.errorEstimateL.data).gradient(at_c<Functional::uSLIdx>(originalEvaluators));
      }
      if( f.considerStateVariable_ )
      {
        yh.value += w * at_c<Functional::yIdx>(f.errorEstimateH.data).value(at_c<Functional::ySHIdx>(extendedEvaluators));
        yl.value += w * at_c<Functional::yIdx>(f.errorEstimateL.data).value(at_c<Functional::ySLIdx>(originalEvaluators));
        yh.gradient += w * at_c<Functional::yIdx>(f.errorEstimateH.data).gradient(at_c<Functional::ySHIdx>(extendedEvaluators));
        yl.gradient += w * at_c<Functional::yIdx>(f.errorEstimateL.data).gradient(at_c<Functional::ySLIdx>(originalEvaluators));
      }
      if( f.considerAdjointVariable_ )
      {
        ph.value += w * at_c<Functional::pIdx>(f.errorEstimateH.data).value(at_c<Functional::pSHIdx>(extendedEvaluators));
        pl.value += w * at_c<Functional::pIdx>(f.errorEstimateL.data).value(at_c<Functional::pSLIdx>(originalEvaluators));
        ph.gradient += w * at_c<Functional::pIdx>(f.errorEstimateH.data).gradient(at_c<Functional::pSHIdx>(extendedEvaluators));
        pl.gradient += w * at_c<Functional::pIdx>(f.errorEstimateL.data).gradient(at_c<Functional::pSLIdx>(extendedEvaluators));
      }
  }

  template <class ArgYL, class ArgYH, class ArgUL, class ArgUH, class ArgPL, class ArgPH>
  void clearVarArgs(ArgYL& yl, ArgYH& yh, ArgUL& ul, ArgUH& uh, ArgPL& pl, ArgPH& ph)
  {
    ul.value = 0;  ul.gradient = 0;
    yl.value = 0;  yl.gradient = 0;
    pl.value = 0;  pl.gradient = 0;
    uh.value = 0;  uh.gradient = 0;
    yh.value = 0;  yh.gradient = 0;
    ph.value = 0;  ph.gradient = 0;
  }

  template <class ArgU, class ArgY, class ArgP>
  double computeL2Error(ArgU const& u, ArgY const& y, ArgP const& p)
  {
    return u.value*u.value + y.value*y.value + p.value*p.value;
  }

  template <class ArgU, class ArgY, class ArgP>
  double computeH1HalfError(ArgU const& u, ArgY const& y, ArgP const& p)
  {
    LinAlg::EuclideanScalarProduct sp;
    return sp(u.gradient,u.gradient) + sp(y.gradient,y.gradient) + sp(p.gradient,p.gradient);
  }

  template <class Functional, class ExtendedAnsatzVars >
  class ErrorDistribution
  {
    typedef typename Functional::AnsatzVars OriginalAnsatzVars;
    typedef typename OriginalAnsatzVars::Grid Grid;
    typedef ShapeFunctionCache<Grid, typename Functional::Scalar> SfCache;
    typedef ShapeFunctionCache<Grid, typename Functional::Scalar> SfCache2;
    typedef typename OriginalAnsatzVars::Spaces OriginalSpaces;
    typedef typename ExtendedAnsatzVars::Spaces ExtendedSpaces;
    typedef typename result_of::as_vector<typename result_of::transform<OriginalSpaces, GetEvaluators<SfCache> >::type>::type OriginalEvaluators;
    typedef typename result_of::as_vector<typename result_of::transform<ExtendedSpaces, GetEvaluators<SfCache2> >::type>::type ExtendedEvaluators;
    typedef typename Grid::ctype CoordType;
    typedef Dune::QuadratureRule<typename Functional::AnsatzVars::Grid::ctype, Functional::AnsatzVars::Grid::dimension> QuadRule;
    typedef Dune::QuadratureRule<typename Functional::AnsatzVars::Grid::ctype, Functional::AnsatzVars::Grid::dimension-1> FaceQuadRule;
  public:
    typedef typename Functional::Scalar Scalar;
    typedef FEFunctionSpace<DiscontinuousLagrangeMapper<Scalar,typename Grid::LeafGridView> > AnsatzSpace;
    typedef boost::fusion::vector<AnsatzSpace const*> AnsatzSpaces;
    typedef Variable<SpaceIndex<0>,Components<1>,VariableId<0> > AnsatzVariableInformation;
    typedef boost::fusion::vector<AnsatzVariableInformation> VariableDescriptions;
    typedef VariableSetDescription<AnsatzSpaces,VariableDescriptions> AnsatzVars;
    typedef typename AnsatzVars::template CoefficientVectorRepresentation<0,1>::type ErrorVector;
    typedef AnsatzVars TestVars;
    typedef OriginalAnsatzVars OriginVars;
    //typedef AnsatzVars OriginVars;
    static int const dim = Grid::dimension;
    static ProblemType const type = Functional::type;

    static constexpr int yIdx = getStateId<Functional>();
    static constexpr int uIdx = getControlId<Functional>();
    static constexpr int pIdx = getAdjointId<Functional>();
    static constexpr int uSLIdx = result_of::value_at_c<typename OriginalAnsatzVars::Variables, uIdx>::type::spaceIndex;
    static constexpr int ySLIdx = result_of::value_at_c<typename OriginalAnsatzVars::Variables, yIdx>::type::spaceIndex;
    static constexpr int pSLIdx = result_of::value_at_c<typename OriginalAnsatzVars::Variables, pIdx>::type::spaceIndex;
    static constexpr int uSHIdx = result_of::value_at_c<typename ExtendedAnsatzVars::Variables, uIdx>::type::spaceIndex;
    static constexpr int ySHIdx = result_of::value_at_c<typename ExtendedAnsatzVars::Variables, yIdx>::type::spaceIndex;
    static constexpr int pSHIdx = result_of::value_at_c<typename ExtendedAnsatzVars::Variables, pIdx>::type::spaceIndex;

    class DomainCache
    {
    public:
      DomainCache(ErrorDistribution const& f_, typename OriginVars::VariableSet const& /*vars*/, int flags=7)
        : f(f_), domainCache(f.functional,f.iterate,flags)
      {
      }

      template <class Entity>
      void moveTo(Entity const &entity)
      {
        e  = &entity;
        domainCache.moveTo(entity);
      }

      template <class Position, class Evaluators>
      void evaluateAt(Position const&, Evaluators const& evaluators)
      {
        clearVarArgs(yl,yh,ul,uh,pl,ph);
        OriginalEvaluators originalEvaluators(transform(f.iterate.descriptions.spaces,GetEvaluators<SfCache>(&sfCache)));
        ExtendedEvaluators extendedEvaluators(transform(f.errorEstimateH.descriptions.spaces,GetEvaluators<SfCache>(&extendedSFCache)));
        QuadRule qr = QuadratureTraits<QuadRule>().rule(e->type(),f.qOrder);
        moveEvaluatorsToCell(originalEvaluators,*e);
        moveEvaluatorsToCell(extendedEvaluators,*e);
        useQuadratureRuleInEvaluators(originalEvaluators,qr,0);
        useQuadratureRuleInEvaluators(extendedEvaluators,qr,0);

        size_t nQuadPos = qr.size();
        for (size_t g=0; g<nQuadPos; ++g)
        {
          // pos of integration point
          Dune::FieldVector<CoordType,dim> quadPos = qr[g].position();
          // for all spaces involved, update the evaluators associated
          // to this quadrature point
          moveEvaluatorsToIntegrationPoint(originalEvaluators,quadPos);
          moveEvaluatorsToIntegrationPoint(extendedEvaluators,quadPos);
          // prepare evaluation of functional
          domainCache.evaluateAt(qr[g].position(),originalEvaluators);

          evaluateData(originalEvaluators, extendedEvaluators, f, yl, yh, ul, uh, pl, ph, qr[g].weight());
        }
      }

      Scalar d0() const { return 0; }

      template<int row, int dim>
      Dune::FieldVector<Scalar, TestVars::template Components<row>::m>
      d1 (VariationalArg<Scalar,dim> const& arg) const
      {
        VariationalArg<Scalar,dim,OriginalAnsatzVars::template Components<yIdx>::m> y(yh);
        VariationalArg<Scalar,dim,OriginalAnsatzVars::template Components<uIdx>::m> u(uh);
        VariationalArg<Scalar,dim,OriginalAnsatzVars::template Components<pIdx>::m> p(ph);

        if( !f.onlyH_ )
        {
          u.value += ul.value;
          u.gradient += ul.gradient;
          y.value += yl.value;
          y.gradient += yl.gradient;
          p.value += pl.value;
          p.gradient += pl.gradient;
        }

        if( f.errorNorm == ErrorNorm::L2 ) return computeL2Error(u,y,p);
        if( f.errorNorm == ErrorNorm::H1_half ) return computeH1HalfError(u,y,p);
        if( f.errorNorm == ErrorNorm::H1 ) return computeL2Error(u,y,p) + computeH1HalfError(u,y,p);

        Scalar result = 0;
        result += Functional::template D2<yIdx,yIdx>::present ? domainCache.template d2_impl<yIdx,yIdx>(y,y) : 0.;
        result += Functional::template D2<yIdx,uIdx>::present ? domainCache.template d2_impl<yIdx,uIdx>(y,u) : 0.;
        result += Functional::template D2<uIdx,yIdx>::present ? domainCache.template d2_impl<uIdx,yIdx>(u,y) : 0.;
        result += Functional::template D2<uIdx,uIdx>::present ? domainCache.template d2_impl<uIdx,uIdx>(u,u) : 0.;
        if( !f.useStateNormForAdjoint_ )
        {
          result += Functional::template D2<yIdx,pIdx>::present ? domainCache.template d2_impl<yIdx,pIdx>(y,p) : 0.;
          result += Functional::template D2<pIdx,yIdx>::present ? domainCache.template d2_impl<pIdx,yIdx>(p,y) : 0.;
          result += Functional::template D2<pIdx,uIdx>::present ? domainCache.template d2_impl<pIdx,uIdx>(p,u) : 0.;
          result += Functional::template D2<uIdx,pIdx>::present ? domainCache.template d2_impl<uIdx,pIdx>(u,p) : 0.;
          result += Functional::template D2<pIdx,pIdx>::present ? domainCache.template d2_impl<pIdx,pIdx>(p,p) : 0.;
        }
        else
        {
          result += Functional::template D2<pIdx,pIdx>::present ? domainCache.template d2_impl<yIdx,yIdx>(p,p) : 0.;
        }
        return result*arg.value[0];
      }

      template<int row, int col, int dim>
      Dune::FieldMatrix<Scalar,1,1>	d2 (VariationalArg<Scalar,dim> const&, VariationalArg<Scalar,dim> const&) const { return Dune::FieldMatrix<Scalar,1,1>(0); }

    private:
      ErrorDistribution const& f;
      typename Functional::DomainCache domainCache;
      typename AnsatzVars::Grid::template Codim<0>::Entity const* e;
      VariationalArg<Scalar,dim,OriginalAnsatzVars::template Components<yIdx>::m> yl;
      VariationalArg<Scalar,dim,ExtendedAnsatzVars::template Components<yIdx>::m> yh;
      VariationalArg<Scalar,dim,OriginalAnsatzVars::template Components<uIdx>::m> ul;
      VariationalArg<Scalar,dim,ExtendedAnsatzVars::template Components<uIdx>::m> uh;
      VariationalArg<Scalar,dim,OriginalAnsatzVars::template Components<pIdx>::m> pl;
      VariationalArg<Scalar,dim,ExtendedAnsatzVars::template Components<pIdx>::m> ph;
      SfCache sfCache;
      SfCache2 extendedSFCache;
    };

    class BoundaryCache
    {
    public:
      BoundaryCache(ErrorDistribution const& f_, typename OriginVars::VariableSet const& vars, int flags=7)
        : f(f_), boundaryCache(f.functional,f.iterate,flags)
      {}

      template <class FaceIterator>
      void moveTo(FaceIterator const& entity)
      {
        e = &entity;
        boundaryCache.moveTo(entity);
      }

      template <class Evaluators>
      void evaluateAt(Dune::FieldVector<typename AnsatzVars::Grid::ctype, AnsatzVars::Grid::dimension-1> const& xlocal, Evaluators const& evaluators)
      {
        clearVarArgs(yl,yh,ul,uh,pl,ph);
        OriginalEvaluators originalEvaluators(transform(f.iterate.descriptions.spaces,GetEvaluators<SfCache>(&sfCache)));
        ExtendedEvaluators extendedEvaluators(transform(f.errorEstimateH.descriptions.spaces,GetEvaluators<SfCache>(&extendedSFCache)));
        FaceQuadRule qr = QuadratureTraits<FaceQuadRule>().rule((*e)->geometryInInside().type(),f.qOrder);
        moveEvaluatorsToCell(originalEvaluators,*at_c<ySLIdx>(evaluators).cell_);
        moveEvaluatorsToCell(extendedEvaluators,*at_c<ySLIdx>(evaluators).cell_);
        moveEvaluatorsToCell(originalEvaluators,*e);
        moveEvaluatorsToCell(extendedEvaluators,*e);
        useQuadratureRuleInEvaluators(originalEvaluators,qr,(*e)->indexInInside());
        useQuadratureRuleInEvaluators(extendedEvaluators,qr,(*e)->indexInInside());

        size_t nQuadPos = qr.size();
        for (size_t g=0; g<nQuadPos; ++g)
        {
          // pos of integration point
          Dune::FieldVector<CoordType,dim> quadPos = (*e)->geometryInInside().global(qr[g].position());
          // for all spaces involved, update the evaluators associated
          // to this quadrature point
          moveEvaluatorsToIntegrationPoint(originalEvaluators,quadPos);
          moveEvaluatorsToIntegrationPoint(extendedEvaluators,quadPos);
          // prepare evaluation of functional
          boundaryCache.evaluateAt(qr[g].position(),originalEvaluators);

          evaluateData(originalEvaluators, extendedEvaluators, f, yl, yh, ul, uh, pl, ph, qr[g].weight());
        }
      }

      Scalar d0() const { return 0; }

      template<int row, int dim>
      Dune::FieldVector<Scalar,1> d1 (VariationalArg<Scalar,dim> const& arg) const
      {
        return Dune::FieldVector<Scalar,1>(0);
        VariationalArg<Scalar,dim,OriginalAnsatzVars::template Components<yIdx>::m> y(yh);
        VariationalArg<Scalar,dim,OriginalAnsatzVars::template Components<uIdx>::m> u(uh);
        VariationalArg<Scalar,dim,OriginalAnsatzVars::template Components<pIdx>::m> p(ph);

        if( !f.onlyH_ )
        {
          u.value += ul.value;
          u.gradient += ul.gradient;
          y.value += yl.value;
          y.gradient += yl.gradient;
          p.value += pl.value;
          p.gradient += pl.gradient;
        }

        if( f.errorNorm != ErrorNorm::Energy ) return computeL2Error(u,y,p);

        Scalar result = Functional::template D2<yIdx,yIdx>::present ? boundaryCache.template d2_impl<yIdx,yIdx>(y,y) : 0.;
        result += Functional::template D2<yIdx,uIdx>::present ? boundaryCache.template d2_impl<yIdx,uIdx>(y,u) : 0.;
        result += Functional::template D2<uIdx,yIdx>::present ? boundaryCache.template d2_impl<uIdx,yIdx>(u,y) : 0.;
        result += Functional::template D2<uIdx,uIdx>::present ? boundaryCache.template d2_impl<uIdx,uIdx>(u,u) : 0.;
        if( !f.useStateNormForAdjoint_ )
        {
          result += Functional::template D2<yIdx,pIdx>::present ? boundaryCache.template d2_impl<yIdx,pIdx>(y,p) : 0.;
          result += Functional::template D2<pIdx,yIdx>::present ? boundaryCache.template d2_impl<pIdx,yIdx>(p,y) : 0.;
          result += Functional::template D2<pIdx,uIdx>::present ? boundaryCache.template d2_impl<pIdx,uIdx>(p,u) : 0.;
          result += Functional::template D2<uIdx,pIdx>::present ? boundaryCache.template d2_impl<uIdx,pIdx>(u,p) : 0.;
          result += Functional::template D2<pIdx,pIdx>::present ? boundaryCache.template d2_impl<pIdx,pIdx>(p,p) : 0.;
        }
        else
        {
          result += Functional::template D2<pIdx,pIdx>::present ? boundaryCache.template d2_impl<yIdx,yIdx>(p,p) : 0.;
        }
        return result*arg.value[0];
      }

      template<int row, int col, int dim>
      Dune::FieldMatrix<Scalar,1,1> d2 (VariationalArg<Scalar,dim> const&, VariationalArg<Scalar,dim> const&) const { return Dune::FieldMatrix<Scalar,1,1>(0); }

    private:
      ErrorDistribution const& f;
      typename Functional::BoundaryCache boundaryCache;
      typename AnsatzVars::Grid::LeafGridView::IntersectionIterator const* e;
      VariationalArg<Scalar,dim,OriginalAnsatzVars::template Components<yIdx>::m> yl;
      VariationalArg<Scalar,dim,ExtendedAnsatzVars::template Components<yIdx>::m> yh;
      VariationalArg<Scalar,dim,OriginalAnsatzVars::template Components<uIdx>::m> ul;
      VariationalArg<Scalar,dim,ExtendedAnsatzVars::template Components<uIdx>::m> uh;
      VariationalArg<Scalar,dim,OriginalAnsatzVars::template Components<pIdx>::m> pl;
      VariationalArg<Scalar,dim,ExtendedAnsatzVars::template Components<pIdx>::m> ph;
      SfCache sfCache;
      SfCache2 extendedSFCache;
    };

    ErrorDistribution(Functional const& functional_, typename OriginalAnsatzVars::VariableSet const& iterate_,
                      typename OriginalAnsatzVars::VariableSet const& errorEstimateL_, typename ExtendedAnsatzVars::VariableSet const& errorEstimateH_)
      : functional(functional_), iterate(iterate_), errorEstimateL(errorEstimateL_), errorEstimateH(errorEstimateH_), ansatzSpace(boost::fusion::at_c<0>(iterate.descriptions.spaces)->gridManager(), iterate.descriptions.gridView,0),
        ansatzSpaces(&ansatzSpace), varName( {"error"} ), ansatzVars(ansatzSpaces,varName)

    {}

    template <int row>
    struct D1
    {
      static bool const present   = true;
      static bool const constant  = false;
    };

    template <int row, int col>
    struct D2
    {
      static bool const present = false;
      static bool const symmetric = false;
      static bool const lumped = false;
    };

    template <class Cell>
    int integrationOrder(Cell const&, int, bool) const { return 0; }

    AnsatzSpaces const& getSpaces() const { return ansatzSpaces; }

    AnsatzVars const& getVariableSetDescription() const { return ansatzVars; }

    void ignoreLowerOrderError(bool ignore) { onlyH_ = ignore; }
    void considerStateVariable(bool consider) { considerStateVariable_ = consider; }
    void considerControlVariable(bool consider) { considerControlVariable_ = consider; }
    void considerAdjointVariable(bool consider) { considerAdjointVariable_ = consider; }
    void setErrorNorm(ErrorNorm errNorm) { errorNorm = errNorm; }
    void useStateNormForAdjoint(bool useStateNorm) { useStateNormForAdjoint_ = useStateNorm; }

    friend class DomainCache;
    friend class BoundaryCache;

    Functional const& functional;
    typename OriginalAnsatzVars::VariableSet const& iterate;
    typename OriginalAnsatzVars::VariableSet const& errorEstimateL;
    typename ExtendedAnsatzVars::VariableSet const& errorEstimateH;
    AnsatzSpace ansatzSpace;
    AnsatzSpaces ansatzSpaces;
    std::string varName[1];
    AnsatzVars ansatzVars;
    ErrorNorm errorNorm = ErrorNorm::Energy;
    bool onlyH_ = false, considerStateVariable_ = true, considerControlVariable_ = true, considerAdjointVariable_ = false, useStateNormForAdjoint_ = false;
    int qOrder = 6;
  };
}
#endif