#include "circuit.h"

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

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

    //STEM_INC = 0;

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

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

    while(true) {

        bool value, propagate; 
        Gate* picked_stem = ls_pick_good_var(value, propagate);

        if(picked_stem == nullptr) {
            ls_update_weight();
            picked_stem = ls_pick_falsified_var(value, propagate);
            printf("[UP] propagate: %lld, stem: %lld, fault:%lld. propagate_cnt: %d, stem_cnt: %d, fault_cnt:%d\n", propagate_total_cost, stem_total_cost, fault_total_cost, propagate_falsified_cnt, stem_falsified_cnt, fault_falsified_cnt);

            int** sa = simulate();
            int cnt = 0;
            static int max_cnt = 0;

            for(Fault* f : faults) {
                if(sa[f->gate->id][f->type]) {
                    cnt++;
                }
            }
            max_cnt = std::max(max_cnt, cnt);
            
            printf("[SOL] detect-faults: %d max-faults: %d\n", cnt, max_cnt);
        } else {
            // printf("pick: %s value: %d propagate: %d\n", picked_stem->name.c_str(), value, propagate);
            // printf("[LS] propagate: %lld, stem: %lld, fault:%lld. propagate_cnt: %d, stem_cnt: %d, fault_cnt:%d\n", propagate_total_cost, stem_total_cost, fault_total_cost, propagate_falsified_cnt, stem_falsified_cnt, fault_falsified_cnt);
        }

        picked_stem->value = value;
        picked_stem->propagate = propagate;

        ls_update(picked_stem);

        //assert(is_valid_circuit());

        if(stem_falsified_cnt == 0 && propagate_falsified_cnt == 0) {
            printf("FIND SOLUTION!\n");
            printf("[SOL] propagate: %lld, stem: %lld, fault:%lld. propagate_cnt: %d, stem_cnt: %d, fault_cnt:%d\n", propagate_total_cost, stem_total_cost, fault_total_cost, propagate_falsified_cnt, stem_falsified_cnt, fault_falsified_cnt);
            break;
        }

        if(fault_total_cost == 6340 && propagate_falsified_cnt == 19 && stem_falsified_cnt == 1 && fault_falsified_cnt == 317) {
            static int cnt = 0;
            cnt++;
            if(cnt == 100) {
                print_gates();
                for(Gate* s: stems) {
                    if(!s->propagate_satisfied) {
                        printf(" > propagate falsified: %s\n", s->name.c_str());

                        printf("s0: %d s1: %d, cal_s0: %d, cal_s1: %d\n", s->fault_satisfied[0], s->fault_satisfied[1], s->cal_sa(0), s->cal_sa(1));

                        printf("score: %lld\n", ls_pick_score(s, 1, 0));

                    }
                }
                exit(0);
            }
        }
    }

    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->fault_satisfied[f->type]) {
            faults.erase(f);
        }
    }

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

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

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

    return true;
}

Gate* Circuit::ls_pick_falsified_var(bool &value, bool &propagate) {
    std::vector<Gate*> candidates;
    for(Gate *g : stems) {
        if(g->po) continue;
        if(g->stem_satisified) continue;
        if(g->propagate_satisfied) continue;
        candidates.push_back(g);
    }

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

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


    value = !pick->value;
    propagate = !propagate;

    return pick;
}

Gate* Circuit::ls_pick_good_var(bool &value, bool &propagate) {

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

    for(Gate* s : stems) {
        if(s->CC) {
            stems_random.push_back(s);
        }
    }

    for(int i=0; i<stems_random.size(); i++) {
        std::swap(stems_random[i], stems_random[rand()%stems_random.size()]);
    }

    const int T = 25;

    for(int i=0; i < stems_random.size() && i < T; i++) {
        Gate* t_stem = stems_random[i];

        for(int j=0; j<4; j++) {
            if((j&1) == t_stem->value && (j&2) == t_stem->propagate) continue;
            ll t_score = ls_pick_score(t_stem, j&1, j&2);
            if(t_score > max_score) {
                max_score = t_score;
                stem = t_stem;
                value = j & 1;
                propagate = j & 2;
            }
        }
    }

    if(max_score > 0) {
        return stem;
    } else {
        return nullptr;
    }
}

