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

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

void init()
{
    int end = size * size;
    data = malloc(sizeof(int) * end);
    for (int i = 0; i < end; i++)
    {
        data[i] = i + 1;
    }
}

int get(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 fitness = 0;
    int wanted = getWanted();
    for (int y = 0; y < size; y++)
    {
        int sum = 0;
        for (int x = 0; x < size; x++)
        {
            sum += get(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(x, y);
        }
        fitness += errorFactor(sum - wanted);
    }
    return fitness;
}

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

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

int getPandiagonalBase()
{
    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((x + y) % size, y);
            sumB += get((x - y + size) % size, y);
        }
        fitness += errorFactor(sumA - wanted);
        fitness += errorFactor(sumB - wanted);
    }
    return fitness;
}

int getAssociativeBase()
{
    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(x, y) + get(size - x - 1, size - y - 1);
            fitness += errorFactor(sum - wanted);
        }
    }
    return fitness;
}

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

int get2x2SquareBase()
{
    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(x, y) + get(x + 1, y) + get(x, y + 1) +
                          get(x + 1, y + 1);
                fitness += errorFactor(sum - wanted);
            }
        }
    }
    return fitness;
}
int getPerfectDiagonalBase()
{
    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(x, y) + get((x + size / 2) % size, (y + size / 2) % size);
            fitness += errorFactor(sum - wanted);
        }
    }
    return fitness;
}

int getMostPerfectFitness()
{
    // printf("ultra: %d\n", getUltraFitness());
    // printf("square: %d\n", get2x2SquareBase());
    // printf("diagonal: %d\n", getPerfectDiagonalBase());
    return getSemiFitness() + getPandiagonalBase() + get2x2SquareBase() +
           getPerfectDiagonalBase();
}

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

void swap(int a, int b)
{
    int tmp = data[a];
    data[a] = data[b];
    data[b] = tmp;
}

void randSwaps(int swaps)
{
    for (int i = 0; i < swaps; i++)
    {
        swap(rand() % (size * size), rand() % (size * size));
    }
}

int main(int argAmount, char **args)
{
    if (argAmount < 5)
    {
        if (argAmount <= 0)
        {
            printf("... <type> <size> <tries> <generations>\n");
        }
        else
        {
            printf("%s <type> <size> <tries> <generations>\n", args[0]);
        }
        return 0;
    }
    int type = atoi(args[1]);
    size = atoi(args[2]);
    int tries = atoi(args[3]);
    int generations = atoi(args[4]);

    switch (type)
    {
    case 1:
        fitness = getNormalFitness;
        break;
    case 2:
        fitness = getUltraFitness;
        break;
    case 3:
        fitness = getMostPerfectFitness;
        break;
    default:
        puts("unknown type");
        return 0;
    }

    init();
    for (int k = 0; k < tries; k++)
    {
        for (int i = 0; i < generations; i++)
        {
            int oldFitness = fitness();

            int a = rand() % (size * size);
            int b = rand() % (size * size);
            swap(a, b);

            int newFitness = fitness();
            if (newFitness == 0)
            {
                break;
            }

            if (newFitness >= oldFitness)
            {
                swap(a, b);
            }
        }
        // if (fitness() < 50)
        if (fitness() == 0)
        {
            print();
        }
        randSwaps(100);
    }
    free(data);
    return 0;
}