########################################################
# author: Grzegorz Sionkowski
# date:   2018.05.06 
# ------------------------------------------------------
# information how to improve the model:
# https://www.kaggle.com/c/trackml-particle-identification/discussion/57180
########################################################

require(data.table)
require(bit64)
require(dbscan)

path='../input/train_1/'

########################################################
# score function by Vicens Gaitan
# https://www.kaggle.com/vicensgaitan/r-scoring-function
# (genius work if you compare it with Python original)
########################################################
score<-function(sub,dft){
  df=merge(sub,dft[,.(hit_id,particle_id,weight)])
  df[,Np:=.N,by=particle_id]# Np = Hits per Particle
  df[,Nt:=.N,by=track_id]   # Nt = Hits per Track
  df[,Ntp:=.N,by=list(track_id,particle_id)]# Hits per Particle per Track
  df[,r1:=Ntp/Nt]
  df[,r2:=Ntp/Np]
  sum(df[r1>.5 & r2>.5,weight])
}
########################################################

trackML <- function(dfh){
  dfh[,r:=sqrt(x*x+y*y+z*z)]
  dfh[,rt:=sqrt(x*x+y*y)]
  dfh[,a0:=atan2(y,x)]
  dfh[,r2:=sqrt(x*x+y*y)]
  dfh[,z1:=z/rt] ###
  dz     <- -0.00012
  stepdz <-  0.00001
  for (ii in 0:12) {
    dz <- dz + ii*stepdz
    dfh[,a1:=a0+dz*z*sign(z)]
    dfh[,x1:=a1/z1]
    dfh[,x2:=(1/z1)]
    dfh[,x3:=x1+x2]
    dfs=scale((dfh[,.(a1,z1,x1,x2,x3)]))
    res=dbscan(dfs,eps=.0040,minPts = 1)
    sub=data.table(hit_id=dfh$hit_id,track_id=res$cluster)
    if (ii==0) {
      dfh[,s1:=res$cluster]
      dfh[,N1:=.N, by=s1]
    }else{
      dfh[,s2:=res$cluster]
      dfh[,N2:=.N, by=s2]
      maxs1 <- max(dfh$s1)
      dfh[,s1:=ifelse(N2>N1,s2+maxs1,s1)]
      dfh[,s1:=as.integer(as.factor(s1))]
      dfh[,N1:=.N, by=s1]
    }
  }
  dz     <-  0.00012
  stepdz <- -0.00001
  for (ii in 0:12) {
    dz <- dz + ii*stepdz
    dfh[,a1:=a0+dz*z*sign(z)]
    dfh[,x1:=a1/z1]
    dfh[,x2:=(1/z1)]
    dfh[,x3:=x1+x2]
    dfs=scale((dfh[,.(a1,z1,x1,x2,x3)]))
    res=dbscan(dfs,eps=.0035,minPts = 1)
    sub=data.table(hit_id=dfh$hit_id,track_id=res$cluster)
    dfh[,s2:=res$cluster]
    dfh[,N2:=.N, by=s2]
    maxs1 <- max(dfh$s1)
    dfh[,s1:=ifelse(N2>N1,s2+maxs1,s1)]
    dfh[,s1:=as.integer(as.factor(s1))]
    dfh[,N1:=.N, by=s1]
    
  }
  return(dfh$s1) 
}


######################## 
###    evaluation    ###
########################

sc=1:1 # 1:10
for(i in 0:0){ # 0:9
  nev=1000+i
  dfh=fread(paste0(path,'event00000',nev,'-hits.csv'))
  dft=fread(paste0(path,'event00000',nev,'-truth.csv'),stringsAsFactors = T)
  
  dfh$s1 <- trackML(dfh)
  sub=data.table(event_id=nev,hit_id=dfh$hit_id,track_id=dfh$s1)
  s=score(sub,dft)
  print(c(i,s))
  sc[i+1]=s
}

mean(sc) 


########################
###    submission    ###
########################
path <- '../input/test/'
for(i in 0:124){
  print(i)
  nev <- i
  if (i<10) {
    path <-'../input/test/event00000000'
  }else{
    if (i<100) {
      path <-'../input/test/event0000000'
    }else{
      path <-'../input/test/event000000'
    }
  }
  dfh <- fread(paste0(path,nev,'-hits.csv'))
  dfh$s1 <- trackML(dfh)
  subN <- data.table(event_id=nev,hit_id=dfh$hit_id,track_id=dfh$s1)
  if (i==0) {
    sub <- subN
  }else{
    sub <- rbind(sub,subN)
  }
}
fwrite(sub, "sub-03472.csv")
print('Finished')