void Circuit::ls_update_weight() {

    //STEM_INC += 5;

    if(false) {

    } else {
        for(Gate* g : gates) {

            if(g->stem && !g->propagate_satisfied && (g->propagate_cost + PROPAGATE_INC <= PROPAGATE_COST_MAX)) {
                g->propagate_cost += PROPAGATE_INC;
                propagate_total_cost += PROPAGATE_INC;

                for(Gate* pre : g->pre_stems) {
                    if(pre->propagate_cost - PROPAGATE_INC >= 1) {
                        pre->propagate_cost -= PROPAGATE_INC;
                        if(!pre->propagate_satisfied) {
                            propagate_total_cost -= PROPAGATE_INC;
                        }
                    }
                }
            }

            if(g->stem && !g->stem_satisified && (g->stem_cost + STEM_INC <= STEM_COST_MAX)) {
                g->stem_cost += STEM_INC;
                stem_total_cost += STEM_INC;

                if(g->stem_satisified)

                for(Gate* suc : g->suc_stems) {
                    if(suc->stem_cost - STEM_INC >= 1) {
                        suc->stem_cost -= STEM_INC;
                        if(!suc->stem_satisified) {
                            stem_total_cost -= STEM_INC;
                        }
                    }
                }
            }

            if(!g->fault_satisfied[0] && g->fault_cost[0] > 0 && (g->fault_cost[0] + FAULT_INC <= FAULT_COST_MAX)) {
                g->fault_cost[0] += FAULT_INC;
                fault_total_cost += FAULT_INC;
            }

            if(!g->fault_satisfied[1] && g->fault_cost[1] > 0 && (g->fault_cost[1] + FAULT_INC <= FAULT_COST_MAX)) {
                g->fault_cost[1] += FAULT_INC;
                fault_total_cost += FAULT_INC;
            }

        }
    }
}


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

void Circuit::ls_update(Gate* stem) {

    stem->CC = 0;

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

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


    ls_block_recal(stem);

    //assert(is_valid_circuit());
    
}
ll Circuit::ls_pick_score(Gate* stem, bool value, bool propagate) {

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

    ll old_score1 = ls_score();
    // print_gates();
    // printf("--------------");
    //printf("[before] propagate: %lld, stem: %lld, fault:%lld\n", Circuit::propagate_total_cost, Circuit::stem_total_cost, Circuit::fault_total_cost);
    //printf("[before] propagate_cnt: %d, stem_cnt: %d, fault_cnt:%d\n", Circuit::propagate_falsified_cnt, Circuit::stem_falsified_cnt, Circuit::fault_falsified_cnt);

    bool old_value = stem->value;
    bool old_propagate = stem->propagate;

    stem->value = value;
    stem->propagate = propagate;

    ls_update(stem);

    ll new_score = ls_score();
    // print_gates();
    // printf("--------------");

    //printf("[now] propagate: %lld, stem: %lld, fault:%lld\n", Circuit::propagate_total_cost, Circuit::stem_total_cost, Circuit::fault_total_cost);
    //printf("[now] propagate_cnt: %d, stem_cnt: %d, fault_cnt:%d\n", Circuit::propagate_falsified_cnt, Circuit::stem_falsified_cnt, Circuit::fault_falsified_cnt);

    stem->value = old_value;
    stem->propagate = old_propagate;

    ls_update(stem);

    ll old_score2 = ls_score();
    // print_gates();
    // printf("--------------");

    //assert(old_score1 == old_score2);
    // if(old_score1 != old_score2) {

    //     printf("old1: %lld old2: %lld\n", old_score1, old_score2);
        //ed_cnt);
    //     exit(-1);
    // }

    return new_score - old_score1;
}

ll Circuit::ls_score() {
    ll score =  - propagate_total_cost - stem_total_cost - fault_total_cost;
    //ll score =  - stem_total_cost - fault_total_cost;
    return score;
}


void Circuit::ls_init_circuit(const std::unordered_set<Fault*> &faults) {

    // fault_falsified_cnt = faults.size();
    // stem_falsified_cnt = stems.size();
    // propagate_falsified_cnt = stems.size();

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

    for(Gate* s : stems) {
        s->stem_cost = STEM_INC;
        stem_total_cost += s->stem_cost;
        stem_falsified_cnt += 1;

        s->propagate_cost = PROPAGATE_INC;
        propagate_total_cost += s->propagate_cost;
        propagate_falsified_cnt += 1;
    }

    for(Gate* g : gates) {
        g->CC = 1;

        fault_total_cost += g->fault_cost[0];
        fault_total_cost += g->fault_cost[1];
        fault_falsified_cnt += 2;
    }
    
    for(Gate* s : stems) {
        s->value = rand() % 2;
        s->propagate = rand() % 2;
    }

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

    if(!is_valid_circuit()) {
        print_gates();
        exit(-1);
    }
    //assert(is_valid_circuit())
    
}

void Circuit::ls_init_data_structs() {

    propagate_total_cost = 0;
    propagate_falsified_cnt = 0;

    stem_total_cost = 0;
    stem_falsified_cnt = 0;

    fault_total_cost = 0;
    fault_falsified_cnt = 0;

    for(Gate* g : gates) {
        g->CC = 0;

        g->propagate = 0;
        g->propagate_cost = 0;
        g->propagate_satisfied = 0;

        g->stem_cost = 0;
        g->stem_satisified = 0;

        g->fault_satisfied[0] = g->fault_satisfied[1] = 0;
        g->fault_cost[0] = g->fault_cost[1] = 0;
    }
}

