#ifndef RFASOLVER_H
#define RFASOLVER_H

#include <dolfin.h>

#include "PrimalSolver.h"
#include "PotentialSolver.h"
#include "Potential.h"
#include "BioSolver.h"
#include "Bioheat.h"
#include "TissueParameters.h"

using namespace dolfin;

class RFASolver
{
public:
    typedef NoSlip SlipPolicy;
    typedef GlobalTimeStep TimeStepPolicy;
    typedef UnitVelocityStabilization StabilizationPolicy;

    typedef PrimalSolver<SlipPolicy,
                         TimeStepPolicy,
                         StabilizationPolicy> NSESolver;

    RFASolver(Mesh& mesh,
            std::vector<SubDomain*> subdomains,
            DirichletBCList dbcs_m,
            DirichletBCList dbcs_c,
            DirichletBCList dbcs_dm,
            SlipPolicy::SlipBCData slipData,
            Function* u_i0,
            OutputQuantityList& outputQuantities,
            Function* f,
            real cfl_target,
            real nu,
            Function* potentialSource,
            SubDomain* electrode,
            DirichletBCList dbcs_p,
            Function* thermalSource,
            Function* TInitial,
            DirichletBCList dbcs_b,
            TissueParameters& tissueParameters
            );
    virtual ~RFASolver ();

    real getTFinal() const {return tFinal;}
    void setTFinal(real t);
    real getAblationTime() const {return ablationTime;}
    void setAblationTime(const real t) {ablationTime = t;}
    real getDt() const {return dt;}
    void setDt(real dt);
    uint getNoSamples() const {return no_samples;}
    void setNoSamples(const uint n);
    real getNSInitTime() const {return nsSolver->getInitTime();}
    void setNSInitTime(const real t) {nsSolver->setInitTime(t);}
    real getTargetPower() const {return potentialProblem->getTargetPower();}
    void setTargetPower(const real p) {potentialProblem->setTargetPower(p);}
    real getTissuePowerFraction() const {return potentialProblem->getTissuePowerFraction();}
    void setTissuePowerFraction(const real f) {potentialProblem->setTissuePowerFraction(f);}
    real getInitialResistance() const {return potentialProblem->getInitialResistance();}
    void setInitialResistance(const real r) {potentialProblem->setInitialResistance(r);}
    real getPowerTolerance() const {return potentialSolver->getPowerTolerance();}
    void setPowerTolerance(const real t) {potentialSolver->setPowerTolerance(t);}
    real getPotential() const {return potentialSolver->getPotential();}
    void setPotential(const real p) {potentialSolver->setPotential(p);}
    real getMinSigmaBoard() const {return potentialSolver->getMinSigmaBoard();}
    void setMinSigmaBoard(const real s) {potentialSolver->setMinSigmaBoard(s);}
    real getMaxSigmaBoard() const {return potentialSolver->getMaxSigmaBoard();}
    void setMaxSigmaBoard(const real s) {potentialSolver->setMaxSigmaBoard(s);}
    real getSigmaBoard() const {return potentialSolver->getSigmaBoard();}
    void setSigmaBoard(const real s);

    real step();
    void run();

private:
    NSESolver* nsSolver;
    PotentialProblem* potentialProblem;
    PotentialSolver* potentialSolver;
    BioProblem* bioProblem;
    BioSolver* bioSolver;

    real tFinal;
    real ablationTime;
    real dt;
    real time;
    real nsTime;

    uint no_samples;
    uint sample;
    uint nsteps;
};

RFASolver::RFASolver(
        Mesh& mesh,
        std::vector<SubDomain*> subdomains,
        DirichletBCList dbcs_m,
        DirichletBCList dbcs_c,
        DirichletBCList dbcs_dm,
        SlipPolicy::SlipBCData slipData,
        Function* u_i0,
        OutputQuantityList& outputQuantities,
        Function* f,
        real cfl_target,
        real nu,
        Function* potentialSource,
        SubDomain* electrode,
        DirichletBCList dbcs_p,
        Function* thermalSource,
        Function* TInitial,
        DirichletBCList dbcs_b,
        TissueParameters& tissueParameters
        ):
    dt(0.01),
    time(0),
    nsTime(0),
    no_samples(100),
    sample(1),
    nsteps(1)
{
    nsSolver = new NSESolver(
            mesh,
            dbcs_m,
            dbcs_c,
            dbcs_dm,
            slipData,
            u_i0,
            outputQuantities,
            NULL,
            NULL,
            NULL,
            f,
            cfl_target,
            nu
            );

    potentialProblem = new PotentialProblem(
            mesh,
            subdomains,
            potentialSource,
            electrode,
            dbcs_p,
            tissueParameters);
    potentialSolver = new PotentialSolver(potentialProblem);

    bioProblem = new BioProblem(
            mesh,
            subdomains,
            nsSolver->getU(),
            potentialSolver->getPhi(),
            TInitial,
            dbcs_b,
            tissueParameters);
    bioSolver = new BioSolver(bioProblem);

    potentialSolver->init(bioSolver->getT());
}

RFASolver::~RFASolver()
{
    delete nsSolver;
    delete potentialProblem;
    delete potentialSolver;
    delete bioProblem;
    delete bioSolver;
}

void RFASolver::setTFinal(const real t)
{
    tFinal = t;
    nsSolver->setT(tFinal);
    bioSolver->setTFinal(tFinal);
}

void RFASolver::setDt(real dt)
{
    this->dt = dt;
    bioSolver->setDt(dt);
}

void RFASolver::setSigmaBoard(const real s)
{
    potentialSolver->setSigmaBoard(s);
    bioSolver->setSigmaBoard(potentialSolver->getSigmaBoard());
}

real RFASolver::step()
{
    while (nsTime < nsteps*dt)
    {
        nsTime = nsSolver->step();
    }
    nsteps++;

    bioSolver->setDt(nsTime-time);

    if (time > getAblationTime())
    {
        cout << "Shutting off potential" << endl;
        if (getTargetPower() != real(0))
        {
            setPotential(0);
            setTargetPower(0);
            setMinSigmaBoard(getSigmaBoard());
            setMaxSigmaBoard(getSigmaBoard());
            potentialSolver->matchInitialPower();
        }
    }
    else
    {
        cout << "Solving potential" << endl;
        potentialSolver->solve();
    }

    real oldTol = dolfin_get("Krylov relative tolerance");
    dolfin_set("Krylov relative tolerance",0.0);
    cout << "Solving bioheat" << endl;
    bioSolver->step();
    dolfin_set("Krylov relative tolerance", oldTol);

    time += nsTime-time;

    real maxT = bioSolver->getT()->vector().norm(linf);
    cout << "Max temperature " << maxT - 273.15 << " at time " << time << endl;

    if (maxT >= 100 + 273.15)
    {
        cout << "Pop temperature reached, aborting simulation." << endl;
        bioSolver->save();
        exit(100);
    }
    
    if (time >= tFinal*(real(sample)/real(no_samples)))
    {
        cout << "saving potential at time " << time << endl;
        potentialSolver->save();
        sample++;
    }
}

void RFASolver::run()
{
    potentialSolver->matchInitialPower();
    bioSolver->setSigmaBoard(potentialSolver->getSigmaBoard());

    while (time < tFinal)
    {
        step();
    }
}

void RFASolver::setNoSamples(const uint n)
{
    no_samples = n;
    nsSolver->setNoSamples(n);
    bioSolver->setNoSamples(n);
}

#endif /* end of include guard: RFASOLVER_H */
