#include <petscsys.h>
#include <petscdm.h>
#include <petscdmplex.h>
#include <petscts.h>
#include <petsc/private/tsimpl.h> /**< Add this one to register the custom TSAdapt */
#include "adapt_dt.h"

double cfl;
double gam;

static PetscErrorCode PetscCalcMaxSpeed(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nf, PetscScalar *u, void *userctx)
{
    PetscReal      *user = (PetscReal*) userctx;
    EulerNode      *uu   = (EulerNode*) u; // See adapt_dt.h
    PetscReal      norm_vel, density, pressure, csound;

    PetscFunctionBegin;

    norm_vel = sqrt((uu->ru[0]/u[0])*(uu->ru[0]/uu->r) + (uu->ru[1]/uu->r)*(uu->ru[1]/uu->r) + (uu->ru[2]/uu->r)*(uu->ru[2]/uu->r));
    density  = uu->r;
    pressure = (gam - 1.0) * (uu->rE - 0.5 * norm_vel*norm_vel * density);
    csound   = sqrt(gam * pressure / density);

    PetscPrintf(PETSC_COMM_WORLD, "Rho = %lf, Vel = (%lf,%lf,%lf), Press = %lf, C = %lf, wsn = %lf.\n", uu->r, uu->ru[0]/uu->r, uu->ru[1]/uu->r, uu->ru[2]/uu->r, pressure, csound, norm_vel + csound);
    if (norm_vel + csound > user[0]) user[0] = norm_vel + csound;

    PetscFunctionReturn(0);
}

PetscErrorCode TSAdaptChoose_MyCFLAdapt(TSAdapt adapt, TS ts, PetscReal h,
                                        PetscInt *next_sc, PetscReal *next_h,
                                        PetscBool *accept, PetscReal *wlte,
                                        PetscReal *wltea, PetscReal *wlter)
{
    PetscErrorCode     ierr;
    DM                 plex, ts_dm;
    PetscErrorCode     (*func[1]) (PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nf, PetscScalar *u, void *ctx);
    void               *ctx[1];
    // Vec                cellGeom = NULL, faceGeom =  NULL;
    // Vec                Y = NULL;
    TSAdapt_MyCFLAdapt *data = (TSAdapt_MyCFLAdapt*) adapt->data;

    PetscReal          user[1];
    PetscReal          globwsn; /* Global maximum wave speed over the whole domain */
    PetscReal          minRadius; /* Global minimum Volume/Surface over the whole mesh */
    PetscReal          new_dt;

    PetscFunctionBegin;

    *next_sc =  0;  /* Reuse the same order scheme */
    *wltea   = -1;  /* Weighted absolute local truncation error is not used */
    *wlter   = -1;  /* Weighted relative local truncation error is not used */
    *wlte    = -1;  /* Weighted local truncation error is not going to be evaluated anyways */

    /* Get the context and the DM through the TS */
    ierr = TSGetDM(ts, &ts_dm); CHKERRQ(ierr);
    ierr = DMConvert(ts_dm, DMPLEX, &plex); CHKERRQ(ierr);

    /* Set the current solution vector as data to the adapter */
    // if (!data->Y) {
    // ierr = VecDuplicate(ts->vec_sol, &data->Y); CHKERRQ(ierr);
    // }
    // ierr = VecDuplicate(ts->vec_sol, &data->Y); CHKERRQ(ierr);
    // ierr = VecCopy(ts->vec_sol, data->Y); CHKERRQ(ierr);
    // ierr = DMPlexGetGeometryFVM(plex, &faceGeom, &cellGeom, NULL); CHKERRQ(ierr);
    // ierr = DMPlexInsertBoundaryValues(plex, PETSC_TRUE, data->Y, 0.0, faceGeom, cellGeom, NULL); CHKERRQ(ierr);
    // ierr = TSGetSolution(ts, &Y); CHKERRQ(ierr);
    ierr = TSGetSolution(ts, &data->Y); CHKERRQ(ierr);

    /* - Then compute the sound speed in the local domain */
    user[0] = 0.0;
    ctx[0]  = (void*) user;
    func[0] = PetscCalcMaxSpeed;
    ierr = DMProjectFunction(plex, 0.0, func, ctx, INSERT_VALUES, data->Y); CHKERRQ(ierr);

    /* In case of several mpi processes, reduce the maxspeed to get the max over the domain */
    ierr = MPI_Allreduce(user, &globwsn, 1, MPIU_REAL, MPIU_MAX, PETSC_COMM_WORLD); CHKERRQ(ierr);

    /* And the minimum inscribed circle radius over the global domain */
    ierr = DMPlexGetGeometryFVM(plex, NULL, NULL, &minRadius); CHKERRQ(ierr);

    /* Evaluate the new DT (inviscid-like) */
    new_dt = ((PetscReal)cfl) * minRadius / globwsn;

    /* Set the new timestep in the TS structure */
    *accept = PETSC_TRUE;
    *next_h = new_dt;

    /* Clean up */
    ierr = DMDestroy(&plex); CHKERRQ(ierr);

    PetscFunctionReturn(0);
}

PetscErrorCode MyTSAdaptRegister(TS ts, TSAdapt adapt, double cgamma, double ccfl)
{
    PetscErrorCode ierr;

    PetscFunctionBegin;

    /* Set the local CFL and gamma values
     * Those are passed from the F90 program
     */
    gam = cgamma;
    cfl = ccfl;

    /* Register the new adaption driver */
    ierr = TSAdaptRegister("my-cfl-adapt", TSAdaptCreate_MyCFLAdapt); CHKERRQ(ierr);
    ierr = TSAdaptSetType(adapt, "my-cfl-adapt"); CHKERRQ(ierr);

    PetscFunctionReturn(ierr);
}

PetscErrorCode TSAdaptCreate_MyCFLAdapt(TSAdapt adapt)
{
    PetscErrorCode ierr;
    TSAdapt_MyCFLAdapt *data;

    PetscFunctionBegin;

    ierr = PetscNewLog(adapt,&data); CHKERRQ(ierr);
    adapt->data         = (void*)data;
    adapt->ops->choose  = TSAdaptChoose_MyCFLAdapt;
    adapt->ops->reset   = TSAdaptReset_MyCFLAdapt;
    adapt->ops->destroy = TSAdaptDestroy_MyCFLAdapt;

    PetscFunctionReturn(0);
}

PetscErrorCode TSAdaptReset_MyCFLAdapt(TSAdapt adapt)
{
    TSAdapt_MyCFLAdapt *data = (TSAdapt_MyCFLAdapt*)adapt->data;
    PetscErrorCode ierr;
    PetscFunctionBegin;
    ierr = VecDestroy(&data->Y); CHKERRQ(ierr);
    PetscFunctionReturn(0);
}

PetscErrorCode TSAdaptDestroy_MyCFLAdapt(TSAdapt adapt)
{
    PetscErrorCode ierr;
    PetscFunctionBegin;
    ierr = TSAdaptReset_MyCFLAdapt(adapt);CHKERRQ(ierr);
    ierr = PetscFree(adapt->data);CHKERRQ(ierr);
    PetscFunctionReturn(0);
}