#### Finite Element Method for Vorticity Equation
### f(x) = (x - 2)^4 * (x + 2)^4 = x^8 - 16 * x^6 + 96 * x^4 - 256 * x^2 + 256
### f'(x) = 8 * x * (x - 2)^3 * (x + 2)^3 = 8 * x^7 - 96 * x^5 + 384 * x^3 - 512 * x
### f''(x) = 8 * (x - 2)^2 * (x + 2)^2 * (7 * x^2 - 4)  = 56 * x^6 - 480 * x^4 + 1152 * x^3 - 512
### F(x) = x^9 / 9 - 16 / 7 * x^7 + 96 / 5 * x^5 - 256 * x^3 / 3 + 256 * x + C
### phi_{mn}(x, y) = f(x) * f(y) in [-2, 2] \\times [-2, 2]
### \\int^{2}_{-2} (f(x))^2 dx = 8589934592 / 109395
### \\int^{2}_{-1} f(x) * f(x - 1) dx = 7172631729 / 170170
### \\int^{2}_{0} f(x) * f(x - 2) dx = 3924361216 / 765765
### \\int^{2}_{1} f(x) * f(x - 3) dx = 7762813 / 218790
### \\int^{2}_{-2} (f'(x))^2 dx = 4294967296 / 45045
### \\int^{2}_{-1} f'(x) * f'(x - 1) dx = -44982216 / 5005
### \\int^{2}_{0} f'(x) * f'(x - 2) dx = -1657536512 / 45045
### \\int^{2}_{1} f'(x) * f'(x - 3) dx = -16712056 / 9009
### f(0) = 256
### f(1) = f(-1) = 81
### f(2) = f(-2) = 0
### f'(0) = 0
### f'(1) = -f'(-1) = -216
### \\int^{.5}_{-.5} f(x) dx = 2 * 19017443 / 161280 = 38034886 / 161280
### \\int^{1.5}_{.5} f(x) dx = 3709761 / 17920 - 19017443 / 161280 
### \\int^{2}_{1.5} f(x) dx = 65536 / 315 - 3709761 / 17920
### \\int^{2}_{-2} f(x) dx = 131072 / 315

library(fields)
library(Matrix)

### Boundary Identification Function
boundary <- function(id_x, id_y, id_boundary){
  value <- 0
  
  if(id_boundary == 1){
    if(id_x <= 3 || id_x > num_x - 3 || id_y <= 3 || id_y > num_y - 3){
      value <- 1
    }
  }else if(id_boundary == 2){
    if((id_y / (num_y + 1) - .5)^2 + (id_x / (num_x + 1) - .5)^2 >= 0.1992985){
      value <- 1
    }
  }
  
  return(value)
}

### Antiderivative of the basis
Integrate_function <- function(z){
  x <- z[1]
  y <- z[2]
  start <- x^9 / 9 - 16 / 7 * x^7 + 96 / 5 * x^5 - 256 * x^3 / 3 + 256 * x
  end <- y^9 / 9 - 16 / 7 * y^7 + 96 / 5 * y^5 - 256 * y^3 / 3 + 256 * y
  
  value <- end - start
  return(value)
}

### Information of the Basis
f_value <- c(8589934592 / 109395, 7172631729 / 170170, 3924361216 / 765765, 7762813 / 218790)
diff_f_value <- c(4294967296 / 45045, -44982216 / 5005, -1657536512 / 45045, -16712056 / 9009)
w_value <- c(256, 81, 0, 0)
w_diff_value <- 216
w_int_value <- c(38034886 / 161280, 3709761 / 17920 - 19017443 / 161280, 65536 / 315 - 3709761 / 17920)
summation_factor <- 174724
integral_factor <- (131072 / 315)^2

### Constants
num_x <- 156
num_y <- 156
num_boundary <- 2
dx <- 1 / num_x
alpha <- 1
Stream_boundary <- c(0, 0)
k_boundary <- Stream_boundary / summation_factor
k_int_boundary <- k_boundary * integral_factor

