#include "grid.hh"
#include "turing_vtk.hh"

#include <vector>
#include <cmath>
#include <cassert>

class turing_problem
{
public:
    // domain
    const structured_grid & grid;
    // data
    const std::vector<double> & a;
    const std::vector<double> & b;
private:
    // parameters
    double Da, Db;
    double eps0, eps1;
    double q, f, m;

public:
    // constructor
    turing_problem(const structured_grid & _g,
        const std::vector<double> & _a,
        const std::vector<double> & _b,
        double _Da, double _Db,
        double _e0, double _e1,
        double _q, double _f, double _m) :
        grid(_g), a(_a), b(_b),
        Da(_Da), Db(_Db),
        eps0(_e0), eps1(_e1),
        q(_q), f(_f), m(_m)
    {}

    // problem description
    void compute_local_update(unsigned int i[3],
        double & da, double & db) const
    {
        unsigned int idx = grid.index(i);

        da = db = 0.0;
        
        // laplace with neumann boundary conditions
        std::vector<unsigned int> neighbors = grid.neighbors(i);
        for (unsigned int n=0; n<neighbors.size(); n++)
        {
            unsigned int nidx = neighbors[n];
            // Da \Delta a[idx]
            da += Da * (a[nidx] - a[idx]);
            // Db \Delta b[idx]
            db += Db * (b[nidx] - b[idx]);
        }

        // reaction terms
        double w0 = (1.0 - m*b[idx])/(1.0 - m*b[idx] + eps1);
        double w1 = f * (q - a[idx])/(q + a[idx]);

        // eq. 0: f(a,b) = 1/eps_0 * [ w_0*a[idx] + w_1*b[idx] - a[idx]^2 ]
        da += 1.0/eps0 * (a[idx]*w0 + b[idx]*w1 - a[idx]*a[idx]);
        // eq. 1: g(a,b) = (w_0*a[idx] - b[idx])
        db += w0*a[idx] - b[idx];
    }
};

void compute_update(const turing_problem & p,
    std::vector<double> & da,
    std::vector<double> & db,
    double & dt)
{
    unsigned int i[3];
    // parallel: nur für eigene Knoten das Update aufrufen
    for (i[0]=0; i[0]<p.grid.size(0); i[0]++)
        for (i[1]=0; i[1]<p.grid.size(1); i[1]++)
            for (i[2]=0; i[2]<p.grid.size(2); i[2]++)
            {
                unsigned int idx = p.grid.index(i);
                p.compute_local_update(i, da[idx], db[idx]);
                if (da[idx] < 0.0)
                {
                    assert(p.a[idx] > 0.0);
                    dt = std::min( dt, std::abs(p.a[idx]/da[idx]) );
                }
                if (db[idx] < 0.0)
                {
                    assert(p.b[idx] > 0.0);
                    dt = std::min( dt, std::abs(p.b[idx]/db[idx]) );
                }
            }
    // parallel: global minimales dt berechnen
}

int main(int argc, char** argv)
{
    // parallel: MPI starten
    
    // read parameter
    int dim = 2;
    int D = 1<<7;   // diameter in cells
    double d = 20;  // diameter in units
    double f = 1.1; // problem param f
    double vt = 1;  // time interval for vtk output
    double max_dt = vt; // max time interval for explicit euler
    double T = 1000; // end time
    // create domain info
    int Z = 2*D;
    if (dim == 2)
        Z = 0;
    // parallel: parallele Information an das Gitter übergeben 
    structured_grid grid(D,D,Z, d/D);
    std::cout << "created a grid with "
              << grid.size() << " vertices" << std::endl;
    // allocate data
    std::vector<double> a(grid.size());
    std::vector<double> b(grid.size());
    // initialize data
    srand48(42);
    for (unsigned int i=0; i<grid.size(); i++)
    {
        a[i] = 1e-2*drand48();
        b[i] = 0.0;
    }
    
    // allocate update vectors
    std::vector<double> da(grid.size());
    std::vector<double> db(grid.size());
    // setup problem description
    turing_problem p(grid, a, b,
        1.0, 10.0,
        2.2, 0.02,
        2e-4, f, 7e-4);

    // vtk writer
    vtk_writer vtk_a("turing_a",grid,a);
    vtk_writer vtk_b("turing_b",grid,b);

    // time loop
    double t = 0.0;
    double v = vt;
    vtk_a.write();
    vtk_b.write();
    while (t < T)
    {
        double sug_dt = T;
        compute_update(p, da, db, sug_dt);
        double dt = std::min(sug_dt * 0.3, max_dt);
        for(unsigned int i=0; i<grid.size(); i++)
        {
            a[i] += dt*da[i];
            b[i] += dt*db[i];
        }
        // parallel: aktualisierte Daten kommunizieren
        if (v <= t)
        {
            std::cout << "from t=" << t
                      << ", dt_sug: " << sug_dt
                      << ", dt: " << dt
                      << " to " << t+dt
                      << std::endl;
            std::cout << "write vtk -- time: " << t << std::endl;
            vtk_a.write_parallel(0,1);
            vtk_b.write();
            v+=vt;
        }
        t += dt;
    }
    vtk_a.write();
    vtk_b.write();

    // parallel: MPI beenden
}