import attrs

from qualtran import Bloq, Signature, BloqBuilder, SoquetT
from qualtran.bloqs.basic_gates import CNOT, TGate
from qualtran.drawing import show_bloq
from qualtran.resource_counting import QubitCount, get_cost_value

@attrs.frozen
class TestManyAlloc(Bloq):
    n: int
    alloc_once: bool

    @property
    def signature(self) -> Signature:
        return Signature.build(x=self.n)

    def build_composite_bloq(self, bb: 'BloqBuilder', *, x: 'SoquetT') -> dict[str, 'SoquetT']:
        x = bb.split(x)
        if self.alloc_once:
            anc = bb.allocate(1)
        for i in range(self.n):
            if not self.alloc_once:
                anc = bb.allocate(1)
            x[i], anc = bb.add(CNOT(), ctrl=x[i], target=anc)
            anc = bb.add(TGate(), q=anc)
            x[i], anc = bb.add(CNOT(), ctrl=x[i], target=anc)
            if not self.alloc_once:
                bb.free(anc)
        if self.alloc_once:
            bb.free(anc)
        return {'x': bb.join(x)}


get_cost_value(TestManyAlloc(10, True), QubitCount()), get_cost_value(TestManyAlloc(10, False), QubitCount())