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

typedef struct {
    int fitness;
    int* data;
} Square;

Square* squares;

long getNanos() {
    struct timespec t;
    clock_gettime(CLOCK_MONOTONIC, &t);
    return t.tv_nsec + t.tv_sec * 1000000000L;
}

int size = 0;
int population = 0;
int (*fitness)(int*) = NULL;

int compare(const void* a, const void* b) {
    int ia = ((Square*)a)->fitness;
    int ib = ((Square*)b)->fitness;
    return ia < ib ? -1 : (ia > ib ? 1 : 0);
}

void sort() {
    qsort(squares, population, sizeof(Square), compare);
}

void mutate(int* data) {
    int a = rand() % (size * size);
    int b = rand() % (size * size);
    while(b == a) {
        b = rand() % (size * size);
    }

    int tmp = data[a];
    data[a] = data[b];
    data[b] = tmp;
}

void reproduce(const int* parent, Square* child) {
    int end = size * size;
    for(int i = 0; i < end; i++) {
        child->data[i] = parent[i];
    }
    mutate(child->data);
    child->fitness = fitness(child->data);
}

void init() {
    squares = malloc(sizeof(Square) * population);
    int end = size * size;
    for(int i = 0; i < population; i++) {
        squares[i].data = malloc(sizeof(int) * end);
        for(int k = 0; k < end; k++) {
            squares[i].data[k] = k + 1;
        }
        for(int k = 0; k < end; k++) {
            mutate(squares[i].data);
        }
        squares[i].fitness = fitness(squares[i].data);
    }
}

void clear() {
    for(int i = 0; i < population; i++) {
        free(squares[i].data);
    }
    free(squares);
}

int get(int* data, int x, int y) {
    return data[y * size + x];
}

int errorFactor(int e) {
    return e * e;
}

int getWanted() {
    return (size * (size * size + 1)) / 2;
}

int getSemiFitness(int* data) {
    int fitness = 0;
    int wanted = getWanted();
    for(int y = 0; y < size; y++) {
        int sum = 0;
        for(int x = 0; x < size; x++) {
            sum += get(data, x, y);
        }
        fitness += errorFactor(sum - wanted);
    }
    for(int x = 0; x < size; x++) {
        int sum = 0;
        for(int y = 0; y < size; y++) {
            sum += get(data, x, y);
        }
        fitness += errorFactor(sum - wanted);
    }
    return fitness;
}

int getDiagonalBase(int* data) {
    int diaA = 0;
    int diaB = 0;
    for(int i = 0; i < size; i++) {
        diaA += get(data, i, i);
        diaB += get(data, i, (size - 1) - i);
    }
    int wanted = getWanted();
    int fitness = errorFactor(diaA - wanted);
    fitness += errorFactor(diaB - wanted);
    return fitness;
}

int getNormalFitness(int* data) {
    return getSemiFitness(data) + getDiagonalBase(data);
}

int getPandiagonalBase(int* data) {
    int fitness = 0;
    int wanted = getWanted();
    for(int x = 0; x < size; x++) {
        int sumA = 0;
        int sumB = 0;
        for(int y = 0; y < size; y++) {
            sumA += get(data, (x + y) % size, y);
            sumB += get(data, (x - y + size) % size, y);
        }
        fitness += errorFactor(sumA - wanted);
        fitness += errorFactor(sumB - wanted);
    }
    return fitness;
}

int getAssociativeBase(int* data) {
    int fitness = 0;
    int wanted = size * size + 1;
    for(int x = 0; x < size; x++) {
        for(int y = 0; y < size; y++) {
            int sum = get(data, x, y) + get(data, size - x - 1, size - y - 1);
            fitness += errorFactor(sum - wanted);
        }
    }
    return fitness;
}

int getUltraFitness(int* data) {
    return getSemiFitness(data) + getPandiagonalBase(data) +
           getAssociativeBase(data);
}

int get2x2SquareBase(int* data) {
    int fitness = 0;
    int wanted = 2 * (size * size + 1);
    for(int x = 0; x < size; x++) {
        for(int y = 0; y < size; y++) {
            if(x % 2 == 0 && y % 2 == 0) {
                int sum = get(data, x, y) + get(data, x + 1, y) +
                          get(data, x, y + 1) + get(data, x + 1, y + 1);
                fitness += errorFactor(sum - wanted);
            }
        }
    }
    return fitness;
}

int getPerfectDiagonalBase(int* data) {
    int fitness = 0;
    int wanted = size * size + 1;
    for(int x = 0; x < size / 2; x++) {
        for(int y = 0; y < size; y++) {
            int sum = get(data, x, y) +
                      get(data, (x + size / 2) % size, (y + size / 2) % size);
            fitness += errorFactor(sum - wanted);
        }
    }
    return fitness;
}

int getMostPerfectFitness(int* data) {
    return getSemiFitness(data) + getPandiagonalBase(data) +
           get2x2SquareBase(data) + getPerfectDiagonalBase(data);
}

void print(Square* s) {
    printf("Fitness: %d\n", s->fitness);
    for(int y = 0; y < size; y++) {
        for(int x = 0; x < size; x++) {
            printf("%2d ", get(s->data, x, y));
        }
        putchar('\n');
    }
    putchar('\n');
}

int main(int argAmount, char** args) {
    if(argAmount < 4) {
        if(argAmount <= 0) {
            printf("... <type> <size> <population>\n");
        } else {
            printf("%s <type> <size> <population>\n", args[0]);
        }
        return 0;
    }
    int type = atoi(args[1]);
    size = atoi(args[2]);
    population = atoi(args[3]);
    if(population < 4) {
        population = 4;
    }
    switch(type) {
        case 1: fitness = getNormalFitness; break;
        case 2: fitness = getUltraFitness; break;
        case 3: fitness = getMostPerfectFitness; break;
        default: puts("unknown type"); return 0;
    }

    long time = -getNanos();

    for(int h = 0; h < 5; h++) {
        init();
        int half = population / 2;
        int run = 1;
        while(run) {
            sort();
            for(int k = half; k < population; k++) {
                int parent = rand() % half;
                reproduce(squares[parent].data, squares + k);
                if(squares[k].fitness == 0) {
                    run = 0;
                    break;
                }
            }
        }
        sort();
        print(squares);
        clear();
    }

    time += getNanos();
    printf("Millis: %lf\n", time / 1000000.0);
    return 0;
}