#ifndef SRC_IO_HANDLER_HH
#define SRC_IO_HANDLER_HH

#include <filesystem>
#include <memory>

#include <dune/common/parametertree.hh>

#include <dune/fufem/boundarypatch.hh>
#include <dune/fufem/geometry/convexpolyhedron.hh>

//#include <dune/tectonic/explicitgrid.hh>

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

#include "../utils/geocoordinate.hh"

#include "../time-stepping/adaptivetimestepper.hh"

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

#include "hdf5-writer.hh"
#include "vtk.hh"

template <class Assembler, class ContactNetwork, class ProgramState>
class IOHandler {
private:
    using Vector = typename Assembler::Vector;
    using LocalVector = typename Vector::block_type;

    using MyVertexBasis = typename Assembler::VertexBasis;
    using MyCellBasis = typename Assembler::CellBasis;
    using HDF5Writer = HDF5Writer<ProgramState, MyVertexBasis, typename Assembler::GV>;

    const size_t bodyCount_;
    std::vector<size_t> nVertices_;

    std::vector<Vector> vertexCoordinates_;
    std::vector<const MyVertexBasis* > vertexBases_;
    std::vector<const MyCellBasis* > cellBases_;
    std::vector<const typename ContactNetwork::BoundaryPatch*> frictionPatches_;
    std::vector<std::vector<const ConvexPolyhedron<LocalVector>*>> weakPatches_;

    std::unique_ptr<HDF5::File> dataFile_;
    std::unique_ptr<HDF5Writer> dataWriter_;

    bool writeVTK_;
    bool writeData_;

    bool writeRestarts_;
    size_t restartIdx_;
    size_t restartSpacing_;

    bool printProgress_;

public:
    IOHandler(const Dune::ParameterTree& parset, const ContactNetwork& contactNetwork)
        : bodyCount_(contactNetwork.nBodies()),
          nVertices_(bodyCount_),
          vertexCoordinates_(bodyCount_),
          vertexBases_(bodyCount_),
          cellBases_(bodyCount_),
          weakPatches_(bodyCount_),
          writeVTK_(parset.get<bool>("vtk.write")),
          writeData_(parset.get<bool>("data.write")),
          writeRestarts_(parset.get<bool>("restarts.write")),
          restartIdx_(parset.get<size_t>("restarts.first")),
          restartSpacing_(parset.get<size_t>("restarts.spacing")),
          printProgress_(parset.get<bool>("printProgress")) {

        for (size_t i=0; i<bodyCount_; i++) {
            nVertices_[i] = contactNetwork.body(i)->nVertices();
        }

        for (size_t i=0; i<bodyCount_; i++) {
          const auto& body = contactNetwork.body(i);
          vertexBases_[i] = &(body->assembler()->vertexBasis);
          cellBases_[i] = &(body->assembler()->cellBasis);

          auto& vertexCoords = vertexCoordinates_[i];
          vertexCoords.resize(nVertices_[i]);

          auto hostLeafView = body->grid()->hostGrid().leafGridView();
          Dune::MultipleCodimMultipleGeomTypeMapper<
              decltype(hostLeafView), Dune::MCMGVertexLayout> const vertexMapper(hostLeafView, Dune::mcmgVertexLayout());
          for (auto &&v : vertices(hostLeafView))
            vertexCoords[vertexMapper.index(v)] = geoToPoint(v.geometry());
        }

        if (writeVTK_) {
            if (!std::filesystem::is_directory("iterates/"))
                std::filesystem::create_directory("iterates");
        }


        contactNetwork.frictionPatches(frictionPatches_);
        dataFile_ = std::make_unique<HDF5::File>("output.h5");
        dataWriter_ = std::make_unique<HDF5Writer>(*dataFile_, vertexCoordinates_, vertexBases_,
                              frictionPatches_);
    }

    template <class GlobalFriction>
    void write(const ProgramState& programState, const ContactNetwork& contactNetwork, const GlobalFriction& friction, const IterationRegister& iterationCount, bool initial = false) {
      if (writeData_) {
        writeData(programState, contactNetwork, friction, iterationCount, initial);
      }

      if (writeVTK_) {
        writeVTK(programState, contactNetwork);
      }

      if (writeRestarts_) {
        writeRestarts(programState);
      }

      if (printProgress_) {
        std::cout << "timeStep = " << std::setw(6) << programState.timeStep
                  << ", time = " << std::setw(12) << programState.relativeTime
                  << ", tau = " << std::setw(12) << programState.relativeTau
                  << std::endl;
      }
    }

    bool read(ProgramState& programState) {
        using ADS = RestartIO<ProgramState>;

        if (restartIdx_ > 0) {// automatically adjusts the time and timestep
            auto restartFile = std::make_unique<HDF5::File>("restarts.h5", HDF5::Access::READONLY);
            auto restartIO = std::make_unique<ADS>(*restartFile, nVertices_);
            restartIO->read(restartIdx_, programState);

            return true;
        } else {
            return false;
        }
    }

private:
    void writeVTK(const ProgramState& programState, const ContactNetwork& contactNetwork) {
        using ScalarVector = typename Assembler::ScalarVector;
        std::vector<ScalarVector> stress(bodyCount_);

        for (size_t i=0; i<bodyCount_; i++) {
          const auto& body = contactNetwork.body(i);
          body->assembler()->assembleVonMisesStress(body->data()->getYoungModulus(),
                                           body->data()->getPoissonRatio(),
                                           programState.u[i], stress[i]);

        }

        const MyVTKWriter<MyVertexBasis, MyCellBasis> vtkWriter(cellBases_, vertexBases_, "iterates/");
        vtkWriter.write(programState.timeStep, programState.u, programState.v,
                        programState.alpha, stress);
    }

    void writeRestarts(const ProgramState& programState) {
        if (programState.timeStep % restartSpacing_ == 0) {
          auto restartFile = std::make_unique<HDF5::File>("restarts.h5", HDF5::Access::READWRITE);
          auto restartIO = std::make_unique<RestartIO<ProgramState>>(*restartFile, nVertices_);

          restartIO->write(programState);
          restartFile->flush();
        }
    }

    template <class GlobalFriction>
    void writeData(const ProgramState& programState, const ContactNetwork& contactNetwork, const GlobalFriction& friction, const IterationRegister& iterationCount, bool initial) {
      dataWriter_->reportSolution(programState, contactNetwork, friction);
      if (!initial) {
        dataWriter_->reportIterations(programState, iterationCount);
      } else {
        dataWriter_->reportWeightedNormalStress(programState);
      }

      dataFile_->flush();
    }

};
#endif