import re

class UniversalTranslator:
    @staticmethod
    def to_qasm(circuit_obj) -> str:
        if isinstance(circuit_obj, str): return circuit_obj
        class_name = circuit_obj.__class__.__name__
        if class_name == 'QuantumCircuit':
            try:
                from qiskit import qasm2
                return qasm2.dumps(circuit_obj)
            except ImportError:
                return circuit_obj.qasm()
        elif class_name == 'Circuit':
            import cirq
            return cirq.qasm(circuit_obj)
        elif class_name == 'QuantumTape':
            import pennylane as qml
            return qml.tape.to_openqasm(circuit_obj)
        raise TypeError(f"Unsupported circuit type: {class_name}.")


class QASMParser:
    def __init__(self):
        self.idx_counter = 0
        self.current_wire_labels = {} 
        self.qubit_names = [] 
        
    def _get_next_index(self):
        idx = self.idx_counter
        self.idx_counter += 1
        return idx

    def _ensure_wire_exists(self, qubit):
        if qubit not in self.current_wire_labels:
            self.current_wire_labels[qubit] = self._get_next_index()

    def parse(self, raw_input):
        qasm_str = UniversalTranslator.to_qasm(raw_input)
        
        clean_lines = []
        in_gate_def = False
        for line in qasm_str.split('\n'):
            line = line.split('//')[0].strip()
            if not line: continue
            
            if line.startswith('gate '):
                in_gate_def = True
                
            if in_gate_def:
                if '}' in line:
                    in_gate_def = False
                continue
                
            clean_lines.append(line)
            
        clean_str = " ".join(clean_lines)
        statements = [s.strip() for s in clean_str.split(';') if s.strip()]
        
        inputs, shapes, gate_info = [], [], []
        tensor_id = 0
        
        for stmt in statements:
            parts = stmt.split(' ', 1)
            if len(parts) < 2: continue
                
            op = parts[0].split('(')[0].lower()
            
            if op == 'qreg' or op.startswith('qubit'):
                if op == 'qreg':
                    args = [arg.strip() for arg in parts[1].split(',')]
                    reg_name = args[0].split('[')[0]
                    num_qubits = int(args[0].split('[')[1].split(']')[0])
                else:
                    num_qubits = int(op.split('[')[1].split(']')[0])
                    reg_name = parts[1].strip()
                for i in range(num_qubits):
                    qubit_id = f"{reg_name}[{i}]"
                    self._ensure_wire_exists(qubit_id)
                    if qubit_id not in self.qubit_names:
                        self.qubit_names.append(qubit_id)
                    
            elif op == 'measure':
                qubit = parts[1].split('->')[0].strip()
                self._ensure_wire_exists(qubit)
                idx_in = self.current_wire_labels[qubit]
                
                inputs.append([idx_in])
                shapes.append((2,))
                gate_info.append({"id": tensor_id, "op": "M", "targets": [qubit], "controls": []})
                tensor_id += 1
                del self.current_wire_labels[qubit]
                
            elif op in ['if', 'reset']:
                print(f"Warning: Operation '{op}' breaks unitarity. It is treated as a wire-reset.")
                if op == 'reset':
                    qubit = parts[1].strip()
                    if qubit in self.current_wire_labels:
                        del self.current_wire_labels[qubit] 
            
            elif op in ['barrier', 'include', 'openqasm']:
                continue
                
            elif op in ['h', 'x', 'y', 'z', 'rx', 'ry', 'rz', 's', 't', 'u', 'u1', 'u2', 'u3']:
                args = [arg.strip() for arg in parts[1].split(',')]
                qubit = args[0]
                self._ensure_wire_exists(qubit)
                idx_in = self.current_wire_labels[qubit]
                idx_out = self._get_next_index()
                self.current_wire_labels[qubit] = idx_out
                
                inputs.append([idx_in, idx_out])
                shapes.append((2, 2))
                # Truncate long parameters like rx(pi/2) to just RX
                display_op = op.upper()
                gate_info.append({"id": tensor_id, "op": display_op, "targets": [qubit], "controls": []})
                tensor_id += 1
                
            elif op in ['cx', 'cz']:
                args = [arg.strip() for arg in parts[1].split(',')]
                control, target = args[0], args[1]
                self._ensure_wire_exists(control)
                self._ensure_wire_exists(target)
                
                ctrl_in, targ_in = self.current_wire_labels[control], self.current_wire_labels[target]
                ctrl_out, targ_out = self._get_next_index(), self._get_next_index()
                
                self.current_wire_labels[control] = ctrl_out
                self.current_wire_labels[target] = targ_out
                
                inputs.append([ctrl_in, targ_in, ctrl_out, targ_out])
                shapes.append((2, 2, 2, 2))
                gate_info.append({"id": tensor_id, "op": op.upper(), "targets": [target], "controls": [control]})
                tensor_id += 1
                
            else:
                args = [arg.strip() for arg in parts[1].split(',')]
                valid_qubits = [q for q in args if q in self.qubit_names]
                if not valid_qubits: continue
                
                idx_in_list, idx_out_list = [], []
                for q in valid_qubits:
                    self._ensure_wire_exists(q)
                    idx_in_list.append(self.current_wire_labels[q])
                    new_idx = self._get_next_index()
                    idx_out_list.append(new_idx)
                    self.current_wire_labels[q] = new_idx
                    
                inputs.append(idx_in_list + idx_out_list)
                shapes.append(tuple([2] * (2 * len(valid_qubits))))
                
                gate_info.append({"id": tensor_id, "op": op.upper(), "targets": valid_qubits, "controls": []})
                tensor_id += 1
                
        output = list(self.current_wire_labels.values())
        return inputs, output, shapes, gate_info

    def generate_svg(self, gate_info):
        y_spacing, margin = 50, 40
        wire_x = {q: margin for q in self.qubit_names}
        wire_y = {q: margin + i * y_spacing for i, q in enumerate(self.qubit_names)}
        
        svg_elements = []
        
        for gate in gate_info:
            involved_qubits = gate["controls"] + gate["targets"]
            if not involved_qubits: continue
            
            involved_indices = [self.qubit_names.index(q) for q in involved_qubits]
            min_idx, max_idx = min(involved_indices), max(involved_indices)
            spanned_qubits = self.qubit_names[min_idx : max_idx + 1]
            
            display_op = gate["op"]
            box_width = max(30, len(display_op) * 8 + 12)
            half_width = box_width / 2
            
            max_prev_right = max(wire_x[q] for q in spanned_qubits)
            gate_x = max_prev_right + 30 + half_width
            
            for q in spanned_qubits: 
                wire_x[q] = gate_x + half_width
                
            g_id = f"gate-T{gate['id']}"
            
            if len(involved_qubits) == 1:
                qy = wire_y[gate["targets"][0]]
                fill_color = "#8b0000" if gate["op"] == "M" else "#2a2a2a"
                
                svg_elements.append(f'<g id="{g_id}" class="svg-gate" style="transition: all 0.2s;">')
                svg_elements.append(f'<rect x="{gate_x - half_width}" y="{qy-15}" width="{box_width}" height="30" fill="{fill_color}" stroke="#666" stroke-width="2" rx="4"/>')
                svg_elements.append(f'<text x="{gate_x}" y="{qy+4}" fill="white" font-family="sans-serif" font-size="12" text-anchor="middle">{display_op}</text>')
                svg_elements.append('</g>')
                
            elif len(gate["controls"]) == 1 and len(gate["targets"]) == 1:
                cy, ty = wire_y[gate["controls"][0]], wire_y[gate["targets"][0]]
                svg_elements.append(f'<g id="{g_id}" class="svg-gate" style="transition: all 0.2s;">')
                svg_elements.append(f'<line x1="{gate_x}" y1="{cy}" x2="{gate_x}" y2="{ty}" stroke="#666" stroke-width="2"/>')
                svg_elements.append(f'<circle cx="{gate_x}" cy="{cy}" r="5" fill="#666"/>')
                svg_elements.append(f'<rect x="{gate_x - half_width}" y="{ty-15}" width="{box_width}" height="30" fill="#2a2a2a" stroke="#666" stroke-width="2" rx="4"/>')
                svg_elements.append(f'<text x="{gate_x}" y="{ty+4}" fill="white" font-family="sans-serif" font-size="12" text-anchor="middle">{display_op}</text>')
                svg_elements.append('</g>')
                
            else:
                qy_min = min(wire_y[q] for q in involved_qubits)
                qy_max = max(wire_y[q] for q in involved_qubits)
                height = qy_max - qy_min + 30
                
                svg_elements.append(f'<g id="{g_id}" class="svg-gate" style="transition: all 0.2s;">')
                svg_elements.append(f'<rect x="{gate_x - half_width}" y="{qy_min-15}" width="{box_width}" height="{height}" fill="#2a2a2a" stroke="#666" stroke-width="2" rx="4"/>')
                
                true_mid_y = (qy_min + qy_max) / 2
                svg_elements.append(f'<text x="{gate_x}" y="{true_mid_y+4}" fill="white" font-family="sans-serif" font-size="12" text-anchor="middle">{display_op}</text>')
                svg_elements.append('</g>')

        max_total_x = max(wire_x.values()) + 30
        height = len(self.qubit_names) * y_spacing + margin
        
        wires_svg = ""
        for q in self.qubit_names:
            y = wire_y[q]
            wires_svg += f'<text x="10" y="{y+4}" fill="#888" font-family="monospace" font-size="12">{q}</text>'
            wires_svg += f'<line x1="{margin}" y1="{y}" x2="{max_total_x}" y2="{y}" stroke="#444" stroke-width="1" z-index="-1"/>'

        svg_str = f'<svg width="{max_total_x + margin}" height="{height}" xmlns="http://www.w3.org/2000/svg" style="background: #1a1a1a;">'
        svg_str += wires_svg + "".join(svg_elements) + "</svg>"
        return svg_str