#' Liang-Kleeman causality measure for 1D time series and 3D array
#'
#' @author Initial code provided by Ángel Muñoz (munoz@ictp.it), modified by Javier Corvillo
#' @description This function calculates the Liang-Kleeman causality measure for a 1D time series
#' and a 3D array. Negative values of the causality measure indicate a causal but stable
#' relationship from ds1 to ds2, while positive values indicate a causal and non-stable
#' relationship. The causality measure can be normalised to the range [0, 1] and significant
#' values can be identified as well.
#'
#' @param ds1 1D array with a named "time" dimension

#' @param ds2 3D array with a named "time" dimension as first dimension, spatial dims as 2nd and 3rd

#' @param dt Number of time steps between observations in the time series for ds1 and ds2

#' @param normalise boolean indicating whether the causality outputs should be normalised

#' @param sig significance level for the causality measure (90, 95 or 99)

#' @return list containing the causality measure and significant values (2D array)

#' @examples
#' # Example 1: Simple causality test
#' ds1 <- rnorm(100)
#' ds2 <- array(rnorm(100 * 10 * 10), dim = c(100, 10, 10))
#' causality <- Causality1D3D(ds1, ds2, dt = 1, normalise = FALSE, sig = 95)


Causality1D3D <- function(ds1, ds2, dt = 1, normalise = FALSE, sig = 95) {
  # Input validation and dimension handling
  if (is.vector(ds1)) {
    n <- length(ds1)
  } else {
    n <- dim(ds1)[1] # time dimension
  }

  if (is.vector(ds2) || length(dim(ds2)) != 3) {
    stop("ds2 should be a 3D array with dimensions [time, spatial1, spatial2]")
  }

  # Check that time dimensions match
  if (dim(ds2)[1] != n) {
    stop("Time dimensions of ds1 and ds2 must match")
  }

  # Get spatial dimensions of ds2
  spatial_dims <- dim(ds2)[-1]

  # Calculate time derivative properly - use central differences where possible
  ds1_diff <- rep(NA, n)
  ds1_diff[1] <- ds1[2] - ds1[1] # forward difference for first point
  if (n > 2) {
    ds1_diff[2:(n - 1)] <- (ds1[3:n] - ds1[1:(n - 2)]) / 2 # central differences
  }
  ds1_diff[n] <- ds1[n] - ds1[n - 1] # backward difference for last point

  # Helper function to calculate covariance between 1D time series and 3D array
  cov_1d_3d <- function(ts1d, arr3d) {
    result <- array(NA, dim = dim(arr3d)[-1])
    for (i in seq_len(dim(arr3d)[2])) {
      for (j in seq_len(dim(arr3d)[3])) {
        result[i, j] <- cov(ts1d, arr3d[, i, j], use = "complete.obs")
      }
    }
    return(result)
  }

  # Helper function to calculate variance within 3D array along time dimension
  var_3d <- function(arr3d) {
    result <- array(NA, dim = dim(arr3d)[-1])
    for (i in seq_len(dim(arr3d)[2])) {
      for (j in seq_len(dim(arr3d)[3])) {
        result[i, j] <- var(arr3d[, i, j], na.rm = TRUE)
      }
    }
    return(result)
  }

  # Calculate covariances
  cov_ds1_ds1 <- var(ds1, na.rm = TRUE) # scalar
  cov_ds1_ds2 <- cov_1d_3d(ds1, ds2) # 2D array (spatial dimensions)
  cov_ds2_ds1 <- cov_ds1_ds2 # same as above (symmetric)
  cov_ds2_ds2 <- var_3d(ds2) # 2D array (spatial dimensions)
  cov_ds1_ds1diff <- cov(ds1, ds1_diff, use = "complete.obs") # scalar
  cov_ds2_ds1diff <- cov_1d_3d(ds1_diff, ds2) # 2D array

  # Calculate numerator and denominator (element-wise for spatial dimensions)
  numer <- cov_ds1_ds1 * cov_ds1_ds2 * cov_ds2_ds1diff - (cov_ds1_ds2^2) * cov_ds1_ds1diff
  denom <- (cov_ds1_ds1^2) * cov_ds2_ds2 - cov_ds1_ds1 * (cov_ds1_ds2^2)

  # Compute cross-variance matrix determinant for each spatial point
  detc <- array(NA, dim = spatial_dims)
  for (i in seq_len(spatial_dims[1])) {
    for (j in seq_len(spatial_dims[2])) {
      cross_var_matrix <- matrix(
        c(
          cov_ds1_ds1, cov_ds1_ds2[i, j],
          cov_ds2_ds1[i, j], cov_ds2_ds2[i, j]
        ),
        nrow = 2, ncol = 2
      )
      detc[i, j] <- det(cross_var_matrix)
    }
  }

  # Calculate p and q (equations 16-18 in Liang 2015)
  p <- (cov_ds2_ds2 * cov_ds1_ds1diff - cov_ds1_ds2 * cov_ds2_ds1diff) / detc
  q <- (cov_ds1_ds1 * cov_ds2_ds1diff - cov_ds1_ds2 * cov_ds1_ds1diff) / detc

  # Information flow (equation 1 in Liang 2015) - element-wise for spatial dimensions
  infflow <- numer / denom

  # Statistical significance testing via Fisher's Information Matrix
  # This needs to be done for each spatial point
  infflow_sig <- array(NA, dim = spatial_dims)

  for (i in seq_len(spatial_dims[1])) {
    for (j in seq_len(spatial_dims[2])) {
      # Extract time series for this spatial point
      ds2_point <- ds2[, i, j]

      # Skip if there are issues with the data at this point
      if (is.na(detc[i, j]) || detc[i, j] == 0 || is.na(infflow[i, j])) {
        infflow_sig[i, j] <- NA
        next
      }

      # Compute additional terms for this point
      f1 <- mean(ds1_diff, na.rm = TRUE) - p[i, j] * mean(ds1) - q[i, j] * mean(ds2_point)
      r1 <- ds1_diff - (f1 + p[i, j] * ds1 + q[i, j] * ds2_point)
      q1 <- sum(r1^2, na.rm = TRUE)
      b1 <- sqrt(q1 * dt / n)

      if (b1 > 0 && is.finite(b1)) {
        # Create Fisher's Information Matrix (n*I) for this point
        ni <- matrix(0, nrow = 4, ncol = 4)

        # Fill diagonal elements
        ni[1, 1] <- n * dt / (b1^2)
        ni[2, 2] <- dt / (b1^2) * sum(ds1^2)
        ni[3, 3] <- dt / (b1^2) * sum(ds2_point^2)
        ni[4, 4] <- 3 * dt / (b1^4) * sum(r1^2, na.rm = TRUE) - n / (b1^2)

        # Fill off-diagonal elements
        ni[1, 2] <- ni[2, 1] <- dt / (b1^2) * sum(ds1)
        ni[1, 3] <- ni[3, 1] <- dt / (b1^2) * sum(ds2_point)
        ni[1, 4] <- ni[4, 1] <- 2 * dt / (b1^3) * sum(r1, na.rm = TRUE)
        ni[2, 3] <- ni[3, 2] <- dt / (b1^2) * sum(ds1 * ds2_point)
        ni[2, 4] <- ni[4, 2] <- 2 * dt / (b1^3) * sum(r1 * ds1, na.rm = TRUE)
        ni[3, 4] <- ni[4, 3] <- 2 * dt / (b1^3) * sum(r1 * ds2_point, na.rm = TRUE)

        # Invert Fisher Information Matrix to find variance
        tryCatch(
          {
            invni <- solve(ni)
            var_a12 <- abs(invni[3, 3]) # key component for variance

            # Variance (uncertainty in the information flow)
            if (cov_ds1_ds1 != 0) {
              varinfflow <- var_a12 * (cov_ds1_ds2[i, j] / cov_ds1_ds1)^2

              # Determine significance threshold based on p-value
              if (sig == 90) {
                ss <- sqrt(varinfflow) * 1.65
              } else if (sig == 95) {
                ss <- sqrt(varinfflow) * 1.96
              } else if (sig == 99) {
                ss <- sqrt(varinfflow) * 2.56
              } else {
                stop("Only 90, 95 and 99 significance levels are supported")
              }

              # Apply significance masking for this point
              infflow_sig[i, j] <- ifelse(abs(infflow[i, j]) > ss, infflow[i, j], NA)
            } else {
              infflow_sig[i, j] <- NA
            }
          },
          error = function(e) {
            # If matrix inversion fails, set to NA
            infflow_sig[i, j] <- NA
          }
        )
      } else {
        infflow_sig[i, j] <- NA
      }
    }
  }

  # Normalization (optional)
  if (normalise == TRUE) {
    dh1noisedt <- array(NA, dim = spatial_dims)

    for (i in 1:spatial_dims[1]) {
      for (j in 1:spatial_dims[2]) {
        ds2_point <- ds2[, i, j]
        if (cov_ds1_ds1 != 0) {
          dh1noisedt[i, j] <- dt / (2 * cov_ds1_ds1) * (
            var(ds1_diff, na.rm = TRUE) +
              p[i, j]^2 * cov_ds1_ds1 +
              q[i, j]^2 * cov_ds2_ds2[i, j] -
              2 * p[i, j] * cov(ds1_diff, ds1, use = "complete.obs") -
              2 * q[i, j] * cov(ds1_diff, ds2_point, use = "complete.obs") +
              2 * p[i, j] * q[i, j] * cov_ds1_ds2[i, j]
          )
        } else {
          dh1noisedt[i, j] <- NA
        }
      }
    }

    norm <- abs(infflow) + abs(p) + abs(dh1noisedt) # eq 20 in Liang (2015)
    # Avoid division by zero
    norm[norm == 0] <- NA
    infflow <- infflow / norm
    infflow_sig <- infflow_sig / norm

    # Add names/attributes for identification
    attr(infflow, "name") <- "Liang-Kleeman causality\n(normalised)"
    attr(infflow_sig, "name") <- "Liang-Kleeman causality\n(normalised, significant)"
  } else {
    attr(infflow, "name") <- "Liang-Kleeman causality\n(not normalised)"
    attr(infflow_sig, "name") <- "Liang-Kleeman causality\n(not normalised, significant)"
  }

  # Return both the full causality measure and the significant values
  return(list(infflow = infflow, infflow_sig = infflow_sig))
}
