/*
 * Author: target0 <target0@geeknode.org>
 * Compilation: gcc -o fractal fractal.c `sdl-config --cflags --libs`
 * Usage: ./fractal -h
 */

#include <stdio.h>
#include <unistd.h>
#include <getopt.h>
#include <SDL.h>
#include <pthread.h>

#define rmask 0x000000ff
#define gmask 0x0000ff00
#define bmask 0x00ff0000
#define amask 0xff000000

#define NTHREADS 16
#define MAX_ITER 4090
#define itoc(x) ((0x00ffffff/MAX_ITER)*(x))

struct threadarg {
    double *coordX, *coordY;
    int from, to;
    int **array;
};

enum { RED, GREEN, BLUE, FULL, };
int color, fsize;
char expfile[128];
double xcoord, ycoord;

void *calc_mandel_mt(void *);

int **get_fractarray(int xsize, int ysize)
{
    int i;
    int **array = malloc(xsize*sizeof(int *));

    for (i = 0; i < xsize; i++)
        array[i] = malloc(ysize*sizeof(int));

    return array;
}

void free_fractarray(int **array, int ysize)
{
    int i;
    for (i = 0; i < ysize; i++)
        free(array[i]);
    free(array);
}

int iter(double zx, double zy, double a, double b, int it)
{
    if (it > MAX_ITER)
        return 0;

    if (zx*zx + zy*zy >= 4.0)
        return it;

    return iter(zx*zx - zy*zy + a, 2*zx*zy + b, a, b, it+1);
}

double iter_julia(double zx, double zy, double a, double b, int it)
{
    if (it > MAX_ITER)
        return zx*zx+zy*zy;

    return iter(zx*zx - zy*zy + a, 2*zx*zy + b, a, b, it+1);
}

int **calc_mandel(double *coordX, double *coordY, int axsize, int aysize, int **array)
{
    int i, j;
    for (i = 0; i < axsize; i++) {
        for (j = 0; j < aysize; j++) {
            array[i][j] = iter(0, 0, coordX[i], coordY[j], 0);
        }
    }

    return array;
}

void *calc_mandel_mt(void *arg)
{
    struct threadarg *t_arg = (struct threadarg *)arg;
    int i, j;
    printf("from %d to %d\n", t_arg->from, t_arg->to);
    for (i = t_arg->from; i < t_arg->to; i++) {
        for (j = 0; j < 800; j++) {
            t_arg->array[i][j] = iter(0, 0, t_arg->coordX[i], t_arg->coordY[j], 0);
        }
    }

    free(t_arg);
    return NULL;
}

int **calc_julia(double *coordX, double *coordY, int axsize, int aysize, double a, double b, int **array)
{
    int i, j;
    for (i = 0; i < axsize; i++) {
        for (j = 0; j < aysize; j++) {
            array[i][j] = iter_julia(coordX[i], coordY[j], a, b, 0);
        }
    }

    return array;
}

double *get_coord(double from, double to, double size, double *array)
{
    double i;
    int ai;
    double gap = abs(from - to)/size;

    for (i = from,  ai = 0; i<=to && ai < size; i+=gap, ai++)
        array[ai] = i;

    return array;
}

void show_fractal(SDL_Surface *screen, int **array, int xsize, int ysize)
{
    SDL_Rect pix;
    int i, j;
    unsigned int tcol, rtcol, gtcol, btcol;
    Uint32 pcol;

    for (i = 0; i < xsize; i++) {
        for (j = 0; j < ysize; j++) {
            tcol = itoc(array[i][j]);
            switch (color) {
            case RED:
                rtcol = tcol & rmask;
                gtcol = 0;
                btcol = 0;
                break;
            case GREEN:
                rtcol = 0;
                gtcol = tcol & rmask;
                btcol = 0;
                break;
            case BLUE:
                rtcol = 0;
                gtcol = 0;
                btcol = tcol & rmask;
                break;
            default:
                rtcol = tcol & rmask;
                gtcol = (tcol & gmask)>>8;
                btcol = (tcol & bmask)>>16;
            }
                
            pcol = SDL_MapRGB(screen->format, rtcol, gtcol, btcol);
            pix.w = pix.h = 1;
            pix.x = i;
            pix.y = j;
            SDL_FillRect(screen, &pix, pcol);
            SDL_UpdateRects(screen, 1, &pix);
        }
    }
}

void fill_array(int **array, int xsize, int ysize)
{
    int i, j;
    for (i = 0; i < xsize; i++) {
        for (j = 0; j < ysize; j++) {
            array[i][j] = random()%0xffffff;
        }
    }
}

SDL_Surface *init_video(int xsize, int ysize)
{
    if (SDL_Init(SDL_INIT_VIDEO) < 0) {
        fprintf(stderr, "SDL error: %s\n", SDL_GetError());
        return NULL;
    }   
    atexit(SDL_Quit);
    
    return SDL_SetVideoMode(xsize, ysize, 32, 0);
}

