deep_causality_algorithms 0.2.12

/*
 * SPDX-License-Identifier: MIT
 * Copyright (c) 2023 - 2026. The DeepCausality Authors and Contributors. All Rights Reserved.
 */
use crate::causal_discovery::surd::surd_utils;
use crate::causal_discovery::surd::surd_utils::surd_utils_cdl;
use crate::causal_discovery::surd::{MaxOrder, SurdResult};
use deep_causality_tensor::{CausalTensor, CausalTensorError, Tensor};
use std::collections::HashMap;

#[cfg(feature = "parallel")]
use rayon::prelude::*;

/// Decomposes mutual information into its Synergistic, Unique, and Redundant components
/// for each state of the target variable, based on the SURD-states algorithm,
/// specifically designed to handle input tensors with `Option<f64>` values.
///
/// This function adapts the core SURD-states algorithm to gracefully manage missing or
/// undefined probability values, represented by `None` in the input `CausalTensor<Option<f64>>`.
/// The handling of `None` values is crucial for robust causal discovery in real-world datasets
/// where data incompleteness is common.
///
/// The algorithm follows these main steps:
/// 1.  **Pre-processing**: The input probability distribution is validated and normalized.
/// 2.  **Information Leak Calculation**: The causality leak, representing the influence of unobserved
///     variables, is calculated as `H(Q_f|Q) / H(Q_f)`.
/// 3.  **Specific Mutual Information**: For every combination of source variables, the specific mutual
///     information `i(Q_f=t; Q_i)` is calculated for each target state `t`.
/// 4.  **Decomposition Loop**: For each target state `t`, the specific information values are sorted.
///     This ordering is crucial for the decomposition. An information-theoretic exclusion principle
///     is applied to prevent double-counting.
/// 5.  **State-Dependent Slice Calculation**: The sorted information values are differenced, and these
///     increments are used to calculate the state-dependent causal maps for redundant, unique, and
///     synergistic components, separating positive (causal) and negative (non-causal) contributions.
/// 6.  **Aggregation**: The state-dependent maps are stacked, and the aggregate SURD values are
///     calculated by summing the increments over all target states.
///
/// # None Value Handling
/// The presence of `Option<f64>` in the input `CausalTensor` necessitates a specific strategy
/// for handling `None` values. The chosen approach aims for "no impact" or "ignore for this
/// specific calculation" for `None` values, aligning with a pairwise deletion strategy. This is
/// crucial for datasets where `None` signifies "not measured" or "unobserved" as in medical
/// data with irregular or on-demand readings.
///
/// -   **Normalization**: When calculating the `total_sum` for normalizing the raw probabilities,
///     `None` values in the input `p_raw` are **filtered out**. This ensures that the sum of
///     observed probabilities is used for normalization, and `None`s do not artificially inflate
///     the total sum.
///
/// -   **Marginalization (`sum_axes_option_f64`)**: This utility function is central to handling `None`s.
///     When summing over axes to compute marginal distributions, `None` values encountered in the
///     `CausalTensor<Option<f64>>` are **ignored** (not treated as `0.0`). The function returns a
///     `CausalTensor<Option<f64>>`. If all values contributing to a particular output element are `None`,
///     then that output element will also be `None`.
///     *   **Rationale**: This approach ensures that `None` values do not artificially influence
///         marginal probabilities by being treated as `0.0`. It allows the algorithm to focus
///         on observed data, which is crucial for datasets where `None` means "not measured"
///         rather than "zero probability." This aligns with a pairwise deletion strategy.
///
/// -   **Information-Theoretic Operations (`entropy_nvars_cdl`, `cond_entropy_cdl`)**: These functions
///     now operate on `CausalTensor<Option<f64>>` inputs. They calculate entropy and conditional
///     entropy by considering only the `Some` values in the marginal distributions, effectively
///     ignoring `None`s in their calculations.
///
/// -   **Arithmetic Operations (`safe_div_cdl`, `surd_log2_cdl`, `sub_cdl`, `mul_cdl`)**: These
///     utility functions operate on `CausalTensor<Option<f64>>` and strictly propagate `None` values.
///     If any operand for an element-wise operation is `None`, the result for that specific element
///     will also be `None`.
///     *   **Rationale**: This strict propagation ensures that any uncertainty or lack of data
///         is carried forward through the calculations, preventing spurious information from being
///         generated from missing values.
///
/// -   **Final Aggregation and Output**: The intermediate state-dependent slices (`causal_slice`,
///     `non_causal_slice`) are computed as `CausalTensor<Option<f64>>`. Only at the very last step,
///     before being stored in the `SurdResult` (which expects `f64` values), are these `Option<f64>`
///     tensors converted to `CausalTensor<f64>`. During this final conversion, any remaining `None`
///     values are explicitly **treated as `0.0`**.
///     *   **Rationale**: This ensures that the final output format is consistent with `SurdResult`
///         while maintaining the "no impact" propagation of `None` throughout the core calculations.
///         The `0.0` at the end signifies that no information could be derived from those missing values.
///
/// # Performance Improvements
/// To make the algorithm practical for high-dimensional data, several performance improvements have been implemented:
/// 1.  **Algorithmic Capping**: The `max_order` parameter allows limiting the analysis to a tractable number
///     of interactions (e.g., pairwise, `k=2`), changing the complexity from exponential `O(2^N)` to polynomial `O(N^k)`.
///     This is the most most significant optimization for datasets with many variables.
/// 2.  **Parallel Execution**: When compiled with the `parallel` feature flag, the main decomposition loop runs in parallel
///     across all available CPU cores using the `rayon` crate. This provides a near-linear speedup for datasets
///     with a large number of target states.
/// 3.  **Memory Optimization**: Unnecessary cloning of large data tensors within the hot loop has been eliminated
///     by using references, reducing memory pressure and improving cache performance.
///
/// # Arguments
/// * `p_raw` - A `CausalTensor<Option<f64>>` representing the joint probability distribution.
///   **Crucially, this must be a joint probability distribution of discrete, binned data.**
///   The first dimension (axis 0) must correspond to the target variable.
/// * `max_order` - An enum specifying the maximum order of interactions to compute.
///
/// # Returns
/// A `Result` containing a `SurdResult` struct with the full decomposition, including
/// separate causal (positive) and non-causal (negative) state-dependent maps.
///
/// # Examples
///
/// ```
/// use deep_causality_algorithms::causal_discovery::surd::{surd_states_cdl, MaxOrder};
/// use deep_causality_tensor::CausalTensor;
///
/// // Create a joint probability distribution for a target and 2 source variables with missing data.
/// // Shape: [target_states, source1_states, source2_states] = [2, 2, 2]
/// let data = vec![
///     Some(0.1), Some(0.2), // P(T=0, S1=0, S2=0), P(T=0, S1=0, S2=1)
///     None,      Some(0.2), // P(T=0, S1=1, S2=0) is missing, P(T=0, S1=1, S2=1)
///     Some(0.3), None,      // P(T=1, S1=0, S2=0), P(T=1, S1=0, S2=1) is missing
///     Some(0.1), Some(0.1), // P(T=1, S1=1, S2=0), P(T=1, S1=1, S2=1)
/// ];
/// let p_raw = CausalTensor::new(data, vec![2, 2, 2]).unwrap();
///
/// // Perform a full decomposition (k=N=2)
/// let full_result = surd_states_cdl(&p_raw, MaxOrder::Max).unwrap();
///
/// // Check some results (synergistic info for pair {1, 2})
/// assert!(full_result.synergistic_info().get(&vec![1, 2]).is_some());
///
/// // Perform a partial, pairwise decomposition (k=2)
/// // This is equivalent to MaxOrder::Max for N=2, but demonstrates the API.
/// let partial_result = surd_states_cdl(&p_raw, MaxOrder::Min).unwrap();
/// assert!(partial_result.synergistic_info().get(&vec![1, 2]).is_some());
/// ```
pub fn surd_states_cdl(
    p_raw: &CausalTensor<Option<f64>>,
    max_order: MaxOrder,
) -> Result<SurdResult<f64>, CausalTensorError> {
    if p_raw.is_empty() {
        return Err(CausalTensorError::EmptyTensor);
    }

    // --- 1. Pre-processing: Normalize the probability distribution ---
    // Sum only Some values for normalization
    let total_sum: f64 = p_raw.as_slice().iter().filter_map(|&x| x).sum();
    if total_sum.abs() < 1e-14 {
        return Err(CausalTensorError::InvalidOperation);
    }

    let p_data: Vec<Option<f64>> = p_raw
        .as_slice()
        .iter()
        .map(|&x| x.map(|val| val / total_sum))
        .collect();
    let p = CausalTensor::new(p_data, p_raw.shape().to_vec())?;

    let n_total_dims = p.num_dim();
    let n_vars = n_total_dims - 1;
    let n_target_states = p.shape()[0];
    let agent_indices: Vec<usize> = (1..n_total_dims).collect();
    let k = max_order.get_k_max_order(n_vars)?;

    // --- 2. Calculate Information Leak ---
    let h = surd_utils_cdl::entropy_nvars_cdl(&p, &[0])?;
    let hc = surd_utils_cdl::cond_entropy_cdl(&p, &[0], &agent_indices)?;

    let info_leak = if h > 1e-14 {
        (hc / h).clamp(0.0, 1.0)
    } else {
        0.0
    };

    // --- 3. Compute Specific and Mutual Information ---
    let mut combs: Vec<Vec<usize>> = Vec::new();
    for i in 1..=k.min(n_vars) {
        let combinations_for_i = surd_utils::combinations(&agent_indices, i);
        combs.extend(combinations_for_i);
    }

    // p_s is CausalTensor<Option<f64>> now
    let p_s = surd_utils_cdl::sum_axes_option_f64(&p, &agent_indices)?;
    let mut is_map: HashMap<Vec<usize>, CausalTensor<Option<f64>>> = HashMap::new();
    for j_comb in &combs {
        let noj: Vec<usize> = agent_indices
            .iter()
            .filter(|&ax| !j_comb.contains(ax))
            .cloned()
            .collect();

        // p(T, j) - marginal distribution of target and the current combination of sources
        let p_as = if noj.is_empty() {
            p.clone()
        } else {
            surd_utils_cdl::sum_axes_option_f64(&p, &noj)?
        };

        // p(j) - marginal distribution of the current combination of sources
        let p_sources = surd_utils_cdl::sum_axes_option_f64(&p, &[0])?;
        let noj_mapped: Vec<usize> = noj.iter().map(|ax| ax - 1).collect();

        let p_j = if noj_mapped.is_empty() {
            p_sources.clone()
        } else {
            surd_utils_cdl::sum_axes_option_f64(&p_sources, &noj_mapped)?
        };

        let p_j_broadcasted = surd_utils_cdl::broadcast_to_cdl(&p_j, p_as.shape())?;
        let p_s_a = surd_utils_cdl::safe_div_cdl(&p_as, &p_j_broadcasted)?;

        let p_s_broadcasted_for_pas = surd_utils_cdl::broadcast_to_cdl(&p_s, p_as.shape())?;
        let p_a_s = surd_utils_cdl::safe_div_cdl(&p_as, &p_s_broadcasted_for_pas)?;

        // Broadcast p_s to the shape of p_s_a for the log difference calculation.
        let p_s_broadcasted_for_log = surd_utils_cdl::broadcast_to_cdl(&p_s, p_s_a.shape())?;

        let log_diff = surd_utils_cdl::sub_cdl(
            &surd_utils_cdl::surd_log2_cdl(&p_s_a)?,
            &surd_utils_cdl::surd_log2_cdl(&p_s_broadcasted_for_log)?,
        )?;
        let specific_info_map = surd_utils_cdl::mul_cdl(&p_a_s, &log_diff)?;

        let dims_to_sum_for_j_comb: Vec<usize> = (1..=j_comb.len()).collect();
        let sum_axes =
            surd_utils_cdl::sum_axes_option_f64(&specific_info_map, &dims_to_sum_for_j_comb)?;

        let ravel = sum_axes.ravel();
        is_map.insert(j_comb.clone(), ravel);
    }

    let mi: HashMap<Vec<usize>, f64> = is_map
        .iter()
        .map(|(k, v)| {
            let multiplied = surd_utils_cdl::mul_cdl(v, &p_s)?;
            Ok((
                k.clone(),
                multiplied.as_slice().iter().filter_map(|&x| x).sum(),
            ))
        })
        .collect::<Result<_, CausalTensorError>>()?;

    // --- 4. Main Decomposition Loop ---
    let results_per_target: Vec<_> = {
        #[cfg(feature = "parallel")]
        {
            (0..n_target_states)
                .into_par_iter()
                .map(|t| {
                    analyze_single_target_state_cdl(
                        t, n_vars, &combs, &is_map, &p_s,
                        &p, // Pass CausalTensor<Option<f64>>
                        k,
                    )
                })
                .collect::<Result<Vec<_>, _>>()?
        }
        #[cfg(not(feature = "parallel"))]
        {
            (0..n_target_states)
                .map(|t| {
                    analyze_single_target_state_cdl(
                        t, n_vars, &combs, &is_map, &p_s,
                        &p, // Pass CausalTensor<Option<f64>>
                        k,
                    )
                })
                .collect::<Result<Vec<_>, _>>()?
        }
    };

    // --- 5. Merge results from all target states ---
    let mut i_r = HashMap::new();
    let mut i_s = HashMap::new();
    let mut temp_causal_rd_states: HashMap<Vec<usize>, Vec<CausalTensor<f64>>> = HashMap::new();
    let mut temp_causal_un_states: HashMap<Vec<usize>, Vec<CausalTensor<f64>>> = HashMap::new();
    let mut temp_causal_sy_states: HashMap<Vec<usize>, Vec<CausalTensor<f64>>> = HashMap::new();
    let mut temp_non_causal_rd_states: HashMap<Vec<usize>, Vec<CausalTensor<f64>>> = HashMap::new();
    let mut temp_non_causal_un_states: HashMap<Vec<usize>, Vec<CausalTensor<f64>>> = HashMap::new();
    let mut temp_non_causal_sy_states: HashMap<Vec<usize>, Vec<CausalTensor<f64>>> = HashMap::new();

    for result in results_per_target {
        for (k, v) in result.i_r {
            *i_r.entry(k).or_insert(0.0) += v;
        }
        for (k, v) in result.i_s {
            *i_s.entry(k).or_insert(0.0) += v;
        }
        for (k, v) in result.causal_rd_states {
            temp_causal_rd_states.entry(k).or_default().push(v);
        }
        for (k, v) in result.causal_un_states {
            temp_causal_un_states.entry(k).or_default().push(v);
        }
        for (k, v) in result.causal_sy_states {
            temp_causal_sy_states.entry(k).or_default().push(v);
        }
        for (k, v) in result.non_causal_rd_states {
            temp_non_causal_rd_states.entry(k).or_default().push(v);
        }
        for (k, v) in result.non_causal_un_states {
            temp_non_causal_un_states.entry(k).or_default().push(v);
        }
        for (k, v) in result.non_causal_sy_states {
            temp_non_causal_sy_states.entry(k).or_default().push(v);
        }
    }

    // --- 6. Finalize State-Dependent Maps by Stacking ---
    let causal_redundant_states = temp_causal_rd_states
        .into_iter()
        .map(|(k, slices)| Ok((k, CausalTensor::stack(&slices, 0)?)))
        .collect::<Result<_, _>>()?;
    let causal_unique_states = temp_causal_un_states
        .into_iter()
        .map(|(k, slices)| Ok((k, CausalTensor::stack(&slices, 0)?)))
        .collect::<Result<_, _>>()?;
    let causal_synergistic_states = temp_causal_sy_states
        .into_iter()
        .map(|(k, slices)| Ok((k, CausalTensor::stack(&slices, 0)?)))
        .collect::<Result<_, _>>()?;
    let non_causal_redundant_states = temp_non_causal_rd_states
        .into_iter()
        .map(|(k, slices)| Ok((k, CausalTensor::stack(&slices, 0)?)))
        .collect::<Result<_, _>>()?;
    let non_causal_unique_states = temp_non_causal_un_states
        .into_iter()
        .map(|(k, slices)| Ok((k, CausalTensor::stack(&slices, 0)?)))
        .collect::<Result<_, _>>()?;
    let non_causal_synergistic_states = temp_non_causal_sy_states
        .into_iter()
        .map(|(k, slices)| Ok((k, CausalTensor::stack(&slices, 0)?)))
        .collect::<Result<_, _>>()?;

    Ok(SurdResult::new(
        i_r,
        i_s,
        mi,
        info_leak,
        causal_redundant_states,
        causal_unique_states,
        causal_synergistic_states,
        non_causal_redundant_states,
        non_causal_unique_states,
        non_causal_synergistic_states,
    ))
}

