"""Quantum Petri net: superposition tokens on places.

Classical SPN tokens are probabilities in [0,1] per place.  Quantum Petri net
encodes tokens as qubit amplitudes -- a transition fires on all branches in
superposition.  Measurement collapses to a single control decision.
"""

from __future__ import annotations

import numpy as np
from qiskit import QuantumCircuit
from qiskit.circuit.library import RYGate
from qiskit.quantum_info import Statevector

from .._constants import WEIGHT_SPARSITY_EPS
from ..bridge.sc_to_quantum import probability_to_angle


class QuantumPetriNet:
    """Quantum Petri net with amplitude-encoded token state.

    Each place is one qubit.  Token density p maps to Ry(2*arcsin(sqrt(p))).
    Transitions apply controlled rotations based on arc weights.
    """

    def __init__(
        self,
        n_places: int,
        n_transitions: int,
        W_in: np.ndarray,
        W_out: np.ndarray,
        thresholds: np.ndarray,
    ):
        """W_in: (n_transitions, n_places), W_out: (n_places, n_transitions)."""
        if n_places <= 0 or n_transitions <= 0:
            raise ValueError(
                f"n_places ({n_places}) and n_transitions ({n_transitions}) must be positive"
            )
        W_in = np.asarray(W_in, dtype=np.float64)
        W_out = np.asarray(W_out, dtype=np.float64)
        thresholds = np.asarray(thresholds, dtype=np.float64)
        if W_in.shape != (n_transitions, n_places):
            raise ValueError(f"W_in shape {W_in.shape} != expected ({n_transitions}, {n_places})")
        if W_out.shape != (n_places, n_transitions):
            raise ValueError(
                f"W_out shape {W_out.shape} != expected ({n_places}, {n_transitions})"
            )
        if len(thresholds) != n_transitions:
            raise ValueError(
                f"thresholds length {len(thresholds)} != n_transitions {n_transitions}"
            )
        self.n_places = n_places
        self.n_transitions = n_transitions
        self.W_in = W_in
        self.W_out = W_out
        self.thresholds = thresholds

    def encode_marking(self, marking: np.ndarray) -> QuantumCircuit:
        """Amplitude-encode token state: p_i -> Ry(theta_i)|0>."""
        qc = QuantumCircuit(self.n_places)
        for i, m in enumerate(marking):
            theta = probability_to_angle(float(np.clip(m, 0.0, 1.0)))
            qc.ry(theta, i)
        return qc

    def apply_transition(self, circuit: QuantumCircuit, t_idx: int):
        """Apply transition t_idx as controlled rotations.

        Input arcs: consume tokens (negative rotation on input places).
        Output arcs: produce tokens controlled by ALL input places (AND gating).
        All input-place qubits must be |1> for the output rotation to fire,
        matching conjunctive Petri net semantics.
        """
        input_places = [
            p for p in range(self.n_places) if abs(self.W_in[t_idx, p]) > WEIGHT_SPARSITY_EPS
        ]

        for p in input_places:
            theta = probability_to_angle(float(abs(self.W_in[t_idx, p]))) * self.thresholds[t_idx]
            circuit.ry(-theta, p)

        for p_out in range(self.n_places):
            w_out = self.W_out[p_out, t_idx]
            if abs(w_out) < WEIGHT_SPARSITY_EPS:
                continue
            theta = probability_to_angle(float(abs(w_out)))
            controls = [p for p in input_places if p != p_out]
            if len(controls) == 0:
                circuit.ry(theta, p_out)
            elif len(controls) == 1:
                circuit.cry(theta, controls[0], p_out)
            else:
                gate = RYGate(theta).control(len(controls))
                circuit.append(gate, controls + [p_out])

    def step(self, marking: np.ndarray) -> np.ndarray:
        """Full step: encode marking, apply all transitions, measure -> new marking."""
        qc = self.encode_marking(marking)
        for t in range(self.n_transitions):
            self.apply_transition(qc, t)

        sv = Statevector.from_instruction(qc)
        new_marking: np.ndarray = np.zeros(self.n_places)
        for p in range(self.n_places):
            probs = sv.probabilities([p])
            new_marking[p] = probs[1]  # P(|1>) = token density
        return new_marking

    @classmethod
    def from_matrices(
        cls,
        W_in: np.ndarray,
        W_out: np.ndarray,
        thresholds: np.ndarray,
    ) -> QuantumPetriNet:
        """Construct from arc weight matrices, inferring n_places and n_transitions."""
        n_t, n_p = W_in.shape
        return cls(n_p, n_t, W_in, W_out, thresholds)
