import numpy as np
import matplotlib.pyplot as plt
from itertools import product



rng = np.random.default_rng()
d = 3
omega = np.exp(2j * np.pi / d)


def single_qutrit_phase_space():
    X = np.zeros((d, d), dtype=complex)
    for x in range(d):
        X[(x + 1) % d, x] = 1

    Z = np.diag([omega**x for x in range(d)])

    def T(i, j):
        return omega ** (-i * j / 2) * np.linalg.matrix_power(Z, i) @ np.linalg.matrix_power(X, j)

    T_ops = [[T(i, j) for j in range(d)] for i in range(d)]
    A_00 = (1 / d) * sum(T_ops[i][j] for i in range(d) for j in range(d))

    A_ops = [[T_ops[i][j] @ A_00 @ T_ops[i][j].conj().T
              for j in range(d)] for i in range(d)]
    return A_ops



def wigner_from_state(rho, A_ops, tol = 1e-10):
    W = np.zeros((d, d), dtype=complex)
    for i in range(d):
        for j in range(d):
            val = np.trace(A_ops[i][j] @ rho) / d
        
            # zero out tiny real or imaginary parts
            real_part = val.real if np.abs(val.real) >= tol else 0.0
            imag_part = val.imag if np.abs(val.imag) >= tol else 0.0
            
            W[i, j] = real_part + 1j * imag_part
    return W



def sample_hidden_variables(W_list):
    n = len(W_list)
    x_vec = np.zeros(n, dtype=int)
    y_vec = np.zeros(n, dtype=int)

    for k in range(n):
        W_flat = np.real(W_list[k]).flatten()
        idx = rng.choice(d**2, p=W_flat)
        x_vec[k], y_vec[k] = divmod(idx, d)

    return x_vec, y_vec


def Z_gate(x_vec, y_vec, q):
    y_old = y_vec[q]
    y_vec[q] = (y_vec[q] + 1) % d
    x_vec[q] = (x_vec[q] + y_old + 1) % d

def Z_gate_batch(x, y, q):
    y_old = y[:, q].copy()
    y[:, q] = (y[:, q] + 1) % d
    x[:, q] = (x[:, q] + y_old + 1) % d



def H_gate(x_vec, y_vec, q):
    x_vec[q], y_vec[q] = y_vec[q], x_vec[q]


def H_gate_batch(x, y, q):
    x[:, q], y[:, q] = y[:, q].copy(), x[:, q].copy()






def CX_gate(x_vec, y_vec, control_q, target_q):
    x_vec[control_q] = (x_vec[control_q] - x_vec[target_q]) % d
    y_vec[target_q] = (y_vec[target_q] + y_vec[control_q]) % d


def CX_gate_batch(x, y, control_q, target_q):
    x[:, control_q] = (x[:, control_q] - x[:, target_q]) % d
    y[:, target_q] = (y[:, target_q] + y[:, control_q]) % d





def Measure(x_vec, y_vec, A_ops):
    n = len(x_vec)
    outcomes = np.zeros(n, dtype=int)

    # computational basis POVM
    E_1 = np.outer(np.array([1,0,0]), np.array([1,0,0]).conj())
    E_2 = np.outer(np.array([0,1,0]), np.array([0,1,0]).conj())
    E_3 = np.outer(np.array([0,0,1]), np.array([0,0,1]).conj())
    E_k = [E_1, E_2, E_3]

    for q in range(n):
        A = A_ops[x_vec[q]][y_vec[q]]
        probs = np.array([np.trace(A @ E_k[k]).real for k in range(d)])
        probs = np.real(probs)
        probs[np.abs(probs) <= 1e-10] = 0.0

        if sum(probs) != 1:
            probs = probs / sum(probs)
            # print("Warning: probabilities do not sum to 1:", probs, " sum=", sum(probs))
        outcomes[q] = rng.choice(d, p=probs)

    return outcomes



def outcomes_to_index(outcomes, d):
    idx = 0
    for o in outcomes:
        idx = idx * d + o
    return idx