/// A private struct to hold the results of analyzing a single target state.
struct PerTargetStateResults {
    i_r: HashMap<Vec<usize>, f64>,
    i_s: HashMap<Vec<usize>, f64>,
    causal_rd_states: HashMap<Vec<usize>, CausalTensor<f64>>,
    causal_un_states: HashMap<Vec<usize>, CausalTensor<f64>>,
    causal_sy_states: HashMap<Vec<usize>, CausalTensor<f64>>,
    non_causal_rd_states: HashMap<Vec<usize>, CausalTensor<f64>>,
    non_causal_un_states: HashMap<Vec<usize>, CausalTensor<f64>>,
    non_causal_sy_states: HashMap<Vec<usize>, CausalTensor<f64>>,
}

fn analyze_single_target_state_cdl(
    t: usize,
    n_vars: usize,
    combs: &[Vec<usize>],
    is_map: &HashMap<Vec<usize>, CausalTensor<Option<f64>>>,
    p_s: &CausalTensor<Option<f64>>,
    p: &CausalTensor<Option<f64>>,
    _k: usize,
) -> Result<PerTargetStateResults, CausalTensorError> {
    let mut i_r = HashMap::new();
    let mut i_s = HashMap::new();
    let mut causal_rd_states = HashMap::new();
    let mut causal_un_states = HashMap::new();
    let mut causal_sy_states = HashMap::new();
    let mut non_causal_rd_states = HashMap::new();
    let mut non_causal_un_states = HashMap::new();
    let mut non_causal_sy_states = HashMap::new();

    // Extract i1_values, filtering out None for sorting, but keeping track of original indices
    let mut i1_values_with_indices: Vec<(f64, usize)> = Vec::new();
    for (idx, c) in combs.iter().enumerate() {
        if let Some(val) = is_map[c].as_slice()[t] {
            i1_values_with_indices.push((val, idx));
        }
    }

    // Sort based on the Some values
    i1_values_with_indices
        .sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));

    let i1_sorted_indices: Vec<usize> =
        i1_values_with_indices.iter().map(|&(_, idx)| idx).collect();
    let i1_sorted_values: Vec<f64> = i1_values_with_indices.iter().map(|&(val, _)| val).collect();

    let lab: Vec<Vec<usize>> = i1_sorted_indices
        .iter()
        .map(|&i| combs[i].clone())
        .collect();
    let mut i1_sorted: Vec<f64> = i1_sorted_values.clone();
    let lens: Vec<usize> = lab.iter().map(|l| l.len()).collect();

    if let Some(&max_len) = lens.iter().max() {
        for l in 1..max_len {
            let max_prev_level = i1_sorted
                .iter()
                .zip(&lens)
                .filter(|&(_, &len)| len == l)
                .map(|(&val, _)| val)
                .fold(f64::NEG_INFINITY, f64::max);
            if max_prev_level.is_finite() {
                i1_sorted
                    .iter_mut()
                    .zip(&lens)
                    .filter(|&(_, &len)| len == l + 1)
                    .for_each(|(val, _)| {
                        if *val < max_prev_level {
                            *val = max_prev_level;
                        }
                    });
            }
        }
    }

    let new_sorted_indices = surd_utils::arg_sort(&i1_sorted);
    let final_i1: Vec<f64> = new_sorted_indices.iter().map(|&i| i1_sorted[i]).collect();
    let final_lab: Vec<Vec<usize>> = new_sorted_indices.iter().map(|&i| lab[i].clone()).collect();
    let di_values = surd_utils::diff(&final_i1);

    // Find the index of the last single-variable term with a significant info contribution.
    let last_single_var_idx = final_lab
        .iter()
        .rposition(|lab| lab.len() == 1)
        .unwrap_or(usize::MAX);

    let mut red_vars: Vec<usize> = (1..=n_vars).collect();

    for (i, ll) in final_lab.iter().enumerate() {
        // info calculation needs to handle Option<f64>
        let info = if let Some(p_s_val) = p_s.as_slice()[t] {
            di_values[i] * p_s_val
        } else {
            0.0 // If p_s is None, info contribution is 0
        };

        if info.abs() < f64::EPSILON {
            continue;
        }

        let prev_ll: &[usize] = if i == 0 { &[] } else { &final_lab[i - 1] };
        let (causal_slice, non_causal_slice) =
            calculate_state_slice_cdl(p, ll, prev_ll, t, n_vars)?;

        if ll.len() > 1
        // Synergistic states
        {
            causal_sy_states.insert(ll.clone(), causal_slice);
            non_causal_sy_states.insert(ll.clone(), non_causal_slice);
            *i_s.entry(ll.clone()).or_insert(0.0) += info;
        } else if ll.len() == 1 {
            // Unique vs Redundant states
            if i == last_single_var_idx
            // This is the highest-ordered single-variable term, so it's UNIQUE.
            {
                causal_un_states.insert(ll.clone(), causal_slice);
                non_causal_un_states.insert(ll.clone(), non_causal_slice);
            } else
            // Lower-ordered single-variable terms contribute to REDUNDANCY.
            // The original logic only stored one redundant slice, so we'll model that.
            // This will capture the state map for the main redundant term.
            {
                causal_rd_states.insert(red_vars.clone(), causal_slice);
                non_causal_rd_states.insert(red_vars.clone(), non_causal_slice);
            }
            // Aggregate calculation for I_R
            *i_r.entry(red_vars.clone()).or_insert(0.0) += info;
            red_vars.retain(|&v| v != ll[0]);
        }
    }

    Ok(PerTargetStateResults {
        i_r,
        i_s,
        causal_rd_states,
        causal_un_states,
        causal_sy_states,
        non_causal_rd_states,
        non_causal_un_states,
        non_causal_sy_states,
    })
}

