#include "circuit.h"

#include <queue>
#include <unordered_set>
#include <unordered_map>
#include <algorithm>
#include "assert.h"
#include <chrono>

#include "clause.h"
#include "CCAnr/ccanr.h"

bool Circuit::local_search() {

    //printf("[start ls]\n");

    ls_reset_data();

    ls_init_stems();

    ls_init_weight();

    CCAnr::module_reset();

    ls_random_circuit();

    //printf("stem-cnt: %d\n", stem_total_cnt);

    for(int i=0; i<MAX_STEPS; i++) {

        auto start = std::chrono::system_clock::now();
        
        //printf("[FLIP] stem: %lld, fault:%lld, stem_cnt: %lld, fault_cnt:%lld, fpl_score: %lld citcuit-score: %lld\n", stem_total_cost, fault_total_weight, stems.size() - stem_total_cnt, fault_total_cnt, fault_propagate_score, ls_circuit_score());

        int id;


        double p = 0.8;

        // if(rand() % 10000 <= p * 10000) {
        //     id = CCAnr::module_pick_var();
        // } else {
            
        // }

        id = ls_pick();

        CCAnr::module_flip_var(id);

        ls_flip_var(id);

        if(stem_total_cnt == stems.size()) {
            printf("FIND SOLUTION!\n");
            break;
        }

        //assert(is_valid_circuit());

        auto end = std::chrono::system_clock::now();
        std::chrono::duration<double> elapsed_seconds = end - start;
        update_cnt++;
        update_time += elapsed_seconds.count();

        //printf("[UP] flip: %lld, stem: %lld, fault:%lld. flip_cnt: %lld, stem_cnt: %lld, fault_cnt:%lld\n", flip_total_weight, stem_total_weight, fault_total_weight, flip_total_cnt, stem_total_cnt, fault_total_cnt);
    }

    ls_statistics();

    return true;
}

void Circuit::ls_statistics() {

    static int pattern = 0;

    int last_undetect = global_fault_undetected_count;

    for(Gate* g : gates) {
        if(g->fault_detected[0] && !g->global_fault_detected[0]) {
            global_fault_undetected_count--;
            g->global_fault_detected[0] = 1;
        }
        if(g->fault_detected[1] && !g->global_fault_detected[1]) {
            global_fault_undetected_count--;
            g->global_fault_detected[1] = 1;
        }
    }

    if(global_fault_undetected_count - last_undetect < 0) {
        pattern++;
    }

    printf("coverage: %.2f%% undected_fault: %d delta: %d total-pt:%d\n", 
        (gates.size() * 2.0 - global_fault_undetected_count) / (gates.size() * 2.0) * 100,
        global_fault_undetected_count, global_fault_undetected_count - last_undetect, pattern);

    printf("flip-cnt: %d flip-time: %.3fs update-cnt: %d update-time: %.3fs\n", flip_cnt, flip_time, update_cnt, update_time);
    printf("time-per-update: %.2fms\n", update_time / update_cnt * 1000);
}

void Circuit::ls_update_weight() {

    if(rand() % 10000 <= SP * 10000) {
        for(auto& clause : ClauseLS::satisfied_clauses) {
            if(clause->weight - CLAUSE_FALSIFIED_INC < 1) continue;
            clause->weight -= CLAUSE_FALSIFIED_INC;
        }

        for(Gate* g : gates) {
            if(g->fault_detected[0] && g->fault_weight[0] - FAULT_INC >= 1) {
                g->fault_weight[0] -= FAULT_INC;
                fault_propagate_score -= FAULT_INC * (g->fault_propagate_length[0]);
            }

            if(g->fault_detected[1] && g->fault_weight[1] - FAULT_INC >= 1) {
                g->fault_weight[1] -= FAULT_INC;
                fault_propagate_score -= FAULT_INC * (g->fault_propagate_length[1]);
            }
        }
    } else {
        for(auto& clause : ClauseLS::falsified_clauses) {
            if(clause->weight + CLAUSE_FALSIFIED_INC > CLAUSE_FALSIFIED_MAX) continue;
            clause->weight += CLAUSE_FALSIFIED_INC;
            clause->total_cost += CLAUSE_FALSIFIED_INC;
        }

        for(Gate* g : gates) {
            if(!g->fault_detected[0] && g->fault_weight[0] > 0 && (g->fault_weight[0] + FAULT_INC < FAULT_WEIGHT_MAX)) {
                g->fault_weight[0] += FAULT_INC;
                fault_propagate_score += FAULT_INC * (g->fault_propagate_length[0]);
            }

            if(!g->fault_detected[1] && g->fault_weight[1] > 0 && (g->fault_weight[1] + FAULT_INC < FAULT_WEIGHT_MAX)) {
                g->fault_weight[1] += FAULT_INC;
                fault_propagate_score += FAULT_INC * (g->fault_propagate_length[1]);
            }
        }
    }
}

