// La reducción de 3-SAT a 3-coloreo. La paleta es un triángulo T, F, B. Cada variable es un par de nodos
// unidos entre sí y con B. Cada cláusula es una cadena de dos compuertas O, cuya salida se une a F y a B.
#pragma once
#include <array>
#include <cstddef>
#include <cstdlib>
#include <vector>
#include "cnf.h"
#include "graph.h"

constexpr int COLOR_T = 0;
constexpr int COLOR_F = 1;
constexpr int COLOR_B = 2;

// Los tres primeros nodos del grafo son la paleta, y luego vienen x1, ¬x1, x2, ¬x2, ...
constexpr int NODE_T = 0;
constexpr int NODE_F = 1;
constexpr int NODE_B = 2;

inline int literal_node(int literal) {
    int v = std::abs(literal);
    return 3 + 2 * (v - 1) + (literal > 0 ? 0 : 1);
}

// Qué aristas poner, para ver qué se rompe si falta alguna.
struct GadgetOptions {
    bool variable_to_b = true;
    bool output_to_f = true;
    bool output_to_b = true;
};

// Compuerta O: agrega los nodos p, q, r con a-p, b-q, p-q, p-r, q-r. Devuelve r, la salida.
inline int add_or_gate(Graph& g, int a, int b) {
    int p = static_cast<int>(g.size());
    int q = p + 1;
    int r = p + 2;
    g.resize(g.size() + 3);
    add_edge(g, a, p);
    add_edge(g, b, q);
    add_edge(g, p, q);
    add_edge(g, p, r);
    add_edge(g, q, r);
    return r;
}

// Cláusula de tres literales: la salida de la segunda compuerta queda obligada a ser T.
inline void add_clause_gadget(Graph& g, int a, int b, int c, const GadgetOptions& options = {}) {
    int first = add_or_gate(g, a, b);
    int output = add_or_gate(g, first, c);
    if (options.output_to_f) add_edge(g, output, NODE_F);
    if (options.output_to_b) add_edge(g, output, NODE_B);
}

inline Graph reduce_to_coloring(const Formula& f, const GadgetOptions& options = {}) {
    Graph g(static_cast<std::size_t>(3 + 2 * f.variables));
    add_edge(g, NODE_T, NODE_F);
    add_edge(g, NODE_F, NODE_B);
    add_edge(g, NODE_T, NODE_B);
    for (int v = 1; v <= f.variables; v++) {
        add_edge(g, literal_node(v), literal_node(-v));
        if (options.variable_to_b) {
            add_edge(g, literal_node(v), NODE_B);
            add_edge(g, literal_node(-v), NODE_B);
        }
    }
    for (const Clause& c : f.clauses) {
        add_clause_gadget(g, literal_node(c[0]), literal_node(c[1]), literal_node(c[2]), options);
    }
    return g;
}

// De un coloreo a una asignación. Los colores no traen nombre: T es el color que recibió el nodo T.
inline Assignment assignment_from_coloring(const Formula& f, const std::vector<int>& color) {
    Assignment a(static_cast<std::size_t>(f.variables) + 1, false);
    for (int v = 1; v <= f.variables; v++) {
        a[static_cast<std::size_t>(v)] = color[static_cast<std::size_t>(literal_node(v))] == color[NODE_T];
    }
    return a;
}

// Colorea una compuerta dados los colores de sus entradas: T en la salida si alguna entrada es T, y si
// no, F. Es la construcción de la demostración.
inline int color_or_gate(std::vector<int>& color, std::size_t p, int a_color, int b_color) {
    int q_color = COLOR_B;
    int p_color = COLOR_B;
    int r_color = COLOR_F;
    if (a_color == COLOR_T) {
        r_color = COLOR_T;
        p_color = b_color == COLOR_T ? COLOR_B : COLOR_F;
        q_color = b_color == COLOR_T ? COLOR_F : COLOR_B;
    } else if (b_color == COLOR_T) {
        r_color = COLOR_T;
        p_color = COLOR_B;
        q_color = COLOR_F;
    } else {
        p_color = COLOR_T;
        q_color = COLOR_B;
    }
    color[p] = p_color;
    color[p + 1] = q_color;
    color[p + 2] = r_color;
    return r_color;
}

// De una asignación que cumple la fórmula a un coloreo, sin buscar.
inline std::vector<int> coloring_from_assignment(const Formula& f, const Assignment& a, const Graph& g) {
    std::vector<int> color(g.size(), -1);
    color[NODE_T] = COLOR_T;
    color[NODE_F] = COLOR_F;
    color[NODE_B] = COLOR_B;
    for (int v = 1; v <= f.variables; v++) {
        bool value = a[static_cast<std::size_t>(v)];
        color[static_cast<std::size_t>(literal_node(v))] = value ? COLOR_T : COLOR_F;
        color[static_cast<std::size_t>(literal_node(-v))] = value ? COLOR_F : COLOR_T;
    }
    std::size_t next = static_cast<std::size_t>(3 + 2 * f.variables);
    for (const Clause& c : f.clauses) {
        int first = color_or_gate(color, next, literal_true(a, c[0]) ? COLOR_T : COLOR_F, literal_true(a, c[1]) ? COLOR_T : COLOR_F);
        color_or_gate(color, next + 3, first, literal_true(a, c[2]) ? COLOR_T : COLOR_F);
        next += 6;
    }
    return color;
}

// Una compuerta sola con la paleta y sus entradas ya coloreadas: qué colores admite la salida.
// Devuelve, para cada color posible de r, si existe un coloreo de p y q que lo permita.
inline std::array<bool, 3> or_gate_outputs(int a_color, int b_color) {
    std::array<bool, 3> possible = {false, false, false};
    for (int p = 0; p < 3; p++) {
        for (int q = 0; q < 3; q++) {
            for (int r = 0; r < 3; r++) {
                if (p != a_color && q != b_color && p != q && p != r && q != r) possible[static_cast<std::size_t>(r)] = true;
            }
        }
    }
    return possible;
}

// Enumera los coloreos de los 6 nodos internos del gadget de cláusula, con T, F, B y las tres entradas
// fijas. Cuenta cuántos son propios. Son 3⁶ posibilidades para cada una de las 8 entradas.
inline int clause_gadget_colorings(bool a, bool b, bool c) {
    Graph g(6);
    add_edge(g, NODE_T, NODE_F);
    add_edge(g, NODE_F, NODE_B);
    add_edge(g, NODE_T, NODE_B);
    add_edge(g, 3, NODE_B);
    add_edge(g, 4, NODE_B);
    add_edge(g, 5, NODE_B);
    add_clause_gadget(g, 3, 4, 5);
    std::vector<int> color(g.size(), 0);
    color[NODE_T] = COLOR_T;
    color[NODE_F] = COLOR_F;
    color[NODE_B] = COLOR_B;
    color[3] = a ? COLOR_T : COLOR_F;
    color[4] = b ? COLOR_T : COLOR_F;
    color[5] = c ? COLOR_T : COLOR_F;
    int total = 0;
    for (int mask = 0; mask < 729; mask++) {
        int rest = mask;
        for (std::size_t v = 6; v < 12; v++) {
            color[v] = rest % 3;
            rest /= 3;
        }
        if (is_proper_coloring(g, color, 3)) total++;
    }
    return total;
}