#ifdef __APPLE__
    #include <GLUT/glut.h>
#else
    #include <GL/glut.h>
#endif

#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <math.h>


#define LATTICE_WIDTH   100
#define LATTICE_HEIGHT  60

#define WINDOW_WIDTH    LATTICE_WIDTH
#define WINDOW_HEIGHT   LATTICE_HEIGHT

#define X_CUTOFF        5
#define FPS             4

#define VECTOR_SPACING  1

#define POINT_SIZE      6
#define SPACING         1

#define Q               9
#define DIM             2

double omega = 1.7;
int step = 0;
int fps = 0;

double lattice[LATTICE_HEIGHT][LATTICE_WIDTH][Q];
double tmp_lattice[LATTICE_HEIGHT][LATTICE_WIDTH][Q];
int boundary[LATTICE_HEIGHT][LATTICE_WIDTH];

int e[Q][DIM] = {{0,0}, {1, 0}, {0, 1}, {-1, 0}, {0, -1},
                        {1, 1}, {-1,1}, {-1,-1}, {1, -1}};
double weights[Q] = {4.0/9.0, 1.0/9.0 , 1.0/9.0 , 1.0/9.0 , 1.0/9.0 ,
                              1.0/36.0, 1.0/36.0, 1.0/36.0, 1.0/36.0};


void resolve_collisions( ) {
    int x,y, i,j;
    int rho;
    double u[DIM];
    double u_squared;
    double u_dot_e[Q];
    double f_eq;

    for(x = 0; x < LATTICE_WIDTH; x++) {
        for(y = 0; y < LATTICE_HEIGHT; y++) {
            
            rho = 0;
            for(i = 0; i < Q; i++) {
                rho += lattice[y][x][i];
            }

            for(i = 0; i < DIM; i++) {
                u[i] = 0;
            }
            for(i = 0; i < DIM; i++) {
                for(j = 0; j < Q; j++) {
                    u[i] += lattice[y][x][j] * e[j][i];
                }
            }

/*            if(rho != 0) {*/
/*                for(i = 0; i < DIM; i++) {*/
/*                    u[i] /= rho;*/
/*                }*/
/*            }*/

            u_squared = 0;
            for(i = 0; i < DIM; i++) {
                u_squared += u[i] * u[i];
            }

            for(i = 0; i < Q; i++) {
                u_dot_e[i] = 0;
            }
            for(i = 0; i < Q; i++) {
                for(j = 0; j < DIM; j++) {
                    u_dot_e[i] += u[j] * e[i][j];
                }
            }

            for(i = 0; i < Q; i++) {
                f_eq = weights[i] * (rho - 1.5 * u_squared + 3 * u_dot_e[i] + 
                                           4.5 * u_dot_e[i] * u_dot_e[i]);
                tmp_lattice[y][x][i] = (1.0 - omega) * lattice[y][x][i] + omega * f_eq;
            }

        }
    }
}

void propagate( ) {
    int x, y, i;
    int e_x, e_y;
    int new_x, new_y, new_i;
    int flip = 1;

    for(x = 0; x < LATTICE_WIDTH; x++) {
        for(y = 0; y < LATTICE_HEIGHT; y++) {
            for(i = 0; i < Q; i++) {
                flip = boundary[y][x] * -2 + 1;
                e_x = e[i][0];
                e_y = e[i][1];
/*                new_x = (x + e_x + LATTICE_WIDTH) % LATTICE_WIDTH;*/
/*                new_y = (y + e_y + LATTICE_HEIGHT) % LATTICE_HEIGHT;*/
                new_x = x + e_x;
                new_y = y + e_y;
                new_i = (i + boundary[y][x]*(2 + (1 & (i*4/Q)*3))) % Q;
                if(new_x >= 0 && new_y >= 0 && new_x < LATTICE_WIDTH && 
                        new_y < LATTICE_HEIGHT) {
                    if(tmp_lattice[y][x][new_i] > 1.0/3.0) {
                        lattice[new_y][new_x][new_i] = 1.0/3.0;
                    }
                    else {
                        lattice[new_y][new_x][new_i] = tmp_lattice[y][x][i];
                    }
                }
                else if(new_x == LATTICE_WIDTH) {
                    lattice[0][0][0] = tmp_lattice[y][x][0];
                    lattice[0][0][1] = tmp_lattice[y][x][1];
                }
            }
        }
    }
}

void next( ) {
    propagate( );
    resolve_collisions( );
}

