// ian henderson <ian@ianhenderson.org>
// published on 2023 february 2
// updated on 2023 march 6 (leave the first rather than the last rectangle unrotated)
// this software belongs to a future without copyright.  please use it however you'd like.

// to build, save this file as main.c, install calcium (https://fredrikj.net/calcium/), and run
// cc -O3 -g -std=c99 -ftrapv -o main main.c -lcalcium -lflint

// see http://ianhenderson.org/similar-rectangles/ for more info about what is happening here.

#include <calcium/ca.h>
#include <calcium/ca_mat.h>
#include <calcium/ca_poly.h>
#include <calcium/ca_vec.h>
#include <stdbool.h>
#include <stdio.h>

#define N 4

static int64_t gcd(int64_t a, int64_t b)
{
    while (b != 0) {
        int64_t t = b;
        b = a % b;
        a = t;
    }
    return a;
}

#define CHECK_DEGREE
#define MAX_DEGREE 50
struct polynomial {
    int64_t coefficient[MAX_DEGREE + 1];
    int degree;
};

static const struct polynomial polynomial_zero;

static bool polynomials_equal(struct polynomial a, struct polynomial b)
{
    if (a.degree != b.degree)
        return false;
    for (int i = 0; i <= a.degree; ++i) {
        if (a.coefficient[i] != b.coefficient[i])
            return false;
    }
    return true;
}
static struct polynomial polynomial_subtract(struct polynomial a, struct polynomial b)
{
    if (b.degree > a.degree)
        a.degree = b.degree;
    for (int i = a.degree; i >= 0; --i) {
        a.coefficient[i] -= b.coefficient[i];
        if (a.coefficient[i] == 0 && i == a.degree && i > 0)
            a.degree--;
    }
    return a;
}
static struct polynomial polynomial_multiply(struct polynomial a, struct polynomial b)
{
    struct polynomial result = { 0 };
    result.degree = a.degree + b.degree;
    if (result.degree > MAX_DEGREE)
        result.degree = MAX_DEGREE;
    for (int i = a.degree + b.degree; i >= 0; --i) {
        int64_t product = 0;
        for (int j = 0; j <= i; ++j) {
            if (j <= a.degree && i - j <= b.degree)
                product += a.coefficient[j] * b.coefficient[i - j];
        }
        if (product == 0) {
            if (i == result.degree && i > 0)
                result.degree--;
            continue;
        }
#ifdef CHECK_DEGREE
        if (i > result.degree) {
            fprintf(stderr, "polynomial product degree %d exceeded max degree\n", i);
            exit(-1);
        }
#endif
        result.coefficient[i] = product;
    }
    return result;
}
struct equation {
    struct polynomial p[N + 1];
};
static void reduce(struct equation *eq)
{
    int power_reduction = MAX_DEGREE + 1;
    int64_t divisor = 0;
    for (int i = 0; i < N + 1; ++i) {
        for (int j = 0; j <= eq->p[i].degree; ++j) {
            if (eq->p[i].coefficient[j] == 0)
                continue;
            if (j < power_reduction)
                power_reduction = j;
            divisor = gcd(divisor, llabs(eq->p[i].coefficient[j]));
        }
    }
    for (int i = 0; i < N + 1; ++i) {
        for (int j = 0; j <= eq->p[i].degree; ++j) {
            if (j + power_reduction <= eq->p[i].degree)
                eq->p[i].coefficient[j] = eq->p[i].coefficient[j + power_reduction] / divisor;
            else
                eq->p[i].coefficient[j] = 0;
        }
        if (power_reduction < eq->p[i].degree)
            eq->p[i].degree -= power_reduction;
        else
            eq->p[i].degree = 0;
    }
}
static void substitute(struct equation *into, struct equation *from, int variable)
{
    struct polynomial a = from->p[variable];
    struct polynomial b = into->p[variable];
    for (int i = 0; i < N + 1; ++i)
        into->p[i] = polynomial_subtract(polynomial_multiply(into->p[i], a), polynomial_multiply(from->p[i], b));
    reduce(into);
}

struct system_of_equations {
    struct equation equation[2 * N];
};
static struct polynomial solve(struct system_of_equations *system)
{
    for (int i = 0; i < 2 * N; ++i) {
        for (int j = 0; j < N + 1; ++j) {
            struct polynomial p = system->equation[i].p[j];
            if (polynomials_equal(p, polynomial_zero))
                continue;
            if (j == N) {
                // the polynomial is already in canonical form except for its
                // sign; ensure the sign is positive before returning it.
                if (p.coefficient[p.degree] < 0) {
                    for (int k = 0; k <= p.degree; ++k)
                        p.coefficient[k] = -p.coefficient[k];
                }
                return p;
            }
            for (int k = i + 1; k < 2 * N; ++k)
                substitute(&system->equation[k], &system->equation[i], j);
            break;
        }
    }
    // see "what if gaussian elimination doesn't produce a single polynomial in
    // the last column?" in http://ianhenderson.org/similar-rectangles/ for more
    // about this situation.
    return polynomial_zero;
}

