/*
 * Josh Holtrop
 * 2008-12-11
 * barebones SDL program
 */

#include <SDL/SDL.h>
#include <mpi.h>
#include <iostream>
#include "Computation.h"
#include "NewtonComputation.h"
#include "FatouComputation.h"
using namespace std;

/* a "task" will be processing TASK_SIZE pixels */
#define TASK_SIZE 100
#define PROGNAME "Josh's CS677 Final : MPI Fractal Generator"
#define getXVirt(x) (((x) - (width >> 1)) * zoom + x_center)
#define getYVirt(y) ((-((y) - (height >> 1))) * zoom + y_center)

bool createWindow(int width, int height,
                  SDL_Surface ** screen, Uint32 ** pixels);
void getSizes(int * rank, int * size, int * nprocs);
void draw(int rank, int world_size, int nprocs, int width, int height,
    Uint32 * pixels, Computation * computation);

static double x_center = 0.0;
static double y_center = 0.0;
static double zoom = 1/300.0;

int main(int argc, char * argv[])
{
    int width = 600;
    int height = 600;
    int my_rank = 0;
    int world_size = 0;
    int nprocs = 0;
    Computation * computation = NULL;
    int fractal_type = 0;

    SDL_Surface * screen;
    Uint32 * pixels;

    MPI_Init(&argc, &argv);
    for (int i = 1; i < argc; i++)
    {
        if (!strncmp(argv[i], "-w", 2))
        {
            width = atoi(strlen(argv[i]) > 2 ? argv[i] + 2 : argv[++i]);
        }
        else if (!strncmp(argv[i], "-h", 2))
        {
            height = atoi(strlen(argv[i]) > 2 ? argv[i] + 2 : argv[++i]);
        }
        else if (!strncmp(argv[i], "-t", 2))
        {
            fractal_type = atoi(strlen(argv[i]) > 2 ? argv[i] + 2 : argv[++i]);
        }
    }
    getSizes(&my_rank, &world_size, &nprocs);

    switch (fractal_type)
    {
        case 0:
        default:
            computation = new NewtonComputation();
            x_center = 0.0;
            y_center = 0.0;
            zoom = 2.0 / width;
            break;
        case 1:
            computation = new FatouComputation();
            x_center = 3.001;
            y_center = 0.075975;
            zoom = 2.0 / width;
            break;
    }

    if (my_rank == 0)
    {
        SDL_Event event;
        bool going = true;
        bool window_success = createWindow(width, height, &screen, &pixels);
        bool redraw = true;

        if (!window_success)
            going = false;

        while (going && SDL_WaitEvent(&event) != 0)
        {
            if (redraw)
            {
                draw(my_rank, world_size, nprocs, width, height,
                    pixels, computation);
                redraw = false;
            }
            SDL_UpdateRect(screen, 0, 0, 0, 0);
            switch (event.type)
            {
                case SDL_QUIT:
                    going = false;
                    break;
                case SDL_KEYDOWN:
                    if (event.key.keysym.sym == SDLK_q)
                        going = false;
                    break;
                case SDL_MOUSEBUTTONDOWN:
                    switch (event.button.button)
                    {
                        case 1:         /* left-click to re-center and zoom in */
                            x_center = getXVirt(event.button.x);
                            y_center = getYVirt(event.button.y);
                            zoom /= 2.0;
                            redraw = true;
                            break;
                        case 2:         /* middle click just to re-center */
                            x_center = getXVirt(event.button.x);
                            y_center = getYVirt(event.button.y);
                            redraw = true;
                            break;
                        case 4:         /* zoom in */
                            zoom /= 2.0;
                            redraw = true;
                            break;
                        case 5:         /* zoom out */
                            zoom *= 2.0;
                            redraw = true;
                            break;
                    }
                    break;
            }
        }
    }

    MPI_Finalize();

    delete computation;
    return 0;
}

bool createWindow(int width, int height,
                  SDL_Surface ** screen, Uint32 ** pixels)
{
    if (SDL_Init(SDL_INIT_VIDEO))
    {
        cerr << "Failed to initialize SDL!" << endl;
        return false;
    }

    atexit(SDL_Quit);

    if (!(*screen = SDL_SetVideoMode(width, height, 32, 0)))
    {
        cerr << "Failed to set video mode!" << endl;
        return false;
    }
    SDL_WM_SetCaption(PROGNAME, PROGNAME);

    *pixels = (Uint32 *) (*screen)->pixels;

    return true;
}

void getSizes(int * rank, int * size, int * nprocs)
{
    MPI_Comm_rank(MPI_COMM_WORLD, rank);
    MPI_Comm_size(MPI_COMM_WORLD, size);

    *nprocs = sysconf(_SC_NPROCESSORS_CONF);

    int displs[*size];
    int counts[*size];
    for (int i = 0; i < *size; i++)
    {
        displs[i] = i;
        counts[i] = 1;
    }
    int all_nprocs[*size];
    MPI_Gatherv(nprocs, 1, MPI_INT,
                &all_nprocs[0], &counts[0], &displs[0], MPI_INT,
                0, MPI_COMM_WORLD);

    if (*rank == 0)
    {
        int total_nprocs = 0;
        cout << "Number of cores on each MPI node:" << endl;
        for (int i = 0; i < *size; i++)
        {
            cout << all_nprocs[i] << " ";
            total_nprocs += all_nprocs[i];
        }
        cout << endl;
        cout << "Total number of cores: " << total_nprocs << endl;
    }
}

void draw(int rank, int world_size, int nprocs, int width, int height,
    Uint32 * pixels, Computation * computation)
{
    if (world_size == 1)
    {
        for (int y = 0; y < height; y++)
        {
            for (int x = 0; x < width; x++)
            {
                double x_virt = getXVirt(x);
                double y_virt = getYVirt(y);
                *pixels++ = computation->compute(x_virt, y_virt);
            }
        }
    }
    else if (rank == 0)
    {
        int num_pixels = width * height;
        for (int to_proc = 1; to_proc < world_size; to_proc++)
        {
        }
    }
    else
    {
        for (;;)
        {
            /* wait to be told what to do */
#if 0
            MPI_Recv();
            if ()   /* exit if we are done */
                break;
            MPI_Send();     /* report results back */
#endif
        }
    }
}