/// Computes the causal (positive) and non-causal (negative) state-dependent information slices.
///
/// This function implements the core calculation of the state-dependent decomposition. It computes
/// the term `p(t, i, j) * log( p(t|i) / p(t|j) )` and separates the result into two tensors
/// based on the sign of the log ratio. This corresponds to the state-dependent causal (`AC`)
/// and non-causal (`AN`) components in (martínezsánchez2025, Supplementary Material, e.g., Eq. S22 & S24).
///
/// # Arguments
/// * `p` - The full joint probability distribution tensor.
/// * `current_vars` - The set of source variables for the current information term (e.g., `{i}`).
/// * `prev_vars` - The set of source variables from the previous step in the sorted hierarchy (e.g., `{j}`).
/// * `target_state_index` - The index `t` of the target variable's state.
/// * `agent_indices` - A slice containing the indices of all source variables.
///
/// # Returns
/// A tuple `(causal_slice, non_causal_slice)` containing the two resulting tensors.
fn calculate_state_slice_cdl(
    p: &CausalTensor<Option<f64>>,
    current_vars: &[usize],
    prev_vars: &[usize],
    target_state_index: usize,
    n_vars: usize,
) -> Result<(CausalTensor<f64>, CausalTensor<f64>), CausalTensorError> {
    // 1. Get the slice corresponding to the current target state.
    let p_slice_option = p.slice(0, target_state_index)?;

    // 2. Remap axes since the target dimension (axis 0) is now gone.
    let current_vars_mapped: Vec<usize> = current_vars.iter().map(|&ax| ax - 1).collect();
    let prev_vars_mapped: Vec<usize> = prev_vars.iter().map(|&ax| ax - 1).collect();
    let all_vars_mapped: Vec<usize> = (0..p_slice_option.num_dim()).collect();

    let p_ti_option = surd_utils_cdl::sum_axes_option_f64(
        &p_slice_option,
        &surd_utils_cdl::set_difference(&all_vars_mapped, &current_vars_mapped),
    )?;

    // To calculate p(i), we marginalize out the target variable and then all other source variables.
    let source_axes: Vec<usize> = (0..n_vars).collect();
    let current_vars_mapped_for_marginal: Vec<usize> =
        current_vars.iter().map(|&ax| ax - 1).collect();
    let axes_to_sum_for_pi =
        surd_utils_cdl::set_difference(&source_axes, &current_vars_mapped_for_marginal);
    let p_i_option = surd_utils_cdl::sum_axes_option_f64(
        &surd_utils_cdl::sum_axes_option_f64(p, &[0])?,
        &axes_to_sum_for_pi,
    )?;

    if p_ti_option.shape() != p_i_option.shape() {
        dbg!("if p_ti_option.shape() != p_i_option.shape() : Tensor ShapeMismatch");

        return Err(CausalTensorError::ShapeMismatch);
    }
    let p_target_given_i_option = surd_utils_cdl::safe_div_cdl(&p_ti_option, &p_i_option)?;

    let p_target_given_j_option = if prev_vars.is_empty() {
        // If prev_vars is empty, p_target_given_j should represent p(T=t).
        // p_slice_option has shape [n_s1_states, n_s2_states, ...]
        // Summing all axes of p_slice_option gives p(T=t) as a scalar tensor.
        let p_t_scalar_option =
            surd_utils_cdl::sum_axes_option_f64(&p_slice_option, &all_vars_mapped)?;
        // This scalar tensor needs to be broadcast to match the shape of p_target_given_i_option.
        surd_utils_cdl::broadcast_to_cdl(&p_t_scalar_option, p_target_given_i_option.shape())?
    } else {
        let p_tj_option = surd_utils_cdl::sum_axes_option_f64(
            &p_slice_option,
            &surd_utils_cdl::set_difference(&all_vars_mapped, &prev_vars_mapped),
        )?;

        let prev_vars_mapped_for_marginal: Vec<usize> =
            prev_vars.iter().map(|&ax| ax - 1).collect();
        let axes_to_sum_for_pj =
            surd_utils_cdl::set_difference(&source_axes, &prev_vars_mapped_for_marginal);
        let p_j_option = surd_utils_cdl::sum_axes_option_f64(
            &surd_utils_cdl::sum_axes_option_f64(p, &[0])?,
            &axes_to_sum_for_pj,
        )?;

        if p_tj_option.shape() != p_j_option.shape() {
            dbg!("if p_tj_option.shape() != p_j_option.shape() : Tensor ShapeMismatch");
            return Err(CausalTensorError::ShapeMismatch);
        }
        surd_utils_cdl::safe_div_cdl(&p_tj_option, &p_j_option)?
    };

    let mut all_involved_vars = current_vars_mapped.to_vec();
    all_involved_vars.extend_from_slice(&prev_vars_mapped);
    all_involved_vars.sort();
    all_involved_vars.dedup();
    let axes_to_sum_out = surd_utils_cdl::set_difference(&all_vars_mapped, &all_involved_vars);
    let p_tij_option = surd_utils_cdl::sum_axes_option_f64(&p_slice_option, &axes_to_sum_out)?;

    // Broadcast p_target_given_i_option to the shape of p_tij_option
    let p_target_given_i_broadcasted_option =
        surd_utils_cdl::broadcast_to_cdl(&p_target_given_i_option, p_tij_option.shape())?;

    // Broadcast p_target_given_j_option to the shape of p_tij_option
    let p_target_given_j_broadcasted_option =
        surd_utils_cdl::broadcast_to_cdl(&p_target_given_j_option, p_tij_option.shape())?;

    let log_ratio_option = surd_utils_cdl::sub_cdl(
        &surd_utils_cdl::surd_log2_cdl(&p_target_given_i_broadcasted_option)?,
        &surd_utils_cdl::surd_log2_cdl(&p_target_given_j_broadcasted_option)?,
    )?;

    // Separate into causal (>0) and non-causal (<0) components
    let causal_log_ratio_option_data: Vec<Option<f64>> = log_ratio_option
        .as_slice()
        .iter()
        .map(|&v_opt| v_opt.map(|v| v.max(0.0)))
        .collect();

    let non_causal_log_ratio_option_data: Vec<Option<f64>> = log_ratio_option
        .as_slice()
        .iter()
        .map(|&v_opt| v_opt.map(|v| v.min(0.0)))
        .collect();

    let causal_log_ratio_tensor = CausalTensor::new(
        causal_log_ratio_option_data,
        log_ratio_option.shape().to_vec(),
    )?;
    let non_causal_log_ratio_tensor = CausalTensor::new(
        non_causal_log_ratio_option_data,
        log_ratio_option.shape().to_vec(),
    )?;

    let causal_slice_option = surd_utils_cdl::mul_cdl(&p_tij_option, &causal_log_ratio_tensor)?;
    let non_causal_slice_option =
        surd_utils_cdl::mul_cdl(&p_tij_option, &non_causal_log_ratio_tensor)?;

    // Convert final Option<f64> slices to f64, treating None as 0.0 for the output
    let causal_slice_data: Vec<f64> = causal_slice_option
        .as_slice()
        .iter()
        .map(|&x| x.unwrap_or(0.0))
        .collect();

    let non_causal_slice_data: Vec<f64> = non_causal_slice_option
        .as_slice()
        .iter()
        .map(|&x| x.unwrap_or(0.0))
        .collect();

    let causal_slice = CausalTensor::new(causal_slice_data, causal_slice_option.shape().to_vec())?;
    let non_causal_slice = CausalTensor::new(
        non_causal_slice_data,
        non_causal_slice_option.shape().to_vec(),
    )?;

    Ok((causal_slice, non_causal_slice))
}