void display( ) {
    int x, y, i, j;
    int v_x, v_y;
    double slope[DIM];
    double intensity;

    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
    glLoadIdentity( );
    glPointSize(POINT_SIZE);

    glColor3f(0.0, 1.0, 0.0);
    glBegin(GL_POINTS);
        for(y = 0; y < LATTICE_HEIGHT; y++) {
            for(x = X_CUTOFF; x < LATTICE_WIDTH; x++) {
                if(boundary[y][x]) {
                    glVertex3f(x, y, 0);
                }
            }
        }
    glEnd( );

    for (x = X_CUTOFF; x < LATTICE_WIDTH; x += VECTOR_SPACING) {
        for (y = 0; y < LATTICE_HEIGHT; y += VECTOR_SPACING) {
            intensity = 0.0;
            for(j = 0; j < DIM; j++) {
                slope[j] = 0.0;
            }
            for(v_x = x; v_x < x + VECTOR_SPACING; v_x++) {
                for(v_y = y; v_y < y + VECTOR_SPACING; v_y++) {
                    for(i = 0; i < Q; i++) {
                        for(j = 0; j < DIM; j++) {
                            intensity += lattice[v_y][v_x][i];
                            slope[j] += e[i][j] * lattice[v_y][v_x][i];
                        }
                    }
                }
            }

            for(j = 0; j < DIM; j++) {
                slope[j] /= 0.5 * VECTOR_SPACING;
            }
            if(slope[0] > 0 && slope[DIM-1] < 0) {
                glColor3f(1.0, intensity + 0.2, 0.0);
            }
            else {
                glColor3f(0.0, intensity + 0.2, 1.0);
            }
            glBegin(GL_LINES);
                glVertex3f(x, y, 0);
                glVertex3f(x + slope[0], y + slope[1], 0);
            glEnd( );

            glPointSize(POINT_SIZE/3);
            glVertex3f(intensity, intensity, intensity);
            glBegin(GL_POINTS);
                glVertex3f(x,y,0);
            glEnd( );
        }
    }

    glutSwapBuffers( );
}

void idle( ) {
    int i, j, k;

    for (i = 2; i < X_CUTOFF; i++) {
        for (j = LATTICE_HEIGHT/2 - X_CUTOFF; j < LATTICE_HEIGHT/2 + X_CUTOFF; j++) {
            lattice[j][i][1] = 0.15;
/*            lattice[j][i][5] = 0.05;*/
/*            lattice[j][i][8] = 0.05;*/

            tmp_lattice[j][i][1] = 0.15;
/*            tmp_lattice[j][i][5] = 0.05;*/
/*            tmp_lattice[j][i][8] = 0.05;*/
        }
    }

    if(step || !step) {
        next( );
        step = 0;
    }
    if(fps == FPS) {
        glutPostRedisplay();
        fps -= FPS;
    }
    fps++;
}

