library(igraph)
setwd('~/Documents/Violence Cascades/')
load('Raw Data/lcc.RData')
load('Raw Data/dag_dat_lcc.RData')
source('criminal_cascades/R Scripts/temporal.R')
source('criminal_cascades/R Scripts/structural.R')
source('criminal_cascades/R Scripts/background.R')

vic_ids = which(lcc_verts$vic==TRUE)
edges = dag_dat_lcc

#### Initialize model parameters
mu0 = 1.1845e-05 #background
alpha = 0.00317 #structural
beta = 0.0039 #temporal

##### Get weights
# p_t = temporal(edges$t1, edges$t2, beta)
# p_s = structural(alpha, edges$dist)
# p = p_s * p_t
# background = background(mu0,t)

##### Find transmission edges
realized = c()
for (u in vic_ids){
  days = as.numeric(vic_times[u,2:7])
  days = days[!is.na(days)]
  for(day in days){
    bkgnd = background(mu0,day)
    Ei = which(edges$to==u & edges$t1<day)
    p_t = temporal(edges$t1[Ei], day, beta)
    p_s = structural(alpha, edges$dist[Ei])
    p = p_s*p_t
    
    pc = sum(p)
    
    if (pc>bkgnd){
      max_prob = -1
      max_edge = NULL
      for(i in 1:length(Ei)){
        if(p[i]>max_prob){
          max_prob = p[i]
          max_edge = i
        }
      }
    }   
    realized = c(realized, max_edge)
  }
  edges[realized,c('from','to')]
}

##### Create forest of cascades
# dag is all of the victim nodes but with only transmission edges, or don't index by vic_ids
dag = graph.data.frame(edges[realized,1:2], directed=TRUE, vertices=lcc_verts[vic_ids,])
dag = set.edge.attribute(dag,'weight',value=weights[realized])
table(clusters(dag)$csize)

t=table(clusters(dag)$csize)
sizes = as.numeric(attr(t,'name'))
counts = as.numeric(t)
plot(sizes,counts,log='xy',type='o',lwd=3)
n.seeds = sum(counts)
n.nodes = sum(sizes*counts) # for vics case is just length(vic_ids)
n.infections = n.nodes-n.seeds

# dag = graph.edgelist(as.matrix(edges[realized,1:2]))