# ********************************************************************************
# * Copyright (c) 2026 Bradan Quantum GmbH
# *
# * 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 qrisp
from qframe.core.qframe_variable import QFrameVariable
from qframe.core.operation_wrapper import OperationWrapper
from qframe.core.qframe_session import QFrameSession

class QFrameUInt(QFrameVariable):
    def __init__(self, size, qfs=None, name=None):
        super().__init__(size, qfs, name)
        self.qv = qrisp.QuantumFloat(size, name=name)

    def __iadd__(self, other):
        from qframe.alg_primitives.adder import recip_adder_gate
        if isinstance(other, QFrameUInt):
            if self == other:
                raise Exception("Incompatible operands: Cannot add a quantum variable to itself")
            if self.qfs != other.qfs:
                self.qfs.merge(other.qfs)
            opw = OperationWrapper()
            def _gate_apply_impl(qfs: QFrameSession):
                self.qv += other.qv
            def _calculate_impl(arg_dict:dict):
                N = 0b1 << self.size  # Addition modulo N
                arg_dict[self] = (arg_dict[self] + arg_dict[other]) % N
            opw.set_gate_apply_impl(_gate_apply_impl)
            opw.set_recip_gate_apply_impl(lambda qfs: recip_adder_gate(other.qv, self.qv, qfs.recip_carry_anc, qfs.phase_anc))
            opw.set_calculate_impl(_calculate_impl)
            self.qfs.append_operation_wrapper(opw)
        elif isinstance(other, int):
            if other.bit_length() > self.size:
                print(f'WARN: Adding constant too long for this register. Auto-truncating constant to {self.size} bits.')
                bit_mask = (0b1 << self.size) - 1
                other &= bit_mask
            opw = OperationWrapper()
            def _gate_apply_impl(qfs: QFrameSession):
                self.qv += other
            def _recip_gate_apply_impl(qfs: QFrameSession):
                from qframe.core.utility import xor_const
                other_qv = qfs._register_anc
                xor_const(other_qv, other)  # other_qv ^= other
                qrisp.h(other_qv)
                recip_adder_gate(other_qv, self.qv, qfs.recip_carry_anc, qfs.phase_anc)
                qrisp.h(other_qv)
                xor_const(other_qv, other)
            def _calculate_impl(arg_dict:dict):
                N = 0b1 << self.size  # Addition modulo N
                arg_dict[self] = (arg_dict[self] + other) % N
            opw.set_gate_apply_impl(_gate_apply_impl)
            opw.set_recip_gate_apply_impl(_recip_gate_apply_impl)
            opw.set_calculate_impl(_calculate_impl)
            self.qfs.append_operation_wrapper(opw)
        elif isinstance(other, OperationWrapper):
            other.check_compatibility(self)
            other.merge_qfs(self.qfs)
            opw = OperationWrapper()
            def _gate_apply_impl(qfs: QFrameSession):
                if other.conjugate_me:
                    with qrisp.conjugate(other.gate_apply)(qfs):
                        self.qv += other.gate_result_qfv.qv
                else:
                    other.gate_apply(qfs)
                    self.qv += other.gate_result_qfv.qv
            def _recip_gate_apply_impl(qfs: QFrameSession):
                if other.conjugate_me:
                    with qrisp.conjugate(other.recip_gate_apply)(qfs):
                        recip_adder_gate(other.gate_result_qfv.qv, self.qv, qfs.recip_carry_anc, qfs.phase_anc)
                else:
                    other.recip_gate_apply(qfs)
                    recip_adder_gate(other.gate_result_qfv.qv, self.qv, qfs.recip_carry_anc, qfs.phase_anc)
            def _calculate_impl(arg_dict:dict):
                N = 0b1 << self.size  # Addition modulo N
                arg_dict[self] = (arg_dict[self] + other.calculate(arg_dict)) % N
            opw.set_gate_apply_impl(_gate_apply_impl)
            opw.set_recip_gate_apply_impl(_recip_gate_apply_impl)
            opw.set_calculate_impl(_calculate_impl)
            self.qfs.append_operation_wrapper(opw)

        return self
