#include "clause.h"

#include <fstream>

ll Clause::total_cost;

void Clause::update_satisfied_lit_count() {

    using namespace ClauseLS;

    ll old_lit_count = satisfied_lit_count;

    satisfied_lit_count = 0;

    for(auto& lit : lits) {
        if(lit_value[abs(lit)]) {
            satisfied_lit_count += (lit > 0);
        } else {
            satisfied_lit_count += (lit < 0);
        }
    }

    if(old_lit_count == 0 && satisfied_lit_count > 0) {
        falsified_clauses.erase(this);
        total_cost -= weight;
        satisfied_clauses.insert(this);
    }

    if(old_lit_count > 0 && satisfied_lit_count == 0) {
        satisfied_clauses.erase(this);
        total_cost += weight;
        falsified_clauses.insert(this);
    }
}

namespace ClauseLS {

int num_vars;
int num_clauses;

std::unordered_set<Clause*> satisfied_clauses;
std::unordered_set<Clause*> falsified_clauses;

std::vector<Clause*> clauses;
std::vector<Clause*> *lit_related_clauses;
int *lit_value;
int *CC;

void write_cnf() {
    std::ofstream f("test.cnf");

    f << "p cnf " << num_vars << " " << num_clauses << std::endl;

    for(auto& c : clauses) {
        for(auto &lit : c->lits) {
            f << lit << " ";
        }
        f << "0" << std::endl;
    }

    f.close();
}

void flip(int var) {

    lit_value[var] = !lit_value[var];

    for(auto& clause : lit_related_clauses[var]) {
        clause->update_satisfied_lit_count();
        for(auto& lit : clause->lits) {
            CC[abs(lit)] = 1;
        }
    }

    CC[var] = 0;
}

void add_to_tmp_clause(int x) {
    static std::vector<int> lits;
    if(x != 0) {
        lits.push_back(x);
    } else {
        Clause* clause = new Clause();
        clause->lits = lits;
        clause->satisfied_lit_count = 0;
        clauses.push_back(clause);
        lits.clear();
    }
}

void init_data_structs(Circuit* circuit) {
    build_clauses(circuit);

    printf("====== Clause Statistics ====== \n");
    printf("num_vars:\t%d\n", num_vars);
    printf("num_clauses:\t%ld\n", num_clauses);
    printf("================================ \n");

    lit_related_clauses = new std::vector<Clause*>[num_vars + 1];
    lit_value = new int[num_vars + 1];
    CC = new int[num_vars + 1];

    for(int i=1; i<=num_vars; i++) {
        lit_value[i] = 0;
        CC[i] = 1;
    }

    for(auto& clause : clauses) {
        clause->weight = 1;
        Clause::total_cost += clause->weight;
        falsified_clauses.insert(clause);

        //printf("fs: %d. ss: %d\n", falsified_clauses.size(), satisfied_clauses.size());
        
        clause->update_satisfied_lit_count();

        for(auto& lit : clause->lits) {
            lit_related_clauses[abs(lit)].push_back(clause);
        }
    }

    write_cnf();
}

void reset_data_structs() {
    for(int i=1; i<=num_vars; i++) {
        lit_value[i] = rand() % 2;
        CC[i] = 1;
    }
}

void build_clauses(Circuit* circuit) {

    int extra_variable = circuit->gates.size();

    int last_id, last_v;

    for(Gate* g : circuit->gates) {
        if(g->pi) continue;

        switch(g->type) {
            case Gate::NOT:
                add_to_tmp_clause(-g->id);
                add_to_tmp_clause(-(g->fan_ins[0]->id));
                add_to_tmp_clause(0);

                add_to_tmp_clause(g->id);
                add_to_tmp_clause(g->fan_ins[0]->id);
                add_to_tmp_clause(0);
                break;
            case Gate::BUF:
                add_to_tmp_clause(-g->id);
                add_to_tmp_clause((g->fan_ins[0]->id));
                add_to_tmp_clause(0);

                add_to_tmp_clause(g->id);
                add_to_tmp_clause(-(g->fan_ins[0]->id));
                add_to_tmp_clause(0);
                break;
            case Gate::AND:
                for(int i=0; i<g->fan_ins.size(); i++) {
                    add_to_tmp_clause(-g->id);
                    add_to_tmp_clause(g->fan_ins[i]->id);
                    add_to_tmp_clause(0);
                }

                for(int i=0; i<g->fan_ins.size(); i++) {
                    add_to_tmp_clause(-g->fan_ins[i]->id);
                }
                add_to_tmp_clause(g->id);
                add_to_tmp_clause(0);
                break;
            case Gate::NAND:
                for(int i=0; i<g->fan_ins.size(); i++) {
                    add_to_tmp_clause(g->id);
                    add_to_tmp_clause(g->fan_ins[i]->id);
                    add_to_tmp_clause(0);
                }

                for(int i=0; i<g->fan_ins.size(); i++) {
                    add_to_tmp_clause(-g->fan_ins[i]->id);
                }
                add_to_tmp_clause(-g->id);
                add_to_tmp_clause(0);
                break;
            case Gate::OR:
                for(int i=0; i<g->fan_ins.size(); i++) {
                    add_to_tmp_clause(g->fan_ins[i]->id);
                }
                add_to_tmp_clause(-g->id);
                add_to_tmp_clause(0);

                for(int i=0; i<g->fan_ins.size(); i++) {
                    add_to_tmp_clause(g->id);
                    add_to_tmp_clause(-g->fan_ins[i]->id);
                    add_to_tmp_clause(0);
                }
                break;
            case Gate::NOR:
                for(int i=0; i<g->fan_ins.size(); i++) {
                    add_to_tmp_clause(g->fan_ins[i]->id);
                }
                add_to_tmp_clause(g->id);
                add_to_tmp_clause(0);

                for(int i=0; i<g->fan_ins.size(); i++) {
                    add_to_tmp_clause(-g->id);
                    add_to_tmp_clause(-g->fan_ins[i]->id);
                    add_to_tmp_clause(0);
                }
                break;
            case Gate::XOR:
                last_v = g->fan_ins[0]->id;
                for(int i=1; i<g->fan_ins.size() - 1; i++) {
                    int new_v = ++extra_variable;
                    add_to_tmp_clause(-g->fan_ins[i]->id); add_to_tmp_clause(-last_v); add_to_tmp_clause(-new_v);
                    add_to_tmp_clause(0);
                    add_to_tmp_clause(-g->fan_ins[i]->id); add_to_tmp_clause(last_v); add_to_tmp_clause(new_v);
                    add_to_tmp_clause(0);
                    add_to_tmp_clause(g->fan_ins[i]->id); add_to_tmp_clause(-last_v); add_to_tmp_clause(new_v);
                    add_to_tmp_clause(0);
                    add_to_tmp_clause(g->fan_ins[i]->id); add_to_tmp_clause(last_v); add_to_tmp_clause(-new_v);
                    add_to_tmp_clause(0);
                    last_v = new_v;
                }

                last_id = g->fan_ins[g->fan_ins.size() - 1]->id;

                add_to_tmp_clause(-g->id); add_to_tmp_clause(-last_v); add_to_tmp_clause(-last_id);
                add_to_tmp_clause(0);
                add_to_tmp_clause(-g->id); add_to_tmp_clause(last_v); add_to_tmp_clause(last_id);
                add_to_tmp_clause(0);
                add_to_tmp_clause(g->id); add_to_tmp_clause(-last_v); add_to_tmp_clause(last_id);
                add_to_tmp_clause(0);
                add_to_tmp_clause(g->id); add_to_tmp_clause(last_v); add_to_tmp_clause(-last_id);
                add_to_tmp_clause(0);

                break;
            case Gate::XNOR:
                last_v = g->fan_ins[0]->id;
                for(int i=1; i<g->fan_ins.size() - 1; i++) {
                    int new_v = ++extra_variable;
                    add_to_tmp_clause(-g->fan_ins[i]->id); add_to_tmp_clause(-last_v); add_to_tmp_clause(-new_v);
                    add_to_tmp_clause(0);
                    add_to_tmp_clause(-g->fan_ins[i]->id); add_to_tmp_clause(last_v); add_to_tmp_clause(new_v);
                    add_to_tmp_clause(0);
                    add_to_tmp_clause(g->fan_ins[i]->id); add_to_tmp_clause(-last_v); add_to_tmp_clause(new_v);
                    add_to_tmp_clause(0);
                    add_to_tmp_clause(g->fan_ins[i]->id); add_to_tmp_clause(last_v); add_to_tmp_clause(-new_v);
                    add_to_tmp_clause(0);
                    last_v = new_v;
                }

                last_id = g->fan_ins[g->fan_ins.size() - 1]->id;

                add_to_tmp_clause(g->id); add_to_tmp_clause(-last_v); add_to_tmp_clause(-last_id);
                add_to_tmp_clause(0);
                add_to_tmp_clause(g->id); add_to_tmp_clause(last_v); add_to_tmp_clause(last_id);
                add_to_tmp_clause(0);
                add_to_tmp_clause(-g->id); add_to_tmp_clause(-last_v); add_to_tmp_clause(last_id);
                add_to_tmp_clause(0);
                add_to_tmp_clause(-g->id); add_to_tmp_clause(last_v); add_to_tmp_clause(-last_id);
                add_to_tmp_clause(0);
                break;
            default:
                exit(-1);
                break;
        }
    }

    num_vars = extra_variable;
    num_clauses = clauses.size();
}

}