void init_vars( )
{
    int i, j, k, l;

    for(i = 0; i < LATTICE_WIDTH; i++) {
        for(j = 0; j < LATTICE_HEIGHT; j++) {
            boundary[i][j] = 0;
            for(k = 0; k < Q; k++) {
                lattice[j][i][k] = 0.00;
                tmp_lattice[j][i][k] = 0.00;
            }
/*            lattice[j][i][1] = 0.01;*/
/*            lattice[j][i][5] = 0.005;*/
/*            lattice[j][i][8] = 0.005;*/
        }
    }

/*    for (i = 0; i < X_CUTOFF; i++) {*/
/*        for (j = 0; j < LATTICE_HEIGHT; j++) {*/
/*            for(k = 0; k < Q; k++) {*/
/*                lattice[j][i][k] = 0.2;*/
/*                tmp_lattice[j][i][k] = 0.2;*/
/*            }*/
/*        }*/
/*    }*/

/*    for(i = LATTICE_HEIGHT/3; i < 3*LATTICE_WIDTH/4; i++) {*/
/*        for(j = LATTICE_WIDTH/3; j < 3*LATTICE_HEIGHT/4; j++) {*/
/*            for(k = 0; k < Q; k++) {*/
/*                lattice[j][i][k] = 0.2;*/
/*                tmp_lattice[j][i][k] = 0.2;*/
/*            }*/
/*            lattice[j][i][1] = 0.2;*/
/*            tmp_lattice[j][i][1] = 0.2;*/
/*        }*/
/*    }*/

/*    for(i = LATTICE_HEIGHT/6; i < LATTICE_WIDTH/3; i++) {*/
/*        for(j = LATTICE_WIDTH/6; j < LATTICE_HEIGHT/3; j++) {*/
/*            for(k = 0; k < Q; k++) {*/
/*                lattice[i][j][k] = 1;*/
/*            }*/
/*        }*/
/*    }*/
    
    int x, y;
    double d_theta = 0.01;
    double theta = d_theta;
    int radius = 1;

/*    for(i = 0; i < 2*M_PI/d_theta; i++, theta += d_theta) {*/
/*        for(j = 0; j < radius; j++) {*/
/*            for(k = 1; k < 1; k++) {*/
/*                x = k*LATTICE_WIDTH/6 + j * cos(theta);*/
/*                y = k*LATTICE_HEIGHT/3 + j * sin(theta);*/
/*                boundary[y][x] = 1;*/
/*                for(l = 0; l < Q; l++) {*/
/*                    lattice[y][x][l] = 0;*/
/*                    tmp_lattice[y][x][l] = 0;*/
/*                }*/
/*            }*/
/*        }*/
/*    }*/

/*    for(i = 0; i < 2*M_PI/d_theta; i++, theta += d_theta) {*/
/*        for(j = 0; j < radius; j++) {*/
/*            x = LATTICE_WIDTH/2 + j * cos(theta);*/
/*            y = LATTICE_HEIGHT/2 + j * sin(theta);*/
/*            boundary[y][x] = 1;*/
/*            for(k = 0; k < Q; k++) {*/
/*                lattice[y][x][k] = 0;*/
/*                tmp_lattice[y][x][k] = 0;*/
/*            }*/
/*        }*/
/*    }*/

    for(i = 0; i < 1; i++) {
        for(j = 0; j < 5; j++) {
            x = LATTICE_WIDTH/2 + i;
            y = LATTICE_HEIGHT/2 + j - 2;
            boundary[y][x]   = 1;
            for(k = 0; k < Q; k++) {
                lattice[y][x][k] = 0;
                tmp_lattice[y][x][k] = 0;
            }
        }
    }
    for(x = 0; x < LATTICE_WIDTH; x++) {
        boundary[0][x] = 1;
        boundary[LATTICE_HEIGHT-1][x] = 1;
        for(k = 0; k < Q; k++) {
            lattice[0][x][k] = 0;
            tmp_lattice[0][x][k] = 0;

            lattice[LATTICE_HEIGHT-1][x][k] = 0;
            tmp_lattice[LATTICE_HEIGHT-1][x][k] = 0;
        }
    }

/*    for(x = 0; x < LATTICE_HEIGHT; x++) {*/
/*        boundary[x][0] = 1;*/
/*        boundary[x][LATTICE_WIDTH-1] = 1;*/
/*        for(k = 0; k < Q; k++) {*/
/*            lattice[x][0][k] = 0;*/
/*            tmp_lattice[x][0][k] = 0;*/

/*            lattice[x][LATTICE_WIDTH-1][k] = 0;*/
/*            tmp_lattice[x][LATTICE_WIDTH-1][k] = 0;*/
/*        }*/
/*    }*/
}

void init_gl()
{
    glClearColor(0.0, 0.0, 0.0, 0.0);

    glMatrixMode(GL_PROJECTION);
    glLoadIdentity();
    glOrtho(0.0, WINDOW_WIDTH, 0.0, WINDOW_HEIGHT, -20.0, 20.0);

    glMatrixMode(GL_MODELVIEW);

    glHint(GL_PERSPECTIVE_CORRECTION_HINT, GL_NICEST);
    glShadeModel(GL_SMOOTH);
    glPointSize(POINT_SIZE);
    glLineWidth(POINT_SIZE);
}

void keyboard(unsigned char key, int x, int y) {
    if(key == 'n') {
        step = 1;
    }
}

int main(int argc, char** argv) {
    glutInit(&argc, argv);
    glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGB | GLUT_DEPTH);
    glutInitWindowSize(WINDOW_WIDTH, WINDOW_HEIGHT);
    glutCreateWindow("D2Q9");
    glEnable(GL_DEPTH_TEST);
    init_vars();
    init_gl();
    glutDisplayFunc(display);
    glutKeyboardUpFunc(keyboard);
    glutIdleFunc(idle);
    glutMainLoop();
    return 0;
}