"""
********************************************************************************
* Copyright (c) 2026 the Qrisp authors
*
* This program and the accompanying materials are made available under the
* terms of the Eclipse Public License 2.0 which is available at
* http://www.eclipse.org/legal/epl-2.0.
*
* This Source Code may also be made available under the following Secondary
* Licenses when the conditions for such availability set forth in the Eclipse
* Public License, v. 2.0 are satisfied: GNU General Public License, version 2
* with the GNU Classpath Exception which is
* available at https://www.gnu.org/software/classpath/license.html.
*
* SPDX-License-Identifier: EPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
********************************************************************************
"""

"""
Boolean Simulation Decorator
============================

This module provides the ``boolean_simulation`` decorator, which enables efficient
classical simulation of quantum programs that contain only boolean/classical logic.

The decorator transforms Jasp functions into pure JAX expressions, allowing them to
be JIT-compiled and executed without any quantum simulation overhead. This is 
particularly useful for:

1. **Verifying uncomputation correctness**: The simulation checks that all deleted
   qubits are properly uncomputed (in the |0⟩ state) and warns otherwise.

2. **Testing classical quantum algorithms at scale**: Algorithms like adders,
   multipliers, and other arithmetic circuits can be tested with large inputs.

3. **Performance benchmarking**: Since the simulation compiles to efficient JAX
   code, it can be used to benchmark the classical logic portion of algorithms.

Supported Operations:
- X, SWAP gates (unconditional bit operations)
- Controlled variants: CX, CCX, CSWAP, multi-controlled X, etc.
- Phase gates (Z, S, T, RZ, P): These are no-ops in classical simulation
- Measurement: Returns the current classical bit values

Unsupported Operations:
- Any gate creating superposition: H, RX, RY, SX, etc.
- These will raise an exception during simulation
"""

from typing import TYPE_CHECKING, Any, Callable

import jax.numpy as jnp
from jax import jit, Array
from jax.tree_util import tree_unflatten

from qrisp.jasp import make_jaspr, AbstractQubitArray, AbstractQubit

from qrisp.jasp.interpreter_tools.interpreters.cl_func_interpreter import (
    jaspr_to_cl_func_jaxpr,
)
from qrisp.jasp.interpreter_tools import Jlist, eval_jaxpr

from qrisp.jasp import Jaspr