void Circuit::ls_block_recal(Gate* stem) {

    ls_block_recal_value(stem);

    std::unordered_set<Gate*> pres;

    for(Gate* suc : stem->suc_stems) {
        ls_block_recal_fault(suc);
    }

    ls_block_recal_fault(stem);
}

void Circuit::ls_block_recal_value(Gate* stem) {

    // 修改当前节点的值信息
    if(stem->cal_value() == stem->value && !stem->stem_satisified) {
        stem->stem_satisified = true;
        stem_total_cost -= stem->stem_cost;
        stem_falsified_cnt -= 1;
    }

    if(stem->cal_value() != stem->value && stem->stem_satisified) {
        stem->stem_satisified = false;
        stem_total_cost += stem->stem_cost;
        stem_falsified_cnt += 1;
    }

    assert((stem->cal_value() == stem->value) == stem->stem_satisified);

    // 更新后继块的值

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

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

    // 更新后继 stem 的关系
    for(Gate* stem : stem->suc_stems) {
        if(stem->cal_value() == stem->value && !stem->stem_satisified) {
            stem->stem_satisified = true;
            stem_total_cost -= stem->stem_cost;
            stem_falsified_cnt -= 1;
        }

        if(stem->cal_value() != stem->value && stem->stem_satisified) {
            stem->stem_satisified = false;
            stem_total_cost += stem->stem_cost;
            stem_falsified_cnt += 1;
        }
    }
}
void Circuit::ls_block_recal_fault(Gate* stem) {
    // po 的 propagte 不可变
    if(stem->po) { stem->propagate = true; }

    // 根据 value 和 propagete 更新当前节点的 fault_satisfied
    for(int i=0; i<=1; i++) {
        int sa[2];
        sa[0] = stem->propagate && stem->value == 1;
        sa[1] = stem->propagate && stem->value == 0;

        if(sa[i] != stem->fault_satisfied[i]) {
            if(stem->fault_satisfied[i]) {
                stem->fault_satisfied[i] = false;
                fault_total_cost += stem->fault_cost[i];
                fault_falsified_cnt++;
            } else {
                stem->fault_satisfied[i] = true;
                fault_total_cost -= stem->fault_cost[i];
                fault_falsified_cnt--;
            }
        }
    }

    // 更新 propagate_satfisfied

    if(stem->cal_sa(0) == stem->fault_satisfied[0] 
            && stem->cal_sa(1) == stem->fault_satisfied[1]
            && !stem->propagate_satisfied) {
        stem->propagate_satisfied = true;
        propagate_total_cost -= stem->propagate_cost;
        propagate_falsified_cnt -= 1;
    }

    if((stem->cal_sa(0) != stem->fault_satisfied[0] || stem->cal_sa(1) != stem->fault_satisfied[1])
            && stem->propagate_satisfied) {
        stem->propagate_satisfied = false;
        propagate_total_cost += stem->propagate_cost;
        propagate_falsified_cnt += 1;
    }

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

    q.push(stem);

    // 修改之前的所有 gate 错误信息直到 stem 边界
    while(!q.empty()) {
        Gate *g = q.front();
        q.pop();
        used[g] = false;

        // 遍历输入迭代
        for(Gate* in : g->inputs) {
            
            // 输入如果是 stem，则更新 propagate 信息是否匹配
            if(in->stem) {
                if(in->cal_sa(0) == in->fault_satisfied[0] 
                    && in->cal_sa(1) == in->fault_satisfied[1]
                    && !in->propagate_satisfied) {
                        in->propagate_satisfied = true;
                        propagate_total_cost -= in->propagate_cost;
                        propagate_falsified_cnt -= 1;
                }

                if((in->cal_sa(0) != in->fault_satisfied[0] || in->cal_sa(1) != in->fault_satisfied[1])
                    && in->propagate_satisfied) {
                        in->propagate_satisfied = false;
                        propagate_total_cost += in->propagate_cost;
                        propagate_falsified_cnt += 1;
                }
            } 
            // 如果是中间节点，直接更新值
            else {

                for(int i=0; i<=1; i++) {
                    int sa[2];
                    sa[0] = in->cal_sa(0);
                    sa[1] = in->cal_sa(1);

                    if(sa[i] != in->fault_satisfied[i]) {
                        if(in->fault_satisfied[i]) {
                            in->fault_satisfied[i] = false;
                            fault_total_cost += in->fault_cost[i];
                            fault_falsified_cnt++;
                        } else {
                            in->fault_satisfied[i] = true;
                            fault_total_cost -= in->fault_cost[i];
                            fault_falsified_cnt--;
                        }
                    }
                }

                in->propagate = in->fault_satisfied[0] || in->fault_satisfied[1];

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

}