void show_cursor(SDL_Surface *screen, int x, int y)
{
    Uint32 pcol;
    SDL_Rect cursor;

    pcol = SDL_MapRGB(screen->format, 255, 255, 255);
    cursor.w = cursor.h = 1;
    cursor.x = x;
    cursor.y = y;
    SDL_FillRect(screen, &cursor, pcol);
    SDL_UpdateRects(screen, 1, &cursor);
}

void fractal(int xsize, int ysize, double xcoord, double ycoord)
{
    SDL_Event event;
    SDL_Surface *screen, *back;
    int **array;
    double *coordX, *coordY;
    pthread_t t[NTHREADS];
    unsigned int i;
    struct threadarg *t_arg;
    int lb, ub, cx, cy;
    Uint32 bpcol;

    printf("Allocating memory... ");
    fflush(stdout);

    array = get_fractarray(xsize, ysize);

    screen = init_video(xsize, ysize);
    if (!screen) {
        fprintf(stderr, "Fail: %s\n", SDL_GetError());
        return;
    }

    printf("done.\nComputing fractal... ");
    fflush(stdout);

    coordX = malloc(xsize*sizeof(double));
    coordY = malloc(ysize*sizeof(double));
//    get_coord(-2, 1, xsize, coordX); // XXX mandel
//    get_coord(-1.5, 2, xsize, coordX); // XXX julia
//    get_coord(-1.5, 2, ysize, coordY);
//    calc_mandel(coordX, coordY, xsize, ysize, array);
//    calc_julia(coordX, coordY, xsize, ysize, -0.6666, 0.3333, array);

    back = SDL_CreateRGBSurface(SDL_SWSURFACE, xsize, ysize, 32, rmask, gmask, bmask, amask);

    if (xcoord == 0 && ycoord == 0) {
        get_coord(-2, 1, xsize, coordX);
        get_coord(-1.5, 2, ysize, coordY);

        lb = 0;
        ub = xsize/NTHREADS;

        for (i = 0; i < NTHREADS; i++) {
            t_arg = malloc(sizeof(*t_arg));
            t_arg->coordX = coordX;
            t_arg->coordY = coordY;
            t_arg->array = array;
            t_arg->from = lb;
            t_arg->to = ub;

            if ((pthread_create(&t[i], NULL, calc_mandel_mt, (void *)t_arg)) != 0) {
                perror("pthread_create");
                exit(0);
            }

            lb = ub;
            ub += xsize/NTHREADS;
        }

        for (i = 0; i < NTHREADS; i++) {
            if ((pthread_join(t[i], NULL)) != 0) {
                perror("pthread_join");
                exit(0);
            }
        }

//        calc_mandel(coordX, coordY, xsize, ysize, array);
    } else {
        get_coord(-1.5, 2, xsize, coordX);
        get_coord(-1.5, 2, ysize, coordY);
        calc_julia(coordX, coordY, xsize, ysize, xcoord, ycoord, array);
    }

    printf("done.\nDisplaying data... ");
    fflush(stdout);

    show_fractal(back, array, xsize, ysize);

    SDL_BlitSurface(back, NULL, screen, NULL);
    SDL_Flip(screen);
    printf("done.\n");

    cx = xsize/2;
    cy = ysize/2;
 //   show_cursor(screen, cx, cy);

    if (*expfile)
        SDL_SaveBMP(screen, expfile);
    while (SDL_WaitEvent(&event)) {
        if (event.type == SDL_QUIT)
            break;
    }

    free(coordX);
    free(coordY);
    free_fractarray(array, ysize);
    SDL_FreeSurface(back);
}

void usage(char *name)
{
    fprintf(stderr, "Usage: %s [-chsxy] [-f filename]\n\n"
                    "\tOptions:\n\n"
                    "\t-c color\tSet color (0: red, 1: green, 2:blue, 3: full colors\n"
                    "\t-s size\tSet fractal size (square)\n"
                    "\t-f file\tExport fractal to specified file (BMP format)\n"
                    "\t-x coord\tSet X coordinate (default: 0)\n"
                    "\t-y coord\tSet Y coordinate (default: 0)\n"
                    "\t-h\tShow this help\n", name);
    _exit(1);
}

int main(int ac, char **av)
{
    char opt;

    color = FULL;
    fsize = 800;
    bzero(expfile, 128);
    xcoord = ycoord = 0;

    while ((opt = getopt(ac, av, "c:s:f:t:x:y:h")) != EOF) {
        switch(opt) {
        case 'c':
            color = atoi(optarg);
            if (color < 0 || color > FULL)
                usage(av[0]);
            break;
        case 's':
            fsize = atoi(optarg);
            break;
        case 'f':
            strncpy(expfile, optarg, 128);
            break;
        case 'x':
            xcoord = strtod(optarg, NULL);
            break;
        case 'y':
            ycoord = strtod(optarg, NULL);
            break;
        case 'h':
        default:
            usage(av[0]);
        }
    }

    fractal(fsize, fsize, xcoord, ycoord);

    return 0;
}