"""Generate bounded circuit-to-Python candidates. This module never executes them."""

import json
import math

import torch
from peft import PeftModel
from transformers import AutoTokenizer,Qwen3_5ForConditionalGeneration

BASE="Qwen/Qwen3.5-9B"
BASE_REVISION="c202236235762e1c871ad0ccb60c8ee5ba337b9a"
ADAPTER="itsalissonsilva/QuantumCoder-9B"


def circuit_prompt(framework,qubits,operations):
    if framework not in {"qiskit","pennylane","cirq"}:
        raise ValueError("framework must be qiskit, pennylane or cirq")
    if type(qubits) is not int or not 2<=qubits<=6:
        raise ValueError("The measured scope is two through six qubits")
    if not isinstance(operations,(list,tuple)) or not 1<=len(operations)<=24:
        raise ValueError("Provide one through24 operations")
    normalized=[]
    for operation in operations:
        if not isinstance(operation,(list,tuple)) or len(operation)!=3:
            raise ValueError("Each operation must be [gate,qubit_indices,angle]")
        gate,indices,angle=operation
        if gate not in {"h","x","z","ry","rz","cx"}:
            raise ValueError("Supported gates: h, x, z, ry, rz, cx")
        if not isinstance(indices,(list,tuple)) or len(indices)!=(2 if gate=="cx" else 1):
            raise ValueError("Gate arity mismatch")
        if any(type(k) is not int or not 0<=k<qubits for k in indices) or len(set(indices))!=len(indices):
            raise ValueError("Invalid/repeated qubit index")
        if type(angle) not in (int,float) or not math.isfinite(angle):
            raise ValueError("Angle must be a finite real number")
        if gate not in {"ry","rz"} and angle!=0:
            raise ValueError("Non-rotation gates require angle0")
        normalized.append([gate,list(indices),float(angle)])
    return [
        dict(role="system",content="Return only Python source code using the requested quantum framework. No prose or markdown."),
        dict(role="user",content=f"Write a {framework} program preparing this circuit from zero on {qubits} qubits. Operations [gate,qubits,angle]: {json.dumps(normalized)}. Angles are radians. Store the final statevector in a variable named state. Use little-endian qubit order."),
    ]


class QuantumCoder:
    def __init__(self,adapter=ADAPTER,revision=None,subfolder=None):
        if not torch.cuda.is_available():
            raise RuntimeError("This measured inference configuration requires a CUDA GPU with enough memory for9B weights")
        self.tokenizer=AutoTokenizer.from_pretrained(BASE,revision=BASE_REVISION)
        self.tokenizer.padding_side="left"
        dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
        base=Qwen3_5ForConditionalGeneration.from_pretrained(BASE,revision=BASE_REVISION,dtype=dtype,
                device_map="cuda",attn_implementation="sdpa")
        kwargs=dict(revision=revision)
        if subfolder is not None:
            kwargs["subfolder"]=subfolder
        self.model=PeftModel.from_pretrained(base,adapter,**kwargs).eval()

    @torch.no_grad()
    def generate(self,circuits):
        """Return program/token metadata for dictionaries with framework/qubits/operations.

        Programs are fallible text, including possible unsupported APIs. For
        measured execution use the separately provided isolated Linux evaluator.
        """
        prompts=[]
        for circuit in circuits:
            messages=circuit_prompt(circuit["framework"],circuit["qubits"],circuit["operations"])
            text=self.tokenizer.apply_chat_template(messages,tokenize=False,add_generation_prompt=True,enable_thinking=False)
            if len(self.tokenizer.encode(text))>2048:
                raise ValueError("Prompt exceeds the measured2048-token bound; no truncation applied")
            prompts.append(text)
        results=[]
        eos=self.model.generation_config.eos_token_id
        eos=set(eos if isinstance(eos,list) else [eos])
        for begin in range(0,len(prompts),8):
            inputs=self.tokenizer(prompts[begin:begin+8],padding=True,return_tensors="pt").to("cuda")
            generated=self.model.generate(**inputs,max_new_tokens=768,do_sample=False,
                                           pad_token_id=self.tokenizer.pad_token_id)
            for tokens in generated[:,inputs["input_ids"].shape[1]:]:
                ids=tokens.tolist();stop=next((i for i,t in enumerate(ids) if t in eos),None)
                used=ids if stop is None else ids[:stop+1]
                results.append(dict(program=self.tokenizer.decode(used,skip_special_tokens=True),generated_tokens=len(used),
                                    eos_observed=stop is not None,at_token_cap=stop is None and len(used)>=768))
        return results