#This code is divided into three parts
#########################################
#Part 1 has code to describe the mechanism that generates the labels
#Part 2 has the code to describe the mechanism that changes the labels
#Part 3 has the main code that runs the two mechanisms on a given graph instance and compares it with the data

##Part 1##
mech_gen_label <- function(i,j,admat,labvec,labelv){
  #Node i is the node with the label
  #Node j is the node without the label
  #admat is the adjacency matrix to figure out if node i and j are connected at time t.
  #labvec is the matrix of labels generated till time t. First row is labels and second row is the time for the corresponding label.
  #labelv is the original number of finite labels. Assume for simplicity labelv = {1,2,3} where 1 is susceptible, 2 is infected and 3 is recovered.
  if(admat[i,j] == 1){#If i and j are in contact at time t
    if(labvec[1,i] == 1){
      #do nothing
    }
    if(labvec[1,i] == 2){
      labvec[1,j] = 2 #infected i infects j
    }
    if(labvec[1,i] == 3){
      #do nothing
    }
  }
   #Infected Time updated for all nodes
  return(labvec)
}


##Part 2###
mech_change_label <- function(i,labvec,labelv,infd){
  #Node i is the vector of nodes with the labels
  #labvec is the vector of labels generated till time t.
  #labelv is the original number of finite labels
  #t is the time step
  #infd is number of days infected
  if(length(which(labvec[3,i] > infd))>0){
    #If node i has exceeded infection threshold
    labvec[1,i] = 3 #Change to recovered
    labvec[3,i] = 0 #Infected time changed to zero
  }
  return(labvec)  
}


###########Main function program##########
r <- 6
c <- 6
m0 <- matrix(0, r, c)
adj_mat <- apply(m0, c(1,2), function(x) sample(c(0,1),1)) 

#Removing self-loops
for (k in 1:nrow(adj_mat)) {
  adj_mat[k,k] <- 0
}

#Making matrix symmetric to remove any directed edges
my_adj_mat <- adj_mat                                        # Duplicate matrix
my_adj_mat[upper.tri(my_adj_mat)] <- t(my_adj_mat)[upper.tri(my_adj_mat)] # Insert lower to upper matrix
my_adj_mat
##      [,1] [,2] [,3] [,4] [,5] [,6]
## [1,]    0    0    1    1    0    0
## [2,]    0    0    0    1    1    1
## [3,]    1    0    0    1    0    1
## [4,]    1    1    1    0    1    1
## [5,]    0    1    0    1    0    1
## [6,]    0    1    1    1    1    0
adj_mat <- my_adj_mat #Updating adjacency matrix

#Parameters and vectors for simulation
time_simul <- 5
infec_thresh <- 3
llmain <- c(1,2,3) #Label vector labelv
llvec <- matrix(data = 0, nrow = length(llmain)+1, ncol = nrow(adj_mat)) ##labvec actual vector
llvec[1,] <- 1 ##Everyone is suspectible at first

#Active edges
act_edg <- which(adj_mat!=0,arr.ind = T)

#Choosing the patient zero#########
p0 <- 1 #Patient zero choice - can be changed
pr0 <- list(0)
llvec[1,p0] <- 2 #Label changed to 2 for patient zero chosen
llvec[3,p0] <- 1 #No. if infected days updated

for (h in 1:time_simul) {
  #Following loop covers all the infected nodes and executes the two functions
  if(length(p0)>0){
    for (v in 1:length(p0)) {
      #following variable selects the nodes connected to patient zero.
      j1 <- unname(act_edg[which(act_edg[,1] == p0[v]),2]) #Choosing the nodes that are on active edges from chosen node
      if(length(j1) > 0){
        numpr0 <- as.vector(unlist(pr0), "numeric")
        if(sum(numpr0) != 0){
          j1 <- j1[! j1 %in% numpr0[numpr0!=0]]
        }
        if(length(j1) >=1){
          for (g in 1:length(j1)) {#Loop updates label of all connected nodes to patient zero
            llvec <- mech_gen_label(p0[v],j1[g],adj_mat,llvec,llmain)
          }
          #Change of labels
          up_vec <- c(p0[v],j1) #List of all infected vectors -- j1 -- along with the infecting node p0[v]
          llvec <- mech_change_label(up_vec,llvec,llmain,infec_thresh)
          #Mech_gen_label updates infected days which is the least we need
          #updating the next node that is infected
          pr0 <- append(pr0,unname(p0[v])) #Previous set of infected nodes
        }
      }
      else{
        #If there are no more nodes eligible for infection, simply check for each infected nodes status
        inf_id <- which(llvec[1,] == 2) #Node ids of those infected
        llvec <- mech_change_label(inf_id,llvec,llmain,infec_thresh)
      }
    }
    #Updating next set of infected nodes
    inf_nod <- which(llvec[1,] == 2) 
    numpr0 <- as.vector(unlist(pr0), "numeric")
    p0 <- as.vector(unname(inf_nod[! inf_nod %in% numpr0[numpr0!=0]])) #Next set of infected nodes
  }
  #Updating time steps for infected
  idx <- which(llvec[1,] == 2)
  llvec[3,idx] <- llvec[3,idx] + 1
}