"""
\********************************************************************************
* Copyright (c) 2025 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
********************************************************************************/
"""

import jax.numpy as jnp

from qrisp.jasp import qache, jrange, AbstractQubit, make_jaspr, Jaspr
from qrisp.core import x, h, cx, t, t_dg, s, measure, cz, mcx, QuantumVariable
from qrisp.qtypes import QuantumBool
from qrisp.environments import control, custom_control
from qrisp.alg_primitives.arithmetic.jasp_arithmetic.jasp_cq_gidney_adder import jasp_cq_gidney_adder

# Addition implementation based on https://arxiv.org/pdf/1709.06648

@custom_control
def jasp_qq_gidney_adder(a, b, ctrl = None):
    
    if isinstance(b, list):
        n = min(len(a), len(b))
        perform_incrementation = n < len(b)
    else:
        n = jnp.min(jnp.array([a.size, b.size]))
        perform_incrementation = n < b.size
    
    if ctrl is not None:
        ctrl_anc = QuantumBool(name = "gidney_anc_2*")
    
    # If the addition is only a single qubit, it can be done with a CX gate (below)
    with control(n > 1):
        
        # Allocate the ancillae        
        gidney_anc = QuantumVariable(n-1, name = "gidney_anc*")
        
        i = 0
        # Perform the initial mcx        
        mcx([a[i], b[i]], gidney_anc[i], method = "gidney")
        
        # Perform the left part of the V-Shape
        for j in jrange(n-2):
            i = j+1
        
            cx(gidney_anc[i-1], a[i])
            cx(gidney_anc[i-1], b[i])
            
            mcx([a[i], b[i]], gidney_anc[i], method = "gidney")
            cx(gidney_anc[i-1], gidney_anc[i])
            
        # This part handles the case that the addition target has more qubit than
        # the control value.
        # We solve this issue by performing a 1-incrementation on the remainder
        # if the carry out value is True
        with control(perform_incrementation):
            
            # Compute the carry out similar to the loop above
            
            carry_out = QuantumBool()
            
            cx(gidney_anc[n-2], a[n-1])
            cx(gidney_anc[n-2], b[n-1])
            
            mcx([a[n-1], b[n-1]], carry_out[0], method = "gidney")
            cx(gidney_anc[n-2], carry_out[0])

            # Perform a controlled incrtementation            
            ctrl_list = [carry_out[0]]
            if ctrl is not None:
                ctrl_list.append(ctrl)
                
            with control(ctrl_list):
                jasp_cq_gidney_adder(1, b[n:])
            
            # Uncompute the carry
            cx(gidney_anc[n-2], carry_out[0])
            mcx([a[n-1], b[n-1]], carry_out[0], method = "gidney_inv")
            
            if ctrl is not None:
                mcx([ctrl, gidney_anc[n-2]], ctrl_anc[0], method = "gidney")
                cx(ctrl_anc[0], b[n-1])
                mcx([ctrl, gidney_anc[n-2]], ctrl_anc[0], method = "gidney_inv")
            else:
                cx(gidney_anc[n-2], b[n-1])
            cx(gidney_anc[n-2], a[n-1])
            
            carry_out.delete()
        
        # This is the CX at the "tip" of the V shape
        cx(gidney_anc[n-2], b[n-1])
        
        # This is the CX at the lower right of the circuit
        if ctrl is not None:
            mcx([ctrl, a[n-1]], ctrl_anc[0], method = "gidney")
            cx(ctrl_anc[0], b[n-1])
            mcx([ctrl, a[n-1]], ctrl_anc[0], method = "gidney_inv")
        else:
            cx(a[n-1], b[n-1])
        
        # Perform the right part of the V shape
        for j in jrange(n-2):
            i = n-j-2
            
            cx(gidney_anc[i-1], gidney_anc[i])
            mcx([a[i], b[i]], gidney_anc[i], method = "gidney_inv")
            
            if ctrl is not None:
                # This is the controlled version described on page 4
                pass
                mcx([ctrl, a[i]], ctrl_anc[0], method = "gidney")
                cx(ctrl_anc[0], b[i])
                mcx([ctrl, a[i]], ctrl_anc[0], method = "gidney_inv")
                cx(gidney_anc[i-1], a[i])
                cx(gidney_anc[i-1], b[i])
            else:
                cx(gidney_anc[i-1], a[i])
                cx(a[i], b[i])
    
        # The final uncomputation
        mcx([a[0], b[0]], gidney_anc[0], method = "gidney_inv")
        
        # Delete the ancilla
        gidney_anc.delete()
    
    # Perform the CX gate at the top right of the circuit
    if ctrl is not None:
        mcx([ctrl, a[0]], ctrl_anc[0], method = "gidney")
        cx(ctrl_anc[0], b[0])
        mcx([ctrl, a[0]], ctrl_anc[0], method = "gidney_inv")
    else:
        cx(a[0], b[0])
    
    if ctrl is not None:
        ctrl_anc.delete()        

