#include <dune/common/bitsetvector.hh>

#include <dune/istl/io.hh>

#include <dune/disc/functions/p1function.hh>
#include <dune/disc/miscoperators/laplace.hh>
#include <dune/disc/operators/p1operator.hh>

#include <dune-solvers/iterationsteps/trustregiongsstep.hh>
#include <dune-solvers/iterationsteps/mmgstep.hh>
#include <dune-solvers/transferoperators/truncatedcompressedmgtransfer.hh>
#include <dune-solvers/transferoperators/mandelobsrestrictor.hh>
#include <dune-solvers/solvers/iterativesolver.hh>

#include <dune-solvers/norms/energynorm.hh>
#include <dune-solvers/norms/h1seminorm.hh>

#include "maxnormtrustregion.hh"

// for debugging
#include "../test/fdcheck.hh"

template <class GridType, class TargetSpace>
void RiemannianTrustRegionSolver<GridType,TargetSpace>::
setup(const GridType& grid,
      const GeodesicFEAssembler<typename GridType::LeafGridView,TargetSpace>* assembler,
      const SolutionType& x,
      const Dune::BitSetVector<blocksize>& dirichletNodes,
      double tolerance,
      int maxTrustRegionSteps,
      double initialTrustRegionRadius,
      int multigridIterations,
      double mgTolerance,
      int mu, 
      int nu1,
      int nu2,
      int baseIterations,
      double baseTolerance,
      bool instrumented)
{
    grid_                     = &grid;
    assembler_                = assembler;
    x_                        = x;
    this->tolerance_          = tolerance;
    maxTrustRegionSteps_      = maxTrustRegionSteps;
    initialTrustRegionRadius_ = initialTrustRegionRadius;
    multigridIterations_      = multigridIterations;
    qpTolerance_              = mgTolerance;
    mu_                       = mu;
    nu1_                      = nu1;
    nu2_                      = nu2;
    baseIt_                   = baseIterations;
    baseTolerance_            = baseTolerance;
    instrumented_             = instrumented;

    int numLevels = grid_->maxLevel()+1;

    // ////////////////////////////////
    //   Create a multigrid solver
    // ////////////////////////////////

    // First create a Gauss-seidel base solver
    TrustRegionGSStep<MatrixType, CorrectionType>* baseSolverStep = new TrustRegionGSStep<MatrixType, CorrectionType>;

    EnergyNorm<MatrixType, CorrectionType>* baseEnergyNorm = new EnergyNorm<MatrixType, CorrectionType>(*baseSolverStep);

    ::LoopSolver<CorrectionType>* baseSolver = new ::LoopSolver<CorrectionType>(baseSolverStep,
                                                                            baseIt_,
                                                                            baseTolerance_,
                                                                            baseEnergyNorm,
                                                                            Solver::QUIET);

    // Make pre and postsmoothers
    TrustRegionGSStep<MatrixType, CorrectionType>* presmoother  = new TrustRegionGSStep<MatrixType, CorrectionType>;
    TrustRegionGSStep<MatrixType, CorrectionType>* postsmoother = new TrustRegionGSStep<MatrixType, CorrectionType>;

    MonotoneMGStep<MatrixType, CorrectionType>* mmgStep = new MonotoneMGStep<MatrixType, CorrectionType>(numLevels);

    mmgStep->setMGType(mu_, nu1_, nu2_);
    mmgStep->ignoreNodes_       = &dirichletNodes;
    mmgStep->basesolver_        = baseSolver;
    mmgStep->presmoother_       = presmoother;
    mmgStep->postsmoother_      = postsmoother; 
    mmgStep->obstacleRestrictor_= new MandelObstacleRestrictor<CorrectionType>();
    mmgStep->hasObstacle_       = &hasObstacle_;
    mmgStep->verbosity_         = Solver::QUIET;

    // //////////////////////////////////////////////////////////////////////////////////////
    //   Assemble a Laplace matrix to create a norm that's equivalent to the H1-norm
    // //////////////////////////////////////////////////////////////////////////////////////
    Dune::LeafP1Function<GridType,double> u(grid),f(grid);
    Dune::LaplaceLocalStiffness<typename GridType::LeafGridView,double> laplaceStiffness;
    Dune::LeafP1OperatorAssembler<GridType,double,1>* A = new Dune::LeafP1OperatorAssembler<GridType,double,1>(grid);
    A->assemble(laplaceStiffness,u,f);

    typedef typename Dune::LeafP1OperatorAssembler<GridType,double,1>::RepresentationType LaplaceMatrixType;

    if (h1SemiNorm_)
        delete h1SemiNorm_;

    h1SemiNorm_ = new H1SemiNorm<CorrectionType>(**A);

    mmgSolver_ = new ::LoopSolver<CorrectionType>(mmgStep,
                                                     multigridIterations_,
                                                     qpTolerance_,
                                                     h1SemiNorm_,
                                                     Solver::QUIET);

    // Write all intermediate solutions, if requested
    if (instrumented_)
        mmgSolver_->historyBuffer_ = "tmp/mgHistory";

    // ////////////////////////////////////////////////////////////
    //    Create Hessian matrix and its occupation structure
    // ////////////////////////////////////////////////////////////
    
    if (hessianMatrix_)
        delete hessianMatrix_;

    hessianMatrix_ = new MatrixType;
    Dune::MatrixIndexSet indices(grid_->size(1), grid_->size(1));
    assembler_->getNeighborsPerVertex(indices);
    indices.exportIdx(*hessianMatrix_);
    
    // //////////////////////////////////////////////////////////
    //   Create obstacles
    // //////////////////////////////////////////////////////////
    
    hasObstacle_.resize(numLevels);
    for (int i=0; i<hasObstacle_.size(); i++)
        hasObstacle_[i].resize(grid_->size(i, gridDim),true);
    
    // ////////////////////////////////////
    //   Create the transfer operators
    // ////////////////////////////////////
    for (int k=0; k<mmgStep->mgTransfer_.size(); k++)
        delete(mmgStep->mgTransfer_[k]);
    
    mmgStep->mgTransfer_.resize(numLevels-1);
    
    for (int i=0; i<mmgStep->mgTransfer_.size(); i++){
        TruncatedCompressedMGTransfer<CorrectionType>* newTransferOp = new TruncatedCompressedMGTransfer<CorrectionType>;
        newTransferOp->setup(*grid_,i,i+1);
        mmgStep->mgTransfer_[i] = newTransferOp;
    }
    
}

