from __future__ import annotations

from typing import Dict, List, Any


SUPPORTED_SINGLE_GATES = {
    "H": "h",
    "X": "x",
    "Y": "y",
    "Z": "z",
    "S": "s",
    "T": "t",
}

SUPPORTED_CONTROL_GATES = {
    "CX": "cx",
    "CY": "cy",
    "CZ": "cz",
    "SWAP": "swap",
}

# Add CCX (Toffoli) to control gates mapping
SUPPORTED_CONTROL_GATES["CCX"] = "ccx"


class TranspileError(Exception):
    pass


def _infer_classical_count(ir: Dict[str, Any]) -> int | None:
    c_decl = ir.get("classical")
    if c_decl is not None:
        return int(c_decl)

    # No explicit classical decl: infer from measurements if any
    ops: List[Dict[str, Any]] = ir.get("ops", [])
    has_measure_all = any(o["op"] == "measure_all" for o in ops)
    if has_measure_all:
        return int(ir.get("qubits") or 0)

    one_measures = [o for o in ops if o["op"] == "measure_one"]
    if one_measures:
        max_c = max(int(o["c"]) for o in one_measures)
        return max_c + 1

    return None


def transpile_to_python(ir: Dict[str, Any]) -> str:
    # Basic validations
    n_qubits = ir.get("qubits")
    if not n_qubits or int(n_qubits) <= 0:
        raise TranspileError("You must declare a positive number of qubits, e.g., 'qubits: 3'")
    n_qubits = int(n_qubits)

    n_classical = _infer_classical_count(ir)
    if n_classical is not None and n_classical < 0:
        raise TranspileError("Classical bit count must be >= 0")

    ops: List[Dict[str, Any]] = ir.get("ops", [])

    lines: List[str] = []
    lines.append("# Auto-generated by Qupac")
    lines.append("from qiskit import QuantumCircuit, transpile")
    lines.append("from qiskit_aer import AerSimulator")
    lines.append("from qiskit.circuit import Parameter")
    lines.append("from math import pi")
    lines.append("try:")
    lines.append("    from qiskit_aer.noise import NoiseModel, depolarizing_error")
    lines.append("except Exception:")
    lines.append("    NoiseModel = None")
    lines.append("    depolarizing_error = None")
    lines.append("")

    if n_classical is None:
        lines.append(f"qc = QuantumCircuit({n_qubits})")
    else:
        lines.append(f"qc = QuantumCircuit({n_qubits}, {n_classical})")

    # Define symbolic parameters if present
    symbols = set()
    # prefer symbols from IR (includes expression symbols)
    for s in sorted(set(ir.get("symbols") or [])):
        symbols.add(str(s))
    # Back-compat: scan ops for param_symbol
    for _op in ops:
        sym = _op.get("param_symbol")
        if sym:
            symbols.add(str(sym))
    if symbols:
        for s in sorted(symbols):
            lines.append(f"{s} = Parameter('{s}')")

    # Emit operations
    for op in ops:
        kind = op["op"]
        if kind == "apply":
            gate = op["gate"].upper()
            controls = op.get("controls") or []
            targets = op.get("targets") or []

            # If there are controls and exactly 1 target
            if controls and len(targets) == 1:
                if len(controls) == 1:
                    # single-control 2-qubit gate
                    c, t = controls[0], targets[0]
                    if gate == "X":
                        lines.append(f"qc.cx({c}, {t})")
                    else:
                        if gate not in SUPPORTED_CONTROL_GATES:
                            raise TranspileError(f"Gate '{gate}' not supported with 'from ... to ...'")
                        method = SUPPORTED_CONTROL_GATES[gate]
                        lines.append(f"qc.{method}({c}, {t})")
                elif len(controls) == 2 and gate == "CCX":
                    # Toffoli (ccx control1 control2 target)
                    method = SUPPORTED_CONTROL_GATES["CCX"]
                    c1, c2 = controls[0], controls[1]
                    t = targets[0]
                    lines.append(f"qc.{method}({c1}, {c2}, {t})")
                elif len(controls) >= 2 and gate == "X":
                    ctrl_list = ", ".join(str(c) for c in controls)
                    t = targets[0]
                    lines.append(f"qc.mcx([{ctrl_list}], {t})")
                else:
                    raise TranspileError("Multi-control gates beyond 2 controls are not supported yet")
            elif not controls and len(targets) == 1:
                # Single-qubit ops
                if gate in {"RX", "RY", "RZ"}:
                    param_value = op.get("param_value")
                    param_symbol = op.get("param_symbol")
                    param_expr = op.get("param_expr")
                    if param_expr is not None:
                        arg = param_expr
                    else:
                        if param_value is None and param_symbol is None:
                            raise TranspileError("Rotation gates RX/RY/RZ require an angle parameter")
                        arg = param_symbol if param_symbol is not None else repr(float(param_value))
                    method = gate.lower()
                    t = targets[0]
                    line = f"qc.{method}({arg}, {t})"
                    if op.get("cond"):
                        c = int(op["cond"]["c"])
                        val = int(op["cond"]["val"])
                        line += f".c_if(qc.clbits[{c}], {val})"
                    lines.append(line)
                    continue
                if gate not in SUPPORTED_SINGLE_GATES:
                    raise TranspileError(f"Single-qubit gate '{gate}' not supported")
                method = SUPPORTED_SINGLE_GATES[gate]
                t = targets[0]
                line = f"qc.{method}({t})"
                if op.get("cond"):
                    c = int(op["cond"]["c"])
                    val = int(op["cond"]["val"])
                    line += f".c_if(qc.clbits[{c}], {val})"
                lines.append(line)
            else:
                # More complex patterns: multiple targets without controls
                if not controls and targets:
                    # broadcast single-qubit gate to all targets
                    if gate in {"RX", "RY", "RZ"}:
                        param_value = op.get("param_value")
                        param_symbol = op.get("param_symbol")
                        param_expr = op.get("param_expr")
                        if param_expr is not None:
                            arg = param_expr
                        else:
                            if param_value is None and param_symbol is None:
                                raise TranspileError("Rotation gates RX/RY/RZ require an angle parameter")
                                
                            arg = param_symbol if param_symbol is not None else repr(float(param_value))
                        method = gate.lower()
                        for t in targets:
                            line = f"qc.{method}({arg}, {t})"
                            if op.get("cond"):
                                c = int(op["cond"]["c"])
                                val = int(op["cond"]["val"])
                                line += f".c_if(qc.clbits[{c}], {val})"
                            lines.append(line)
                    elif gate in SUPPORTED_SINGLE_GATES:
                        method = SUPPORTED_SINGLE_GATES[gate]
                        for t in targets:
                            line = f"qc.{method}({t})"
                            if op.get("cond"):
                                c = int(op["cond"]["c"])
                                val = int(op["cond"]["val"])
                                line += f".c_if(qc.clbits[{c}], {val})"
                            lines.append(line)
                    else:
                        raise TranspileError(f"Gate '{gate}' cannot be broadcast to multiple targets")
                else:
                    raise TranspileError("Unsupported apply pattern; complex control/target combinations are not supported yet")

        elif kind == "measure_all":
            if n_classical is None:
                # If not declared, implicitly make classical == qubits in circuit creation
                n_classical_eff = n_qubits
            else:
                n_classical_eff = n_classical
            if n_classical_eff < n_qubits:
                raise TranspileError("Not enough classical bits to measure all qubits; declare 'classical: N' >= qubits")
            lines.append(f"qc.measure(range({n_qubits}), range({n_classical_eff}))")

        elif kind == "measure_one":
            q = int(op["q"])
            c = int(op["c"])
            if n_classical is None:
                # If not declared, ensure circuit has enough classical bits by recreating qc is hard.
                # We require explicit classical or a prior measure_all to infer. Keep simple:
                raise TranspileError("Measuring a specific qubit requires 'classical: N' declaration in this version")
            if c >= n_classical:
                raise TranspileError("Classical index out of range; increase 'classical:'")
            lines.append(f"qc.measure({q}, {c})")

        elif kind == "reset":
            q = int(op["q"]) if isinstance(op["q"], (int,)) else op["q"]
            # allow q to be int
            lines.append(f"qc.reset({q})")

        elif kind == "call":
            name = op["name"]
            args = op.get("args") or []
            subs = ir.get("subcircuits") or []
            subdef = next((s for s in subs if s.get("name") == name), None)
            if not subdef:
                raise TranspileError(f"Unknown subcircuit: {name}")
            params = list(subdef.get("params") or [])
            if len(args) != len(params):
                raise TranspileError(f"call {name}: expected {len(params)} args, got {len(args)}")
            subvar = f"sub_{name}"
            lines.append(f"{subvar} = QuantumCircuit({len(params)})")
            lines.append(f"# build subcircuit: {name}")
            lines.extend(_emit_sub_ops(subvar, subdef))
            arg_list_src = ", ".join(str(a) for a in args)
            lines.append(f"qc.compose({subvar}, qubits=[{arg_list_src}], inplace=True)")

        else:
            raise TranspileError(f"Unknown operation: {kind}")

    # Optional draw-to-file request
    draw_file = ir.get("draw_file")
    if draw_file:
        lines.append("")
        path = draw_file.get("path")
        # Resolve relative path against the source .qu file directory if available
        base_dir = ir.get("_source_dir")
        import os as _os
        if path and not (_os.path.isabs(path)) and base_dir:
            try:
                path = str((_os.path.normpath(_os.path.join(base_dir, path))))
            except Exception:
                pass
        mode = (draw_file.get("mode") or "mpl").lower()
        if mode == "text":
            lines.append("import os as _os")
            lines.append(f"_p = {repr(path)}")
            lines.append("_os.makedirs(_os.path.dirname(_p) or '.', exist_ok=True)")
            lines.append("with open(_p, 'w', encoding='utf-8') as _f: _f.write(str(qc.draw()))")
        else:
            lines.append("try:")
            lines.append("    import os as _os")
            lines.append(f"    _p = {repr(path)}")
            lines.append("    _os.makedirs(_os.path.dirname(_p) or '.', exist_ok=True)")
            lines.append("    fig = qc.draw('mpl')")
            lines.append("    fig.savefig(_p, bbox_inches='tight')")
            lines.append("except Exception as e:")
            lines.append("    print('Failed to save draw (mpl):', e)")

    # Simulation
    # Optional draw request
    draw_mode = ir.get("draw")
    if draw_mode:
        lines.append("")
        if draw_mode == "text":
            lines.append("# Draw circuit (text)")
            lines.append("print(qc.draw())")
        else:
            lines.append("# Draw circuit (mpl) - requires matplotlib and qiskit visualization support")
            lines.append("try:")
            lines.append("    qc.draw('mpl')")
            lines.append("    import matplotlib.pyplot as plt")
            lines.append("    plt.show()")
            lines.append("except Exception as e:")
            lines.append("    print('Failed to draw with mpl:', e)")

    if ir.get("simulate"):
        lines.append("")
        lines.append("# Prepare simulator and run")
        sim_type_val = ir.get('simulator') or 'default'
        lines.append(f"sim_type = {repr(sim_type_val)}")
        lines.append("if sim_type == 'statevector':")
        lines.append("    sim = AerSimulator(method='statevector')")
        lines.append("elif sim_type == 'unitary':")
        lines.append("    sim = AerSimulator(method='unitary')")
        lines.append("else:")
        lines.append("    sim = AerSimulator()")

        opt_level_val = ir.get('optimize_level')
        lines.append(f"opt_level = {('None' if opt_level_val is None else int(opt_level_val))}")
        shots_val = int(ir.get('shots') or 1024)
        lines.append(f"shots = {shots_val}")
        noise_ir = ir.get('noise')
        if noise_ir:
            lines.append(f"noise = {repr(noise_ir)}")
        else:
            lines.append("noise = None")

        lines.append("# Optional noise model (depolarizing)")
        lines.append("if noise and noise.get('kind') == 'depol':")
        lines.append("    if NoiseModel is None or depolarizing_error is None:")
        lines.append("        print('NoiseModel not available; skipping noise')")
        lines.append("        noise_model=None")
        lines.append("    else:")
        lines.append("        noise_model = NoiseModel()")
        lines.append("        p = float(noise.get('p', 0.0))")
        lines.append("        try:")
        lines.append("            noise_model.add_all_qubit_quantum_error(depolarizing_error(p, 1), ['x','y','z','h','s','t'])")
        lines.append("            noise_model.add_all_qubit_quantum_error(depolarizing_error(min(2*p,1.0), 2), ['cx'])")
        lines.append("        except Exception as e:")
        lines.append("            print('Failed to configure noise model:', e); noise_model=None")
        lines.append("else:")
        lines.append("    noise_model = None")

        layout = ir.get('initial_layout')
        cmap = ir.get('coupling_map')
        lines.append("transpile_kwargs = {}")
        lines.append("if opt_level is not None: transpile_kwargs['optimization_level'] = int(opt_level)")
        if layout:
            lines.append(f"transpile_kwargs['initial_layout'] = {repr(layout)}")
        if cmap:
            lines.append(f"transpile_kwargs['coupling_map'] = {repr(cmap)}")
        lines.append("compiled = transpile(qc, sim, **{k:v for k,v in transpile_kwargs.items() if v is not None})")

        lines.append("result = sim.run(compiled, shots=int(shots), noise_model=noise_model).result()")
        lines.append("# Print results depending on simulator type")
        lines.append("if sim_type == 'statevector':")
        lines.append("    try:")
        lines.append("        print(result.get_statevector(qc))")
        lines.append("    except Exception as e:")
        lines.append("        print('Statevector unavailable:', e); print(result)")
        lines.append("elif sim_type == 'unitary':")
        lines.append("    try:")
        lines.append("        print(result.get_unitary(qc))")
        lines.append("    except Exception as e:")
        lines.append("        print('Unitary unavailable:', e); print(result)")
        lines.append("else:")
        lines.append("    try:")
        lines.append("        print(result.get_counts())")
        lines.append("    except Exception:")
        lines.append("        print(result)")

    return "\n".join(lines) + "\n"


