import numpy as np

def print_mat(P_mat):
    age = ["X_a", "X_b", "X_c", "X_d"]
    # Print column headers
    print("\tY_a\tY_b\tY_c")
    i = 0
    for row in P_mat:
        # Print row header
        print("{}".format(age[i]), end="")
        for element in row:
            # Print elements formatted to two decimal places
            print(f"\t{element:.2f}", end="  ")
        print()
        i += 1


def iterative_proportional_fitting_fixed_cell(
    P_A, P_B, fixed_coords, fixed_value, max_iterations=100, tolerance=1e-6
):
    m = len(P_A)
    n = len(P_B)
    P_A_target = P_A.reshape(-1, 1)
    P_B_target = P_B.reshape(1, -1)

    # 1. Set the fixed value in the initial distribution P^{(0)} and fill the rest based on independence.
    P_initial = np.outer(P_A, P_B)
    P_initial[fixed_coords] = fixed_value

    # 2. Adjustment for the remaining cells
    # Subtract the fixed cell P(X_a, Y_a) value from the marginals
    P_A_adj = P_A.copy()
    P_B_adj = P_B.copy()

    P_A_adj[fixed_coords[0]] -= fixed_value
    P_B_adj[fixed_coords[1]] -= fixed_value

    # 3. Check that the remaining cells are non-negative (Axiom of Probability)
    if P_A_adj[fixed_coords[0]] < 0 or P_B_adj[fixed_coords[1]] < 0:
        raise ValueError(
            "The fixed value exceeds the marginal probability. IPF cannot be applied."
        )

    # 4. Matrix to mask the location of the fixed cell
    mask = np.ones((m, n))
    mask[fixed_coords] = 0

    # 5. Set the initial joint distribution P^{(0)} (fixed cell remains, normalize the rest using independence)
    # Calculate the product of the remaining marginal sums and allocate proportionally to the remaining cells
    P_rest_sum = np.sum(P_A_adj)

    # Calculate the remaining initial distribution (starting here as the outer product of adjusted marginals)
    P_current = np.outer(P_A_adj, P_B_adj)
    # Normalize so the remaining marginal sum is 1.0
    P_current /= P_current.sum()  # Normalize to the sum of P_rest_sum

    # 6. Add the fixed value to P_current to form the overall initial distribution
    P_current *= mask  # Apply the mask
    P_current[fixed_coords] = fixed_value  # Set the fixed value

    # Final normalization of the initial distribution
    P_current /= P_current.sum()

    for i in range(max_iterations):
        P_prev = P_current.copy()

        # === Match Row Constraint (P_A) ===
        # Only adjust the part that does not include the fixed cell
        # Use the approach of applying standard IPF and finally overwriting the fixed cell

        # Calculate row sums
        row_sums = P_current.sum(axis=1).reshape(-1, 1)
        adjustment_factor_row = P_A_target / (row_sums + 1e-18)
        P_current *= adjustment_factor_row

        # Overwrite the fixed cell (to satisfy the constraint)
        P_current[fixed_coords] = fixed_value

        # === Match Column Constraint (P_B) ===
        col_sums = P_current.sum(axis=0).reshape(1, -1)
        adjustment_factor_col = P_B_target / (col_sums + 1e-18)
        P_current *= adjustment_factor_col

        # Overwrite the fixed cell (to satisfy the constraint)
        P_current[fixed_coords] = fixed_value

        # === Convergence Check ===
        if np.sum(np.abs(P_current - P_prev)) < tolerance:
            # print(f"--- Converged. Iterations: {i+1} ---")
            break
    else:
        print(f"--- Maximum iterations ({max_iterations}) reached ---")

    return P_current


