
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

# ─────────────────────────────────────────────────────────────────
#  PKTron Advanced Visualization & Utilities
#
#  draw_circuit_text   → text-based circuit diagram (no matplotlib)
#  draw_circuit_mpl    → matplotlib circuit diagram
#  plot_city           → city plot of density matrix
#  plot_hinton         → Hinton diagram of a matrix
#  plot_bloch_multi    → Bloch spheres for all qubits
#  EnhancedResult      → Result object with metadata
# ─────────────────────────────────────────────────────────────────


def draw_circuit_text(circuit):
    """
    Draw a simple text-based circuit diagram.
    Works in any terminal or notebook without matplotlib.

    Usage:
        draw_circuit_text(qc)
    """
    n = circuit.num_qubits
    lines = [f"q{i}: ─" for i in range(n)]

    for gate_str in circuit.gate_history:
        parts = gate_str.split()
        gname = parts[0]

        if gname in ("H","X","Y","Z","S","T"):
            try:
                q = int(parts[-1])
                box = f"[{gname}]─"
                for i in range(n):
                    lines[i] += box if i == q else "─" * len(box)
            except Exception:
                pass

        elif gname in ("Rx","Ry","Rz"):
            try:
                q = int(parts[-1])
                box = f"[{gname}]─"
                for i in range(n):
                    lines[i] += box if i == q else "─" * len(box)
            except Exception:
                pass

        elif gname == "CNOT":
            try:
                c = int(parts[2])
                t = int(parts[4])
                box_c = "[●]─"
                box_t = "[⊕]─"
                pad   = "─" * len(box_c)
                for i in range(n):
                    if i == c:   lines[i] += box_c
                    elif i == t: lines[i] += box_t
                    else:        lines[i] += pad
            except Exception:
                pass

        elif gname == "CZ":
            try:
                c = int(parts[2])
                t = int(parts[4])
                box = "[Z]─"
                pad = "─" * len(box)
                for i in range(n):
                    if i == c or i == t: lines[i] += box
                    else:                lines[i] += pad
            except Exception:
                pass

        elif gname == "SWAP":
            try:
                q1 = int(parts[2])
                q2 = int(parts[4])
                box = "[×]─"
                pad = "─" * len(box)
                for i in range(n):
                    if i == q1 or i == q2: lines[i] += box
                    else:                  lines[i] += pad
            except Exception:
                pass

    print("\n".join(lines))
    return lines


def draw_circuit_mpl(circuit, figsize=None):
    """
    Draw a matplotlib circuit diagram.

    Usage:
        draw_circuit_mpl(qc)
    """
    n = circuit.num_qubits
    gates = circuit.gate_history
    num_layers = max(len(gates), 1)
    width = max(10, num_layers * 1.2 + 2)
    height = max(3, n * 1.0 + 1)
    fig, ax = plt.subplots(figsize=figsize or (width, height))

    # Draw qubit lines
    for q in range(n):
        ax.plot([0, num_layers + 1], [q, q], color="black", lw=1.5, zorder=1)
        ax.text(-0.3, q, f"q{q}", fontsize=11, ha="right", va="center",
                fontweight="bold")

    colors = {
        "H": "#4FC3F7", "X": "#EF9A9A", "Y": "#CE93D8",
        "Z": "#80CBC4", "S": "#FFE082", "T": "#A5D6A7",
        "Rx": "#FFCC80", "Ry": "#FFCC80", "Rz": "#FFCC80",
        "CNOT": "#F48FB1", "CZ": "#B39DDB", "SWAP": "#90CAF9",
    }

    for layer, gate_str in enumerate(gates):
        x = layer + 0.5
        parts = gate_str.split()
        gname = parts[0]
        color = colors.get(gname, "#E0E0E0")

        try:
            if gname in ("H","X","Y","Z","S","T","Rx","Ry","Rz"):
                q = int(parts[-1])
                rect = plt.Rectangle((x - 0.3, q - 0.25), 0.6, 0.5,
                                     color=color, ec="black", lw=1.2, zorder=3)
                ax.add_patch(rect)
                ax.text(x, q, gname, ha="center", va="center",
                        fontsize=9, fontweight="bold", zorder=4)

            elif gname == "CNOT":
                c, t = int(parts[2]), int(parts[4])
                ax.plot([x, x], [c, t], color="black", lw=1.5, zorder=2)
                ax.add_patch(plt.Circle((x, c), 0.12, color="black", zorder=4))
                ax.add_patch(plt.Circle((x, t), 0.22, color="white",
                                        ec="black", lw=1.5, zorder=4))
                ax.text(x, t, "⊕", ha="center", va="center",
                        fontsize=12, zorder=5)

            elif gname == "CZ":
                c, t = int(parts[2]), int(parts[4])
                ax.plot([x, x], [c, t], color="black", lw=1.5, zorder=2)
                for q in [c, t]:
                    ax.add_patch(plt.Circle((x, q), 0.12, color="black", zorder=4))

            elif gname == "SWAP":
                q1, q2 = int(parts[2]), int(parts[4])
                ax.plot([x, x], [q1, q2], color="black", lw=1.5, zorder=2)
                for q in [q1, q2]:
                    ax.text(x, q, "×", ha="center", va="center",
                            fontsize=14, fontweight="bold", zorder=4)

        except Exception:
            pass

    ax.set_xlim(-0.6, num_layers + 0.8)
    ax.set_ylim(-0.7, n - 0.3)
    ax.set_axis_off()
    ax.set_title("PKTron Circuit Diagram", fontsize=13,
                 color="darkblue", fontweight="bold", pad=12)
    plt.tight_layout()
    plt.show()
    return fig


