import argparse
import ast
from qiskit import QuantumCircuit
from .architecture import square_sparse_layout, compact_layout
from .dascot import (
    extract_gates_from_file,
    extract_qubits_from_gates,
    dump,
    run_dascot,
    run_sat_scmr,
)
from .guoq import run_guoq, print_help, CLIFFORDT, FAULT_TOLERANT_OPTIMIZATION_OBJECTIVE
from .utils import create_scratch_dir
import os
import shutil
import json
from importlib.metadata import version

from rich.console import Console
from rich.panel import Panel
from rich.rule import Rule
from rich.text import Text
from rich.table import Table
from rich import box

console = Console()

OPT_MODE = "opt"
FULL_FT_MODE = "full_ft"
SCMR_MODE = "scmr"

DEFAULT_EXT = {
    OPT_MODE: "qasm",
    FULL_FT_MODE: "json",
    SCMR_MODE: "json",
}

_VERSION = version("wisq")

MODE_LABELS = {
    OPT_MODE: "Circuit Optimization",
    FULL_FT_MODE: "Full Fault-Tolerant Compilation",
    SCMR_MODE: "Surface Code Mapping & Routing",
}


_LOGO = Text(
    "        \u2580\u2580              \n"
    "\u2588\u2588   \u2588\u2588 \u2588\u2588  \u2584\u2588\u2580\u2580\u2580 \u2584\u2588\u2588\u2588\u2588 \n"
    "\u2588\u2588 \u2588 \u2588\u2588 \u2588\u2588  \u2580\u2588\u2588\u2588\u2584 \u2588\u2588 \u2588\u2588 \n"
    " \u2588\u2588\u2580\u2588\u2588  \u2588\u2588\u2584 \u2584\u2584\u2584\u2588\u2580 \u2580\u2588\u2588\u2588\u2588 \n"
    "                     \u2588\u2588 \n"
    "                     \u2580\u2580 ",
    style="bold #7C3AED",
)


def _print_banner():
    subtitle = Text.assemble(
        ("\n\n", ""),
        ("Quantum Circuit Compiler", "#7C3AED"),
        ("  \u00b7  ", "dim #7C3AED"),
        (f"v{_VERSION}", "dim"),
    )
    layout = Table.grid(padding=(0, 3))
    layout.add_row(_LOGO, subtitle)
    console.print(Panel(layout, border_style="#7C3AED", padding=(0, 2), box=box.ROUNDED))
    console.print()


def _print_job_summary(mode: str, input_path: str, output_path: str):
    console.print(
        f"  [dim]Mode  [/dim]  [bold]{MODE_LABELS[mode]}[/bold]\n"
        f"  [dim]Input [/dim]  {input_path}\n"
        f"  [dim]Output[/dim]  {output_path}"
    )
    console.print()


def _step(n: int, total: int, label: str, detail: str = ""):
    detail_str = f"  [dim]{detail}[/dim]" if detail else ""
    console.print(
        f"[bold #7C3AED]  [{n}/{total}][/bold #7C3AED]  [bold]{label}[/bold]{detail_str}"
    )


def _done(label: str = ""):
    if label:
        console.print(f"  [bold green]\u2713[/bold green]  {label}")
    else:
        console.print(f"  [bold green]\u2713[/bold green]  Done")


def _print_formatted_help(parser):
    console.print(f"  [bold]Usage:[/bold]  wisq [dim][options][/dim] [bold]input_path[/bold]\n")
    for group in parser._action_groups:
        actions = [a for a in group._group_actions]
        if not actions:
            continue
        console.print(Rule(f"  {group.title}  ", style="dim #7C3AED", align="left"))
        console.print()
        for action in actions:
            if action.option_strings:
                flags = ", ".join(action.option_strings)
                if action.choices:
                    flags += f"  [dim]{{{', '.join(str(c) for c in action.choices)}}}[/dim]"
                elif action.type:
                    flags += f"  [dim]{action.type.__name__.upper()}[/dim]"
            else:
                flags = action.dest
            console.print(f"  [bold #7C3AED]{flags}[/bold #7C3AED]")
            if action.help:
                help_text = " ".join(action.help.split())
                console.print(f"    [dim]{help_text}[/dim]")
            console.print()


class WisqArgumentParser(argparse.ArgumentParser):
    def error(self, message):
        _print_banner()
        console.print(f"  [bold red]\u2717[/bold red]  [bold]{message}[/bold]\n")
        console.print(f"  Run [bold #7C3AED]wisq --help[/bold #7C3AED] for usage.\n")
        self.exit(2)

    def print_help(self, file=None):
        _print_banner()
        _print_formatted_help(self)


class Guoq_Help_Action(argparse.Action):
    def __init__(
        self,
        option_strings,
        dest=argparse.SUPPRESS,
        default=argparse.SUPPRESS,
        help=None,
    ):
        super(Guoq_Help_Action, self).__init__(
            option_strings=option_strings,
            dest=dest,
            default=default,
            nargs=0,
            help=help,
        )

    def __call__(self, parser, namespace, values, option_string=None):
        print_help()
        parser.exit()