### Main Matrices
A <- matrix(0, nrow = num_x * num_y, ncol = num_x * num_y)
B <- matrix(0, nrow = num_x * num_y, ncol = num_x * num_y)
W <- matrix(0, nrow = num_x * num_y, ncol = num_x * num_y)
W_int <- matrix(0, nrow = num_x * num_y, ncol = num_x * num_y)
domain <- c()
boundary_elements <- list()
for(id_boundary in 1:num_boundary){
  boundary_elements[[id_boundary]] <- NA
}

neighbor_dir_y <- c(-3: 3)
neighbor_dir_x <- c(-3: 3)

for(id_y in 1:num_y){
  for(id_x in 1:num_x){
    
    ### Check Index of Current Position in the Large Matrix
    pos_current <- (id_y - 1) * num_x + id_x
    
    ### Check if Current Position is in the Boundary
    in_boundary <- 0
    for(id_boundary in 1:num_boundary){
      if(boundary(id_x, id_y, id_boundary)){
        boundary_elements[[id_boundary]] <- c(boundary_elements[[id_boundary]], pos_current)
        in_boundary <- 1
        break
      }
    }
    if(in_boundary){
      ### Skip the process below if in the boundary
      next
    }
    
    ### Record Current Position in the Domain
    domain <- c(domain, pos_current)
    
    ### Scan all the neighbors
    for(neighbor_id_x in neighbor_dir_x){
      for(neighbor_id_y in neighbor_dir_y){
        pos_neighbor <- pos_current + neighbor_id_y * num_x + neighbor_id_x
        
        if(abs(neighbor_id_x) + abs(neighbor_id_y) == 6){
          A[pos_current, pos_neighbor] <- (f_value[4])^2
          B[pos_current, pos_neighbor] <- 2 * diff_f_value[4] * f_value[4]
          W[pos_current, pos_neighbor] <- (w_value[4])^2
        }else if(abs(neighbor_id_x) + abs(neighbor_id_y) == 5){
          A[pos_current, pos_neighbor] <- f_value[4] * f_value[3]
          B[pos_current, pos_neighbor] <- diff_f_value[4] * f_value[3] + diff_f_value[3] * f_value[4]
          W[pos_current, pos_neighbor] <- w_value[4] * w_value[3]
        }else if(abs(neighbor_id_x) + abs(neighbor_id_y) == 4){
          if(abs(neighbor_id_x) != abs(neighbor_id_y)){
            A[pos_current, pos_neighbor] <- f_value[4] * f_value[2]
            B[pos_current, pos_neighbor] <- diff_f_value[4] * f_value[2] + diff_f_value[2] * f_value[4]
            W[pos_current, pos_neighbor] <- w_value[4] * w_value[2]
          }else{
            A[pos_current, pos_neighbor] <- f_value[3] * f_value[3]
            B[pos_current, pos_neighbor] <- 2 * diff_f_value[3] * f_value[3]
            W[pos_current, pos_neighbor] <- w_value[3] * w_value[3]
            W_int[pos_current, pos_neighbor] <- w_int_value[3] * w_int_value[3]
          }
        }else if(abs(neighbor_id_x) + abs(neighbor_id_y) == 3){
          if(abs(neighbor_id_x) == 3 || abs(neighbor_id_y) == 3){
            A[pos_current, pos_neighbor] <- f_value[4] * f_value[1]
            B[pos_current, pos_neighbor] <- diff_f_value[4] * f_value[1] + diff_f_value[1] * f_value[4]
            W[pos_current, pos_neighbor] <- w_value[4] * w_value[1]
          }else{
            A[pos_current, pos_neighbor] <- f_value[3] * f_value[2]
            B[pos_current, pos_neighbor] <- diff_f_value[3] * f_value[2] + diff_f_value[2] * f_value[3]
            W[pos_current, pos_neighbor] <- w_value[3] * w_value[2]
            W_int[pos_current, pos_neighbor] <- w_int_value[3] * w_int_value[2]
          }
        }else if(abs(neighbor_id_x) + abs(neighbor_id_y) == 2){
          if(abs(neighbor_id_x) != abs(neighbor_id_y)){
            A[pos_current, pos_neighbor] <- f_value[3] * f_value[1]
            B[pos_current, pos_neighbor] <- diff_f_value[3] * f_value[1] + diff_f_value[1] * f_value[3]
            W[pos_current, pos_neighbor] <- w_value[3] * w_value[1]
            W_int[pos_current, pos_neighbor] <- w_int_value[3] * w_int_value[1]
          }else{
            A[pos_current, pos_neighbor] <- f_value[2] * f_value[2]
            B[pos_current, pos_neighbor] <- 2 * diff_f_value[2] * f_value[2]
            W[pos_current, pos_neighbor] <- w_value[2] * w_value[2]
            W_int[pos_current, pos_neighbor] <- w_int_value[2] * w_int_value[2]
          }
        }else if(abs(neighbor_id_x) + abs(neighbor_id_y) == 1){
          A[pos_current, pos_neighbor] <- f_value[2] * f_value[1]
          B[pos_current, pos_neighbor] <- diff_f_value[2] * f_value[1] + diff_f_value[1] * f_value[2]
          W[pos_current, pos_neighbor] <- w_value[2] * w_value[1]
          W_int[pos_current, pos_neighbor] <- w_int_value[2] * w_int_value[1]
        }else{
          A[pos_current, pos_neighbor] <- (f_value[1])^2
          B[pos_current, pos_neighbor] <- 2 * diff_f_value[1] * f_value[1]
          W[pos_current, pos_neighbor] <- w_value[1] * w_value[1]
          W_int[pos_current, pos_neighbor] <- w_int_value[1] * w_int_value[1]
        }
      }
    }
  }
}