int Circuit::ls_pick() {
    int var = -1;
    ll max_score = 0;

    for(int i=0; i<SAMPLING_COUNT; i++) {
        int t_var = rand() % ClauseLS::num_vars + 1;
        if(!ClauseLS::CC[t_var]) continue;

        ll t_score = ls_pick_score(t_var);
        
        if(t_score > max_score) {
            max_score = t_score;
            var = t_var;
        }
    }

    if(var == -1) {
        ls_update_weight();

        //printf("[UP] stem: %lld, fault:%lld, stem_cnt: %lld, fault_cnt:%lld, fpl_score: %lld citcuit-score: %lld\n", stem_total_cost, fault_total_weight, stems.size() - stem_total_cnt, fault_total_cnt, fault_propagate_score, ls_circuit_score());

        auto it = std::next(ClauseLS::falsified_clauses.begin(), rand() % ClauseLS::falsified_clauses.size());

        auto& lits = (*it)->lits;

        var = abs(lits[rand()%lits.size()]);
    }

    int var_score = ls_pick_score(var);

    assert(var != -1);;

    return var;
}

void Circuit::ls_init_stems() {
    stems.clear();

    for(Gate* g : gates) {
        if(g->pi || g->po) {
            g->stem = true;
        }
        
        if(!g->global_fault_detected[0] || !g->global_fault_detected[0]) {
            g->stem = true;
        }

        if(g->fan_outs.size() >= 2) {
            g->stem = true;
        }

        if(g->stem) {
            stems.push_back(g);
        }
    }

    for(Gate *g : gates) {
        if(g->pi) continue;
        std::queue<Gate*> q;
        std::unordered_map<Gate*, bool> used;
        q.push(g);

        while(!q.empty()) {
            Gate* now = q.front();
            q.pop();
            for(Gate* in : now->fan_ins) {
                if(in->stem) {
                    g->pre_stems.push_back(in);
                } else if(!used[in]) {
                    used[in] = true;
                    q.push(in);
                }
            }
        }
    }

    for(Gate *g : gates) {
        if(g->po) continue;
        std::queue<Gate*> q;
        std::unordered_map<Gate*, bool> used;
        q.push(g);

        while(!q.empty()) {
            Gate* now = q.front();
            q.pop();
            for(Gate* out : now->fan_outs) {
                if(out->stem) {
                    g->suc_stems.push_back(out);
                } else if(!used[out]) {
                    used[out] = true;
                    q.push(out);
                }
            }
        }
    }
}

void Circuit::ls_flip_var(int var) {
    ClauseLS::flip(var);
    if(id2gate.count(var) && id2gate[var]->stem) {
        ls_flip_stem(id2gate[var]);
    }
}

ll Circuit::ls_pick_score(int var) {

    ll old_score = ls_circuit_score();

    ls_flip_var(var);

    ll new_score = ls_circuit_score();

    ls_flip_var(var);

    assert(old_score == ls_circuit_score());

    return new_score - old_score;
}

ll Circuit::ls_circuit_score() {
    ll score = - Clause::total_cost + fault_propagate_score + fault_total_weight;
    //ll score = - Clause::total_cost;
    return score;
}

void Circuit::ls_init_weight() {
    for(Gate* s : stems) {
        s->stem_weight = 1;
        stem_total_cost += s->stem_weight;
    }

    for(Gate* g : gates) {
        g->fault_weight[0] = !g->global_fault_detected[0];
        g->fault_weight[1] = !g->global_fault_detected[1];
    }
}

void Circuit::ls_random_circuit() {

    // random init clauseLS
    for(int i=1; i<=ClauseLS::num_vars; i++) {
        int new_v = CCAnr::module_cur_soln()[i];
        if(new_v != ClauseLS::lit_value[i]) {
            //ls_flip_var(i);
            ClauseLS::flip(i);
        }
    }

    // init assignment
    for(Gate* s : stems) {
        s->value = ClauseLS::lit_value[s->id];
    }

    // recal value by topo
    for(Gate *g : gates) {
        if(g->stem) {
            g->stem_satisfied = (g->recal_value() == g->value);
            if(g->stem_satisfied) {
                stem_total_cost -= g->stem_weight;
                stem_total_cnt++;
            }
        } else {
            g->value = g->recal_value();
        }
    }

    // recal fault by rtopo
    for(Gate* g : rtopo_gates) {
        g->recal_fault(g->fault_detected);
        if(g->fault_detected[0]) {
            fault_total_weight += g->fault_weight[0];
            fault_total_cnt++;
        }
        if(g->fault_detected[1]) {
            fault_total_weight += g->fault_weight[1];
            fault_total_cnt++;
        }

        g->propagate = (g->fault_detected[0] || g->fault_detected[1]);

        g->recal_propagate_len(g->fault_propagate_length);

        fault_propagate_score += g->fault_weight[0] * g->fault_propagate_length[0];
        fault_propagate_score += g->fault_weight[1] * g->fault_propagate_length[1];
    }

    assert(is_valid_circuit());
}

