{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: Method definition (::Type{QEDcore.ParticleStateful{DIR, SPECIES, ELEMENT} where ELEMENT<:QEDbase.AbstractFourMomentum})(QEDbase.AbstractFourMomentum) where {DIR<:QEDbase.ParticleDirection, SPECIES<:QEDbase.AbstractParticleType} in module QEDcore at /home/antonr/.julia/packages/QEDcore/uVldP/src/phase_spaces/create.jl:7 overwritten in module MetagraphOptimization at /home/antonr/.julia/packages/MetagraphOptimization/kQaEJ/src/QEDprocesses_patch.jl:15.\n",
      "ERROR: Method overwriting is not permitted during Module precompilation. Use `__precompile__(false)` to opt-out of precompilation.\n",
      "WARNING: Method definition (::Type{QEDcore.ParticleStateful{DIR, SPECIES, ELEMENT} where ELEMENT<:QEDbase.AbstractFourMomentum})(QEDbase.AbstractFourMomentum) where {DIR<:QEDbase.ParticleDirection, SPECIES<:QEDbase.AbstractParticleType} in module QEDcore at /home/antonr/.julia/packages/QEDcore/uVldP/src/phase_spaces/create.jl:7 overwritten in module MetagraphOptimization at /home/antonr/.julia/packages/MetagraphOptimization/kQaEJ/src/QEDprocesses_patch.jl:15.\n",
      "ERROR: Method overwriting is not permitted during Module precompilation. Use `__precompile__(false)` to opt-out of precompilation.\n"
     ]
    }
   ],
   "source": [
    "using Revise\n",
    "using FeynmanDiagramGenerator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "42-element Vector{VirtualParticle{QEDProcess{Tuple{Photon, Electron}, Tuple{Electron, Electron, Electron, Positron, Positron}, Tuple{AllPolarization, AllSpin}, NTuple{5, AllSpin}}, PT, 2, 5} where PT<:AbstractParticleType}:\n",
       " photon: \t00 | 00101\n",
       " photon: \t00 | 00110\n",
       " photon: \t00 | 01001\n",
       " photon: \t00 | 01010\n",
       " photon: \t00 | 10001\n",
       " photon: \t00 | 10010\n",
       " photon: \t01 | 00100\n",
       " photon: \t01 | 01000\n",
       " photon: \t01 | 10000\n",
       " electron: \t10 | 00001\n",
       " ⋮\n",
       " photon: \t10 | 00101\n",
       " photon: \t10 | 00110\n",
       " photon: \t10 | 01001\n",
       " photon: \t10 | 01010\n",
       " photon: \t10 | 10001\n",
       " photon: \t10 | 10010\n",
       " photon: \t11 | 00100\n",
       " photon: \t11 | 01000\n",
       " photon: \t11 | 10000"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: both QEDcore and QEDbase export \"mul\"; uses of it in module FeynmanDiagramGenerator must be qualified\n"
     ]
    }
   ],
   "source": [
    "proc = QEDProcess(1, 0, 1, 3, 0, 2)\n",
    "all_particles = Set()\n",
    "for fd in feynman_diagrams(proc)\n",
    "    push!(all_particles, virtual_particles(proc, fd)...)\n",
    "end\n",
    "all_particles = sort([all_particles...])\n",
    "display(all_particles)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OrderedCollections.OrderedDict{VirtualParticle, Vector{Tuple{VirtualParticle, VirtualParticle}}} with 42 entries:\n",
       "  photon: \t00 | 00101   => [(electron: \t00 | 00001, positron: \t00 | 00100)]\n",
       "  photon: \t00 | 00110   => [(electron: \t00 | 00010, positron: \t00 | 00100)]\n",
       "  photon: \t00 | 01001   => [(electron: \t00 | 00001, positron: \t00 | 01000)]\n",
       "  photon: \t00 | 01010   => [(electron: \t00 | 00010, positron: \t00 | 01000)]\n",
       "  photon: \t00 | 10001   => [(electron: \t00 | 00001, positron: \t00 | 10000)]\n",
       "  photon: \t00 | 10010   => [(electron: \t00 | 00010, positron: \t00 | 10000)]\n",
       "  photon: \t01 | 00100   => [(positron: \t00 | 00100, electron: \t01 | 00000)]\n",
       "  photon: \t01 | 01000   => [(positron: \t00 | 01000, electron: \t01 | 00000)]\n",
       "  photon: \t01 | 10000   => [(positron: \t00 | 10000, electron: \t01 | 00000)]\n",
       "  electron: \t10 | 00001 => [(electron: \t00 | 00001, photon: \t10 | 00000)]\n",
       "  electron: \t10 | 00010 => [(electron: \t00 | 00010, photon: \t10 | 00000)]\n",
       "  positron: \t10 | 00100 => [(positron: \t00 | 00100, photon: \t10 | 00000)]\n",
       "  positron: \t10 | 01000 => [(positron: \t00 | 01000, photon: \t10 | 00000)]\n",
       "  positron: \t10 | 10000 => [(positron: \t00 | 10000, photon: \t10 | 00000)]\n",
       "  electron: \t11 | 00000 => [(electron: \t01 | 00000, photon: \t10 | 00000)]\n",
       "  electron: \t00 | 00111 => [(electron: \t00 | 00001, photon: \t00 | 00110), (electro…\n",
       "  electron: \t00 | 01011 => [(electron: \t00 | 00001, photon: \t00 | 01010), (electro…\n",
       "  positron: \t00 | 01101 => [(positron: \t00 | 00100, photon: \t00 | 01001), (positro…\n",
       "  positron: \t00 | 01110 => [(positron: \t00 | 00100, photon: \t00 | 01010), (positro…\n",
       "  ⋮                    => ⋮"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pairs = sort(FeynmanDiagramGenerator.particle_pairs(all_particles))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "81-element Vector{Tuple{VirtualParticle, VirtualParticle, VirtualParticle}}:\n",
       " (photon: \t00 | 00101, electron: \t01 | 10010, positron: \t10 | 01000)\n",
       " (photon: \t00 | 00101, electron: \t11 | 00000, positron: \t00 | 11010)\n",
       " (photon: \t00 | 00101, electron: \t10 | 00010, positron: \t01 | 11000)\n",
       " (photon: \t00 | 00101, electron: \t01 | 01010, positron: \t10 | 10000)\n",
       " (photon: \t00 | 00110, electron: \t11 | 00000, positron: \t00 | 11001)\n",
       " (photon: \t00 | 00110, electron: \t10 | 00001, positron: \t01 | 11000)\n",
       " (photon: \t00 | 00110, electron: \t01 | 10001, positron: \t10 | 01000)\n",
       " (photon: \t00 | 00110, electron: \t01 | 01001, positron: \t10 | 10000)\n",
       " (photon: \t00 | 01001, electron: \t01 | 10010, positron: \t10 | 00100)\n",
       " (photon: \t00 | 01001, electron: \t11 | 00000, positron: \t00 | 10110)\n",
       " ⋮\n",
       " (photon: \t10 | 00000, electron: \t01 | 10010, positron: \t00 | 01101)\n",
       " (photon: \t10 | 00000, electron: \t01 | 00110, positron: \t00 | 11001)\n",
       " (photon: \t10 | 00000, electron: \t00 | 00111, positron: \t01 | 11000)\n",
       " (photon: \t10 | 00000, electron: \t00 | 01011, positron: \t01 | 10100)\n",
       " (photon: \t10 | 00000, electron: \t01 | 10001, positron: \t00 | 01110)\n",
       " (photon: \t10 | 00000, electron: \t00 | 10011, positron: \t01 | 01100)\n",
       " (photon: \t10 | 00000, electron: \t01 | 01001, positron: \t00 | 10110)\n",
       " (photon: \t10 | 00000, electron: \t01 | 00101, positron: \t00 | 11010)\n",
       " (photon: \t10 | 00000, electron: \t01 | 01010, positron: \t00 | 10101)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "triples = FeynmanDiagramGenerator.total_particle_triples(all_particles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "s: 252, should be: 252\n",
      "number of triples: 81\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "81-element Vector{Tuple{VirtualParticle, VirtualParticle, VirtualParticle}}:\n",
       " (photon: \t10 | 00000, electron: \t00 | 00111, positron: \t01 | 11000)\n",
       " (photon: \t10 | 00000, electron: \t00 | 01011, positron: \t01 | 10100)\n",
       " (photon: \t10 | 00000, electron: \t00 | 10011, positron: \t01 | 01100)\n",
       " (photon: \t10 | 00000, electron: \t01 | 00101, positron: \t00 | 11010)\n",
       " (photon: \t10 | 00000, electron: \t01 | 00110, positron: \t00 | 11001)\n",
       " (photon: \t10 | 00000, electron: \t01 | 01001, positron: \t00 | 10110)\n",
       " (photon: \t10 | 00000, electron: \t01 | 01010, positron: \t00 | 10101)\n",
       " (photon: \t10 | 00000, electron: \t01 | 10001, positron: \t00 | 01110)\n",
       " (photon: \t10 | 00000, electron: \t01 | 10010, positron: \t00 | 01101)\n",
       " (photon: \t00 | 00101, electron: \t10 | 00010, positron: \t01 | 11000)\n",
       " ⋮\n",
       " (photon: \t11 | 00100, electron: \t00 | 10011, positron: \t00 | 01000)\n",
       " (photon: \t11 | 01000, electron: \t00 | 00001, positron: \t00 | 10110)\n",
       " (photon: \t11 | 01000, electron: \t00 | 00010, positron: \t00 | 10101)\n",
       " (photon: \t11 | 01000, electron: \t00 | 00111, positron: \t00 | 10000)\n",
       " (photon: \t11 | 01000, electron: \t00 | 10011, positron: \t00 | 00100)\n",
       " (photon: \t11 | 10000, electron: \t00 | 00001, positron: \t00 | 01110)\n",
       " (photon: \t11 | 10000, electron: \t00 | 00010, positron: \t00 | 01101)\n",
       " (photon: \t11 | 10000, electron: \t00 | 00111, positron: \t00 | 01000)\n",
       " (photon: \t11 | 10000, electron: \t00 | 01011, positron: \t00 | 00100)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "function n(vp::VirtualParticle)\n",
    "    if !haskey(pairs, vp)\n",
    "        return 1\n",
    "    end\n",
    "    s = 0\n",
    "    for (l, r) in pairs[vp]\n",
    "        s += n(l) * n(r)\n",
    "    end\n",
    "    return s\n",
    "end\n",
    "\n",
    "s = 0\n",
    "for (ph, e, p) in triples\n",
    "    s += n(ph) * n(e) * n(p)\n",
    "end\n",
    "\n",
    "println(\"s: $s, should be: $(length(feynman_diagrams(proc)))\")\n",
    "println(\"number of triples: $(length(triples))\")\n",
    "\n",
    "sort(triples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Graph:\n",
       "  Nodes: Total: 23250, FeynmanDiagramGenerator.ComputeTask_CollectTriples: 128, FeynmanDiagramGenerator.ComputeTask_Pair: 492, \n",
       "         FeynmanDiagramGenerator.ComputeTask_BaseState: 14, MetagraphOptimization.DataTask: 11653, FeynmanDiagramGenerator.ComputeTask_Triple: 10368, \n",
       "         FeynmanDiagramGenerator.ComputeTask_PropagatePairs: 276, FeynmanDiagramGenerator.ComputeTask_CollectPairs: 276, FeynmanDiagramGenerator.ComputeTask_Propagator: 42, \n",
       "         FeynmanDiagramGenerator.ComputeTask_SpinPolCumulation: 1\n",
       "  Edges: 55281\n",
       "  Total Compute Effort: 0.0\n",
       "  Total Data Transfer: 0.0\n",
       "  Total Compute Intensity: 0.0\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "graph = generate_DAG(proc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "mock_machine (generic function with 1 method)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "using MetagraphOptimization\n",
    "using UUIDs\n",
    "\n",
    "function mock_machine()\n",
    "    return Machine(\n",
    "        [\n",
    "            MetagraphOptimization.NumaNode(\n",
    "                0,\n",
    "                1,\n",
    "                MetagraphOptimization.default_strategy(MetagraphOptimization.NumaNode),\n",
    "                -1.0,\n",
    "                UUIDs.uuid1(),\n",
    "            ),\n",
    "        ],\n",
    "        [-1.0;;],\n",
    "    )\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PhaseSpacePoint{QEDProcess{Tuple{Photon, Electron}, Tuple{Electron, Electron, Electron, Positron, Positron}, Tuple{AllPolarization, AllSpin}, NTuple{5, AllSpin}}, PerturbativeQED, PhasespaceDefinition{SphericalCoordinateSystem, ElectronRestFrame}, Tuple{ParticleStateful{Incoming, Photon, SFourMomentum}, ParticleStateful{Incoming, Electron, SFourMomentum}}, Tuple{ParticleStateful{Outgoing, Electron, SFourMomentum}, ParticleStateful{Outgoing, Electron, SFourMomentum}, ParticleStateful{Outgoing, Electron, SFourMomentum}, ParticleStateful{Outgoing, Positron, SFourMomentum}, ParticleStateful{Outgoing, Positron, SFourMomentum}}, SFourMomentum}"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "psp = PhaseSpacePoint(\n",
    "    proc, \n",
    "    PerturbativeQED(), \n",
    "    PhasespaceDefinition(SphericalCoordinateSystem(), ElectronRestFrame()), \n",
    "    tuple((rand(SFourMomentum) for _ in 1:number_incoming_particles(proc))...),\n",
    "    tuple((rand(SFourMomentum) for _ in 1:number_outgoing_particles(proc))...)\n",
    ")\n",
    "typeof(psp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "using MetagraphOptimization: unpack_identity\n",
    "\n",
    "func = eval(get_compute_function(graph, proc, mock_machine()))\n",
    "\n",
    "func(psp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using BenchmarkTools\n",
    "\n",
    "@benchmark func($psp)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.10.4",
   "language": "julia",
   "name": "julia-1.10"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}