struct dissection {
    unsigned char colors[N][N];
};

static void system_for_oriented_dissection(struct system_of_equations *system, unsigned orientation, struct dissection dissection)
{
    for (int row = 0; row < N; ++row) {
        for (int column = 0; column < N; ++column) {
            int color = dissection.colors[row][column];
            if ((orientation >> (N - color - 1)) & 1) {
                system->equation[row].p[color].coefficient[1] = 1;
                system->equation[row].p[color].degree = 1;
                system->equation[N + column].p[color].coefficient[0] = 1;
            } else {
                system->equation[row].p[color].coefficient[0] = 1;
                system->equation[N + column].p[color].coefficient[1] = 1;
                system->equation[N + column].p[color].degree = 1;
            }
        }
        system->equation[row].p[N].coefficient[0] = -1;
    }
    for (int column = 0; column < N; ++column)
        system->equation[N + column].p[N].coefficient[0] = -1;
}

struct dissection_process {
    struct dissection result;
    unsigned char row;
    unsigned char column;
    unsigned char largest_used_color;
};

static void generate_dissection(struct dissection d);

static void generate_dissections(void)
{
    struct dissection_process *stack = calloc(100000, sizeof(struct dissection_process));
    int depth = 1;
    while (depth > 0) {
        if (depth > 90000) {
            fprintf(stderr, "stack overflow in generate_dissections()\n");
            exit(-1);
        }
        struct dissection_process d = stack[--depth];
        d.column++;
        if (d.column == N) {
            d.column = 0;
            d.row++;
            if (d.row == N) {
                if (d.largest_used_color == N - 1)
                    generate_dissection(d.result);
                continue;
            }
            if (d.row > 1) {
                bool matches_previous_row = true;
                for (int column = 0; column < N && matches_previous_row; ++column) {
                    if (d.result.colors[d.row - 1][column] != d.result.colors[d.row - 2][column])
                        matches_previous_row = false;
                }
                if (matches_previous_row) {
                    if (d.largest_used_color == N - 1) {
                        for (int row = d.row; row < N; ++row) {
                            for (int column = 0; column < N; ++column)
                                d.result.colors[row][column] = d.result.colors[row - 1][column];
                        }
                        generate_dissection(d.result);
                    }
                    continue;
                }
            }
        }
        bool allow_row = true;
        bool allow_column = true;
        if (d.row > 0 && d.column > 0) {
            unsigned char a = d.result.colors[d.row - 1][d.column - 1];
            unsigned char b = d.result.colors[d.row - 1][d.column];
            unsigned char c = d.result.colors[d.row][d.column - 1];
            if (b == c) {
                d.result.colors[d.row][d.column] = c;
                stack[depth++] = d;
                continue;
            }
            if (a == b)
                allow_row = false;
            if (a == c)
                allow_column = false;
        }
        if (allow_row && d.row > 0) {
            d.result.colors[d.row][d.column] = d.result.colors[d.row - 1][d.column];
            stack[depth++] = d;
        }
        if (allow_column && d.column > 0) {
            d.result.colors[d.row][d.column] = d.result.colors[d.row][d.column - 1];
            stack[depth++] = d;
        }
        if (d.largest_used_color < N - 1) {
            d.result.colors[d.row][d.column] = ++d.largest_used_color;
            stack[depth++] = d;
        }
    }
}

static void ca_poly_from_polynomial(ca_poly_t poly, const struct polynomial *p, ca_ctx_t ctx)
{
    ca_t coeff;
    ca_init(coeff, ctx);
    for (int i = 0; i <= p->degree; ++i) {
        ca_set_si(coeff, p->coefficient[i], ctx);
        ca_poly_set_coeff_ca(poly, i, coeff, ctx);
    }
    ca_clear(coeff, ctx);
}

static void polynomial_evaluate_to_ca(ca_ptr result, const struct polynomial *p, ca_t x, ca_ctx_t ctx)
{
    ca_t xn;
    ca_t prod;
    ca_init(xn, ctx);
    ca_init(prod, ctx);
    ca_one(xn, ctx);
    for (int i = 0; i <= p->degree; ++i) {
        ca_set_si(prod, p->coefficient[i], ctx);
        ca_mul(prod, prod, xn, ctx);
        ca_add(result, result, prod, ctx);
        ca_mul(xn, xn, x, ctx);
    }
    ca_clear(prod, ctx);
    ca_clear(xn, ctx);
}