### Clean up of boundary elements
for(id_boundary in 1:num_boundary){
  boundary_elements[[id_boundary]] <- boundary_elements[[id_boundary]][-1]
}
A <- A[domain, domain]
B <- -B
B <- B[domain, ]
BC_2 <- B[, -domain]
B <- B[, domain]
W <- W[domain, ]
BC_0 <- W[, -domain]
W_int <- W_int[domain, ]
BC_int <- W_int[, -domain]
W <- W[, domain]
W_int <- W_int[, domain]
A <- as.spam(A)
B <- as.spam(B)
BC_2 <- as.spam(BC_2)
W <- as.spam(W)
BC_0 <- as.spam(BC_0)
W_int <- as.spam(W_int)
BC_int <- as.spam(BC_int)
A <- as.dgCMatrix.spam(A)
B <- as.dgCMatrix.spam(B)
BC_2 <- as.dgCMatrix.spam(BC_2)
W <- as.dgCMatrix.spam(W)
BC_0 <- as.dgCMatrix.spam(BC_0)
W_int <- as.dgCMatrix.spam(W_int)
BC_int <- as.dgCMatrix.spam(BC_int)

B_inv <- solve(B)
W_int_inv <- solve(W_int)
M <- -B_inv %*% A * (6 * dx)^2
M <- as.matrix(M)

### Set Boundary Conditions and Homogeneous Solution
bc <- c()
for(id_boundary in 1:num_boundary){
  bc[boundary_elements[[id_boundary]]] <- k_boundary[id_boundary]
}
bc <- bc[-domain]
k_eq <- -B_inv %*% BC_2 %*% bc
s_eq <- W%*%k_eq[domain] + BC_0 %*% bc
k_eq_actual <- rep(0, num_x * num_y)
s_eq_actual <- rep(0, num_x * num_y)
k_eq_actual[domain] <- k_eq
for(id_boundary in 1:num_boundary){
  k_eq_actual[boundary_elements[[id_boundary]]] <- k_boundary[id_boundary]
  s_eq_actual[boundary_elements[[id_boundary]]] <- Stream_boundary[id_boundary]
}
image.plot(matrix(s_eq_actual, nrow = num_x), asp = 1)