def sample_hidden_variables_batch(W_flat_list, num_runs):
    n = len(W_flat_list)
    x = np.zeros((num_runs, n), dtype=int)
    y = np.zeros((num_runs, n), dtype=int)

    for q in range(n):
        idx = rng.choice(d**2, size=num_runs, p=W_flat_list[q])
        x[:, q] = idx // d
        y[:, q] = idx % d

    return x, y



def precompute_measurement_probs(A_ops):
    probs = np.zeros((d, d, d))
    E = np.eye(d)

    for x in range(d):
        for y in range(d):
            for k in range(d):
                probs[x, y, k] = np.trace(A_ops[x][y] @ np.outer(E[k], E[k])).real
    probs[probs < 1e-12] = 0
    probs /= probs.sum(axis=2, keepdims=True)
    return probs





def Measure_batch(x, y, meas_probs):
    num_runs, n = x.shape
    outcomes = np.zeros((num_runs, n), dtype=int)

    for q in range(n):
        p = meas_probs[x[:, q], y[:, q]]  # shape (num_runs, d)

        # cumulative distribution
        cdf = np.cumsum(p, axis=1)
        r = rng.random(num_runs)[:, None]
        outcomes[:, q] = (r < cdf).argmax(axis=1)

    return outcomes


def outcomes_to_index_batch(outcomes, d):
    idx = np.zeros(outcomes.shape[0], dtype=int)
    for q in range(outcomes.shape[1]):
        idx = idx * d + outcomes[:, q]
    return idx








def plot_joint_probs(outcomes, d=3, plot_type='bar', title=None):
    """
    Plot joint measurement probabilities for qutrits (or qudits) from outcomes.

    Parameters
    ----------
    outcomes : np.ndarray
        Array of shape (num_runs, n_qutrits) containing measured outcomes.
    d : int
        Local dimension (default 3 for qutrits).
    plot_type : str
        Type of plot: 'bar', 'heatmap', or 'stem'.
    title : str, optional
        Plot title. If None, a default is generated.
    """

    n = outcomes.shape[1]

    # Convert outcomes to joint indices
    idx = np.zeros(outcomes.shape[0], dtype=int)
    for q in range(n):
        idx = idx * d + outcomes[:, q]

    # Histogram
    joint_counts = np.bincount(idx, minlength=d**n)
    joint_probs = joint_counts / outcomes.shape[0]

    # Labels for bar/stem plots
    labels = [''.join(map(str, tup)) for tup in product(range(d), repeat=n)]

    if title is None:
        title = f"Joint measurement probabilities for {n} qutrits"

    # ---- Plotting ----
    if plot_type == 'bar':
        plt.figure(figsize=(max(6, 0.5 * d**n), 4))
        plt.bar(labels, joint_probs)
        plt.xlabel("Joint measurement outcome")
        plt.ylabel("Probability")
        plt.title(title)
        plt.xticks(rotation=90)
        plt.tight_layout()
        plt.show()

    elif plot_type == 'topk':
        k = min(10, len(joint_probs))
        topk_indices = np.argsort(joint_probs)[-k:][::-1]
        topk_labels = [labels[i] for i in topk_indices]
        topk_probs = joint_probs[topk_indices]

        plt.figure(figsize=(max(6, 0.5 * k), 4))
        plt.bar(topk_labels, topk_probs)
        plt.xlabel("Joint measurement outcome")
        plt.ylabel("Probability")
        plt.title(f"Top {k} joint measurement probabilities")
        plt.xticks(rotation=90)
        plt.tight_layout()
        plt.show()

    elif plot_type == 'heatmap':
        if n != 2:
            raise ValueError("Heatmap plot only supported for 2 qutrits/qudits")
        plt.figure(figsize=(5, 4))
        plt.imshow(joint_probs.reshape(d, d), origin='lower', cmap='viridis', interpolation='nearest')
        plt.colorbar(label='Probability')
        plt.xlabel("Qutrit 1 outcome")
        plt.ylabel("Qutrit 0 outcome")
        plt.title(title)
        plt.show()

    else:
        raise ValueError(f"Unknown plot_type '{plot_type}'. Use 'bar', 'stem', or 'heatmap'.")
