#include "sweep.hpp"
#include "src/cadical.hpp"
#include "circuit.hpp"
  
extern "C" {
    #include "aiger.h"
}

int Sweep::get_time() {
    return std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::high_resolution_clock::now() - clk_st).count();
}

void Sweep::cal_topo_index() {

    int cnt = 0;

    std::queue<int> q;
    for(int i=0; i<aig->num_inputs; i++) {
        int input = aiger_lit2var(aig->inputs[i].lit);
        topo_index[input] = ++cnt;
        q.push(input);
    }

    

    static int* topo_counter = new int[maxvar + 1];
    memset(topo_counter, 0, (maxvar + 1) * sizeof(int));

    while(!q.empty()) {
        int cur = q.front();
        q.pop();
        for(auto& edge : graph.edge[cur]) {
            int v = edge.to;
            topo_counter[v]++;
            if(topo_counter[v] == 2) {
                q.push(v);
                topo_index[v] = ++cnt;
            }
        }
    }

    for(int i=1; i<=maxvar; i++) {
        rtopo_index[i] = maxvar - topo_index[i] + 1;
        //printf("topo[%d]=%d\n", i, rtopo_index[i]);
    }
}


void Sweep::miter() {

    std::queue<int> q;
    
    int x = aiger_lit2var(output);

    if (!used[x])
        q.push(x), used[x] = true;

    while(!q.empty()) {
        int u = q.front();
        q.pop();

        if(graph.redge[u].size() == 0) {
            //cout << "oh FUCK" << endl;
            continue;
        }

        int a = graph.redge[u][0].to;
        int b = graph.redge[u][1].to;

        if(graph.redge[u][0].neg) a = -a;
        if(graph.redge[u][1].neg) b = -b;
        

        solver->add(-u); solver->add(a); solver->add(0);
        solver->add(-u); solver->add(b); solver->add(0);
        solver->add(-a); solver->add(-b); solver->add(u); solver->add(0);

        a = abs(a);
        b = abs(b);

        if(!used[a]) {
            q.push(a);
            used[a] = true;
        }

        if(!used[b]) {
            q.push(b);
            used[b] = true;
        }

    }

    if(output & 1) {
        solver->add(-x); solver->add(0);
    } else {
        solver->add(x); solver->add(0);
    }
    

    int result = solver->solve ();

    delete solver;

    printf("c final: %d\n", result);

    return;
}

void Sweep::aiger_preprocess() {
    aiger_read_from_file(aig, file);
    maxvar = aig->maxvar;
    output = aig->outputs[0].lit;

    used = new int[maxvar + 1];
    C = new circuit[maxvar + 1];
    memset(used, 0, (maxvar + 1) * sizeof(int));
    sum_pos = new int[2 * maxvar + 2];
    sum_neg = new int[2 * maxvar + 2];
    memset(sum_pos, 0, sizeof(int) * (2 * maxvar + 2));
    memset(sum_neg, 0, sizeof(int) * (2 * maxvar + 2));

    for(int i=0; i<aig->num_ands; i++) {
        auto gate = aig->ands[i];
        int o  = aiger_lit2var(gate.lhs);
        int i0 = aiger_lit2var(gate.rhs0);
        int i1 = aiger_lit2var(gate.rhs1);
        graph.add(i0, o, gate.rhs0 & 1);
        graph.add(i1, o, gate.rhs1 & 1);

        assert((gate.lhs & 1) == 0);
        C[o].push(i0 * (gate.rhs0 & 1 ? -1 : 1));
        C[o].push(i1 * (gate.rhs1 & 1 ? -1 : 1));
    }
    
    // for (int i = 1; i <= maxvar; i++) {
    //     circuit *c = &C[i];
    //     if (!c->sz) continue; 
    //     printf("%d = %s ( ", i, gate_type[c->type].c_str());
    //     for (int j = 0; j < c->sz; j++) {
    //         printf("%d ", c->to[j]);
    //     }
    //     puts(")");
    // }
    cal_topo_index();
    return;

    printf("c inputs: \nc ");
    for(int i=0; i<aig->num_inputs; i++) {
        printf("%d ", aig->inputs[i].lit);
    }

    printf("\nc outputs: \nc ");

    for(int i=0; i<aig->num_outputs; i++) {
        printf("%d ", aig->outputs[i].lit);
    }
    printf("\nc maxvar: \nc %d\n", maxvar);
    // int x = det_v > 0 ? det_v : -det_v;
    // if (x > maxvar) det_v = 0;
    printf("assume:");
    for (int i = 0; i < det_v_sz; i++) {
        printf(" %d", det_v[i]);
    }
    puts("");
}

