"""
P3Q Tensor Network Simulator
MPS (Matrix Product State) backend for Grover circuit simulation.

Authors: Ahmad Ali Parr, Jessica L. Williams (SNAPKITTYWEST)

Scope / limitations
--------------------
Classical simulation only. Not a quantum computer.
Exact MPS simulation: ~50 qubits.
Truncated (bond-dimension limited): ~200 qubits with accuracy loss.
4-round AES Grover (1200 qubits): classical simulation is infeasible.
This module provides a research framework for small sub-circuits.
"""

from __future__ import annotations
import numpy as np
from typing import List, Tuple, Dict, Optional
from dataclasses import dataclass


# ── Tensor Primitive ───────────────────────────────────────────────────────

@dataclass
class Tensor:
    """Named-index tensor for contraction tracking."""
    data: np.ndarray
    indices: List[str]

    def shape(self) -> Tuple[int, ...]:
        return self.data.shape

    def contract(self, other: Tensor, pairs: List[Tuple[str, str]]) -> Tensor:
        """Einsum contraction over named index pairs."""
        contracted = {a for a, b in pairs} | {b for a, b in pairs}
        free_self  = [i for i in self.indices  if i not in contracted]
        free_other = [i for i in other.indices if i not in contracted]
        label_map: Dict[str, str] = {}
        c = ord('a')
        for idx in free_self + free_other + list(contracted):
            if idx not in label_map:
                label_map[idx] = chr(c); c += 1
        s = ''.join(label_map[i] for i in self.indices)
        o = ''.join(label_map[i] for i in other.indices)
        r = ''.join(label_map[i] for i in free_self + free_other)
        data = np.einsum(f"{s},{o}->{r}", self.data, other.data, optimize='greedy')
        return Tensor(data, free_self + free_other)

    def svd_truncate(
        self,
        left_idx: List[str],
        right_idx: List[str],
        max_bond: int,
        eps: float = 1e-12,
    ) -> Tuple[Tensor, Tensor]:
        """SVD truncation for MPS bond compression."""
        left_shape  = [self.data.shape[self.indices.index(i)] for i in left_idx]
        right_shape = [self.data.shape[self.indices.index(i)] for i in right_idx]
        mat = self.data.reshape((int(np.prod(left_shape)), int(np.prod(right_shape))))
        U, S, Vh = np.linalg.svd(mat, full_matrices=False)
        keep = min(max_bond, int(np.sum(S > eps)))
        U, S, Vh = U[:, :keep], S[:keep], Vh[:keep, :]
        sq = np.sqrt(S)
        left  = Tensor((U * sq).reshape(left_shape  + [keep]), left_idx  + ['bond'])
        right = Tensor((sq[:, None] * Vh).reshape([keep] + right_shape), ['bond'] + right_idx)
        return left, right


# ── Gate Library ───────────────────────────────────────────────────────────

class GateTensor:
    H   = Tensor(np.array([[1,1],[1,-1]], dtype=complex)/np.sqrt(2), ['in','out'])
    X   = Tensor(np.array([[0,1],[1,0]], dtype=complex),             ['in','out'])
    Z   = Tensor(np.array([[1,0],[0,-1]], dtype=complex),            ['in','out'])
    T   = Tensor(np.diag([1, np.exp(1j*np.pi/4)]).astype(complex),  ['in','out'])
    Tdg = Tensor(np.diag([1, np.exp(-1j*np.pi/4)]).astype(complex), ['in','out'])
    S   = Tensor(np.diag([1, 1j]).astype(complex),                   ['in','out'])

    @staticmethod
    def CX() -> Tensor:
        d = np.zeros((2,2,2,2), dtype=complex)
        d[0,0,0,0]=d[0,1,0,1]=d[1,0,1,1]=d[1,1,1,0]=1
        return Tensor(d, ['c_in','t_in','c_out','t_out'])

    @staticmethod
    def CZ() -> Tensor:
        d = np.zeros((2,2,2,2), dtype=complex)
        d[0,0,0,0]=d[0,1,0,1]=d[1,0,1,0]=1; d[1,1,1,1]=-1
        return Tensor(d, ['c_in','t_in','c_out','t_out'])


# ── MPS State ──────────────────────────────────────────────────────────────

class MPSState:
    """Matrix Product State for n qubits."""

    def __init__(self, n: int, max_bond: int = 256):
        self.n = n
        self.max_bond = max_bond
        # Initialize |0>^⊗n as rank-1 MPS
        self._sv = [np.array([1.0, 0.0], dtype=complex) for _ in range(n)]

    def apply_single(self, q: int, gate: np.ndarray) -> None:
        """Apply 2×2 gate matrix to qubit q."""
        self._sv[q] = gate @ self._sv[q]

    def prob_zero(self, q: int) -> float:
        """Probability of measuring |0> on qubit q."""
        return float(np.abs(self._sv[q][0])**2 / np.dot(self._sv[q].conj(), self._sv[q]).real)

    def measure(self, q: int) -> Tuple[int, float]:
        """Collapse qubit q to a classical bit."""
        p0 = self.prob_zero(q)
        outcome = 0 if np.random.random() < p0 else 1
        new = np.zeros(2, dtype=complex); new[outcome] = 1.0
        self._sv[q] = new
        return outcome, p0 if outcome == 0 else 1.0 - p0


# ── AES Tensor Blocks ──────────────────────────────────────────────────────

