"""
\********************************************************************************
* Copyright (c) 2023 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 numpy as np

from qrisp.qtypes.quantum_float import QuantumFloat
from qrisp.alg_primitives.arithmetic import multi_controlled_U_g, hybrid_mult
from qrisp.alg_primitives import QFT, h, cx, swap, mcx
from qrisp.environments import conjugate, control
from qrisp.alg_primitives.arithmetic.modular_arithmetic.mod_tools import modinv, montgomery_encoder

def QREDC(t, N, m):
    
    u = QuantumFloat(0, qs = t.qs, name = "u*")
    
    for i in range(m):
        
        h(t[0])
        
        transfer_lsb(t, u)
        
        t.exp_shift(-1)
        multi_controlled_U_g(t, [u[-1]], -N/2)

    
    QFT(t, inv = True, exec_swap = False)
    
    with conjugate(QFT)(t[:-1], exec_swap = False):
        
        multi_controlled_U_g(t[:-1], [t.sign()], N)
    
    cx(t[0], t.sign())
    
    sgn = t.reg.pop(-1)
    t.size -= 1
    t.signed = False
    
    u.reg.insert(len(u), sgn)
    u.size += 1
    u.mshape[1] += 1
    
    
    t.m = t.m - m
    
    return t, u


def find_optimal_m(b, N):
    m = 0
    n = int(np.ceil(np.log2(N)))
    
    while True:
        b_trial = montgomery_encoder(b, 2**m, N)%N
        if (N-1)*b_trial < N*2**(m):
            return m
        m += 1

def montgomery_mod_semi_mul(a, b, output_qg = None, permeable_if_zero = False):
    
    
    N = a.modulus
    
    b = b%N
    
    m = find_optimal_m(b, N)
    
    if b == 0:
        if output_qg is None:
            return a.duplicate()
        else:
            return output_qg
    if b == 1:
        if output_qg is None:
            return a.duplicate(init = True)
        else:
            output_qg += a
            return output_qg
    
    m = find_optimal_m(b, N)
    b = montgomery_encoder(b, 2**m, N)%N
    
    
    if output_qg is None:
        t = QuantumFloat(a.size+m, signed = True)
        
        h(t)
    else:
        if output_qg.modulus != N:
            raise Exception("Output QuantumModulus has incompatible modulus")
        
        output_qg.extend(m, 0)
        output_qg.add_sign()
        
        QFT(output_qg, exec_swap = False)
        
        t = output_qg
    
    
    for i in range(a.size):
        multi_controlled_U_g(t, [a[i]], (2**i*b))
    
    from qrisp import QuantumModulus
    t.__class__ = QuantumModulus
    t.modulus = a.modulus
    t.m = a.m + m
    t.inpl_adder = a.inpl_adder
    
    
    
    return montgomery_red(t, a, b, N, m, permeable_if_zero = permeable_if_zero)


def montgomery_red(t, a, b, N, m, permeable_if_zero = False):
    
    t, u = QREDC(t, N, m)
    
    with conjugate(QFT)(u, exec_swap = False):
        if isinstance(b, QuantumFloat):
            
            hybrid_mult(a, 
                        b, 
                        output_qf=u, 
                        init_op = None, 
                        terminal_op = None, 
                        cl_factor = -modinv(N, 2**(m+1)))
            
        else:
            for i in range(len(a)):
                multi_controlled_U_g(u, [a[i]], -((2**i*b))*modinv(N, 2**(m+1)))

    if permeable_if_zero:    
        mcx(list(a) + [t[0]],  
            u[-1], ctrl_state = "0"*len(a) + "1",
            method = "balauca")
    
    u.delete(verify = False)
    
    return t

from qrisp import merge
def montgomery_mod_mul(a, b, output_qg = None):
    
    m = int(np.ceil(np.log2((a.modulus-1)**2)+1)) - a.size
    
    if a.modulus != b.modulus:
        raise Exception("Tried to multiply two QuantumModulus with differing modulus")
        
    if output_qg is None:
        t = QuantumFloat(a.size + m, signed = True)
        h(t)
        
    else:
        if output_qg.modulus != a.modulus:
            raise Exception("Output QuantumModulus has incompatible modulus")
        
        merge(output_qg.qs, a.qs)
        output_qg.extend(m, 0)
        output_qg.add_sign()
        output_qg.reg.insert(0, output_qg.reg.pop(-1))
        
        QFT(output_qg, exec_swap = False)
        
        t = output_qg
    
    t = hybrid_mult(a, b, init_op = None, terminal_op = None, output_qf = t)
    
    from qrisp import QuantumModulus
    t.__class__ = QuantumModulus
    t.modulus = a.modulus
    t.m = (a.m + b.m)
    t.inpl_adder = a.inpl_adder
    
    t = montgomery_red(t, a, b, a.modulus, m)
    
    return t

from qrisp import invert
from qrisp.environments import custom_control

@custom_control
def qft_semi_cl_inpl_mult(a, X, ctrl = None, treat_invalid = False):
    
    X = X%a.modulus
    
    if X == 0:
        raise Exception("Tried to perform in-place multiplication with 0 (not invertible)")
    if X == 1:
        return a
    
    tmp = a.duplicate(qs = a.qs)
    
    from qrisp import multi_measurement, less_than, redirect_qfunction, QuantumModulus
    
    if treat_invalid:
        a.__class__ = QuantumFloat
        reduced = a < a.modulus
        a.__class__ = QuantumModulus
        ctrl = [ctrl, reduced[0]]
    
    if ctrl is not None:
        with control(ctrl, invert = True):
            swap(tmp, a)
            
    tmp = montgomery_mod_semi_mul(a, 
                                  X, 
                                  output_qg = tmp,
                                  permeable_if_zero = ctrl is not None)
    
    if ctrl is not None:
        with control(ctrl, invert = True):
            swap(tmp, a)
    
    with invert():
        a = montgomery_mod_semi_mul(tmp, 
                                    modinv(X, a.modulus)%a.modulus, 
                                    output_qg = a, 
                                    permeable_if_zero = ctrl is not None)
        
    if ctrl is not None:
        with control(ctrl, invert = False):
            swap(tmp, a)
    else:
        swap(tmp, a)
    
    tmp.delete(verify = False)
    
    if treat_invalid:
        a.__class__ = QuantumFloat
        redirect_qfunction(less_than)(a,a.modulus,target= reduced)
        a.__class__ = QuantumModulus
        reduced.delete(verify = False)
    
    
    return a



            
    
    

def transfer_lsb(from_qv, to_qv):
    lsb = from_qv.reg.pop(0)
    from_qv.size -= 1
    from_qv.msize -= 1
    from_qv.mshape[0] += 1
    from_qv.exponent += 1

    to_qv.reg.insert(len(to_qv), lsb)
    to_qv.size += 1
    to_qv.msize += 1
    to_qv.mshape[1] += 1