library(s2dverification)
library(multiApply)
anom2regime <- function(ref, target, method = "distance", lat) {
  posdim <- which(names(dim(ref)) == "nclust")
  poslat <- which(names(dim(ref)) == "lat")
  poslon <- which(names(dim(ref)) == "lon")
  print(dim(ref))
  print(dim(target))
  nclust <- dim(ref)[posdim]
  
  if (all(dim(ref)[-posdim] != dim(target))) {
    stop("The target should have the same dimensions [lat,lon] that
         the reference ")
  }
  if (is.null(names(dim(ref))) | is.null(names(dim(target)))) {
    stop(
      "The arrays should include dimensions names ref[nclust,lat,lon]
      and target [lat,lon]"
    )
  }
  
  if (length(lat) != dim(ref)[poslat]) {
    stop("latitudes do not match with the maps")
  }
  
  # This dimensions are reorganized
  ref <- aperm(ref, c(posdim, poslat, poslon))
  target <-
    aperm(target, c(which(names(dim(
      target
    )) == "lat"), which(names(dim(
      target
    )) == "lon")))
  
  # weights are defined
  latWeights <- InsertDim(sqrt(cos(lat * pi / 180)), 2, dim(ref)[3]) #nolint
  
  
  rmsdiff <- function(x, y) {
    dims <- dim(x)
    ndims <- length(dims)
    if (ndims != 2 | ndims != length(dim(y))) {
      stop("x and y should be maps")
    }
    map_diff <- NA * x
    for (i in 1 : dims[1]) {
      for (j in 1 : dims[2]) {
        map_diff[i, j] <- (x[i, j] - y[i, j]) ^ 2
      }
    }
    rmsdiff <- sqrt(mean(map_diff))
    return(rmsdiff)
  }
  
  if (method == "ACC") {
    corr <- rep(NA, nclust)
    for (i in 1:nclust) {
      corr[i] <-
        ACC(InsertDim(InsertDim( #nolint
          InsertDim(ref[i, , ] * latWeights, 1, 1), 2, 1 #nolint
        ), 3, 1),
        InsertDim(InsertDim( #nolint
          InsertDim(target * latWeights, 1, 1), 2, 1 #nolint
        ), 3, 1))$ACC[2]
    }
    assign <- which(corr == max(corr))
  }
  
  if (method == "distance") {
    rms <- rep(NA, nclust)
    for (i in 1 : nclust) {
      rms[i] <- rmsdiff(ref[i, , ] * latWeights, target * latWeights)#nolint
    }
    assign <- which(rms == min(rms))
  }
  
  return(assign)
}

ref <- rnorm(2 * 3 * 5)
dim(ref) <- c(nclust = 2, lat = 3, lon = 5)
target <- 1 : (3 * 5)
dim(target) <- c(lat = 3, lon = 5)
res <-  anom2regime(ref, target, method = "distance", lat = 1 : 3)
str(res) # funciona y da 1

dim(ref) <- c(lat = 3, lon = 5, nclust = 2)
res <-  anom2regime(ref, target, method = "distance", lat = 1 : 3)
str(res)  # funciona y da 1

load("/home/Earth/nperez/git/ESMValTool/output_dir/recipe_modes_of_variability_wp4_20190131_135233/plots/weather_regime/main/lat.RData")
load("/home/Earth/nperez/git/ESMValTool/output_dir/recipe_modes_of_variability_wp4_20190131_135233/plots/weather_regime/main/lon.RData")
load("/home/Earth/nperez/git/ESMValTool/output_dir/recipe_modes_of_variability_wp4_20190131_135233/plots/weather_regime/main/anom_exp.RData")
PlotEquiMap(anom_exp[1, 1, 1, 1, , ], lon = lon, lat = lat, filled.continents = FALSE)
load("/home/Earth/nperez/git/ESMValTool/output_dir/recipe_modes_of_variability_wp4_20190131_135233/plots/weather_regime/main/ref.RData")
PlotEquiMap(reference[1, , ], lon = lon, lat = lat, filled.continents = FALSE)

res <- anom2regime(ref = reference, target = anom_exp[1, 1, 1, 1, , ],  method = "distance", lat = lat)

str(res) # dice que ese mapa pertenece al cluster 2

RegimesAssign <- function(var_ano, ref_maps, lats, #nolint
                          method = "distance") {
  posdim <- which(names(dim(ref_maps)) == "nclust")
  poslat <- which(names(dim(ref_maps)) == "lat")
  poslon <- which(names(dim(ref_maps)) == "lon")
  poslat_ano <- which(names(dim(var_ano)) == "lat")
  poslon_ano <- which(names(dim(var_ano)) == "lon")
  
  nclust <- dim(ref_maps)[posdim]
  nlat <- dim(ref_maps)[poslat]
  nlon <- dim(ref_maps)[poslon]
  
  
  if (is.null(names(dim(ref_maps))) | is.null(names(dim(var_ano)))) {
    stop(
      "The arrays should include dimensions names ref[nclust,lat,lon]
      and target [lat,lon]"
    )
  }
  
  if (length(lats) != dim(ref_maps)[poslat]) {
    stop("latitudes do not match with the maps")
  }
  print("HERE")
  print(dim(var_ano))
  print(dim(ref_maps))
  assign <-
    Apply(
      data = list(target = var_ano),
      margins = c( (1 : length(dim(var_ano)) )[-c(poslat_ano, poslon_ano)]),
      fun = "anom2regime",
      ref = ref_maps,
      lat = lats,
      method = method
    )
  if (poslat_ano < poslon_ano) {
    dim_order = c(nlat, nlon)
  } else {
    dim_order = c(nlon, nlat)
  }
  
  anom_array <-
    array(var_ano, dim = c(prod(dim(var_ano)[-c(poslat_ano, poslon_ano)]),
                           dim_order))
                           #nlat, nlon))
  rm(var_ano)
  index <- as.vector(assign$output1)
  recon <- Composite(var = aperm(anom_array, c(3, 2, 1)), occ = index)
  freqs <- rep(NA, nclust)
  for (n in 1 : nclust) {
    freqs[n] <- (length(which(index == n)) / length(index)) * 100
  }
  output <-
    list(
      composite = recon$composite,
      pvalue = recon$pvalue,
      cluster = assign$output1,
      frequency = freqs
    )
  return(output)
  }


res <- RegimesAssign(var_ano = anom_exp, ref_maps = reference, lats = lat, #nolint
                          method = "distance")
str(res)
PlotEquiMap(res$composite[ , , 1], lon = lon, lat = lat)