if __name__ == "__main__":
    # --- Example Usage ---
    # Use data from the question example and modify the marginal distribution to incorporate the constraint P(X_a, Y_a) = 0.2

    # P_A: [X_a, X_b, X_c, X_d] (4 items)
    P_A_example = np.array([0.2, 0.25, 0.35, 0.2])
    # P_B: [Y_a, Y_b, Y_c] (3 items)
    P_B_example = np.array([0.4, 0.55, 0.05])

    print("==== Initial Conditions ====")
    print("P_A: [X_a, X_b, X_c, X_d]")
    print(P_A_example)
    print("P_B: [Y_a, Y_b, Y_c]")
    print(P_B_example)
    print("")

    P_current = np.outer(P_A_example, P_B_example)

    # Row marginal distribution (constraint) shaped as (m, 1)
    P_A_target = P_A_example.reshape(-1, 1)
    # Column marginal distribution (constraint)
    P_B_target = P_B_example.reshape(1, -1)

    print("==== Joint Distribution (Independence Assumption) ====")

    print_mat(P_current)
    print("")

    # Case from the question: Fix P(X_a, Y_a) = 0.2 and satisfy marginal constraints
    fixed_coords = (0, 0)  # X_a (0th row), Y_a (0th column)
    fixed_value = 0.2

    P_fixed = iterative_proportional_fitting_fixed_cell(
        P_A_example, P_B_example, fixed_coords, fixed_value
    )

    print("\n==== Final Result (P(X_a, Y_a)=0.2 Fixed, KL Minimization) ====")
    print_mat(P_fixed)
    print("\nRow Marginal Sums (compared to P_A_example):", P_fixed.sum(axis=1))
    print("Column Marginal Sums (compared to P_B_example):", P_fixed.sum(axis=0))

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np


# NN Model Definition
class ProbabilityModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(NUM_A + NUM_B, 128),
            nn.ReLU(),
            nn.Linear(128, 1),  # Change output size to 1
        )

    def forward(self, coords):
        # Simulate one-hot encoding
        coords_int = coords.long()
        features = []
        for i in range(coords.size(0)):
            one_hot_A = torch.zeros(NUM_A, device=coords.device)
            one_hot_A[coords_int[i, 0]] = 1.0
            one_hot_B = torch.zeros(NUM_B, device=coords.device)
            one_hot_B[coords_int[i, 1]] = 1.0
            features.append(torch.cat([one_hot_A, one_hot_B]))

        features_tensor = torch.stack(features)
        logits = self.fc(features_tensor)
        return logits.squeeze(1)


def display_joint_probability(
    model, coords_tensor, num_a, num_b, title="Joint Probability Distribution P(A, B)"
):
    """
    Retrieves and displays the joint probability distribution from the trained NN model.
    """
    print(f"\n--- {title} ---")

    with torch.no_grad():
        # 1. Get logits for all bins from NN (NN forward pass)
        individual_logits = model(coords_tensor)
        # 2. Calculate joint probability distribution P_NN (shape [80])
        P_NN_flat = torch.softmax(individual_logits, dim=0)
        # 3. Reshape into a 2D matrix [NUM_A, NUM_B]
        P_NN_matrix = P_NN_flat.view(num_a, num_b).cpu().numpy()

    # Header (Indices of H_B)
    header = "A\\B | " + " | ".join([f"B[{j}]" for j in range(num_b)])
    print("-" * (len(header) + 4))
    print(header)
    print("-" * (len(header) + 4))

    for i in range(num_a):
        row_str = f"A[{i}] | "
        for j in range(num_b):
            prob = P_NN_matrix[i, j]
            row_str += f"{prob:.3f} | "
        print(row_str.strip())
    print("-" * (len(header) + 4))

    # Verification: Check if the area sum is close to the target
    area_sum_check = P_NN_matrix[:N_CONSTRAINT, :M_CONSTRAINT].sum()
    P_A_check = P_NN_matrix.sum(axis=1)
    P_B_check = P_NN_matrix.sum(axis=0)

    # Calculate marginal error
    L_Margin = (
        nn.MSELoss()(torch.tensor(P_A_check), H_A).item()
        + nn.MSELoss()(torch.tensor(P_B_check), H_B).item()
    )
    print(
        f"Verification: Sum of area P(A[0:{N_CONSTRAINT-1}], B[0:{M_CONSTRAINT-1}]): {area_sum_check:.6f} (Target: {C_TRUE:.6f})"
    )
    print(f"Verification: Marginal Error (MSE): {L_Margin:.8f}")