def _emit_sub_ops(subvar: str, subdef: Dict[str, Any]) -> List[str]:
    """Emit Python lines to build a subcircuit in variable `subvar`.
    subdef contains: name, params (list of names), ops (list of dicts with *_names fields).
    """
    out: List[str] = []
    params = list(subdef.get("params") or [])
    idx_of = {name: i for i, name in enumerate(params)}
    for op in subdef.get("ops", []):
        kind = op.get("op")
        if kind == "apply":
            gate = op.get("gate", "").upper()
            ctrls_n = op.get("controls_names") or []
            tgs_n = op.get("targets_names") or []
            ctrls = [idx_of[n] for n in ctrls_n]
            tgs = [idx_of[n] for n in tgs_n]
            if ctrls and len(tgs) == 1:
                if len(ctrls) == 1:
                    c, t = ctrls[0], tgs[0]
                    if gate == 'X':
                        out.append(f"{subvar}.cx({c}, {t})")
                    else:
                        method = SUPPORTED_CONTROL_GATES.get(gate)
                        if not method:
                            raise TranspileError(f'subcircuit gate {gate} not supported with controls')
                        out.append(f"{subvar}.{method}({c}, {t})")
                elif len(ctrls) >= 2 and gate == 'X':
                    ctrl_list = ", ".join(str(c) for c in ctrls)
                    t = tgs[0]
                    out.append(f"{subvar}.mcx([{ctrl_list}], {t})")
                else:
                    raise TranspileError('subcircuit multi-control not supported for this gate')
            elif (not ctrls) and len(tgs) == 1:
                if gate in {'RX','RY','RZ'}:
                    expr = op.get('param_expr')
                    if not expr:
                        raise TranspileError('rotation in subcircuit requires param')
                    method = gate.lower()
                    t = tgs[0]
                    out.append(f"{subvar}.{method}({expr}, {t})")
                else:
                    method = SUPPORTED_SINGLE_GATES.get(gate)
                    if not method:
                        raise TranspileError(f'subcircuit single gate {gate} not supported')
                    t = tgs[0]
                    out.append(f"{subvar}.{method}({t})")
            elif (not ctrls) and len(tgs) > 1:
                method = SUPPORTED_SINGLE_GATES.get(gate)
                if gate in {'RX','RY','RZ'} and op.get('param_expr'):
                    expr = op['param_expr']
                    for t in tgs:
                        out.append(f"{subvar}.{gate.lower()}({expr}, {t})")
                elif method:
                    for t in tgs:
                        out.append(f"{subvar}.{method}({t})")
                else:
                    raise TranspileError('subcircuit broadcast not supported for this gate')
            else:
                raise TranspileError('unsupported subcircuit apply pattern')
        elif kind == 'reset':
            qn = op.get('q_name')
            if qn not in idx_of:
                raise TranspileError(f'subcircuit reset: unknown name {qn}')
            out.append(f"{subvar}.reset({idx_of[qn]})")
        elif kind == 'measure_one':
            # not yet supported inside subcircuits
            raise TranspileError('measure inside subcircuits not supported yet')
        else:
            raise TranspileError(f'unsupported subcircuit op: {kind}')
    return out