template <class GridType, class TargetSpace>
void RiemannianTrustRegionSolver<GridType,TargetSpace>::solve()
{
    MaxNormTrustRegion<blocksize> trustRegion(x_.size(), initialTrustRegionRadius_);

    std::vector<std::vector<BoxConstraint<field_type,blocksize> > > trustRegionObstacles(dynamic_cast<MultigridStep<MatrixType,CorrectionType>*>(mmgSolver_->iterationStep_)->numLevels_);
    
    // /////////////////////////////////////////////////////
    //   Set up the log file, if requested
    // /////////////////////////////////////////////////////
    FILE* fp;
    if (instrumented_) {

        fp = fopen("statistics", "w");
        if (!fp)
            DUNE_THROW(Dune::IOError, "Couldn't open statistics file for writing!");
    
    }

    // /////////////////////////////////////////////////////
    //   Trust-Region Solver
    // /////////////////////////////////////////////////////
    for (int i=0; i<maxTrustRegionSteps_; i++) {
        
        if (this->verbosity_ == Solver::FULL) {
            std::cout << "----------------------------------------------------" << std::endl;
            std::cout << "      Trust-Region Step Number: " << i 
                      << ",     radius: " << trustRegion.radius()
                      << ",     energy: " << assembler_->computeEnergy(x_) << std::endl;
            std::cout << "----------------------------------------------------" << std::endl;
        }

        CorrectionType rhs;
        CorrectionType corr(x_.size());
        corr = 0;

        assembler_->assembleGradient(x_, rhs);
        assembler_->assembleMatrix(x_, *hessianMatrix_);

        //gradientFDCheck(x_, rhs, *rodAssembler_);
        //hessianFDCheck(x_, *hessianMatrix_, *rodAssembler_);

        rhs *= -1;

        dynamic_cast<MultigridStep<MatrixType,CorrectionType>*>(mmgSolver_->iterationStep_)->setProblem(*hessianMatrix_, corr, rhs, grid_->maxLevel()+1);

        trustRegionObstacles.back() = trustRegion.obstacles();
        dynamic_cast<MonotoneMGStep<MatrixType, CorrectionType>*>(mmgSolver_->iterationStep_)->obstacles_ = &trustRegionObstacles;

        
        mmgSolver_->preprocess();
        
        dynamic_cast<MultigridStep<MatrixType,CorrectionType>*>(mmgSolver_->iterationStep_)->preprocess();
        
        
        // /////////////////////////////
        //    Solve !
        // /////////////////////////////
        mmgSolver_->solve();
        
        corr = dynamic_cast<MultigridStep<MatrixType,CorrectionType>*>(mmgSolver_->iterationStep_)->getSol();
        
        //std::cout << "Correction: " << std::endl << corr << std::endl;
        
        if (instrumented_) {

            fprintf(fp, "Trust-region step: %d, trust-region radius: %g\n",
                    i, trustRegion.radius());
                    
            // ///////////////////////////////////////////////////////////////
            //   Compute and measure progress against the exact solution
            //   for each trust region step
            // ///////////////////////////////////////////////////////////////
    
            CorrectionType exactSolution = corr;

            // Start from 0
            double oldError = 0;
            double totalConvRate = 1;
            double convRate = 1;
    
            // Write statistics of the initial solution
            // Compute the energy norm
            oldError = h1SemiNorm_->operator()(exactSolution);
    
            for (int j=0; j<multigridIterations_; j++) {
        
                // read iteration from file
                CorrectionType intermediateSol(grid_->size(1));
                intermediateSol = 0;
                char iSolFilename[100];
                sprintf(iSolFilename, "tmp/mgHistory/intermediatesolution_%04d", j);
                
                FILE* fpInt = fopen(iSolFilename, "rb");
                if (!fpInt)
                    DUNE_THROW(Dune::IOError, "Couldn't open intermediate solution");
                for (int k=0; k<intermediateSol.size(); k++)
                    for (int l=0; l<blocksize; l++)
                        fread(&intermediateSol[k][l], sizeof(double), 1, fpInt);
                
                fclose(fpInt);
                //std::cout << "intermediateSol\n" << intermediateSol << std::endl;

                // Compute errors
                intermediateSol -= exactSolution;
      
                //std::cout << "error\n" << intermediateSol << std::endl;

                // Compute the H1 norm
                double error = h1SemiNorm_->operator()(intermediateSol);

                convRate = error / oldError;
                totalConvRate *= convRate;
                
                if (error < 1e-12)
                    break;
                
                std::cout << "Iteration: " << j << "  ";
                std::cout << "Errors:  error " << error << ", convergence rate: " << convRate
                          << ",  total conv rate " << pow(totalConvRate, 1/((double)j+1)) << std::endl;
                

                fprintf(fp, "%d %g %g %g\n", j+1, error, convRate, pow(totalConvRate, 1/((double)j+1)));
                
                
                oldError = error;
                
            }
            
            
        }

        if (this->verbosity_ == NumProc::FULL) {
            double translationMax = 0;
            double rotationMax    = 0;
            for (size_t j=0; j<corr.size(); j++) {
                for (int k=0; k<3; k++) {
                    translationMax = std::max(translationMax, corr[j][k]);
                    rotationMax    = std::max(rotationMax, corr[j][k+3]);
                }
            }
            printf("infinity norm of the correction: %g %g\n", translationMax, rotationMax);
        }

        if (corr.infinity_norm() < this->tolerance_) {
            if (this->verbosity_ == NumProc::FULL)
                std::cout << "CORRECTION IS SMALL ENOUGH" << std::endl;

            if (this->verbosity_ != NumProc::QUIET)
                std::cout << i+1 << " trust-region steps were taken." << std::endl;
            break;
        }
        
        // ////////////////////////////////////////////////////
        //   Check whether trust-region step can be accepted
        // ////////////////////////////////////////////////////
        
        SolutionType newIterate = x_;
        for (int j=0; j<newIterate.size(); j++) 
            newIterate[j] = TargetSpace::exp(newIterate[j], corr[j]);
        
        /** \todo Don't always recompute oldEnergy */
        double oldEnergy = assembler_->computeEnergy(x_);
        double energy    = assembler_->computeEnergy(newIterate); 
        
        // compute the model decrease
        // It is $ m(x) - m(x+s) = -<g,s> - 0.5 <s, Hs>
        // Note that rhs = -g
        CorrectionType tmp(corr.size());
        tmp = 0;
        hessianMatrix_->umv(corr, tmp);
        double modelDecrease = (rhs*corr) - 0.5 * (corr*tmp);
        
        if (this->verbosity_ == NumProc::FULL) {
            std::cout << "Absolute model decrease: " << modelDecrease 
                      << ",  functional decrease: " << oldEnergy - energy << std::endl;
            std::cout << "Relative model decrease: " << modelDecrease / energy
                      << ",  functional decrease: " << (oldEnergy - energy)/energy << std::endl;
        }            

        assert(modelDecrease >= 0);
        
        if (energy >= oldEnergy) {
            if (this->verbosity_ == NumProc::FULL)
                printf("Richtung ist keine Abstiegsrichtung!\n");
        }

        if (energy >= oldEnergy &&
            (std::abs(oldEnergy-energy)/energy < 1e-9 || modelDecrease/energy < 1e-9)) {
            if (this->verbosity_ == NumProc::FULL)
                std::cout << "Suspecting rounding problems" << std::endl;

            if (this->verbosity_ != NumProc::QUIET)
                std::cout << i+1 << " trust-region steps were taken." << std::endl;

            x_ = newIterate;
            break;
        }

        // //////////////////////////////////////////////
        //   Check for acceptance of the step
        // //////////////////////////////////////////////
        if ( (oldEnergy-energy) / modelDecrease > 0.9) {
            // very successful iteration
            
            x_ = newIterate;
            trustRegion.scale(2);
            
        } else if ( (oldEnergy-energy) / modelDecrease > 0.01
                    || std::abs(oldEnergy-energy) < 1e-12) {
            // successful iteration
            x_ = newIterate;
            
        } else {
            // unsuccessful iteration
            trustRegion.scale(0.5);
            if (this->verbosity_ == NumProc::FULL)
                std::cout << "Unsuccessful iteration!" << std::endl;
        }
        
        //  Write current energy
        if (this->verbosity_ == NumProc::FULL)
            std::cout << "--- Current energy: " << energy << " ---" << std::endl;

        // /////////////////////////////////////////////////////////////////////
        //   Write the iterate to disk for later convergence rate measurement
        // /////////////////////////////////////////////////////////////////////

        if (instrumented_) {

            char iRodFilename[100];
            sprintf(iRodFilename, "tmp/intermediateSolution_%04d", i);

            FILE* fpRod = fopen(iRodFilename, "wb");
            if (!fpRod)
                DUNE_THROW(SolverError, "Couldn't open file " << iRodFilename << " for writing");
            
            for (int j=0; j<x_.size(); j++) {

                for (int k=0; k<3; k++)
                    fwrite(&x_[j].r[k], sizeof(double), 1, fpRod);

                for (int k=0; k<4; k++)  // 3d hardwired here!
                    fwrite(&x_[j].q[k], sizeof(double), 1, fpRod);

            }

            fclose(fpRod);

        }

    }

    // //////////////////////////////////////////////
    //   Close logfile
    // //////////////////////////////////////////////
    if (instrumented_)
        fclose(fp);
}