#define TABLE_SIZE (1<<22)
#define TABLE_MASK (TABLE_SIZE-1)

struct table {
    struct polynomial key[TABLE_SIZE];
    bool could_have_further_roots[TABLE_SIZE];
};

static struct table *table;

// this is the 32-bit FNV-1a hash diffusion algorithm.
// http://www.isthe.com/chongo/tech/comp/fnv/index.html
static uint32_t fnv(const void *dataPointer, size_t length)
{
    uint32_t hash = 0x811c9dc5;
    const unsigned char *data = dataPointer;
    for (size_t i = 0; i < length; ++i) {
        hash ^= data[i];
        hash *= 0x01000193;
    }
    return hash;
}

static bool polynomial_could_have_further_roots(struct polynomial p, int *index_for_polynomial)
{
    if (polynomials_equal(p, polynomial_zero))
        return false;
    uint32_t hash = fnv(&p, sizeof(p));
    uint32_t index = hash & TABLE_MASK;
    while (true) {
        if (polynomials_equal(table->key[index], p)) {
            *index_for_polynomial = index;
            return table->could_have_further_roots[index];
        }
        if (polynomials_equal(table->key[index], polynomial_zero)) {
            *index_for_polynomial = index;
            return true;
        }
        index = (index + 1) & TABLE_MASK;
        if (index == (hash & TABLE_MASK)) {
            fprintf(stderr, "hash table overflow\n");
            exit(-1);
        }
    }
}

#define MAX_TOTAL_ROOTS 1000000
struct root {
    ca_struct root;
    int table_index;
    unsigned orientation;
    struct dissection dissection;
    char *scale[N];
};

static ca_ctx_t ctx;
static ca_t zero;
static ca_t one;
static ca_vec_t roots;
static struct root all_roots[MAX_TOTAL_ROOTS];
static int number_of_roots;
static ulong multiplicities[MAX_DEGREE + 1];
static ca_vec_t rectangle_scales;

static int compare_root_indexes(const void *aa, const void *bb)
{
    ca_ptr a = &all_roots[*(const int *)aa].root;
    ca_ptr b = &all_roots[*(const int *)bb].root;
    if (ca_check_equal(a, b, ctx) == T_TRUE)
        return 0;
    else if (ca_check_lt(a, b, ctx) == T_TRUE)
        return -1;
    else if (ca_check_gt(a, b, ctx) == T_TRUE)
        return 1;
    fprintf(stderr, "calcium couldn't decide comparison\n");
    exit(-1);
    return 0;
}