# Loss Function and Fine-tuning
def fine_tune_constrained(
    model,
    optimizer,
    c_true,
    area_indices,
    H_A_target,
    H_B_target,
    lambda_area=10.0,
    lambda_margin=1.0,
    num_steps=500,
):

    criterion_mse = nn.MSELoss()

    for step in range(num_steps):
        optimizer.zero_grad()

        individual_logits = model(all_coords_tensor)
        P_NN_flat = torch.softmax(individual_logits, dim=0)

        # Loss 1: Area Sum Constraint (L_area)
        c_nn = P_NN_flat[area_indices].sum()
        c_true_tensor = torch.tensor(c_true).float().to(c_nn.device)
        loss_area = criterion_mse(c_nn, c_true_tensor)

        # Loss 2: Marginal Constraints (L_margin)
        P_NN_matrix = P_NN_flat.view(NUM_A, NUM_B)
        P_A_NN = P_NN_matrix.sum(dim=1)
        P_B_NN = P_NN_matrix.sum(dim=0)

        loss_margin_A = criterion_mse(P_A_NN, H_A_target)
        loss_margin_B = criterion_mse(P_B_NN, H_B_target)
        loss_margin = loss_margin_A + loss_margin_B

        # Total Loss
        loss_total = lambda_area * loss_area + lambda_margin * loss_margin
        loss_total.backward()
        optimizer.step()

        if (step + 1) % 50 == 0:
            print(
                f"Step {step+1}: L_Total={loss_total.item():.6f}, L_Area={loss_area.item():.6f}, L_Margin={loss_margin.item():.6f}, C_NN={c_nn.item():.6f}"
            )


if __name__ == "__main__":

    # --- 1. Initial Conditions and Setup ---

    # Distribution Histograms (Fixed as Marginal Constraints)
    H_A = torch.tensor(
        [0.085, 0.112, 0.063, 0.141, 0.039, 0.187, 0.095, 0.051, 0.155, 0.072],
        dtype=torch.float32,
    )
    H_B = torch.tensor(
        [0.021, 0.058, 0.150, 0.245, 0.280, 0.148, 0.065, 0.033], dtype=torch.float32
    )

    NUM_A = H_A.size(0)  # 10
    NUM_B = H_B.size(0)  # 8
    TOTAL_BINS = NUM_A * NUM_B  # 80

    # Target Value for Area Constraint (Fitting Target)
    C_TRUE = 0.3

    # Area Constraint Definition: H_A[0:n] and H_B[0:m] (Assume n=5, m=4)
    N_CONSTRAINT = 5
    M_CONSTRAINT = 4

    # --- Generate coordinates (i, j) for all bins (Global scope) ---
    all_coords = []
    for i in range(NUM_A):
        for j in range(NUM_B):
            all_coords.append([i, j])
    all_coords_tensor = torch.tensor(all_coords).float()  # Shape [80, 2]

    # --- Identify flat indices of bins included in the area constraint ---
    AREA_INDICES = []
    for i in range(N_CONSTRAINT):
        for j in range(M_CONSTRAINT):
            AREA_INDICES.append(i * NUM_B + j)
    AREA_INDICES = torch.tensor(AREA_INDICES)  # Shape [20]

    # --- Execution ---
    model = ProbabilityModel()
    optimizer = optim.Adam(model.parameters(), lr=5e-4)

    # Display before fine-tuning
    display_joint_probability(
        model,
        all_coords_tensor,
        NUM_A,
        NUM_B,
        title="Before Fine-tuning (Initial Random State)",
    )

    print("\n--- Starting Fine-tuning ---")
    fine_tune_constrained(model, optimizer, C_TRUE, AREA_INDICES, H_A, H_B)
    print("--- Fine-tuning Complete ---")

    # Display after fine-tuning
    display_joint_probability(
        model,
        all_coords_tensor,
        NUM_A,
        NUM_B,
        title="After Fine-tuning (Constraints Applied)",
    )

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import copy


# NN Model Definition
class ProbabilityModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(NUM_A + NUM_B, 128),
            nn.ReLU(),
            nn.Linear(128, 1),  # Change output size to 1
        )

    def forward(self, coords):
        # Simulate one-hot encoding
        coords_int = coords.long()
        features = []
        for i in range(coords.size(0)):
            one_hot_A = torch.zeros(NUM_A, device=coords.device)
            one_hot_A[coords_int[i, 0]] = 1.0
            one_hot_B = torch.zeros(NUM_B, device=coords.device)
            one_hot_B[coords_int[i, 1]] = 1.0
            features.append(torch.cat([one_hot_A, one_hot_B]))

        features_tensor = torch.stack(features)
        logits = self.fc(features_tensor)
        return logits.squeeze(1)