def plot_city(density_matrix_data, title="City Plot"):
    """
    City plot: 3D bar chart of the density matrix.
    Real part = blue bars, Imaginary part = red bars.

    Usage:
        plot_city(dm.data)
        plot_city(circuit)  # also accepts a QuantumCircuit
    """
    from .quantum import QuantumCircuit
    if isinstance(density_matrix_data, QuantumCircuit):
        state = density_matrix_data.get_state()
        rho = np.outer(state, state.conj())
    else:
        rho = np.array(density_matrix_data)

    n = rho.shape[0]
    fig = plt.figure(figsize=(12, 5))

    for idx, (part, part_label, color) in enumerate([
        (np.real(rho), "Real", "#2196F3"),
        (np.imag(rho), "Imaginary", "#F44336"),
    ]):
        ax = fig.add_subplot(1, 2, idx + 1, projection="3d")
        xpos, ypos = np.meshgrid(range(n), range(n))
        xpos = xpos.flatten()
        ypos = ypos.flatten()
        zpos = np.zeros_like(xpos)
        dz = part.flatten()

        colors = [color if v >= 0 else "#9E9E9E" for v in dz]
        ax.bar3d(xpos, ypos, zpos, 0.6, 0.6, dz,
                 color=colors, alpha=0.85, edgecolor="white", lw=0.3)
        ax.set_title(f"{part_label} Part", fontsize=11, fontweight="bold")
        ax.set_xlabel("Column"); ax.set_ylabel("Row"); ax.set_zlabel("Value")
        labels = [f"|{i:0{int(np.log2(n))}b}⟩" for i in range(n)]
        ax.set_xticks(range(n)); ax.set_xticklabels(labels, fontsize=7)
        ax.set_yticks(range(n)); ax.set_yticklabels(labels, fontsize=7)

    fig.suptitle(f"PKTron {title}", fontsize=13,
                 color="darkblue", fontweight="bold")
    plt.tight_layout()
    plt.show()
    return fig


def plot_hinton(matrix_data, title="Hinton Diagram"):
    """
    Hinton diagram: squares whose size = magnitude of matrix element.
    White = positive, Black = negative.

    Usage:
        plot_hinton(dm.data)
        plot_hinton(operator.data)
    """
    matrix = np.real(np.array(matrix_data))
    n, m = matrix.shape
    fig, ax = plt.subplots(figsize=(max(6, m), max(6, n)))
    ax.set_facecolor("#607D8B")
    ax.set_aspect("equal")

    max_val = np.max(np.abs(matrix))
    if max_val == 0:
        max_val = 1

    for (row, col), val in np.ndenumerate(matrix):
        size = np.sqrt(np.abs(val) / max_val) * 0.9
        color = "white" if val >= 0 else "#212121"
        rect = mpatches.FancyBboxPatch(
            [col - size/2, row - size/2], size, size,
            boxstyle="round,pad=0.05",
            facecolor=color, edgecolor="#37474F", lw=0.8
        )
        ax.add_patch(rect)

    ax.set_xlim(-0.6, m - 0.4)
    ax.set_ylim(-0.6, n - 0.4)
    ax.set_xticks(range(m))
    ax.set_yticks(range(n))
    labels = [f"|{i:0{int(np.ceil(np.log2(max(n,2))))}b}⟩" for i in range(n)]
    ax.set_xticklabels(labels, fontsize=8)
    ax.set_yticklabels(labels, fontsize=8)
    ax.invert_yaxis()
    ax.set_title(f"PKTron {title}\n(white=positive, black=negative)",
                 fontsize=12, color="darkblue", fontweight="bold")
    plt.tight_layout()
    plt.show()
    return fig