### Set Initial Conditions
zeta_0 <- matrix(0, nrow = num_x, ncol = num_y)
zeta_0[domain] <- .1
for(id_y in 1:num_y){
  for(id_x in 1:num_x){
    # dist <- sqrt((id_y / (num_y + 1) - .5)^2 + (id_x / (num_x + 1) - .5)^2)
    # theta <- atan2((id_y / (num_y + 1) - .5), (id_x / (num_x + 1) - .5))
    # if(dist <= 0.4464286){
    #   zeta_0[id_x, id_y] <- zeta_0[id_x, id_y] + cos(2.5 * pi * (sqrt(dist) / .4464286)^.8) - .2 * sin(6 * theta)
    # }
    
    dist_1 <- sqrt((id_y / (num_y + 1) - .35)^2 + (id_x / (num_x + 1) - .35)^2)
    dist_2 <- sqrt((id_y / (num_y + 1) - .65)^2 + (id_x / (num_x + 1) - .65)^2)
    dist_3 <- sqrt((id_y / (num_y + 1) - .35)^2 + (id_x / (num_x + 1) - .65)^2)
    dist_4 <- sqrt((id_y / (num_y + 1) - .65)^2 + (id_x / (num_x + 1) - .35)^2)
    if(dist_1 <= 0.1){
      zeta_0[id_x, id_y] <- zeta_0[id_x, id_y] + 1 * (cos(pi * dist_1 / .1) + 1)
    }
    if(dist_2 <= 0.1){
      zeta_0[id_x, id_y] <- zeta_0[id_x, id_y] + 1 * (cos(pi * dist_2 / .1) + 1)
    }
    if(dist_3 <= 0.1){
      zeta_0[id_x, id_y] <- zeta_0[id_x, id_y] - 1 * (cos(pi * dist_3 / .1) + 1)
    }
    if(dist_4 <= 0.1){
      zeta_0[id_x, id_y] <- zeta_0[id_x, id_y] - 1 * (cos(pi * dist_4 / .1) + 1)
    }
  }
}
zeta_0 <- -zeta_0 / summation_factor
image.plot(zeta_0, asp = 1)
lines(c(0.4464286 * cos(0:100/100 * 2 * pi) + .5), c(0.4464286 * sin(0:100/100 * 2 * pi) + .5), lwd = 3)
zeta_int_0 <- rep(0, num_x * num_y)
zeta_int_0[domain] <- W_int %*% zeta_0[domain]
zeta_int_0 <- matrix(zeta_int_0, nrow <- num_x)
image.plot(zeta_int_0, asp = 1, xlim = c(2/53, 51/53), ylim = c(2/53, 51/53))
lines(c(0.4464286 * cos(0:100/100 * 2 * pi) + .5), c(0.4464286 * sin(0:100/100 * 2 * pi) + .5), lwd = 3)

k_0 <- rep(0, num_x * num_y)
k_0[domain] <- M %*% c(zeta_0)[domain] + k_eq
for(id_boundary in 1:num_boundary){
  k_0[boundary_elements[[id_boundary]]] <- k_boundary[id_boundary]
}
k_0 <- matrix(k_0, nrow <- num_x)
s_0 <- rep(0, num_x * num_y)
s_0[domain] <- W%*%k_0[domain] + BC_0 %*% bc
for(id_boundary in 1:num_boundary){
  s_0[boundary_elements[[id_boundary]]] <- Stream_boundary[id_boundary]
}
s_0 <- matrix(s_0, nrow <- num_x)
image.plot(k_0, asp = 1, xlim = c(2/53, 51/53), ylim = c(2/53, 51/53))
lines(c(0.4464286 * cos(0:100/100 * 2 * pi) + .5), c(0.4464286 * sin(0:100/100 * 2 * pi) + .5), lwd = 3)
image.plot(s_0, asp = 1, xlim = c(2/53, 51/53), ylim = c(2/53, 51/53))
lines(c(0.4464286 * cos(0:100/100 * 2 * pi) + .5), c(0.4464286 * sin(0:100/100 * 2 * pi) + .5), lwd = 3)

### Run Dynamic Loop
dt <- .5 * 10^(-5)
zeta <- zeta_0
k <- k_0
s <- s_0
u <- matrix(0, nrow = num_x, ncol = num_y)
v <- matrix(0, nrow = num_x, ncol = num_y)

image.plot(zeta_0, asp = 1, xlim = c(2/53, 51/53), ylim = c(2/53, 51/53))
lines(c(0.4464286 * cos(0:100/100 * 2 * pi) + .5), c(0.4464286 * sin(0:100/100 * 2 * pi) + .5), lwd = 3)