void Sweep::check_with_SAT() {
    printf("c =================================== check =================================\n");

    for(auto& clas : classes) {
        //if(clas.size() != 2) continue;
        int lit = clas[0];
        if (sum_pos[lit] == 0 || sum_neg[lit] == 0) {
            // printf("c Find ALWAYS NEG SET\n");
            for(int i=0; i<clas.size(); i++) {
                if (clas[i] & 1) continue;
                int vara =  aiger_lit2var(clas[i]);
                pairs.push_back(std::make_pair(topo_index[vara], std::make_pair(clas[i], !sum_pos[lit] ? 0 : -1)));      
            }
            continue;
        }

        for(int i=0; i<clas.size(); i++) {
            for(int j=i+1; j<clas.size(); j++) {
                if ((clas[i] & 1) & (clas[j] & 1)) continue;    
                int vara =  aiger_lit2var(clas[i]);
                int varb =  aiger_lit2var(clas[j]);
                if (abs(vara - varb) <= 5) {printf("%d %d\n", clas[i], clas[j]); continue;}
                pairs.push_back(std::make_pair(std::max(topo_index[vara], topo_index[varb]), std::make_pair(clas[i], clas[j])));      
            }
        }
    }
    sort(pairs.begin(), pairs.end());

    if (det_v_sz) {
        for (int i = 0; i < det_v_sz; i++) {
            solver->add(det_v[i]);
            solver->add(0);
        }
    }

    int s = 0, pair_cnt = 0, nv = maxvar;
    auto clk_ori_st = clk_st;
    printf("c Number of pairs detected: %d\n", pairs.size());
    for(auto& pr : pairs) {
        clk_st = std::chrono::high_resolution_clock::now();
        int a = pr.second.first;
        int b = pr.second.second;
        printf("[%4d/%4d]   check-eq %6d %6d: ", ++pair_cnt, pairs.size(), a, abs(b));
        ++nv;
        if (b > 0) { //check a = b
            if(check_var_equal(a, b, -nv)) {
                printf("  YES (%7.2lf s)\n", 0.001 * get_time());
                ++s;
                
                int x = a & 1 ? -aiger_lit2var(a) : aiger_lit2var(a);
                int y = b & 1 ? -aiger_lit2var(b) : aiger_lit2var(b);
                    
                solver->add(-x), solver->add( y), solver->add(0);
                solver->add( x), solver->add(-y), solver->add(0);

            }
            else printf("  \t\t\t NO (%7.2lf s)\n", 0.001 * get_time());
        }
        else { //check a = 1 or a = 0
            int val = abs(b);
            if(check_constant(a, val)) {
                printf("  YES (%7.2lf s)\n", 0.001 * get_time());
                ++s;
                
                int x = a & 1 ? -aiger_lit2var(a) : aiger_lit2var(a);
                
                solver->add(x * (val ? 1 : -1));
                solver->add(0);
            }
            else printf("  \t\t\t NO (%7.2lf s)\n", 0.001 * get_time());
        }
    }
    printf("c Number of equivalent pairs: %d\n", s);
    clk_st = clk_ori_st;
}

void Sweep::solve() {
    aiger_preprocess();
    //printf("c preprocess finish: %.2lf s\n", 0.001 * get_time());
    identify();
    //simulation();
    return;
    printf("c simulation finish: %.2lf s\n", 0.001 * get_time());

    check_with_SAT();
    printf("c check finish: %.2lf s\n", 0.001 * get_time());
    
    miter();
    printf("c miter finish: %.2lf s\n", 0.001 * get_time());
}



// int main(int argv, char* args[]) {
//     printf("c =================================== input =================================\n");
//     Sweep *S = sweep_init(argv, args);
//     S->solve();
    
//     return 0;
// }