#include "circuit.h"

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


bool Circuit::local_search(std::unordered_set<Fault*> &faults) {

    // 初始化并重置所有 ls 数据结构
    ls_init_data_structs();

    // 赋值初始权重
    ls_init_weight(faults);

    // 随机生成初始电路
    ls_init_circuit();
    
    //printf("local search!\n");

    while(true) {

        // int ustem = 0;
        // for(Gate* stem : stems) {
        //     ustem += !stem_satisfied[stem->id];
        // }
        // printf("ustem: %d %d\n", ustem, stem_unsatisfied_set.size());
        // assert(ustem == stem_unsatisfied_set.size());

        auto start = std::chrono::system_clock::now();

        Gate* stem = nullptr;
        ll max_score = 0;
        
        const int T = 20;

        for(int i=0; i<T; i++) {
            Gate* t_stem = stems[rand()%stems.size()];
            if(!CC[t_stem->id]) continue;

            ll t_score = ls_pick_score(t_stem);
            if(t_score > max_score) {
                max_score = t_score;
                stem = t_stem;
            }
        }

        if(max_score > 0) {
            // printf("FLIP: %s (+%lld)\n", stem->name.c_str(), max_score);
            // printf("[LS] flip: %lld, stem: %lld, fault:%lld. flip_cnt: %d, stem_cnt: %d, fault_cnt:%d\n", flip_total_weight, stem_total_weight, fault_total_weight, flip_total_cnt, stem_total_cnt, fault_total_cnt);
            ls_flip(stem);

            CC[stem->id] = 0;

            for(Gate* pre : stem->pre_stems) {
                CC[pre->id] = 1;
            }

            for(Gate* suc : stem->suc_stems) {
                CC[suc->id] = 1;
            }

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

        } else {
            ls_update_weight();

            if(stem_total_cnt == stems.size()) {
                while(!flip_update_queue.empty()) {
                    Gate* g = flip_update_queue.back();
                    flip_update_queue.pop_back();
                    if(!flip_need_update[g->id]) continue;
                    flip_need_update[g->id] = false;
                    flip_total_weight -= flip_weight[g->id];
                    flip_total_cnt -= 1;
                    ls_update(g);
                }
                //printf("FIND SOLUTION!\n");
                printf("[SOL] flip: %lld, stem: %lld, fault:%lld. flip_cnt: %d, stem_cnt: %d, fault_cnt:%d\n", flip_total_weight, stem_total_weight, fault_total_weight, flip_total_cnt, stem_total_cnt, fault_total_cnt);
                break;
            }

            std::vector<Gate*> candidates;
            for(Gate *g : stems) {
                if(g->isPO) continue;
                if(stem_satisfied[g->id]) continue;
                candidates.push_back(g);
            }

            if(candidates.size() == 0) {
                candidates.push_back(stems[rand()%stems.size()]);
            }

            Gate* pick = candidates[rand()%candidates.size()];

            ls_flip(pick);

            CC[pick->id] = 0;

            for(Gate* pre : pick->pre_stems) {
                CC[pre->id] = 1;
            }

            for(Gate* suc : pick->suc_stems) {
                CC[suc->id] = 1;
            }

            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);
        }
    }

    static int original_faults = -1;
    if(original_faults == - 1) {
        original_faults = faults.size();
    }
    static int pattern = 0;

    std::unordered_set<Fault*> tmp = faults;

    for(Fault* f : tmp) {
        if(f->gate->sa[f->type]) {
            faults.erase(f);
        }
    }

    if(tmp.size() == faults.size()) pattern--;

    printf("coverage: %.3f%%\tpattern: %d\tbefore: %d\tnow: %d\n", (double)(original_faults - faults.size()) / (original_faults) * 100, ++pattern, tmp.size(), faults.size());

    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-flip: %.2fms time-per-update: %.2fms\n",  flip_time / flip_cnt * 1000, update_time / update_cnt * 1000);

    //if(tmp.size() == faults.size()) return false;

    return true;
}