def map_and_route(
    input_path: str, arch_name: str, output_path: str, timeout: int, mode="dascot"
):
    """
    Apply a surface code mapping and routing pass to the given circuit

    Args:
        input_path: Path to the input circuit file.
        output_path: Path to the output circuit file.
        arch_name: A description of the target architecture. Valid options are the
        strings "square_sparse_layout" or "compact_layout" representing built-in architectures,
        or the path to a text file containing the description of a custom architecture
        timeout: Total timeout in seconds for both mapping and routing.


    Writes a JSON representing
    the scheduled circuit after mapping and routing to output_path.
    """
    circ = QuantumCircuit.from_qasm_file(input_path)
    gates, ops = extract_gates_from_file(input_path)
    id_to_op = {i: ops[i] for i in range(len(ops))}
    total_qubits = len(extract_qubits_from_gates(gates))
    circ = QuantumCircuit.from_qasm_file(input_path)
    if arch_name == "square_sparse_layout":
        arch = square_sparse_layout(total_qubits, magic_states="all_sides")
    elif arch_name == "compact_layout":
        arch = compact_layout(total_qubits, magic_states="all_sides")
    else:
        with open(arch_name) as f:
            arch = ast.literal_eval(f.read())
    if mode == "dascot":
        map, steps = run_dascot(circ, gates, arch, output_path, timeout)
    elif mode == "sat":
        map, steps = run_sat_scmr(circ, gates, arch, output_path, timeout)
    dump(arch, map, steps, id_to_op, output_path, gates)


def optimize(
    input_path: str,
    output_path: str,
    target_gateset: str,
    optimization_objective: str,
    timeout: int,
    approximation_epsilon: float = 0,
    advanced_args: dict = None,
    verbose: bool = False,
    path_to_synthetiq: str = None,
) -> None:
    """
    Use the default GUOQ parameters to optimize a circuit. Recommended for most users. Advanced users can use `advanced_args` to override default values.

    Args:
        input_path: Path to the input circuit file.
        output_path: Path to the output circuit file.
        target_gateset: Target gateset to optimize the circuit to.
        timeout: Timeout in seconds. If set to 0, only transpiles and performs no optimization.
        approximation_epsilon: Approximation epsilon to use.
        advanced_args: Dictionary containing advanced arguments to pass to GUOQ, overriding default values except `-out` and `-job`. `guoq.print_help` displays available options.
        For example, if we want to override the default for `--rules` and use the `--remove-size-preserving-rules` flag, the dictionary would be `{"--rules": "file.txt", "--remove-size-preserving-rules": None}`.
        verbose: Whether to print verbose output.
    """
    run_guoq(
        input_path,
        output_path,
        target_gateset,
        optimization_objective,
        timeout,
        approximation_epsilon=approximation_epsilon,
        args=advanced_args,
        verbose=verbose,
        path_to_synthetiq=path_to_synthetiq,
    )


def compile_fault_tolerant(
    input_path,
    output_path,
    opt_timeout,
    arch_name,
    approximation_epsilon=1e-10,
    verbose=False,
    mr_timeout=1800,
    mr_solver="dascot",
    path_to_synthetiq=None,
):
    """
    Compiles a circuit to a fault-tolerant architecture using the Clifford + T gate set.
    The input is a QASM circuit and architecture and the output is a JSON representing
    the scheduled circuit after optimizing, mapping, and routing.
    """
    scratch_dir_path, _ = create_scratch_dir(output_path)

    try:
        transpiled_and_optimized_path = os.path.join(
            scratch_dir_path, "after_guoq.qasm"
        )
        _step(1, 2, "Decompose to Clifford + T (if needed) & Optimize", f"timeout: {opt_timeout}s")
        optimize(
            input_path,
            transpiled_and_optimized_path,
            CLIFFORDT,
            FAULT_TOLERANT_OPTIMIZATION_OBJECTIVE,
            opt_timeout,
            approximation_epsilon,
            verbose=verbose,
            path_to_synthetiq=path_to_synthetiq,
        )
        _done("Optimization complete")
        console.print()
        _step(2, 2, "Map & Route", f"timeout: {mr_timeout}s")
        map_and_route(
            transpiled_and_optimized_path,
            arch_name,
            output_path,
            mr_timeout,
            mode=mr_solver,
        )
        _done("Mapping and routing complete")
    finally:
        if os.path.exists(scratch_dir_path):
            shutil.rmtree(scratch_dir_path)