def display_joint_probability(
    model,
    coords_tensor,
    num_a,
    num_b,
    H_A_target,
    H_B_target,
    C_true,
    area_indices,
    title="Joint Probability Distribution P(A, B)",
):
    """
    Retrieves the joint probability distribution from the trained NN model and displays it using nested loops.
    """
    print(f"\n--- {title} ---")

    with torch.no_grad():
        # 1. Get logits for all bins from NN (NN forward pass)
        individual_logits = model(coords_tensor)
        # 2. Calculate joint probability distribution P_NN (shape [80])
        P_NN_flat = torch.softmax(individual_logits, dim=0)
        # 3. Reshape into a 2D matrix [NUM_A, NUM_B]
        P_NN_matrix = P_NN_flat.view(num_a, num_b).cpu().numpy()

    # Header (Indices of H_B)
    header = "A\\B | " + " | ".join([f"B[{j}]" for j in range(num_b)])
    print("-" * (len(header) + 4))
    print(header)
    print("-" * (len(header) + 4))

    for i in range(num_a):
        row_str = f"A[{i}] | "
        for j in range(num_b):
            prob = P_NN_matrix[i, j]
            row_str += f"{prob:.3f} | "
        print(row_str.strip())
    print("-" * (len(header) + 4))

    # Verification
    area_sum_check = P_NN_matrix[:N_CONSTRAINT, :M_CONSTRAINT].sum()
    P_A_check = P_NN_matrix.sum(axis=1)
    P_B_check = P_NN_matrix.sum(axis=0)

    L_Margin = (
        nn.MSELoss()(torch.tensor(P_A_check), H_A_target).item()
        + nn.MSELoss()(torch.tensor(P_B_check), H_B_target).item()
    )

    print(
        f"Verification: Sum of area P(A[0:{N_CONSTRAINT-1}], B[0:{M_CONSTRAINT-1}]): {area_sum_check:.6f} (Target: {C_true:.6f})"
    )
    print(f"Verification: Marginal Error (MSE): {L_Margin:.8f}")


# Function to Calculate P* (Constrained IPF Principle)
def calculate_p_star_with_constraints(
    P_initial,
    H_A_target,
    H_B_target,
    C_true,
    area_indices,
    max_iter=100,
    tolerance=1e-6,
):

    P_current = P_initial.view(NUM_A, NUM_B).clone()

    # Convert marginal sums to NumPy arrays
    H_A_np = H_A_target.cpu().numpy()
    H_B_np = H_B_target.cpu().numpy()

    # Area indices (2D)
    area_rows, area_cols = np.unravel_index(area_indices.cpu().numpy(), (NUM_A, NUM_B))

    for i in range(max_iter):
        P_prev = P_current.clone()

        # === Step 1: Adjust Area Constraint ===
        P_np = P_current.cpu().numpy()
        current_area_sum = P_np[area_rows, area_cols].sum()

        if abs(current_area_sum - C_true) > 1e-18:
            # Calculate ratios for the area and the rest
            ratio_area = C_true / (current_area_sum + 1e-18)
            P_rest_sum = 1.0 - current_area_sum
            P_rest_target = 1.0 - C_true
            ratio_rest = P_rest_target / (P_rest_sum + 1e-18)

            P_temp_np = P_np.copy()

            # Adjust inside the constrained area
            P_temp_np[area_rows, area_cols] *= ratio_area

            # Adjust outside the constrained area
            mask_rest = np.ones_like(P_np, dtype=bool)
            mask_rest[area_rows, area_cols] = False
            P_temp_np[mask_rest] *= ratio_rest

            P_current = torch.tensor(P_temp_np).float()
            P_current /= P_current.sum()  # Ensure normalization after adjustment

        # === Step 2: Adjust Marginal Constraints (Standard IPF Steps) ===

        # Adjustment for A (Row) direction
        row_sums = P_current.sum(dim=1).unsqueeze(1)
        adjustment_factor_row = H_A_target.unsqueeze(1) / (row_sums + 1e-18)
        P_current *= adjustment_factor_row

        # Adjustment for B (Column) direction
        col_sums = P_current.sum(dim=0).unsqueeze(0)
        adjustment_factor_col = H_B_target.unsqueeze(0) / (col_sums + 1e-18)
        P_current *= adjustment_factor_col

        # === Convergence Check ===
        if torch.sum(torch.abs(P_current - P_prev)) < tolerance:
            break

    return P_current.view(1, TOTAL_BINS)  # Shape [1, 80]


# Learning Step Using KL Divergence as Teacher
def fine_tune_to_p_star(model, optimizer, P_star, num_steps=500):

    # KL Divergence Loss: Minimize D_KL(P* || P_NN)
    criterion_kl = nn.KLDivLoss(reduction="batchmean", log_target=False)

    for step in range(num_steps):
        optimizer.zero_grad()

        individual_logits = model(all_coords_tensor)
        log_probs_nn = torch.log_softmax(individual_logits, dim=0).unsqueeze(0)

        # Calculate loss: loss_kl = D_KL(P* || P_NN)
        loss_kl = criterion_kl(log_probs_nn, P_star)

        loss_total = loss_kl
        loss_total.backward()
        optimizer.step()

        if (step + 1) % 50 == 0:
            print(f"Step {step+1}: KL_Loss = {loss_kl.item():.6f}")


