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

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

#include <mpi.h>

const int TAG_OVERLAP_A = 42;
const int TAG_OVERLAP_B = 43;

class turing_problem
{
public:
    // domain
    const parallel_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 parallel_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]=p.grid.overlap(0,0); i[0]<p.grid.size(0)-p.grid.overlap(0,1); i[0]++)
        for (i[1]=p.grid.overlap(1,0); i[1]<p.grid.size(1)-p.grid.overlap(1,1); i[1]++)
            for (i[2]=p.grid.overlap(2,0); i[2]<p.grid.size(2)-p.grid.overlap(2,1); 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
    double dt_global_min;
    MPI_Allreduce (&dt, &dt_global_min, 1, MPI_DOUBLE, MPI_MIN, MPI_COMM_WORLD );
    dt = dt_global_min;
}

void getOverlap(const turing_problem & p, unsigned int d, bool posDir, const std::vector<double> & data, double *overlap) {
    unsigned int dims[2];
    switch (d) {
        case 0: dims[0]=1; dims[1]=2; break;
        case 1: dims[0]=0; dims[1]=2; break;
        case 2: dims[0]=0; dims[1]=1; break;
    }
    unsigned int i[3];
    i[d] = p.grid.overlap(d,posDir);
    if (posDir)
        i[d] = p.grid.size(d)-1-i[d];
    for (i[dims[0]]=0; i[dims[0]]<p.grid.size(dims[0]); i[dims[0]]++)
        for (i[dims[1]]=0; i[dims[1]]<p.grid.size(dims[1]); i[dims[1]]++)
            overlap[i[dims[0]]*p.grid.size(dims[1])+i[dims[1]]]=data[p.grid.index(i)];
}

void setOverlap(const turing_problem & p, unsigned int d, bool posDir, std::vector<double> & data, double *overlap) {
    unsigned int dims[2];
    switch (d) {
        case 0: dims[0]=1; dims[1]=2; break;
        case 1: dims[0]=0; dims[1]=2; break;
        case 2: dims[0]=0; dims[1]=1; break;
    }
    unsigned int i[3];
    i[d] = posDir? p.grid.size(d)-p.grid.overlap(d,posDir) : 0;
    for (i[dims[0]]=0; i[dims[0]]<p.grid.size(dims[0]); i[dims[0]]++)
        for (i[dims[1]]=0; i[dims[1]]<p.grid.size(dims[1]); i[dims[1]]++)
            data[p.grid.index(i)] = overlap[i[dims[0]]*p.grid.size(dims[1])+i[dims[1]]];
}

void communicateOverlap(const turing_problem & p, const structured_grid & proc_grid, std::vector<double> & data, unsigned int rank, int TAG) {
    unsigned int i[3], n[3];
    for (int d=0; d<3; d++) {
        i[d] = proc_grid.index(rank,d);
        n[d] = i[d];
    }
    for (int k=0; k<6; k++) {
        int d=k/2;
        bool posDir=k%2;
        
        if (p.grid.overlap(d,posDir)) {
            n[d]+=posDir?1:-1;
            // retrieve own overlap
            unsigned int count=(p.grid.size()/p.grid.size(d))*p.grid.overlap(d,posDir);
            double soverlap[count];
            getOverlap(p,d,posDir,data,soverlap);
            unsigned int to = proc_grid.index(n);
            // send it to neighbor
            if (MPI_SUCCESS!=MPI_Send(soverlap,count,MPI_DOUBLE, to, TAG, MPI_COMM_WORLD))
                std::cout<<"sending failed"<<std::endl;
            // recieve overlap from neighbor
            double roverlap[count];
            if (MPI_SUCCESS!=MPI_Recv(roverlap,count,MPI_DOUBLE, to ,TAG, MPI_COMM_WORLD,NULL))
                std::cout<<"receiving failed"<<std::endl;
            // store it in data
            setOverlap(p,d,posDir,data,roverlap);
            n[d]+=posDir?-1:1;
        }
    }
}

int main(int argc, char** argv)
{
    // 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;
    // initialize MPI
    MPI_Init(&argc ,&argv );
    int rank;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank );
    int size;
    MPI_Comm_size(MPI_COMM_WORLD, &size);
	
    // create process grid
    // try to factorize s...for now only in 2D
    unsigned int s = std::sqrt(size);
    while ((size/s)*s!=size) s--;
    structured_grid proc_grid(size/s-1,s-1,0);
    std::cout << rank <<" uses a process grid: "<<proc_grid.size(0)<<"x"<<proc_grid.size(1)<<"x"<<proc_grid.size(2)<<std::endl;
    
    //create parallel grid
    parallel_structured_grid grid(D,D,Z, d/D,proc_grid,rank);
    
    std::cout << rank << " created a grid with "
              << grid.size() << " vertices " 
              << "of size "<<grid.size(0)<<"x"<<grid.size(1)<<"x"<<grid.size(2)<< std::endl;
    // allocate data
    std::vector<double> a(grid.size());
    std::vector<double> b(grid.size());
    // initialize data
    srand48(42*rank);
    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,proc_grid,a,rank);
    vtk_writer vtk_b("turing_b",grid,proc_grid,b,rank);

    // time loop
    double t = 0.0;
    double v = vt;
    vtk_a.write();
    vtk_b.write();
//    std::cout<<rank<<" finished writing"<<std::endl;
    while (t < T)
    {
        double sug_dt = T;
        compute_update(p, da, db, sug_dt);
        double dt = std::min(sug_dt * 0.3, max_dt);
//        std::cout<<"####### changing b #####"<<std::endl;
        for(unsigned int i=0; i<grid.size(); i++)
        {
            a[i] += dt*da[i];
            b[i] += dt*db[i];
        }
        // parallel: aktualisierte Daten kommunizieren
        communicateOverlap(p,proc_grid,a,rank,TAG_OVERLAP_A);
        communicateOverlap(p,proc_grid,b,rank,TAG_OVERLAP_B);
        if (v <= t)
        {
            if (rank==0) {
                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();
//            std::cout<<rank<<" finished writing"<<std::endl;
            v+=vt;
        }
        t += dt;
    }
    vtk_a.write();
    vtk_b.write();
    // parallel: MPI beenden
    MPI_Finalize();
    return 0;
}