# Study of nonlinear dynamics

# R setup
#library(Rmpi)
library(snow); library(rlecuyer)

# make parallel cluster
# cl = makeCluster(2, type="MPI")
# cl = getMPIcluster()

# initialize random number generator
# clusterSetupRNG(cl, seed=rep(11, 6))

# output file
#Out.File.Dir = "/data/stancari/impact-parameter-study/"
Out.File.Dir = "./"
Out.File.Base = "tmp_octu_Q0618"
Out.File = paste(Out.File.Dir, Out.File.Base, ".RData", sep="")

# number of particles per case
Np = 64
# number of turns
Nt = 256

## MACHINE DESCRIPTION
# fractional tune
Q = 0.5*(sqrt(5)-1)
# Courant-Snyder parameter
alpha = 0

## BEAM
# minimum and maximum normalized amplitudes
Am = 0
AM = 12
# initial conditions
#A0 = runif(Np, Am, AM)
A0 = seq(Am, AM, length=Np)
phi0 = runif(Np, -pi, pi)
# positive beam?
posi = TRUE

Nk = 12 # number of different cases
#k = c(0, 10^seq(-3, 1, length=(Nk-1))) # normalized maximum kicks
#k = c(0, 5e-2, 1e-1, 3e-1, 5e-1, 1, 3, 5, 10, 30, 50, 100)
k = c(0, 10^seq(-4, 0, length=(Nk-1)))

# hollow electron lens parameters
Rm = 4 # minimum radius
RM = 7.6 # maximum radius


# list of cases
fwi = 1+floor(log10(Nk))
Par = data.frame(
  name = paste("K", formatC(1:Nk, width=fwi, flag="0"), sep=""),
  Np = rep(Np, Nk),
  Nt = rep(Nt, Nk),
  Q = rep(Q, Nk),
  alpha = rep(alpha, Nk),
  k = k,
  Rm = rep(Rm, Nk),
  RM = rep(RM, Nk),
  posi = rep(posi, Nk),
  stringsAsFactors=FALSE)
CaseList = list()
CaseList = apply(Par, 1, FUN=as.list)

# track particles
# save initial conditions, impact parameters and impact turn
Track = function(X, A0=0:15, phi0=seq(-pi, pi, length=16),
                 type='sext'){

  # get system info
  node = Sys.info()["nodename"]
  # get case info
  name = X["name"]
  Np = as.integer(X["Np"])
  Nt = as.integer(X["Nt"])
  Q = as.numeric(X["Q"])
  twopiQ = 2*pi*Q
  alpha = as.numeric(X["alpha"])
  sqrt1a = sqrt(1+alpha^2)
  k = as.numeric(X["k"])
  Rm = as.numeric(X["Rm"])
  RM = as.numeric(X["RM"])
  posi = as.logical(X["posi"])

  # sextupole
  skick = function(x=1:10,
                   kmax=k){
    stopifnot(is.numeric(x), is.numeric(kmax))
    k = kmax*x^2
    return(k)
  }

  # octupole
  okick = function(x=1:10,
                   kmax=k){
    stopifnot(is.numeric(x), is.numeric(kmax))
    k = kmax*x^3
    return(k)
  }

  # McMillan
  mcmkick = function(x=1:10,
                     kmax=k){
    stopifnot(is.numeric(x), is.numeric(kmax))
    k = 2*kmax*x/(1+x^2)
    return(k)
  }

  # e-lens kick vs. amplitude
  ekick = function(x=1:10, # normalized position
                   rmin=Rm, # normalized inner radius
                   rmax=RM, # normalized outer radius
                   kmax=k, # normalized max kick
                   po=posi){ # positive beam?
    stopifnot(is.numeric(x), rmin>=0, rmax>rmin, kmax>=0, is.logical(po))
    k = numeric(length(x))
    # inside hole
    k[abs(x)<=rmin] = 0
    # inside electron beam
    sub = rmin<=abs(x) & abs(x)<=rmax
    k[sub] = kmax*(x[sub]^2-rmin^2)/(rmax^2-rmin^2)/abs(x[sub])
    # outside electron beam
    sub = rmax<=abs(x)
    k[sub] = kmax/abs(x[sub])
    # reverse appropriate kicks
    if(po) k[x>rmin] = -k[x>rmin] else k[x<(-rmin)] = -k[x<(-rmin)]
    # returns normalized kick
    return(k)
  }

  # generate particle coordinates
  # amplitude
  A = matrix(nrow=Np, ncol=Nt)
  A[, 1] = A0
  # phase
  phi = matrix(nrow=Np, ncol=Nt)
  phi[, 1] = phi0
  # initial coordinates
  X = matrix(nrow=Np, ncol=Nt)
  XP = matrix(nrow=Np, ncol=Nt)
  X[, 1] =  A[, 1] * cos(phi[, 1])
  p = -A[, 1] * sin(phi[, 1]) # normalized momentum
  XP[, 1] = (p - alpha*X[, 1]) / sqrt1a # normalized angle
  # track
  for(i in 2:Nt){
    X[, i] = (cos(twopiQ)+alpha*sin(twopiQ)) * X[, (i-1)] +
             sqrt1a * sin(twopiQ) * XP[, (i-1)]
    if(type=='sext') K = skick(x=X[, i])
    if(type=='octu') K = okick(x=X[, i])
    if(type=='mcmi') K = mcmkick(x=X[, i])
    if(type=='hebc') K = ekick(x=X[, i])
    XP[, i] = -sqrt1a * sin(twopiQ) * X[, (i-1)] +
             (cos(twopiQ)-alpha*sin(twopiQ)) * XP[, (i-1)] + K
    p = alpha*X[, i] + sqrt1a*XP[, i]
    A[, i] = sqrt(X[, i]^2 + p^2)
    phi[, i] = atan2(p, X[, i])
  }
  # assemble results
  TR = list(name=name, node=node,
           Np=Np, Nt=Nt, Q=Q, alpha=alpha, k=k, Rm=Rm, RM=RM, po=posi,
           A=A, phi=phi)
  return(TR)
}

# baseline
T = lapply(CaseList, FUN=Track, type='octu')

# snow package
#message("Starting parallel calculation...")
#T = clusterApplyLB(cl, CaseList, fun=Track, A0=A0, phi0=phi0)

message("Saving results to file ", Out.File)
save(T, file=Out.File)

# clean up
#stopCluster(cl)