def boolean_simulation(*func: Callable, bit_array_padding: int = 2**16) -> Callable:
    """
    Decorator to simulate Jasp functions containing only classical logic (like X, CX, CCX etc.).

    This decorator transforms the function into a JAX expression without any
    quantum primitives and leverages the JAX compilation pipeline to compile
    a highly efficient simulation.

    .. note::

        The ``boolean_simulation`` decorator will check if deleted
        :ref:`QuantumVariables <QuantumVariable>`
        have been properly uncomputed and submit a warning otherwise.
        It therefore provides a valuable tool for verifying the correctness
        of your algorithms at scale.

    Parameters
    ----------
    func : Callable
        A Python function performing Jasp logic.
    bit_array_padding : int, optional
        An integer specifying the size of the classical array containing the
        (classical) bits which are simulated. Since JAX doesn't allow dynamically
        sized arrays but Jasp supports dynamically sized QuantumVariables, the
        array has to be "padded". The padding therefore indicates an upper boundary
        for how many qubits are required to execute ``func``. A large padding
        slows down the simulation but prevents overflow errors.The default is
        ``2**16``. The minimum value is 64. This threshold describes the maximum
        OVERALL amount of qubits that can appear in the simulation. The maximum
        amount per QuantumVariable/per QuantumArray is tied to this number
        and will always be ``1/64`` of the ``bit_array_padding``.
        The default here is therefore ``2**16/2**6 = 1024``.

    Returns
    -------
    Callable
        A function performing the simulation for the given input parameters.

    Raises
    ------
    Exception
        If ``bit_array_padding`` is less than 64.
    Exception
        If the function returns a quantum value (must measure before returning).

    Examples
    --------
    We create a simple script that demonstrates the functionality:

    ::

        from qrisp import QuantumFloat, measure
        from qrisp.jasp import boolean_simulation, jrange

        @boolean_simulation
        def main(i, j):

            a = QuantumFloat(10)

            b = QuantumFloat(10)

            a[:] = i
            b[:] = j

            c = QuantumFloat(30)

            for i in jrange(150):
                c += a*b

            return measure(c)

    This script evaluates the multiplication of the two inputs 150 times and adds
    them into the same QuantumFloat. The respected result is therefore ``i*j*150``.

    >>> main(1,2)
    Array(300., dtype=float64)
    >>> main(3,4)
    Array(1800., dtype=float64)

    Next we demonstrate the behavior under a faulty uncomputation:

    ::

        @boolean_simulation
        def main(i):

            a = QuantumFloat(10)
            a[:] = i
            a.delete()
            return


    >>> main(0)
    >>> main(1)
    WARNING: Faulty uncomputation found during simulation.
    >>> main(3)
    WARNING: Faulty uncomputation found during simulation.
    WARNING: Faulty uncomputation found during simulation.

    For the first case, the deletion is valid, because ``a`` is initialized in the
    $\\ket{0}$ state. For the second case, the first qubit is in the $\\ket{1}$
    state, so the deletion is not valid. The third case has both the first and
    the second qubit in the $\\ket{1}$ state (because 3 = ``11`` in binary) so
    there are two warnings.

    **Padding**

    We demonstrate the effects of the padding feature. For this we recreate the
    above script but with different padding selections.

    ::

        @boolean_simulation(bit_array_padding = 64)
        def main(i, j):

            a = QuantumFloat(10)

            b = QuantumFloat(10)

            a[:] = i
            b[:] = j

            c = QuantumFloat(30)

            for i in jrange(150):
                c += a*b

            return measure(c)

    >>> main(1,2)
    Array(8.92323439e+08, dtype=float64)

    A faulty result because the script needs more than 64 qubits.

    Increasing the padding ensures that enough qubits are available at the cost
    of simulation speed.
    """
    # Handle both @boolean_simulation and @boolean_simulation(...) syntax
    if len(func) == 0:
        # Called with arguments: @boolean_simulation(bit_array_padding=...)
        return lambda x: boolean_simulation(x, bit_array_padding=bit_array_padding)
    else:
        # Called without arguments: @boolean_simulation
        func = func[0]

    if bit_array_padding < 64:
        raise Exception(
            "Tried to initialize boolean_simulation with less than 512 bits"
        )

    @jit
    def return_function(*args: Any) -> Any:
        """
        JIT-compiled simulation function that transforms and executes the quantum program.

        This inner function:
        1. Creates a Jaspr from the decorated function (with output tree structure)
        2. Validates that no quantum values are returned (must be measured first)
        3. Transforms the Jaspr to a classical Jaxpr
        4. Initializes the boolean quantum circuit state (bit array + free qubit pool)
        5. Executes the classical Jaxpr and reconstructs the PyTree structure

        Parameters
        ----------
        *args : Any
            Arguments to pass to the original function.

        Returns
        -------
        Any
            The result of the simulation, with the original PyTree structure preserved:
            - None if the function has no return value
            - dict, list, tuple, or nested structure matching the original return type
        """
        # Create the Jaspr representation of the quantum program
        # Use return_shape=True to capture the output PyTree structure
        jaspr: Jaspr
        jaspr, out_tree = make_jaspr(func, return_shape=True)(*args)

        # Validate that no quantum values are returned
        # (quantum values must be measured before returning from boolean simulation)
        for var in jaspr.outvars:
            if isinstance(var.aval, (AbstractQubitArray, AbstractQubit)):
                raise Exception(
                    "Tried to perform boolean simulation of a function returning "
                    "a quantum value (please measure before returning)"
                )

        # Transform the Jaspr to a classical Jaxpr
        # flatten_environments() resolves any nested control flow structures
        cl_func_jaxpr = jaspr_to_cl_func_jaxpr(
            jaspr.flatten_environments(), bit_array_padding
        )

        # Initialize the boolean quantum circuit representation:
        # - bit_array: packed uint64 array storing qubit states (all zeros initially)
        # - free_qubit_list: Jlist containing all available qubit indices
        aval = cl_func_jaxpr.jaxpr.invars[-3].aval
        bit_array: Array = jnp.zeros(aval.shape, dtype=aval.dtype)

        # Create the free qubit pool as a flattened Jlist
        free_qubit_list = Jlist(
            jnp.arange(bit_array_padding), max_size=bit_array_padding
        ).flatten()[0]

        # The boolean quantum circuit is represented as a tuple:
        # (bit_array, jlist_array, jlist_counter)
        boolean_quantum_circuit: tuple[Array, ...] = (bit_array, *free_qubit_list)

        # Combine user arguments with the quantum circuit state
        ammended_args: list[Any] = list(args) + list(boolean_quantum_circuit)

        # Execute the classical Jaxpr
        res = eval_jaxpr(cl_func_jaxpr)(*ammended_args)

        # Extract the return values from the result
        # The result always contains 3 trailing elements for quantum circuit state:
        # (bit_array, jlist_array, jlist_counter)
        # The leading elements are the actual return values from the user's function.
        num_output_values = len(res) - 3

        if num_output_values == 0:
            # Function returns nothing (only quantum circuit state)
            return None
        else:
            # Extract user return values and reconstruct the PyTree structure
            flat_results = res[:num_output_values]
            return tree_unflatten(out_tree, flat_results)

    return return_function