void Circuit::ls_update_weight() {

    STEM_INC += 1;

    if(rand() % 10000 <= SP * 10000) {
        for(Gate* g : gates) {
            if(g->stem && stem_satisfied[g->id] && (stem_weight[g->id] - STEM_INC >= 1)) {
                stem_weight[g->id] -= STEM_INC;
                for(Gate* suc : g->suc_stems) {
                    if(stem_weight[suc->id] + STEM_INC <= STEM_WEIGHT_MAX) {
                        stem_weight[suc->id] += STEM_INC;
                        if(!stem_satisfied[suc->id]) {
                            stem_total_weight += STEM_INC;
                        }
                    }
                }
            }
        }
    } else {
        for(Gate* g : gates) {
            if(flip_need_update[g->id] && (flip_weight[g->id] + FLIP_INC < FLIP_WEIGHT_MAX)) {
                flip_weight[g->id] += FLIP_INC;
                flip_total_weight += FLIP_INC;
            }

            if(g->stem && !stem_satisfied[g->id] && (stem_weight[g->id] + STEM_INC < STEM_WEIGHT_MAX)) {
                stem_weight[g->id] += STEM_INC;
                stem_total_weight += STEM_INC;

                for(Gate* suc : g->suc_stems) {
                    if(stem_weight[suc->id] - STEM_INC > 1) {
                        stem_weight[suc->id] -= STEM_INC;
                        if(!stem_satisfied[suc->id]) {
                            stem_total_weight -= STEM_INC;
                        }
                    }
                }
            }

            if(!g->sa[0] && fault_weight[g->id][0] > 0 && (fault_weight[g->id][0] + FAULT_INC < FAULT_WEIGHT_MAX)) {
                fault_weight[g->id][0] += FAULT_INC;
                fault_propagate_score += FAULT_INC * (g->fault_propagate_len[0]);
            }

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


bool cmp(Gate* a, Gate *b) {
    return a->id > b->id;
}

void Circuit::ls_flip(Gate* stem) {
    stem->value = !stem->value;
    ls_block_recal(stem);
}

void Circuit::ls_update(Gate* stem) {
    ls_block_recal(stem);
}

ll Circuit::ls_pick_score(Gate* stem) {

    ll old_score = ls_score();

    ls_flip(stem);

    ll new_score = ls_score();

    ls_flip(stem);

    old_score = std::max(old_score, ls_score());

    return new_score - old_score;
}

ll Circuit::ls_score() {
    //ll score =  - flip_total_weight - stem_total_weight + fault_total_weight + fault_propagate_tatal_len;
    ll score =  - flip_total_weight - stem_total_weight + fault_propagate_score + fault_total_weight;
    return score;
}

void Circuit::ls_init_weight(const std::unordered_set<Fault*> &faults) {
    for(Gate* s : stems) {
        stem_weight[s->id] = 1;
        stem_total_weight += stem_weight[s->id];
    }

    for(Fault* f : faults) {
        fault_weight[f->gate->id][f->type] = 1;
    }

    // int r = rand() % faults.size();
    // auto it = faults.begin();
    // std::advance(it, r);

    // fault_weight[(*it)->gate->id][(*it)->type] = 1000000;

    for(Gate* s: stems) {
        flip_weight[s->id] = 1;
    }
}

void Circuit::ls_init_circuit() {

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

    for(Gate* s : stems) {
        s->value = rand() % 2;
    }

    for(int i=stems.size()-1; i>=0; i--) {
        ls_update(stems[i]);
    }

    while(!flip_update_queue.empty()) {
        Gate* g = flip_update_queue.back();
        flip_update_queue.pop_back();
        if(!flip_need_update[g->id]) continue;
        flip_need_update[g->id] = false;
        flip_total_weight -= flip_weight[g->id];
        flip_total_cnt -= 1;
        ls_update(g);
    }
}

void Circuit::ls_init_data_structs() {
    const int MAX_LEN = gates.size() + 1;

    if(flip_weight == nullptr) {
        CC = new int[MAX_LEN];

        flip_weight = new int[MAX_LEN];
        flip_need_update = new int[MAX_LEN];

        stem_weight = new int[MAX_LEN];
        stem_satisfied = new int[MAX_LEN];

        fault_weight = new int*[MAX_LEN];
        for(int i=0; i<MAX_LEN; i++) {
            fault_weight[i] = new int[2];
        }
        fault_detected = new int*[MAX_LEN];
        for(int i=0; i<MAX_LEN; i++) {
            fault_detected[i] = new int[2];
        }
    }

    STEM_INC = 1;

    fault_propagate_score = 0;

    fault_propagate_tatal_len = 0;

    flip_total_weight = 0;
    flip_total_cnt = 0;

    stem_total_weight = 0;
    stem_total_cnt = 0;

    fault_total_weight = 0;
    fault_total_cnt = 0;

    for(int i=0; i<MAX_LEN; i++) {
        CC[i] = 1;
        flip_weight[i] = 0;
        flip_need_update[i] = 0;
        stem_weight[i] = 0;
        stem_satisfied[i] = 0;
        fault_weight[i][0] = 0;
        fault_weight[i][1] = 0;
        fault_detected[i][0] = 0;
        fault_detected[i][1] = 0;
    }

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


void Circuit::ls_block_recal(Gate* stem) {
    if(flip_need_update[stem->id]) {
        flip_need_update[stem->id] = false;
        flip_total_weight -= flip_weight[stem->id];
        flip_total_cnt -= 1;
    }

    if(stem->cal_value() == stem->value && !stem_satisfied[stem->id]){
        stem_unsatisfied_set.erase(stem);

        stem_satisfied[stem->id] = true;
        stem_total_weight -= stem_weight[stem->id];
        stem_total_cnt += 1;
        for(Gate* pre : stem->pre_stems) {
            if(flip_need_update[pre->id]) continue;

            flip_need_update[pre->id] = true;
            flip_update_queue.push_back(pre);

            flip_total_weight += flip_weight[pre->id];
            flip_total_cnt += 1;
        }
    }
    
    if(stem->cal_value() != stem->value && stem_satisfied[stem->id]) {
        stem_unsatisfied_set.insert(stem);

        stem_satisfied[stem->id] = false;
        stem_total_weight += stem_weight[stem->id];
        stem_total_cnt -= 1;
        for(Gate* pre : stem->pre_stems) {
            if(flip_need_update[pre->id]) continue;

            flip_need_update[pre->id] = true;
            flip_update_queue.push_back(pre);

            flip_total_weight += flip_weight[pre->id];
            flip_total_cnt += 1;
        }
    }

    if(stem->isPO) {
        if(stem->sa[!stem->value] == false) {
            fault_total_weight += fault_weight[stem->id][!stem->value];
            fault_total_cnt += 1;
            stem->sa[!stem->value] = true;

            for(Gate* pre : stem->pre_stems) {
                if(flip_need_update[pre->id]) continue;

                flip_need_update[pre->id] = true;
                flip_update_queue.push_back(pre);

                flip_total_weight += flip_weight[pre->id];
                flip_total_cnt += 1;
            }
        }

        if(stem->sa[stem->value] == true) {
            fault_total_weight -= fault_weight[stem->id][stem->value];
            fault_total_cnt -= 1;
            stem->sa[stem->value] = false;

            for(Gate* pre : stem->pre_stems) {
                if(flip_need_update[pre->id]) continue;

                flip_need_update[pre->id] = true;
                flip_update_queue.push_back(pre);

                flip_total_weight += flip_weight[pre->id];
                flip_total_cnt += 1;
            }
        }
    }

    static std::queue<Gate*> q;
    static std::unordered_map<Gate*, int> used;
    assert(q.empty());
    used.clear();

    q.push(stem);
    
    while(!q.empty()) {
        Gate* g = q.front();
        q.pop();
        used[g] = false;
        for(Gate* out : g->outputs) {
            if(out->stem) {
                continue;
            }
            out->value = out->cal_value();
            if(!used[out]) {
                used[out] = true;
                q.push(out);
            }
        }
    }

    assert(q.empty());
    used.clear();

    for(Gate* stem : stem->suc_stems) {
        q.push(stem);

        int fpl0 = stem->cal_propagate_len(0);
        int fpl1 = stem->cal_propagate_len(1);

        if(fault_weight[stem->id][0]) {
            fault_propagate_tatal_len += (fpl0 - stem->fault_propagate_len[0]);
        }

        if(fault_weight[stem->id][1]) {
            fault_propagate_tatal_len += (fpl1 - stem->fault_propagate_len[1]);
        }

        fault_propagate_score += (fault_weight[stem->id][0] * (fpl0 - stem->fault_propagate_len[0]));
        fault_propagate_score += (fault_weight[stem->id][1] * (fpl1 - stem->fault_propagate_len[1]));        
            
        stem->fault_propagate_len[0] = fpl0;
        stem->fault_propagate_len[1] = fpl1;

        if(stem->cal_value() == stem->value && !stem_satisfied[stem->id]){
            stem_unsatisfied_set.erase(stem);
            stem_satisfied[stem->id] = true;
            stem_total_weight -= stem_weight[stem->id];
            stem_total_cnt += 1;
        }
    
        if(stem->cal_value() != stem->value && stem_satisfied[stem->id]) {
            stem_unsatisfied_set.insert(stem);
            stem_satisfied[stem->id] = false;
            stem_total_weight += stem_weight[stem->id];
            stem_total_cnt -= 1;
        }
    }

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

        used[g] = false;

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

            bool old_sa[2];
            old_sa[0] = in->sa[0];
            old_sa[1] = in->sa[1];

            in->sa[0] = in->cal_sa(0);
            in->sa[1] = in->cal_sa(1);

            if(in->stem && !in->isPI && (in->sa[0] != old_sa[0] || in->sa[1] != old_sa[1])) {
                for(Gate* pre : in->pre_stems) {
                    if(flip_need_update[pre->id]) continue;

                    flip_need_update[pre->id] = true;
                    flip_update_queue.push_back(pre);

                    flip_total_weight += flip_weight[pre->id];
                    flip_total_cnt += 1;
                }
            }

            int fpl0 = in->cal_propagate_len(0);
            int fpl1 = in->cal_propagate_len(1);

            // if(in->name == "422") {
            //     printf("%s changed: %d fpl0: %d fpl1: %d \n", in->name.c_str(), (in->fault_propagate_len[0] != fpl0 || in->fault_propagate_len[1] != fpl1), fpl0, fpl1);   
            // }

            if(in->stem && !in->isPI && (in->fault_propagate_len[0] != fpl0 || in->fault_propagate_len[1] != fpl1)) {
                
                for(Gate* pre : in->pre_stems) {
                    if(flip_need_update[pre->id]) continue;

                    flip_need_update[pre->id] = true;
                    flip_update_queue.push_back(pre);

                    flip_total_weight += flip_weight[pre->id];
                    flip_total_cnt += 1;
                }
            }

            if(fault_weight[in->id][0]) {
                fault_propagate_tatal_len += fpl0 - in->fault_propagate_len[0];
            }

            if(fault_weight[in->id][1]) {
                fault_propagate_tatal_len += fpl1 - in->fault_propagate_len[1];
            }

            fault_propagate_score += fault_weight[in->id][0] * (fpl0 - in->fault_propagate_len[0]);
            fault_propagate_score += fault_weight[in->id][1] * (fpl1 - in->fault_propagate_len[1]);
            
            in->fault_propagate_len[0] = fpl0;
            in->fault_propagate_len[1] = fpl1;

            if(old_sa[0] != in->sa[0]) {
                if(in->sa[0]) {
                    fault_total_weight += fault_weight[in->id][0];
                    fault_total_cnt += 1;
                } else {
                    fault_total_weight -= fault_weight[in->id][0];
                    fault_total_cnt -= 1;
                }
            }

            if(old_sa[1] != in->sa[1]) {
                if(in->sa[1]) {
                    fault_total_weight += fault_weight[in->id][1];
                    fault_total_cnt += 1;
                } else {
                    fault_total_weight -= fault_weight[in->id][1];
                    fault_total_cnt -= 1;
                }
            }
                
            if(!in->stem && !used[in]) {
                used[in] = true;
                q.push(in);
            }
        }
    }
}