def main():
    parser = WisqArgumentParser(
        prog="wisq",
        description="A compiler for quantum circuits. Optimize your circuits and/or map them to a surface code architecture. See README for example usage and documentation.",
    )
    opt = parser.add_argument_group(title="optimization config")
    scmr = parser.add_argument_group(title="mapping and routing config")
    parser.add_argument("input_path", help="path to the input circuit")
    parser.add_argument(
        "--output_path",
        "-op",
        help="path to write the output. Default is out.qasm or out.json",
    )
    parser.add_argument(
        "--mode",
        "-m",
        default=FULL_FT_MODE,
        help="""
        control which compilation passes to apply. |
        opt: circuit optimization only |
        scmr: surface code mapping and routing only |
        full_ft (default): circuit optimization, then surface code mapping and routing
        """,
        choices=[OPT_MODE, FULL_FT_MODE, SCMR_MODE],
    )
    opt.add_argument(
        "--target_gateset",
        "-tg",
        help="target gateset for circuit optimization (default: Clifford + T)",
        default=CLIFFORDT,
        choices=guoq.GATE_SETS.keys(),
    )
    opt.add_argument(
        "--optimization_objective",
        "-obj",
        help="objective function used to guide the optimization (default: fault-tolerant objective function)",
        default=FAULT_TOLERANT_OPTIMIZATION_OBJECTIVE,
        choices=["TWO_Q", "FIDELITY", "FT", "TOTAL", "T"],
    )
    opt.add_argument(
        "--opt_timeout",
        "-ot",
        help="integer representing timeout for optimization in seconds (default: 3600)",
        type=int,
        default=3600,
    )
    opt.add_argument(
        "--approx_epsilon",
        "-ap",
        help="the approximation epsilon for optimization (represented as a plain decimal or in scientific notation, e.g., 1e-8); output of optimization pass is equivalent to input circuit up to epsilon error (default: 0)",
        type=float,
        default=0,
    )
    parser.add_argument(
        "--verbose", "-v", help="print verbose output", action="store_true"
    )
    scmr.add_argument(
        "--architecture",
        "-arch",
        help="target architecture for mapping and routing, can be one of the strings {'square_sparse_layout', 'compact_layout'} or the path to a file specifying a custom architecture. See README for format. (default: 'square_sparse')",
        default="square_sparse_layout",
    )
    scmr.add_argument(
        "--mr_timeout",
        "-tmr",
        type=int,
        help="integer representing timeout for mapping and routing in seconds (default: 1800)",
        default=1800,
    )
    scmr.add_argument(
        "--mr_solver",
        "-smr",
        help="solver to use for mapping and routing (default: 'dascot')",
        default="dascot",
    )
    parser.add_argument(
        "--guoq_help", "-gh", help="print GUOQ options", action=Guoq_Help_Action
    )
    parser.add_argument(
        "--advanced_args",
        "-aa",
        help="file path to JSON with advanced GUOQ args. See `optimize` for example",
    )
    parser.add_argument(
        "--abs_path_to_synthetiq",
        "-apts",
        help="absolute path to Synthetiq `main` binary",
    )
    args = parser.parse_args()

    _print_banner()

    if not os.path.exists(args.input_path) and "wisq-circuits" in args.input_path:
        args.input_path = os.path.join(
            os.path.dirname(__file__),
            args.input_path,
        )

    if args.output_path is None:
        args.output_path = f"out.{DEFAULT_EXT[args.mode]}"

    if args.advanced_args is not None:
        with open(args.advanced_args, "r") as f:
            args.advanced_args = json.load(f)

    _print_job_summary(args.mode, args.input_path, args.output_path)

    if args.mode == OPT_MODE:
        _step(1, 1, "Optimize", f"timeout: {args.opt_timeout}s")
        optimize(
            input_path=args.input_path,
            output_path=args.output_path,
            target_gateset=args.target_gateset,
            optimization_objective=args.optimization_objective,
            timeout=args.opt_timeout,
            approximation_epsilon=args.approx_epsilon,
            advanced_args=args.advanced_args,
            verbose=args.verbose,
            path_to_synthetiq=args.abs_path_to_synthetiq,
        )
        _done("Optimization complete")
    elif args.mode == FULL_FT_MODE:
        compile_fault_tolerant(
            input_path=args.input_path,
            output_path=args.output_path,
            opt_timeout=args.opt_timeout,
            arch_name=args.architecture,
            approximation_epsilon=args.approx_epsilon,
            verbose=args.verbose,
            mr_timeout=args.mr_timeout,
            mr_solver=args.mr_solver,
            path_to_synthetiq=args.abs_path_to_synthetiq,
        )
    elif args.mode == SCMR_MODE:
        _step(1, 1, "Map & Route", f"timeout: {args.mr_timeout}s")
        map_and_route(
            input_path=args.input_path,
            output_path=args.output_path,
            arch_name=args.architecture,
            timeout=args.mr_timeout,
            mode=args.mr_solver,
        )
        _done("Mapping and routing complete")

    console.print()
    console.print(Rule(style="dim #7C3AED"))
    console.print(
        f"  [bold green]\u2713[/bold green]  Output written to [bold]{args.output_path}[/bold]"
    )
    console.print()


if __name__ == "__main__":
    main()
