{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Revise; using QEDbase; using QEDprocesses; using MetagraphOptimization; using BenchmarkTools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "QEDParticleTypeIn = Union{Type{FermionStateful{Incoming}},Type{AntiFermionStateful{Incoming}},Type{PhotonStateful{Incoming}}};\n",
    "QEDParticleTypeOut = Union{Type{FermionStateful{Outgoing}}, Type{AntiFermionStateful{Outgoing}},Type{PhotonStateful{Outgoing}}};\n",
    "\n",
    "PhotonTypes = Type{PhotonStateful}\n",
    "FermionTypes = Union{Type{FermionStateful{Incoming}},Type{AntiFermionStateful{Outgoing}}};\n",
    "AntifermionTypes = Union{Type{FermionStateful{Outgoing}},Type{AntiFermionStateful{Incoming}}};\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "gen_diagram_helper (generic function with 4 methods)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "function gen_diagram_helper(photons::Int, fermionsIn::Int, fermionsOut::Int, antifermionsIn::Int, antifermionsOut::Int, trace::String = \"\", fermionsN::Int = 1, photonsN::Int = 1)\n",
    "    subtrace = \"ph\"\n",
    "    if (fermionsIn == 1)\n",
    "        subtrace *= \" + f(in)\"\n",
    "    end\n",
    "    if (antifermionsIn == 1)\n",
    "        subtrace *= \" + af(in)\"\n",
    "    end\n",
    "    if (fermionsOut == 1)\n",
    "        subtrace *= \" + f(out)\"\n",
    "    end\n",
    "    if (antifermionsOut == 1)\n",
    "        subtrace *= \" + af(out)\"\n",
    "    end\n",
    "    if photons == 1 && (fermionsIn + antifermionsOut == 1) && (antifermionsIn + fermionsOut == 1)\n",
    "        println(\"Average over $fermionsN fermion perturbations and sum over $photonsN photon perturbations for diagrams of form: $(trace)$(subtrace)\")\n",
    "        return nothing\n",
    "    end\n",
    "    \n",
    "    #println(\"trace: $trace  photons: $(photons)  fermions: $(fermionsIn) / $fermionsOut  antifermions: $antifermionsIn / $antifermionsOut\")\n",
    "\n",
    "    if (photons != 0)\n",
    "        if (fermionsIn != 0)\n",
    "            # photon + fermionin\n",
    "            gen_diagram_helper(photons - 1, fermionsIn, fermionsOut, antifermionsIn, antifermionsOut, trace * \"ph + f(in) = f(in) -- \", fermionsN * fermionsIn, photonsN * photons)\n",
    "        end\n",
    "        if (fermionsOut != 0)\n",
    "            # photon + fermionout\n",
    "            gen_diagram_helper(photons - 1, fermionsIn + 1, fermionsOut, antifermionsIn, antifermionsOut, trace * \"ph + f(out) = f(out) -- \", fermionsN * fermionsOut, photonsN * photons)\n",
    "        end\n",
    "\n",
    "        if (antifermionsIn != 0)\n",
    "            # photon + antifermionin\n",
    "            gen_diagram_helper(photons - 1, fermionsIn, fermionsOut, antifermionsIn, antifermionsOut, trace * \"ph + af(in) = af(in) -- \", fermionsN * antifermionsIn, photonsN * photons)\n",
    "        end\n",
    "        if (antifermionsOut != 0)\n",
    "            # photon + antifermionOut\n",
    "            gen_diagram_helper(photons - 1, fermionsIn, fermionsOut, antifermionsIn, antifermionsOut, trace * \"ph + af(out) = af(out) -- \", fermionsN * antifermionsOut, photonsN * photons)\n",
    "        end\n",
    "\n",
    "        if fermionsIn + fermionsOut + antifermionsIn + antifermionsOut != 0\n",
    "            return nothing\n",
    "        end\n",
    "    end\n",
    "\n",
    "    if (fermionsIn != 0)\n",
    "        if (fermionsOut != 0)\n",
    "            # fermionIn + fermionOut\n",
    "            gen_diagram_helper(photons + 1, fermionsIn - 1, fermionsOut - 1, antifermionsIn, antifermionsOut, trace * \"f(in) + f(out) = ph -- \", fermionsN * fermionsIn * fermionsOut, photonsN)\n",
    "            return nothing\n",
    "        end\n",
    "        if (antifermionsIn != 0)\n",
    "            #fermionIn + antfermionIn\n",
    "            gen_diagram_helper(photons + 1, fermionsIn - 1, fermionsOut, antifermionsIn - 1, antifermionsOut, trace * \"f(in) + af(in) = ph -- \", fermionsN * fermionsIn * antifermionsIn, photonsN)\n",
    "            return nothing\n",
    "        end\n",
    "        #fermionIn + antifermionOut is impossible\n",
    "    end\n",
    "\n",
    "    if (fermionsOut != 0)\n",
    "        if (antifermionsOut != 0)\n",
    "            #fermionOut + antifermionOut\n",
    "            gen_diagram_helper(photons + 1, fermionsIn, fermionsOut - 1, antifermionsIn, antifermionsOut - 1, trace * \"f(out) + af(out) = ph -- \", fermionsN * fermionsOut * antifermionsOut, photonsN)\n",
    "            return nothing\n",
    "        end\n",
    "        #fermionOut + antifermionIn is impossible\n",
    "    end\n",
    "\n",
    "    if (antifermionsIn != 0)\n",
    "        if (antifermionsOut != 0)\n",
    "            #antifermionIn + antifermionOut\n",
    "            gen_diagram_helper(photons + 1, fermionsIn, fermionsOut, antifermionsIn - 1, antifermionsOut - 1, trace * \"af(in) + af(out) = ph -- \", fermionsN * antifermionsIn * antifermionsOut, photonsN)\n",
    "            return nothing\n",
    "        end\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average over 1 fermion perturbations and sum over 2 photon perturbations for diagrams of form: ph + f(in) = f(in) -- ph + f(in) + f(out)\n"
     ]
    }
   ],
   "source": [
    "# Compton:\n",
    "gen_diagram_helper(2, 1, 1, 0, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average over 1 fermion perturbations and sum over 2 photon perturbations for diagrams of form: ph + f(in) = f(in) -- ph + f(in) + af(in)\n",
      "Average over 1 fermion perturbations and sum over 2 photon perturbations for diagrams of form: ph + af(in) = af(in) -- ph + f(in) + af(in)\n"
     ]
    }
   ],
   "source": [
    "# Pair annihilation\n",
    "gen_diagram_helper(2, 1, 0, 1, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average over 1 fermion perturbations and sum over 2 photon perturbations for diagrams of form: ph + af(out) = af(out) -- ph + f(out) + af(out)\n"
     ]
    }
   ],
   "source": [
    "# Pair production\n",
    "gen_diagram_helper(2, 0, 1, 0, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average over 2 fermion perturbations and sum over 1 photon perturbations for diagrams of form: ph + f(in) = f(in) -- f(in) + f(out) = ph -- ph + f(out) + af(out)\n",
      "Average over 2 fermion perturbations and sum over 1 photon perturbations for diagrams of form: ph + af(out) = af(out) -- f(in) + f(out) = ph -- ph + f(out) + af(out)\n"
     ]
    }
   ],
   "source": [
    "# Trident\n",
    "gen_diagram_helper(1, 1, 2, 0, 1)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.9.4",
   "language": "julia",
   "name": "julia-1.9"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.9.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}