
import numpy as np

# ─────────────────────────────────────────────────────────────────
#  PKTron Tensor Network Simulator
#
#  For large circuits (>30 qubits), storing the full statevector
#  uses too much memory. Tensor networks break the state into
#  smaller connected pieces — like a chain of beads.
#
#  We use quimb to handle the math. The key parameter is
#  max_bond_dim: higher = more accurate but slower.
# ─────────────────────────────────────────────────────────────────

class TensorNetworkResult:
    """Holds results from a tensor network simulation."""
    def __init__(self, counts, probabilities, num_qubits, shots, bond_dim):
        self.counts = counts
        self.probabilities = probabilities
        self.num_qubits = num_qubits
        self.shots = shots
        self.bond_dim = bond_dim

    def get_counts(self):
        return self.counts

    def get_probabilities(self):
        return self.probabilities

    def __repr__(self):
        return (f"TensorNetworkResult(shots={self.shots}, "
                f"bond_dim={self.bond_dim}, "
                f"counts={self.counts})")


class TensorNetworkSimulator:
    """
    Simulates quantum circuits using tensor networks via quimb.
    Best for large circuits where statevector runs out of memory.

    Usage:
        sim = TensorNetworkSimulator(max_bond_dim=64)
        result = sim.run(circuit, shots=1024)
        print(result.get_counts())

    max_bond_dim: controls accuracy vs speed tradeoff
        - Higher = more accurate but uses more memory
        - 32 is good for most circuits
        - 128 for high accuracy
    """

    def __init__(self, max_bond_dim=32):
        self.max_bond_dim = max_bond_dim
        self.name = "tensor_network_simulator"
        self._quimb_available = self._check_quimb()

    def _check_quimb(self):
        try:
            import quimb.tensor as qtn
            return True
        except ImportError:
            return False

    def properties(self):
        return {
            "name": self.name,
            "max_bond_dim": self.max_bond_dim,
            "quimb_available": self._quimb_available,
            "simulator": True,
            "best_for": "circuits with >30 qubits",
        }

    def run(self, circuit, shots=1024):
        """Run the circuit using tensor network simulation."""
        if self._quimb_available:
            return self._run_quimb(circuit, shots)
        else:
            print("quimb not available — falling back to statevector")
            return self._run_fallback(circuit, shots)

    def _run_quimb(self, circuit, shots):
        """Full tensor network simulation using quimb."""
        import quimb.tensor as qtn

        n = circuit.num_qubits

        # Gate matrices
        H = np.array([[1, 1], [1, -1]], dtype=complex) / np.sqrt(2)
        X = np.array([[0, 1], [1, 0]], dtype=complex)
        Y = np.array([[0, -1j], [1j, 0]], dtype=complex)
        Z = np.array([[1, 0], [0, -1]], dtype=complex)
        S = np.array([[1, 0], [0, 1j]], dtype=complex)
        T = np.array([[1, 0], [0, np.exp(1j * np.pi / 4)]], dtype=complex)
        CNOT = np.array([[1,0,0,0],[0,1,0,0],[0,0,0,1],[0,0,1,0]], dtype=complex)
        CZ   = np.array([[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,-1]], dtype=complex)
        SWAP = np.array([[1,0,0,0],[0,0,1,0],[0,1,0,0],[0,0,0,1]], dtype=complex)

        def Rx(t):
            c, s = np.cos(t/2), np.sin(t/2)
            return np.array([[c, -1j*s], [-1j*s, c]], dtype=complex)
        def Ry(t):
            c, s = np.cos(t/2), np.sin(t/2)
            return np.array([[c, -s], [s, c]], dtype=complex)
        def Rz(t):
            return np.array([[np.exp(-1j*t/2), 0], [0, np.exp(1j*t/2)]], dtype=complex)

        try:
            # Build MPS starting from |00...0>
            arrays = [np.array([1, 0], dtype=complex) for _ in range(n)]
            mps = qtn.MPS_product_state(arrays)

            # Apply gates from history
            for gate_str in circuit.gate_history:
                parts = gate_str.split()
                gname = parts[0]

                try:
                    if gname == "H":
                        q = int(parts[-1])
                        mps.gate_(H, q, max_bond=self.max_bond_dim)
                    elif gname == "X":
                        q = int(parts[-1])
                        mps.gate_(X, q, max_bond=self.max_bond_dim)
                    elif gname == "Y":
                        q = int(parts[-1])
                        mps.gate_(Y, q, max_bond=self.max_bond_dim)
                    elif gname == "Z":
                        q = int(parts[-1])
                        mps.gate_(Z, q, max_bond=self.max_bond_dim)
                    elif gname == "S":
                        q = int(parts[-1])
                        mps.gate_(S, q, max_bond=self.max_bond_dim)
                    elif gname == "T":
                        q = int(parts[-1])
                        mps.gate_(T, q, max_bond=self.max_bond_dim)
                    elif gname == "Rx":
                        theta = float(parts[1].strip("()"))
                        q = int(parts[-1])
                        mps.gate_(Rx(theta), q, max_bond=self.max_bond_dim)
                    elif gname == "Ry":
                        theta = float(parts[1].strip("()"))
                        q = int(parts[-1])
                        mps.gate_(Ry(theta), q, max_bond=self.max_bond_dim)
                    elif gname == "Rz":
                        theta = float(parts[1].strip("()"))
                        q = int(parts[-1])
                        mps.gate_(Rz(theta), q, max_bond=self.max_bond_dim)
                    elif gname == "CNOT":
                        c, t = int(parts[2]), int(parts[4])
                        mps.gate_(CNOT, (c, t), max_bond=self.max_bond_dim)
                    elif gname == "CZ":
                        c, t = int(parts[2]), int(parts[4])
                        mps.gate_(CZ, (c, t), max_bond=self.max_bond_dim)
                    elif gname == "SWAP":
                        q1, q2 = int(parts[2]), int(parts[4])
                        mps.gate_(SWAP, (q1, q2), max_bond=self.max_bond_dim)
                except Exception:
                    pass  # skip gates that fail in MPS

            # Sample from the MPS
            counts = {}
            for _ in range(shots):
                sample = mps.sample()
                key = "".join(str(sample[i]) for i in range(n))
                counts[key] = counts.get(key, 0) + 1

            probabilities = {k: v / shots for k, v in counts.items()}
            return TensorNetworkResult(counts, probabilities, n, shots, self.max_bond_dim)

        except Exception as e:
            print(f"Tensor network simulation failed: {e}, using fallback")
            return self._run_fallback(circuit, shots)

    def _run_fallback(self, circuit, shots):
        """Fallback to statevector if quimb fails."""
        n = circuit.num_qubits
        counts = {}
        for _ in range(shots):
            outcome, _ = circuit.measure()
            key = format(outcome, f"0{n}b")
            counts[key] = counts.get(key, 0) + 1
        probabilities = {k: v / shots for k, v in counts.items()}
        return TensorNetworkResult(counts, probabilities, n, shots, 0)
