"""Adaptive: `let h = n` aliases `n` after store pruning, so `set n = 7` changes `h`

Compiles the Q# program below to QIR for each target profile, runs that QIR with
qir-runner, and runs the same program in the Q# simulator for comparison.
Exits 0 while the bug reproduces: Adaptive returns something other than 3.

Requires: pip install qdk qirrunner
"""

import os
import re
import sys
import tempfile

import qirrunner
from qdk import Context, TargetProfile

SOURCE = """
operation Main() : Int {
    use q = Qubit();
    X(q);
    mutable n = 0;
    if MResetZ(q) == One { set n = 2; }
    set n = n + 1;
    let h = n;
    set n = 7;
    h
}
"""
ENTRY = "Main()"
EXPECTED = "3"  # the correct output
AFFECTED = ['Adaptive']  # runs that show the bug
SHOTS = 10
PROFILES = ["Base", "Adaptive_RI", "Adaptive_RIF", "Adaptive"]


def fmt(value):
    """Formats a simulator value the same way as qir-runner output below."""
    if value is None:
        return "()"
    if type(value).__name__ == "Result":
        return "1" if str(value).endswith("One") else "0"
    if isinstance(value, bool):
        return str(value).lower()
    if isinstance(value, float):
        return repr(value)
    if isinstance(value, (tuple, list)):
        inner = ", ".join(fmt(v) for v in value)
        return f"({inner})" if isinstance(value, tuple) else f"[{inner}]"
    return str(value)


def parse_shots(text):
    """Turns qir-runner's output records into one formatted value per shot."""
    shots, records = [], None
    for line in text.splitlines():
        fields = line.split("\t")
        if fields[0] == "START":
            records = []
        elif fields[0] == "OUTPUT" and records is not None:
            records.append((fields[1], fields[2]))
        elif fields[0] == "END" and records is not None:
            shots.append(to_text(iter(records)) if records else "()")
            if fields[1] != "0":
                shots[-1] = f"exit code {fields[1]}"
            records = None
    return shots


def to_text(records):
    kind, value = next(records)
    if kind in ("TUPLE", "ARRAY"):
        inner = ", ".join(to_text(records) for _ in range(int(value)))
        return f"({inner})" if kind == "TUPLE" else f"[{inner}]"
    if kind == "DOUBLE":
        return repr(float(value))
    if kind == "BOOL":
        return value.lower()
    return value


def run_qir(qir):
    with tempfile.NamedTemporaryFile("w", suffix=".ll", delete=False) as file:
        file.write(qir)
    handler = qirrunner.OutputHandler()
    try:
        qirrunner.run(file.name, shots=SHOTS, rng_seed=1, output_fn=handler.handle)
    except Exception as error:  # qir-runner rejects the module
        message = str(error).strip().splitlines()[0]
        if message.startswith("Failed to link"):  # valid QIR calling functions qir-runner lacks
            return {"qir-runner can't link: " + message.split(": ", 1)[-1]}
        return {"invalid QIR: " + re.sub(r"^\S+\.ll:\d+:\d+: (error: )?", "", message)}
    finally:
        os.unlink(file.name)
    return set(parse_shots(handler.get_output()))


def outputs():
    rows = {}
    ctx = Context(target_profile=TargetProfile.Unrestricted)
    ctx.eval(SOURCE)
    rows["simulator"] = {fmt(v) for v in ctx.run(ENTRY, SHOTS, seed=1)}
    for profile in PROFILES:
        ctx = Context(target_profile=getattr(TargetProfile, profile))
        try:
            ctx.eval(SOURCE)
            qir = str(ctx.compile(ENTRY))
        except Exception as error:  # the profile rejects the program
            rows[profile] = {"does not compile: " + str(error).strip().splitlines()[0]}
            continue
        rows[profile] = run_qir(qir)
    return rows


def main():
    rows = outputs()
    width = max(map(len, rows)) + 2
    for name, values in rows.items():
        print(f"{name:<{width}}{' | '.join(sorted(values))}")
    print(f"{'expected':<{width}}{EXPECTED}")
    wrong = [name for name in AFFECTED if name in rows and rows[name] != {EXPECTED}]
    if wrong:
        print(f"REPRODUCED: {', '.join(wrong)} differ from {EXPECTED}")
        return 0
    print("NOT REPRODUCED")
    return 1


if __name__ == "__main__":
    sys.exit(main())