for(Time in 1:2000){
  ## Stream Function
  k[domain] <- M %*% c(zeta)[domain] + k_eq
  s[domain] <- W%*%k[domain] + BC_0 %*% bc
  
  jpeg(filename = paste("ve_stream_poly_", Time, ".jpeg", sep = ""), width = 1152 * 2, height = 1080 * 2,
       pointsize = 12, quality = 100, bg = "white", res = 300)
  image.plot(s, asp = 1, xlim = c(3/num_x, 1 - 3/num_x), ylim = c(3/num_y, 1 - 3/num_y), main = paste("Time = ", Time * dt, " seconds", sep = ""))
  lines(c(0.4464286 * cos(0:100/100 * 2 * pi) + .5), c(0.4464286 * sin(0:100/100 * 2 * pi) + .5), lwd = 3)
  dev.off()
  
  ## Velocity
  for(id_y in 2:(num_y - 1)){
    for(id_x in 2:(num_x - 1)){
      u[id_x, id_y] <- (s[id_x, id_y + 1] - s[id_x, id_y - 1]) * (num_y + 1) * w_diff_value  ## Seems wrong should be subtracting k[,]
      v[id_x, id_y] <- -(s[id_x + 1, id_y] - s[id_x - 1, id_y]) * (num_x + 1) * w_diff_value
    }
  }
  
  ## Vorticity Transport
  zeta_int_temp <- matrix(0, nrow = num_x, ncol = num_y)
  for(id_y in 4:(num_y - 3)){
    for(id_x in 4:(num_x - 3)){
      # Scan all the reachable neighbors
      for(neighbor_id_x in neighbor_dir_x[2:(length(neighbor_dir_x) - 1)]){
        for(neighbor_id_y in neighbor_dir_y[2:(length(neighbor_dir_y) - 1)]){
          pos_diff <- c(neighbor_id_x, neighbor_id_y)
          pos_diff[1] <- pos_diff[1] + dt * (num_x + 1) * u[id_x + neighbor_id_x, id_y + neighbor_id_y]
          pos_diff[2] <- pos_diff[2] + dt * (num_y + 1) * v[id_x + neighbor_id_x, id_y + neighbor_id_y]
          
          x_int_lim <- c(max(-2, pos_diff[1] - .5), min(2, pos_diff[1] + .5))
          y_int_lim <- c(max(-2, pos_diff[2] - .5), min(2, pos_diff[2] + .5))
          
          zeta_int_temp[id_x, id_y] <- zeta_int_temp[id_x, id_y] + zeta[id_x + neighbor_id_x, id_y + neighbor_id_y] * Integrate_function(x_int_lim) * Integrate_function(y_int_lim)
        }
      }
    }
  }
  
  zeta_int <- rep(0, num_x * num_y)
  zeta_int[domain] <- zeta_int_temp[domain]
  zeta_int <- matrix(zeta_int, nrow <- num_x)
  
  zeta <- rep(0, num_x * num_y)
  zeta[domain] <- W_int_inv %*% (zeta_int[domain] - BC_int %*% bc)
  zeta <- matrix(zeta, nrow <- num_x)
  
  zeta_smooth <- rep(0, num_x * num_y)
  zeta_smooth[domain] <- W%*%zeta[domain]
  zeta_smooth <- matrix(zeta_smooth, nrow <- num_x)
  
  jpeg(filename = paste("ve_vorticity_poly_", Time, ".jpeg", sep = ""), width = 1152 * 2, height = 1080 * 2,
       pointsize = 12, quality = 100, bg = "white", res = 300)
  image.plot(zeta_smooth, asp = 1, xlim = c(3/num_x, 1 - 3/num_x), ylim = c(3/num_y, 1 - 3/num_y), main = paste("Time = ", Time * dt, " seconds", sep = ""))
  lines(c(0.4464286 * cos(0:100/100 * 2 * pi) + .5), c(0.4464286 * sin(0:100/100 * 2 * pi) + .5), lwd = 3)
  dev.off()
  print(Time)
}