#include "infection.h"
#include <vector>
#include <utility>
#include <queue>
#include <iostream>
#include  <cassert>
#define pb push_back
const int NMAX = 5010;
const int INF = 1e9;
std::vector <int> adj[NMAX + 1];
int tata[NMAX + 1][NMAX + 1];

bool infected[NMAX + 1];
int coada[10 * NMAX + 1], viz[NMAX], N, M;
void bfs(int start, int last){
    int st =1, dr = 1;
    coada[st] = start;
    viz[start] = 1;
    while(st <= dr){
        for(int i = 0;i<adj[coada[st]].size();i++){
            int vec = adj[coada[st]][i];
            if(!viz[vec]){
                coada[++dr] = vec;
                viz[vec] = 1;
                tata[last][vec] = coada[st];
            }
        }
        st++;
    }
}
void reset(){
    for(int i = 0;i<N;i++){
        viz[i] = 0;
    }
}
bool check(int nod){
    if(infected[nod]){
       //  std::cout<<"In ziua: "<< ziua<<", am fost infectat, eu "<<crt<<" in orasul "<<nod<<std::endl;
        return true;
    }
    return false;
}
std::vector<int> solve(
    std::vector<std::pair<int, int>> R,
    std::vector<long long> D,
    std::vector<int> S,
    std::vector<int> T,
    std::vector<bool> I
) {
    N = R.size() + 1;
    M = D.size();
    std::vector<int> rez(M);
    for(int i = 0;i<N - 1;i++){
        adj[R[i].first].pb(R[i].second);
        adj[R[i].second].pb(R[i].first);
    }

    int ziua = INF;
    for(int i = 0;i<M;i++){
//        std::cout<<"Am inceput bfs ul de la"<<S[i]<<" la "<< T[i]<<std::endl;
        bfs(T[i], i);
        reset();
        if(ziua > D[i]){
            ziua = D[i];
        }
    }

//    for(int i = 0;i<M;i++){
//        std::cout<<"Pentru nodul "<<i<<". Printam drumul de la "<<S[i]<<" la "<<T[i]<<":\n";
//        int cur = S[i];
//        while(cur != T[i]){
//            std::cout<<cur<<" ";
//            cur = tata[i][cur];
//        }
//    }
    ///parcurgerea efectiva pe zile
    bool ok = true;
    while(ok){
        ok = false; /// doar cand mai are trasee continua
        ///ne apucam mai intai de cei infectati
        for(int i = 0;i<M;i++){
            if(D[i] <= ziua && I[i] && S[i] != T[i]){
                ok = true;
                if(D[i] == ziua){
                    infected[S[i]] = true;
                    rez[i] = ziua; /// marcam ziua in care au inceput sa fie infectati
                }else{
                    S[i] = tata[i][S[i]];
                    infected[S[i]] = true;
                }
                ///ziua in care au terminat a fi fost infectati
                if(S[i] == T[i]){
                    rez[i] = ziua - rez[i] + 1;
                }
           //     std::cout<<"Azi am infectata: "<<S[i]<<"\n";
            }
        }
        /// acum pentru cei neinfectati
        for(int i = 0;i<M;i++){
            if(D[i] <= ziua && !I[i] && S[i] != T[i]){
                ok = true;
                if(D[i] == ziua){
                    if(check(S[i])){
                        I[i] =1;
                        rez[i] = ziua;
                    }
                }else{
                    S[i] = tata[i][S[i]];
                    if(check(S[i])){
                        I[i] =1;
                        rez[i] = ziua;
                    }
                }
            }
        }

        ///resetam infectatii
        for(int i = 0;i<M;i++){
            if(D[i] <= ziua && I[i]){
                infected[S[i]] = false;
            }
        }
        ziua++;
    }
    for(int i = 0;i<M;i++){
        assert(rez[i] >= 0);
    }
    return {rez};
}