static void generate_dissection(struct dissection d)
{
    // dissections with repeating columns can be canonicalized by putting the
    // repeating columns at the end.  skip any dissections which aren't
    // canonical in this sense.
    bool should_repeat = false;
    for (int column = 1; column < N; ++column) {
        bool matches_previous_column = true;
        for (int row = 0; row < N && matches_previous_column; ++row) {
            if (d.colors[row][column] != d.colors[row][column - 1])
                matches_previous_column = false;
        }
        if (!matches_previous_column && should_repeat)
            return;
        if (matches_previous_column)
            should_repeat = true;
    }

    // iterate over all orientations for the rectangles.  orientations are
    // represented as a bitset of the rectangles that are rotated by 90 degrees.
    // the loop goes up to 2^(N-1) instead of 2^N because, due to symmetry, the
    // first rectangle can always be left unrotated.
    for (unsigned i = 0; i < (1 << (N - 1)); ++i) {
        struct system_of_equations system = { 0 };
        system_for_oriented_dissection(&system, i, d);
        struct polynomial p = solve(&system);
        int index_for_polynomial = 0;
        if (!polynomial_could_have_further_roots(p, &index_for_polynomial))
            continue;
        ca_poly_t poly;
        ca_poly_init(poly, ctx);
        ca_poly_from_polynomial(poly, &p, ctx);
        if (ca_poly_roots(roots, multiplicities, poly, ctx) != 1) {
            fprintf(stderr, "couldn't find roots for polynomial\n");
            exit(-1);
        }
        bool could_have_further_roots = false;
        for (int j = 0; j < ca_vec_length(roots, ctx); ++j) {
            ca_ptr root = ca_vec_entry(roots, j);
            if (ca_check_is_zero(root, ctx) == T_TRUE)
                continue;
            if (ca_check_is_real(root, ctx) == T_FALSE)
                continue;
            if (ca_check_le(root, one, ctx) == T_FALSE)
                continue;
            ca_vec_zero(rectangle_scales, N, ctx);
            ca_t divisor;
            ca_init(divisor, ctx);
            ca_t term;
            ca_init(term, ctx);
            for (int k = 2 * N - 1; k >= 0; --k) {
                for (int l = 0; l < N; ++l) {
                    struct polynomial p = system.equation[k].p[l];
                    if (polynomials_equal(p, polynomial_zero))
                        continue;
                    ca_zero(divisor, ctx);
                    polynomial_evaluate_to_ca(divisor, &system.equation[k].p[l], root, ctx);
                    if (ca_check_is_zero(divisor, ctx) == T_TRUE)
                        break;
                    ca_ptr scale = ca_vec_entry(rectangle_scales, l);
                    if (ca_check_is_zero(scale, ctx) != T_TRUE) {
                        fprintf(stderr, "invalid system of equations\n");
                        exit(-1);
                    }
                    for (int m = l + 1; m <= N; ++m) {
                        ca_zero(term, ctx);
                        polynomial_evaluate_to_ca(term, &system.equation[k].p[m], root, ctx);
                        if (m < N)
                            ca_mul(term, term, ca_vec_entry(rectangle_scales, m), ctx);
                        ca_add(scale, scale, term, ctx);
                    }
                    ca_div(scale, scale, divisor, ctx);
                    ca_neg(scale, scale, ctx);
                    break;
                }
            }
            ca_clear(divisor, ctx);
            ca_clear(term, ctx);
            bool has_valid_scales = true;
            for (int k = 0; k < N && has_valid_scales; ++k) {
                truth_t is_positive = ca_check_gt(ca_vec_entry(rectangle_scales, k), zero, ctx);
                if (is_positive == T_FALSE) {
                    has_valid_scales = false;
                    break;
                } else if (is_positive != T_TRUE) {
                    fprintf(stderr, "couldn't figure out whether a rectangle had a positive scale or not\n");
                    exit(-1);
                }
            }
            if (!has_valid_scales) {
                could_have_further_roots = true;
                continue;
            }
            for (int k = 0; k < N; ++k)
                all_roots[number_of_roots].scale[k] = ca_get_decimal_str(ca_vec_entry(rectangle_scales, k), 8, 0, ctx);
            ca_init(&all_roots[number_of_roots].root, ctx);
            ca_set(&all_roots[number_of_roots].root, ca_vec_entry(roots, j), ctx);
            all_roots[number_of_roots].table_index = index_for_polynomial;
            all_roots[number_of_roots].orientation = i;
            all_roots[number_of_roots].dissection = d;
            if ((number_of_roots % 100) == 0)
                fprintf(stderr, "%d\n", number_of_roots);
            number_of_roots++;
        }
        table->key[index_for_polynomial] = p;
        table->could_have_further_roots[index_for_polynomial] = could_have_further_roots;
        ca_poly_clear(poly, ctx);
    }
}

int main(void)
{
    table = calloc(1, sizeof(struct table));
    ca_ctx_init(ctx);
    ca_init(one, ctx);
    ca_one(one, ctx);
    ca_init(zero, ctx);
    ca_vec_init(rectangle_scales, N, ctx);

    generate_dissections();
    int *indexes = calloc(number_of_roots, sizeof(int));
    for (int i = 0; i < number_of_roots; ++i)
        indexes[i] = i;
    qsort(indexes, number_of_roots, sizeof(int), compare_root_indexes);
    int number_of_distinct_roots = 0;
    printf("[\n");
    for (int i = 0; i < number_of_roots; ++i) {
        if (i > 0 && ca_check_equal(&all_roots[indexes[i]].root, &all_roots[indexes[i - 1]].root, ctx) == T_TRUE)
            continue;
        number_of_distinct_roots++;
        if (i != 0)
            printf(",");
        printf("{\"rectangles\":[");
        unsigned char next_color = 0;
        for (int row = 0; row < N; ++row) {
            for (int column = 0; column < N; ++column) {
                if (all_roots[indexes[i]].dissection.colors[row][column] == next_color) {
                    if (next_color != 0)
                        printf(",");
                    printf("{");
                    if (row > 0)
                        printf("\"below\":%d,", all_roots[indexes[i]].dissection.colors[row - 1][column]);
                    if (column > 0)
                        printf("\"beside\":%d,", all_roots[indexes[i]].dissection.colors[row][column - 1]);
                    printf("\"orientation\":%s,", (all_roots[indexes[i]].orientation & (1 << (N - next_color - 1))) ? "true" : "false");
                    printf("\"scale\":%s}", all_roots[indexes[i]].scale[next_color]);
                    next_color++;
                }
            }
        }
        printf("],\"root\":\"");
        ca_print(&all_roots[indexes[i]].root, ctx);
        printf("\"}\n");
    }
    printf("]\n");
    fprintf(stderr, "number_of_distinct_roots = %d\n", number_of_distinct_roots);
    ca_vec_clear(roots, ctx);
    ca_clear(one, ctx);
    ca_ctx_clear(ctx);
    flint_cleanup();
}