def plot_bloch_multi(circuit, figsize=None):
    """
    Bloch sphere for every qubit in the circuit, all in one figure.

    Usage:
        plot_bloch_multi(qc)
    """
    state = circuit.get_state()
    n = circuit.num_qubits
    rho_full = np.outer(state, state.conj())

    def partial_trace_single(rho, keep, total):
        r = rho.reshape([2] * (2 * total))
        traced = []
        for i in range(total):
            if i != keep:
                ri = i - len([j for j in traced if j < i])
                ci = (total + i) - len([j for j in traced if j < total + i])
                r = np.trace(r, axis1=ri, axis2=ci)
                traced += [i, total + i]
        return r

    def bloch_vector(rho_q):
        px = np.array([[0,1],[1,0]], dtype=complex)
        py = np.array([[0,-1j],[1j,0]], dtype=complex)
        pz = np.array([[1,0],[0,-1]], dtype=complex)
        return (float(np.real(np.trace(rho_q @ px))),
                float(np.real(np.trace(rho_q @ py))),
                float(np.real(np.trace(rho_q @ pz))))

    cols = min(n, 4)
    rows = (n + cols - 1) // cols
    fig = plt.figure(figsize=figsize or (4 * cols, 4 * rows))

    u = np.linspace(0, 2 * np.pi, 25)
    v = np.linspace(0, np.pi, 25)
    xs = np.outer(np.cos(u), np.sin(v))
    ys = np.outer(np.sin(u), np.sin(v))
    zs = np.outer(np.ones(25), np.cos(v))

    for q in range(n):
        ax = fig.add_subplot(rows, cols, q + 1, projection="3d")
        ax.plot_surface(xs, ys, zs, alpha=0.08, color="#1976D2")

        # Axes lines
        for start, end in [
            ([-1.3,0,0],[1.3,0,0]), ([0,-1.3,0],[0,1.3,0]), ([0,0,-1.3],[0,0,1.3])
        ]:
            ax.plot(*zip(start, end), color="gray", lw=0.8, ls="--", alpha=0.5)

        try:
            rho_q = partial_trace_single(rho_full, q, n)
            xv, yv, zv = bloch_vector(rho_q)
            ax.quiver(0, 0, 0, xv, yv, zv,
                      color="#E53935", lw=2, arrow_length_ratio=0.15)
            ax.scatter([xv], [yv], [zv], color="#E53935", s=60, zorder=5)
            ax.text(xv * 1.15, yv * 1.15, zv * 1.15,
                    f"({xv:.2f},{yv:.2f},{zv:.2f})", fontsize=6)
        except Exception:
            pass

        ax.set_xlim([-1.3, 1.3]); ax.set_ylim([-1.3, 1.3]); ax.set_zlim([-1.3, 1.3])
        ax.set_title(f"Qubit {q}", fontsize=10, fontweight="bold")
        ax.text(1.4, 0, 0, "x", fontsize=9)
        ax.text(0, 1.4, 0, "y", fontsize=9)
        ax.text(0, 0, 1.4, "|0⟩", fontsize=9)
        ax.text(0, 0, -1.5, "|1⟩", fontsize=9)
        ax.set_axis_off()

    fig.suptitle("PKTron Multi-Qubit Bloch Spheres",
                 fontsize=13, color="darkblue", fontweight="bold")
    plt.tight_layout()
    plt.show()
    return fig


class EnhancedResult:
    """
    An enhanced result object with metadata and extra analysis.

    Usage:
        result = EnhancedResult(counts, circuit, shots=1024)
        print(result.get_counts())
        print(result.get_probabilities())
        print(result.most_likely())
        print(result.metadata)
    """

    def __init__(self, counts, circuit=None, shots=1024, backend_name="unknown"):
        self.counts = counts
        self.shots = shots
        self._circuit = circuit
        self.metadata = {
            "shots": shots,
            "backend": backend_name,
            "num_qubits": circuit.num_qubits if circuit else None,
            "gate_count": len(circuit.gate_history) if circuit else None,
            "total_counts": sum(counts.values()),
        }

    def get_counts(self):
        return self.counts

    def get_probabilities(self):
        """Returns probabilities dict from counts."""
        total = sum(self.counts.values())
        return {k: v / total for k, v in self.counts.items()}

    def most_likely(self):
        """Returns the most frequently measured state."""
        return max(self.counts, key=self.counts.get)

    def least_likely(self):
        """Returns the least frequently measured state."""
        return min(self.counts, key=self.counts.get)

    def marginal_counts(self, qubits):
        """
        Marginalize counts over specific qubits.
        qubits: list of qubit indices to keep

        Example:
            result.marginal_counts([0])  # only look at qubit 0
        """
        marginal = {}
        for bitstring, count in self.counts.items():
            key = "".join(bitstring[q] for q in qubits)
            marginal[key] = marginal.get(key, 0) + count
        return marginal

    def plot(self):
        """Quick bar chart of results."""
        import matplotlib.pyplot as plt
        probs = self.get_probabilities()
        states = sorted(probs.keys())
        vals = [probs[s] for s in states]
        fig, ax = plt.subplots(figsize=(max(6, len(states)), 4))
        bars = ax.bar(states, vals, color="#42A5F5", edgecolor="black", lw=0.8)
        ax.set_xlabel("State", fontsize=11)
        ax.set_ylabel("Probability", fontsize=11)
        ax.set_title("PKTron EnhancedResult", fontsize=12,
                     color="darkblue", fontweight="bold")
        for bar, val in zip(bars, vals):
            ax.text(bar.get_x() + bar.get_width()/2,
                    bar.get_height() + 0.005,
                    f"{val:.3f}", ha="center", fontsize=8)
        plt.xticks(rotation=45, ha="right")
        plt.tight_layout()
        plt.show()
        return fig

    def __repr__(self):
        return (f"EnhancedResult(shots={self.shots}, "
                f"most_likely={self.most_likely()}, "
                f"metadata={self.metadata})")
