#include "circuit.h"

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

std::vector<int> Circuit::local_search(const std::vector<Fault*> &faults) {

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

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

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

    while(true) {

        if(stem_total_cnt == stems.size() && flip_total_cnt == 0) {
            printf("FIND SOLUTION!\n");
            printf("[SOL] 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);
            break;
        }

        Gate* stem = nullptr;
        ll max_score = 0;
        
        std::vector<Gate*> stems_random;
        std::vector<Gate*> candidates;

        for(int i=0; i<stems.size(); i++) {
            if(CC[stems[i]->id]) {
                stems_random.push_back(stems[i]);
            }
        }
        for(int i=0; i<stems_random.size(); i++) {
            std::swap(stems_random[i], stems_random[rand()%stems_random.size()]);
        }

        const int T = 50;
        int t = 0;

        for(int i=0; i<stems_random.size(); i++) {
            Gate* t_stem = stems_random[i];
            ll t_score = ls_pick_score(t_stem);
            if(t_score > max_score) {
                max_score = t_score;
                stem = t_stem;
            }
            if(t_score > 0) t++;
            if(t >= T) break;
        }

        if(max_score > 0) {
            printf("FLIP: %s\n", stem->name.c_str());
            printf("[LS] 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_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;
            }

        } else {
            ls_update_weight();

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

            std::queue<Gate*> q;
            std::unordered_map<Gate*, bool> used;
            for(Gate* pi : PIs) {
                used[pi] = true;
                q.push(pi);
            }

            // while(!q.empty()) {

            // }

            // assert(flip_total_cnt == 0);

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

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

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

    //print_gates();

    return std::vector<int>();
}

void Circuit::ls_update_weight() {
    if(rand() % 10 < 3) {

        for(Gate* g : gates) {

            if(g->stem && stem_satisfied[g->id] && (stem_weight[g->id] - STEM_INC > 1)) {
                stem_weight[g->id] -= STEM_INC;
                stem_total_weight -= STEM_INC;

                for(Gate* suc : g->suc_stems) {
                    if(!stem_satisfied[suc->id]) {
                        stem_weight[suc->id] += STEM_INC;
                    }
                }
                //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;

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

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


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) {

    ls_flip(stem);

    ll new_score = ls_score();

    ls_flip(stem);

    ll old_score = ls_score();

    return new_score - old_score;
}

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

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_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_satisfied[stem->id] = false;
        stem_total_weight -= stem_weight[stem->id];
        stem_total_cnt -= 1;
    }

    //printf("flip: %s\n", stem->name.c_str());

    //stem->value = !stem->value;

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

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

    std::queue<Gate*> q;
    std::unordered_map<Gate*, int> used;
    std::vector<Gate*> suc_stems;

    //printf("suc: %d %d\n", suc_stems.size(), stem->suc_stems.size());

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

    // sort(suc_stems.begin(), suc_stems.end(), cmp);
    // sort(stem->suc_stems.begin(), stem->suc_stems.end(), cmp);
    // assert(suc_stems == stem->suc_stems);

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

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

        if(stem->cal_value() == stem->value && !stem_satisfied[stem->id]){
            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_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;

        bool right_value = (g->cal_value() == g->value);
        for(Gate* in : g->inputs) {
            in->value = !in->value;
            bool input_detected = (g->cal_value() != g->value);
            in->value = !in->value;
                
            bool sa0 = right_value && input_detected && g->sa[!g->value] && in->value;
            bool sa1 = right_value && input_detected && g->sa[!g->value] && !in->value;

            //printf("gate: %s -> %s rv: %d id: %d p:%d sa0: %d sa1: %d\n", in->name.c_str(), g->name.c_str(), right_value, input_detected, g->is_propagated(), sa0, sa1);

            in->sa_by_out[g] = std::make_pair(sa0, sa1);

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

            in->sa[0] = in->sa[1] = 0;
            for(Gate* out : in->outputs) {
                auto &p = in->sa_by_out[out];
                //printf("%d %d\n", p.first, p.second);
                in->sa[0] |= p.first;
                in->sa[1] |= p.second;
            }

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

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

            //printf("gate: %s -> %s rv: %d id: %d p:%d sa0: %d sa1: %d\n", in->name.c_str(), g->name.c_str(), right_value, input_detected, g->is_propagated(), in->sa[0], in->sa[1]);
                
            if(!in->stem && !used[in]) {
                used[in] = true;
                q.push(in);
            }
        }
    }
}

void Circuit::ls_init_weight(const std::vector<Fault*> &faults) {
    for(Gate* s : stems) {
        stem_weight[s->id] = 1;
    }
    for(Fault* f : faults) {
        fault_weight[f->gate->id][f->type] = 1;
    }
    for(Gate* s: stems) {
        flip_weight[s->id] = 1;
    }
}

void Circuit::ls_init_circuit() {
    // for(auto pi : PIs) {
    //     pi->value = rand() % 2;
    // }

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

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

    //printf("flip: %lld, stem: %lld, fault:%lld\n", flip_total_weight, stem_total_weight, fault_total_weight);

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

    }

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