if __name__ == "__main__":

    # --- 1. Initial Conditions and Setup ---

    # Distribution Histograms (Fixed as Marginal Constraints)
    H_A = torch.tensor(
        [0.085, 0.112, 0.063, 0.141, 0.039, 0.187, 0.095, 0.051, 0.155, 0.072],
        dtype=torch.float32,
    )
    H_B = torch.tensor(
        [0.021, 0.058, 0.150, 0.245, 0.280, 0.148, 0.065, 0.033], dtype=torch.float32
    )

    NUM_A = H_A.size(0)  # 10
    NUM_B = H_B.size(0)  # 8
    TOTAL_BINS = NUM_A * NUM_B  # 80

    # Target Value for Area Constraint (Fitting Target)
    C_TRUE = 0.3

    # Area Constraint Definition: H_A[0:n] and H_B[0:m] (Assume n=5, m=4)
    N_CONSTRAINT = 5
    M_CONSTRAINT = 4

    # --- Generate coordinates (i, j) for all bins (Global scope) ---
    all_coords = []
    for i in range(NUM_A):
        for j in range(NUM_B):
            all_coords.append([i, j])
    all_coords_tensor = torch.tensor(all_coords).float()  # Shape [80, 2]

    # --- Identify flat indices of bins included in the area constraint ---
    AREA_INDICES = []
    for i in range(N_CONSTRAINT):
        for j in range(M_CONSTRAINT):
            AREA_INDICES.append(i * NUM_B + j)
    AREA_INDICES = torch.tensor(AREA_INDICES)  # Shape [20]

    # --- Execution ---
    model_hybrid = ProbabilityModel()
    optimizer_hybrid = optim.Adam(model_hybrid.parameters(), lr=1e-3)

    # ----------------------------------------------------
    # Display Before Fine-tuning
    # ----------------------------------------------------
    display_joint_probability(
        model_hybrid,
        all_coords_tensor,
        NUM_A,
        NUM_B,
        H_A,
        H_B,
        C_TRUE,
        AREA_INDICES,
        title="Before Fine-tuning (Initial Random State)",
    )

    print("\n--- Calculating P* Target Distribution (Applying IPF Principle) ---")

    # 1. Get the initial output of the NN (P_initial)
    with torch.no_grad():
        initial_logits = model_hybrid(all_coords_tensor)
        P_initial = torch.softmax(initial_logits, dim=0)

    # 2. Modify P_initial and calculate the target distribution P*
    P_star = calculate_p_star_with_constraints(
        P_initial, H_A, H_B, C_TRUE, AREA_INDICES
    )

    # Check P* constraint achievement
    P_star_matrix_check = P_star.squeeze().view(NUM_A, NUM_B)
    P_star_area_sum_check = (
        P_star_matrix_check[:N_CONSTRAINT, :M_CONSTRAINT].sum().item()
    )

    print(f"P* Target Area Sum: {P_star_area_sum_check:.6f} (Target: {C_TRUE:.6f})")

    # 3. Fine-tune the NN to approach P*
    print("\n--- Starting NN Fine-tuning (KL Minimization) ---")
    fine_tune_to_p_star(model_hybrid, optimizer_hybrid, P_star)
    print("--- NN Fine-tuning Complete ---")

    # ----------------------------------------------------
    # Display After Fine-tuning
    # ----------------------------------------------------
    display_joint_probability(
        model_hybrid,
        all_coords_tensor,
        NUM_A,
        NUM_B,
        H_A,
        H_B,
        C_TRUE,
        AREA_INDICES,
        title="After Fine-tuning (Approximation to P* Target)",
    )

import numpy as np
import math
import hashlib


