{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Grover Search"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "using Yao\n",
    "using Yao.EasyBuild: variational_circuit\n",
    "using LinearAlgebra"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Grover Step\n",
    "A single grover step is consist of applying oracle circuit and reflection circuit.\n",
    "The `reflection_circuit` function takes the wave function generator `U` as the input and returns `U|0><0|U'`."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "function grover_step!(reg::AbstractRegister, oracle, U::AbstractBlock)\n",
    "    apply!(reg |> oracle, reflect_circuit(U))\n",
    "end\n",
    "\n",
    "function reflect_circuit(gen::AbstractBlock)\n",
    "    N = nqubits(gen)\n",
    "    reflect0 = control(N, -collect(1:N-1), N=>-Z)\n",
    "    chain(gen', reflect0, gen)\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Compute the propotion of target states to estimate the number of iterations,\n",
    "which requires computing the output state."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "function solution_state(oracle, gen::AbstractBlock)\n",
    "    N = nqubits(gen)\n",
    "    reg= zero_state(N) |> gen\n",
    "    reg.state[real.(statevec(ArrayReg(ones(ComplexF64, 1<<N)) |> oracle)) .> 0] .= 0\n",
    "    normalize!(reg)\n",
    "end\n",
    "\n",
    "function num_grover_step(oracle, gen::AbstractBlock)\n",
    "    N = nqubits(gen)\n",
    "    reg = zero_state(N) |> gen\n",
    "    ratio = abs2(solution_state(oracle, gen)'*reg)\n",
    "    Int(round(pi/4/sqrt(ratio)))-1\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Run\n",
    "First, we define the problem by an oracle, it finds bit string `bit\"000001100100\"`."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "num_bit = 12\n",
    "oracle = matblock(Diagonal((v = ones(ComplexF64, 1<<num_bit); v[Int(bit\"000001100100\")+1]*=-1; v)))"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "then solve the above problem"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "gen = repeat(num_bit, H, 1:num_bit)\n",
    "reg = zero_state(num_bit) |> gen\n",
    "\n",
    "target_state = solution_state(oracle, gen)\n",
    "\n",
    "for i = 1:num_grover_step(oracle, gen)\n",
    "    grover_step!(reg, oracle, gen)\n",
    "    overlap = abs(reg'*target_state)\n",
    "    println(\"step $(i-1), overlap = $overlap\")\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Rejection Sampling\n",
    "In practise, it is often not possible to determine the number of iterations before actual running.\n",
    "we can use rejection sampling technique to avoid estimating the number of grover steps."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "In a single try, we `apply` the grover algorithm for `nstep` times."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "function single_try(oracle, gen::AbstractBlock, nstep::Int; nbatch::Int)\n",
    "    N = nqubits(gen)\n",
    "    reg = zero_state(N+1; nbatch)\n",
    "    focus(reg, (1:N...,)) do r\n",
    "        r |> gen\n",
    "        for i = 1:nstep\n",
    "            grover_step!(r, oracle, gen)\n",
    "        end\n",
    "        return r\n",
    "    end\n",
    "    reg |> checker\n",
    "    res = measure!(RemoveMeasured(), reg, (N+1))\n",
    "    return res, reg\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "After running the grover search, we have a checker program that flips the ancilla qubit\n",
    "if the output is the desired value, we assume the checker program can be implemented in polynomial time.\n",
    "to gaurante the output is correct.\n",
    "We contruct a checker \"program\", if the result is correct, flip the ancilla qubit"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "ctrl = -collect(1:num_bit); ctrl[[3,6,7]] *= -1\n",
    "checker = control(num_bit+1,ctrl, num_bit+1=>X)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The register is batched, with batch dimension `nshot`.\n",
    "`focus!` views the first 1-N qubts as system.\n",
    "For a batched register, `measure!`\n",
    "returns a vector of bitstring as output."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Run"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "maxtry = 100\n",
    "nshot = 3\n",
    "\n",
    "for nstep = 0:maxtry\n",
    "    println(\"number of iter = $nstep\")\n",
    "    res, regi = single_try(oracle, gen, nstep; nbatch=3)\n",
    "\n",
    "    # success!\n",
    "    if any(==(1), res)\n",
    "        overlap_final = viewbatch(regi, findfirst(==(1), res))'*target_state\n",
    "        println(\"success, overlap = $(overlap_final)\")\n",
    "        break\n",
    "    end\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The final state has an overlap of `1` with the target state."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Amplitude Amplification\n",
    "Given a circuit to generate a state,\n",
    "now we want to project out the subspace with [1,3,5,8,9,11,12] fixed to 1 and [4,6] fixed to 0.\n",
    "We can construct an oracle"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "evidense = [1, 3, -4, 5, -6, 8, 9, 11, 12]\n",
    "function inference_oracle(nbit::Int, locs::Vector{Int})\n",
    "    control(nbit, locs[1:end-1], abs(locs[end]) => (locs[end]>0 ? Z : -Z))\n",
    "end\n",
    "oracle = inference_oracle(nqubits(reg), evidense)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "We use a variational circuit generator defined in `Yao.EasyBuild`"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "gen = dispatch!(variational_circuit(num_bit), :random)\n",
    "reg = zero_state(num_bit) |> gen"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Run"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "solution = solution_state(oracle, gen)\n",
    "for i = 1:num_grover_step(oracle, gen)\n",
    "    grover_step!(reg, oracle, gen)\n",
    "    println(\"step $(i-1), overlap = $(abs(reg'*solution))\")\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "---\n",
    "\n",
    "*This notebook was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).*"
   ],
   "metadata": {}
  }
 ],
 "nbformat_minor": 3,
 "metadata": {
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.9.2"
  },
  "kernelspec": {
   "name": "julia-1.9",
   "display_name": "Julia 1.9.2",
   "language": "julia"
  }
 },
 "nbformat": 4
}