void Circuit::ls_reset_data() {

    stems.clear();

    fault_propagate_score = 0;

    stem_total_cost = 0;
    stem_total_cnt = 0;

    fault_total_weight = 0;
    fault_total_cnt = 0;

    for(Gate *g : gates) {
        g->CC = 1;
        g->stem = 0;
        g->propagate = 0;
        g->stem_weight = 0;
        g->stem_satisfied = 0;

        g->fault_weight[0] = g->fault_weight[1] = 0;
        g->fault_detected[0] = g->fault_detected[1] = 0;

        g->fault_propagate_length[0] = 0;
        g->fault_propagate_length[1] = 0;        
    }
}


void Circuit::ls_flip_stem(Gate* stem) {

    stem->value = !stem->value;

    // update CC
    stem->CC = 0;
    for(Gate* pre : stem->pre_stems) {
        pre->CC = 1;
    }
    for(Gate* suc : stem->suc_stems) {
        suc->CC = 1;
    }

    // update value
    bool new_stem_satisfied = (stem->recal_value() == stem->value);

    if(new_stem_satisfied && !stem->stem_satisfied){
        stem->stem_satisfied = true;
        stem_total_cost -= stem->stem_weight;
        stem_total_cnt += 1;
    }

    if(!new_stem_satisfied && stem->stem_satisfied){
        stem->stem_satisfied = false;
        stem_total_cost += stem->stem_weight;
        stem_total_cnt -= 1;
    }

    // update po fault
    if(stem->po) {
        stem->propagate = true;
        if(!stem->fault_detected[!stem->value]) {
            fault_total_weight += stem->fault_weight[!stem->value];
            fault_total_cnt += 1;
            stem->fault_detected[!stem->value] = true;
        }

        if(stem->fault_detected[stem->value]) {
            fault_total_weight -= stem->fault_weight[stem->value];
            fault_total_cnt -= 1;
            stem->fault_detected[stem->value] = false;
        }
    }

    static std::queue<Gate*> q1;
    static std::unordered_map<Gate*, int> used1;
    static std::queue<Gate*> q2;
    static std::unordered_map<Gate*, int> used2;
    used1.clear();
    used2.clear();

    q1.push(stem);
    
    while(!q1.empty()) {
        Gate* g = q1.front();
        q1.pop();
        used1[g] = false;

        for(Gate* suc : g->suc_stems) {
            if(!used2[suc]) {
                used2[suc] = true;
                q2.push(suc);
            }
        }

        for(Gate* out : g->fan_outs) {
            if(out->stem) {
                Gate* stem = out;
                bool new_stem_satisfied = (stem->recal_value() == stem->value);
                if(new_stem_satisfied && !stem->stem_satisfied){
                    stem->stem_satisfied = true;
                    stem_total_cost -= stem->stem_weight;
                    stem_total_cnt += 1;
                }

                if(!new_stem_satisfied && stem->stem_satisfied){
                    stem->stem_satisfied = false;
                    stem_total_cost += stem->stem_weight;
                    stem_total_cnt -= 1;
                }
                continue;
            }
            int new_value = out->recal_value();
            if(new_value == out->value) continue;
            out->value = new_value;
            if(!used1[out]) {
                used1[out] = true;
                q1.push(out);
            }
        }
    }

    q2.push(stem);
    used2[stem] = true;

    while(!q2.empty()) {
        Gate *g = q2.front();
        q2.pop();

        used2[g] = false;

        for(Gate* in : g->fan_ins) {

            bool update = false;

            bool fd[2];
            in->recal_fault(fd);

            if(fd[0] != in->fault_detected[0]) {
                update = true;
                if(fd[0]) {
                    fault_total_weight += in->fault_weight[0];
                    fault_total_cnt += 1;
                } else {
                    fault_total_weight -= in->fault_weight[0];
                    fault_total_cnt -= 1;
                }
                in->fault_detected[0] = fd[0];
            }

            if(fd[1] != in->fault_detected[1]) {
                update = true;
                if(fd[1]) {
                    fault_total_weight += in->fault_weight[1];
                    fault_total_cnt += 1;
                } else {
                    fault_total_weight -= in->fault_weight[1];
                    fault_total_cnt -= 1;
                }
                in->fault_detected[1] = fd[1];
            }

            in->propagate = (in->fault_detected[0] || in->fault_detected[1]);

            int fpl[2];
            in->recal_propagate_len(fpl);

            if(fpl[0] != in->fault_propagate_length[0] || fpl[1] != in->fault_propagate_length[1]) {
                update = true;
                fault_propagate_score += in->fault_weight[0] * (fpl[0] - in->fault_propagate_length[0]);
                fault_propagate_score += in->fault_weight[1] * (fpl[1] - in->fault_propagate_length[1]);

                in->fault_propagate_length[0] = fpl[0];
                in->fault_propagate_length[1] = fpl[1];
            }
                
            if(update && !used2[in]) {
                used2[in] = true;
                q2.push(in);
            }
        }
    }
}