class CountMinSketch:
    def __init__(self, d_rows, w_cols):
        # Initialize the sketch matrix with zeros
        self.sketch = np.zeros((d_rows, w_cols), dtype=np.int32)
        self.d = d_rows
        self.w = w_cols

        # Seeds to simulate different hash functions
        # Using NumPy random seeds (in practice, cryptographic hash functions are often used)
        self.seeds = np.random.randint(0, 1000000, d_rows)

    def _hash(self, i, j, seed):
        """
        Calculates a hash value within w_cols using 2D indices (i, j) and a seed.
        """
        # Combine indices and hash (simple concatenation and hashing)
        # Assuming i and j are integers
        combined_key = str(i) + "-" + str(j)

        # Combine Python's standard hash with the seed
        # Append the seed to ensure diversity of hash values
        h = int(
            hashlib.sha256((combined_key + str(seed)).encode("utf-8")).hexdigest(), 16
        )

        # Map to the range of w_cols
        return h % self.w

    def add(self, i: int, j: int, k: int):
        """
        Adds value k to the array cell (i, j).
        Args:
            i (int): Index of the first dimension
            j (int): Index of the second dimension
            k (int): Value (frequency or weight) to add
        """
        if i >= N_DIM or j >= M_DIM:
            raise IndexError("Index out of bounds for the assumed array size.")

        for row in range(self.d):
            # Use a different hash function for each row
            col_index = self._hash(i, j, self.seeds[row])
            self.sketch[row, col_index] += k

    def get(self, i: int, j: int) -> int:
        """
        Returns the estimated value of the array cell (i, j).
        Args:
            i (int): Index of the first dimension
            j (int): Index of the second dimension
        Returns:
            int: Estimated value of cell (i, j)
        """
        if i >= N_DIM or j >= M_DIM:
            raise IndexError("Index out of bounds for the assumed array size.")

        min_value = np.inf

        for row in range(self.d):
            col_index = self._hash(i, j, self.seeds[row])
            min_value = min(min_value, self.sketch[row, col_index])

        return int(min_value)

    def show(self, max_dim=10):
        """
        Displays the estimated data for the array (only a subset is shown due to computation cost).
        Args:
            max_dim (int): Maximum dimension to display (e.g., shows the top-left max_dim x max_dim)
        """
        print(f"\n--- Estimated Joint Probability Distribution (CM Sketch) ---")
        print(
            f" (Displaying estimate for top-left {min(N_DIM, max_dim)}x{min(M_DIM, max_dim)})"
        )

        # Array to store results
        estimated_array = np.zeros(
            (min(N_DIM, max_dim), min(M_DIM, max_dim)), dtype=np.int32
        )

        # Estimate all elements in the nested loop and display
        for i in range(min(N_DIM, max_dim)):
            for j in range(min(M_DIM, max_dim)):
                estimated_array[i, j] = self.get(i, j)

        # Print the formatted NumPy array
        print(estimated_array)
        print("-" * 40)


if __name__ == "__main__":

    # --- 1. Initial Parameter Setup ---
    # Corresponds to user input: error epsilon and failure probability delta
    EPSILON = 0.01  # Tolerance for error (1%)
    DELTA = 0.01  # Tolerance for failure probability (1%)

    # Array dimensions (assuming 400x400, but doesn't directly affect CM Sketch size)
    N_DIM = 400
    M_DIM = 400

    # --- 2. CM Sketch Parameter Calculation ---
    # Number of rows d (number of hash functions)
    D_ROWS = math.ceil(math.log(1 / DELTA))

    # Number of columns w (width of buckets)
    W_COLS = math.ceil(math.e / EPSILON)

    # --- Example Execution ---
    # 1. Initialize the Sketch
    cm_sketch = CountMinSketch(D_ROWS, W_COLS)

    # 2. Add Data (Add operation)
    # i, j are in the range 0 to 399
    print("Adding data (Add Operation)...")
    cm_sketch.add(1, 1, 100)
    cm_sketch.add(1, 1, 50)  # Total 150
    cm_sketch.add(1, 2, 20)  # Total 20
    cm_sketch.add(200, 300, 5)  # Far index
    cm_sketch.add(2, 1, 100)  # Data that might cause hash collision with (1,1)

    # 3. Get value of specific cells (Get operation)
    val_1_1 = cm_sketch.get(1, 1)
    val_1_2 = cm_sketch.get(1, 2)
    val_2_1 = cm_sketch.get(2, 1)
    val_200_300 = cm_sketch.get(200, 300)

    print(f"Estimated value (1, 1): {val_1_1} (True value: 150)")
    print(f"Estimated value (1, 2): {val_1_2} (True value: 20)")
    print(f"Estimated value (2, 1): {val_2_1} (True value: 100)")
    print(f"Estimated value (200, 300): {val_200_300} (True value: 5)")

    # 4. Display the entire array (Show operation)
    cm_sketch.show(max_dim=5)