#ifndef SRC_HDF5_WRITER_HH
#define SRC_HDF5_WRITER_HH

#include <vector>

#include "hdf5/iteration-writer.hh"
#include "hdf5/time-writer.hh"
#include "hdf5-bodywriter.hh"

#include <dune/fufem/hdf5/file.hh>

template <class ProgramState, class VertexBasis, class GridView>
class HDF5Writer {
public:
    using HDF5BodyWriter = HDF5BodyWriter<ProgramState, VertexBasis, GridView>;
    using VertexCoordinates = std::vector<typename HDF5BodyWriter::VertexCoordinates>;
    using VertexBases = std::vector<const VertexBasis*>;
    using FrictionPatches = std::vector<const typename HDF5BodyWriter::Patch*>;
    using ScalarVector = typename ProgramState::ScalarVector;
    //using WeakPatches = std::vector<typename HDF5BodyWriter::WeakPatches>;

    //friend class HDF5NetworkWriter<ProgramState, VertexBasis, GridView>;

    HDF5Writer(HDF5::File& file,
             const VertexCoordinates& vertexCoordinates,
             const VertexBases& vertexBases,
             const FrictionPatches& frictionPatches)
             //const WeakPatches& weakPatches)
      : file_(file),
        frictionPatches_(frictionPatches),
        iterationWriter_(file_),
        timeWriter_(file_) {

        for (size_t i=0; i<vertexCoordinates.size(); i++) {
            if (frictionPatches_[i]->numVertices() > 0)
                bodyWriters_.push_back(std::make_unique<HDF5BodyWriter>(file_, i, vertexCoordinates[i], *vertexBases[i], *frictionPatches_[i])); //, weakPatches[i]));
        }
    }

    template <class ContactNetwork, class Friction>
    void reportSolution(const ProgramState& programState, const ContactNetwork& contactNetwork, const Friction& friction) {

        timeWriter_.write(programState);

        friction.updateAlpha(programState.alpha);

        // extract relative velocities
        using Vector = typename ProgramState::Vector;
        Vector mortarV;
        contactNetwork.nBodyAssembler().nodalToTransformed(programState.v, mortarV);

        std::vector<Vector> v_rel;
        split(mortarV, v_rel);

        using ScalarVector = typename ProgramState::ScalarVector;
        const auto frictionCoeff = friction.coefficientOfFriction(mortarV);

        /*double norm = 0;
        const auto& bodyNodes = *frictionPatches_[0]->getVertices();
        for (size_t i=bodyNodes.size(); i<frictionCoeff.size(); i++) {
            norm += frictionCoeff[i].two_norm();
        } */
        //std::cout << std::setprecision(10) << "friction coefficients norm: " << norm << std::endl;

        std::vector<ScalarVector> splitCoeff;
        split(frictionCoeff, splitCoeff);

        /*if (std::isnan(norm) or std::isinf(norm)) {
            print(programState.alpha, "alpha");
            print(frictionCoeff, "frictionCoeff");
            print(splitCoeff, "splitCoeff");
            print(v_rel, "v_rel: ");
            DUNE_THROW(Dune::Exception, "invalid state");
        }*/

        for (size_t i=0; i<bodyWriters_.size(); i++) {
            auto bodyID = bodyWriters_[i]->id();
            bodyWriters_[i]->reportSolution(programState, v_rel[bodyID], splitCoeff[bodyID]);
        }
    }

    void reportIterations(const ProgramState& programState, const IterationRegister& iterationCount) {
        iterationWriter_.write(programState.timeStep, iterationCount);
    }


    void reportWeightedNormalStress(const ProgramState& programState) {
        for (size_t i=0; i<bodyWriters_.size(); i++) {
            auto bodyID = bodyWriters_[i]->id();
            bodyWriters_[i]->reportWeightedNormalStress(programState);
        }
    }

    template <class VectorType>
    void split(const VectorType& v, std::vector<VectorType>& splitV) const {
        const auto nBodies = frictionPatches_.size();

        size_t globalIdx = 0;
        splitV.resize(nBodies);

        for (size_t bodyIdx=0; bodyIdx<nBodies; bodyIdx++) {
            const auto& bodyNodes = *frictionPatches_[bodyIdx]->getVertices();

            auto& splitV_ = splitV[bodyIdx];
            splitV_.resize(bodyNodes.size());

            for (size_t i=0; i<splitV_.size(); i++) {
                if (toBool(bodyNodes[i])) {
                    splitV_[i] = v[globalIdx];
                }
                globalIdx++;
            }
        }
    }

private:
  HDF5::File& file_;
  const FrictionPatches& frictionPatches_;

  IterationWriter iterationWriter_;
  TimeWriter<ProgramState> timeWriter_;

  std::vector<std::unique_ptr<HDF5BodyWriter>> bodyWriters_;
};
#endif