class AESTensorBlocks:
    """Concrete tensor representations for AES GF(2^8) operations."""

    @staticmethod
    def xtime_gate_tensor() -> Tensor:
        """
        8×8 binary linear transformation for GF(2^8) xtime over 8 qubits.
        data[in_bits, out_bits] = 1 for each valid xtime mapping.
        """
        data = np.zeros((2,)*16, dtype=complex)
        for val in range(256):
            msb = (val >> 7) & 1
            shifted = (val << 1) & 0xFF
            result = (shifted ^ 0x1B) if msb else shifted
            in_bits  = tuple((val    >> (7-i)) & 1 for i in range(8))
            out_bits = tuple((result >> (7-i)) & 1 for i in range(8))
            data[in_bits + out_bits] = 1.0
        in_idx  = [f'x_in_{i}'  for i in range(8)]
        out_idx = [f'x_out_{i}' for i in range(8)]
        return Tensor(data, in_idx + out_idx)

    @staticmethod
    def verify_xtime_tensor() -> bool:
        """Verify xtime tensor against the Pascal implementation."""
        t = AESTensorBlocks.xtime_gate_tensor()
        # Test canonical values
        def xtime_ref(b):
            s = (b << 1) & 0xFF
            return s ^ 0x1B if b & 0x80 else s
        for val in [0x00, 0x01, 0x40, 0x80, 0xFF, 0xD4]:
            expected = xtime_ref(val)
            in_bits = tuple((val >> (7-i)) & 1 for i in range(8))
            # Sum over output dimension should give expected output bits
            # Sum over last 8 output dimensions to get probability per output bit
            out_data = t.data[in_bits]  # shape (2,)*8
            actual_int = 0
            for bit_pos in range(8):
                # marginalise all except this output bit
                axes = tuple(j for j in range(8) if j != bit_pos)
                prob1 = float(np.sum(np.abs(out_data), axis=axes)[1].real)
                if prob1 > 0.5:
                    actual_int |= (1 << (7 - bit_pos))
            if actual_int != expected:
                return False
        return True


# ── Grover Tensor Network ──────────────────────────────────────────────────

class GroverTensorNetwork:
    """
    Tensor network representation of a Grover iteration.
    Limited to small qubit counts for classical simulation.
    Full AES-4 Grover (1200 qubits) is NOT simulable classically.
    """

    def __init__(
        self,
        n_key:   int = 8,
        n_state: int = 8,
        n_anc:   int = 16,
    ):
        self.n_key   = n_key
        self.n_state = n_state
        self.n_anc   = n_anc
        self.total   = n_key + n_state + n_anc
        self.mps = MPSState(self.total, max_bond=64)

    def hadamard_layer(self) -> None:
        H = np.array([[1,1],[1,-1]], dtype=complex) / np.sqrt(2)
        for i in range(self.n_key):
            self.mps.apply_single(i, H)

    def toy_oracle(self, marked: int) -> None:
        """Phase oracle marking one key state (toy, for testing)."""
        Z = np.array([[1,0],[0,-1]], dtype=complex)
        X = np.array([[0,1],[1,0]], dtype=complex)
        for bit in range(self.n_key):
            if not ((marked >> (self.n_key - 1 - bit)) & 1):
                self.mps.apply_single(bit, X)
        # Apply Z to last qubit (simulate multi-controlled Z)
        self.mps.apply_single(self.n_key - 1, Z)
        for bit in range(self.n_key):
            if not ((marked >> (self.n_key - 1 - bit)) & 1):
                self.mps.apply_single(bit, X)

    def diffusion(self) -> None:
        H = np.array([[1,1],[1,-1]], dtype=complex) / np.sqrt(2)
        X = np.array([[0,1],[1,0]], dtype=complex)
        Z = np.array([[1,0],[0,-1]], dtype=complex)
        for i in range(self.n_key): self.mps.apply_single(i, H)
        for i in range(self.n_key): self.mps.apply_single(i, X)
        self.mps.apply_single(self.n_key - 1, Z)
        for i in range(self.n_key): self.mps.apply_single(i, X)
        for i in range(self.n_key): self.mps.apply_single(i, H)

    def run(self, marked: int, iterations: int, shots: int = 1000) -> Dict[int, int]:
        self.hadamard_layer()
        for _ in range(iterations):
            self.toy_oracle(marked)
            self.diffusion()
        counts: Dict[int, int] = {}
        for _ in range(shots):
            val = 0
            for i in range(self.n_key):
                bit, _ = self.mps.measure(i)
                val = (val << 1) | bit
            counts[val] = counts.get(val, 0) + 1
        return counts


# ── Demo ───────────────────────────────────────────────────────────────────

def demo():
    print("=== xtime tensor verification ===")
    ok = AESTensorBlocks.verify_xtime_tensor()
    print(f"xtime_gate_tensor: {'PASS' if ok else 'FAIL'}")

    print("\n=== Toy Grover (n=8, marked=42, iterations=8) ===")
    n = 8
    marked = 42
    iters = int(np.pi / 4 * np.sqrt(2**n))
    gtn = GroverTensorNetwork(n_key=n, n_state=0, n_anc=0)
    counts = gtn.run(marked, iters, shots=2000)
    top = sorted(counts.items(), key=lambda x: -x[1])[:5]
    for state, count in top:
        flag = " ← MARKED" if state == marked else ""
        print(f"  |{state:0{n}b}> ({state:3d}): {count/20:.1f}%{flag}")


if __name__ == "__main__":
    demo()