(*
 * # Tune Mathematica CLI
 *)

SetOptions[$Output, PageWidth -> 65];
$HistoryLength=2;

(*
 * # General parameter. Because we can.
 *)

NLOOPS = 2;

(*
 * # Generate the diagrams and load them
 *)

Run["./mkdia.py dia-A-A-" <> ToString[NLOOPS] <> ".m"];

diagrams = Get["dia-A-A-" <> ToString[NLOOPS] <> ".m"];
Print["Loaded ", diagrams//Length, " diagrams"];

(*
 * # Insert Feynman rules
 *)

(* Shortcut for index deltas. *)
deltaflv[a_, b_] := deltaf[flv[a], flv[b]]
deltaflvt[a_, b_] := deltaft[flvt[a], flvt[b]]
deltafun[a_, b_] := delta[fun[a], fun[b]]
deltaadj[a_, b_] := delta[adj[a], adj[b]]
deltalor[a_, b_] := delta[lor[a], lor[b]]

ClearAll[Amplitude];
Amplitude[Diagram[id_, factor_, ifields_List, ofields_List, propagators_List, vertices_List]] :=
  Times[
    factor,
    propagators // Map[Amplitude] // Apply[Times],
    vertices // Map[Amplitude] // Apply[Times]
  ]
(* (Light) quark propagator *)
Amplitude[P["q", fi1_, fi2_, _, _, p_]] :=
  I deltaflv[fi1, fi2] deltafun[fi2, fi1] \
  gammachain[slash[p], spn[fi2], spn[fi1]] den[p]
(* Heavy quark propagator *)
Amplitude[P["t", fi1_, fi2_, _, _, p_]] :=
  I deltaflvt[fi1, fi2] deltafun[fi2, fi1] \
  (gammachain[slash[p], spn[fi2], spn[fi1]] + mt1 gammachain[spn[fi2], spn[fi1]]) den[p, mt2]
(* Gluon propagator *)
Amplitude[P["g", fi1_, fi2_, _, _, p_]] :=
  -I deltaadj[fi1, fi2] (deltalor[fi1, fi2] den[p] -
    Xi momentum[p, lor[fi1]] momentum[p, lor[fi2]] den[p]^2)
(* Ghost propagator *)
Amplitude[P["c", fi1_, fi2_, _, _, p_]] := I deltaadj[fi1, fi2] den[p]
(* Light-quark-gluon vertex *)
Amplitude[V[_, "Qqg", fi1_, p1_, fi2_, p2_, fi3_, p3_]] :=
  I gs deltaflv[fi2, fi1] gammachain[gamma[lor[fi3]], spn[fi1],
  spn[fi2]] colorT[adj[fi3], fun[fi1], fun[fi2]]
(* Heavy-quark-gluon vertex *)
Amplitude[V[_, "Ttg", fi1_, p1_, fi2_, p2_, fi3_, p3_]] :=
  I gs deltaflvt[fi2, fi1] gammachain[gamma[lor[fi3]], spn[fi1],
  spn[fi2]] colorT[adj[fi3], fun[fi1], fun[fi2]]
(* Ghost-gluon vertex *)
Amplitude[V[_, "Ccg", fi1_, p1_, fi2_, p2_, fi3_, p3_]] :=
  gs colorf[adj[fi3], adj[fi2], adj[fi1]] momentum[p1, lor[fi3]]
(* Quark-photon vertex *)
Amplitude[V[_, "QqA", fi1_, p1_, fi2_, p2_, fi3_, p3_]] :=
  (* ge *) I deltaflv[fi2, fi1] deltafun[fi1, fi2] chargeQ[flv[fi1]] \
  gammachain[gamma[lor[fi3]], spn[fi1], spn[fi2]]
(* Heavy quark-photon vertex *)
Amplitude[V[_, "TtA", fi1_, p1_, fi2_, p2_, fi3_, p3_]] :=
  (* ge *) I deltaflvt[fi2, fi1] deltafun[fi1, fi2] chargeQt[flvt[fi1]] \
  gammachain[gamma[lor[fi3]], spn[fi1], spn[fi2]]
(* 3-gluon vertex *)
Amplitude[V[_, "ggg", fi1_, p1_, fi2_, p2_, fi3_, p3_]] :=
  gs colorf[adj[fi1], adj[fi2], adj[fi3]] (
    deltalor[fi1, fi2] momentum[p1 - p2, lor[fi3]] +
    deltalor[fi2, fi3] momentum[p2 - p3, lor[fi1]] +
    deltalor[fi3, fi1] momentum[p3 - p1, lor[fi2]]
  )
(* 4-gluon vertex *)
Amplitude[V[vi_, "gggg", fi1_, p1_, fi2_, p2_, fi3_, p3_, fi4_, p4_]] :=
  Module[{adjX = adj[1000 + vi]},
    -I gs^2 (
      colorf[adjX, adj[fi1], adj[fi2]] colorf[adjX, adj[fi3], adj[fi4]]
        (deltalor[fi1, fi3] deltalor[fi2, fi4] -
          deltalor[fi1, fi4] deltalor[fi2, fi3]) +
      colorf[adjX, adj[fi1], adj[fi3]] colorf[adjX, adj[fi4], adj[fi2]]
        (deltalor[fi1, fi4] deltalor[fi2, fi3] -
          deltalor[fi1, fi2] deltalor[fi3, fi4]) +
      colorf[adjX, adj[fi1], adj[fi4]] colorf[adjX, adj[fi2], adj[fi3]]
        (deltalor[fi1, fi2] deltalor[fi3, fi4] -
          deltalor[fi1, fi3] deltalor[fi2, fi4])
    )
  ]
(* Default case, just to be safe *)
Amplitude[x__] := Error["Don't know an amplitude for: ", x]

(* Print a message and fail *)
Error[msg__] := (Print["ERROR: ", msg]; Throw[$Failed];)

(*
 * # Apply Feynman rules to get amplitudes from the diagrams.
 *)

projector = deltalor[-1,-2] / (d-1);

amplitudes = diagrams // Map[Amplitude[#] * projector&];

(* Note the den[0] and momentum[0, ...] in the last amplitude: *)
amplitudes[[-1]]
(* Look at diagrams.nb to figure out the source. *)

(*
 * # Cleanup scaleless integrals
 *)

amplitudes = amplitudes /. den[0] -> 0 /. momentum[0,_] -> 0;
Print["Non-zero amplitudes: ", amplitudes//Count[Except[0]], " of ", amplitudes//Length];

(*
 * # Tensor summation through FORM
 *)

(* Flatten, join, and convert the arguments to a string. *)
MkString[args__] := args // List // Flatten // Map[ToString] // StringJoin

(* Return a name of a non-existing temporary file. *)
MkTemp[prefix_, suffix_] := Module[{i, fn, alphabet},
  alphabet = Characters["abcdefghijklmnopqrstuvwxyz0123456789"];
  While[True,
    i = RandomSample[alphabet, 8];
    fn = FileNameJoin[{$TemporaryDirectory, MkString[prefix, ".", Environment["USER"], ".", $ProcessID, ".", i, suffix]}];
    If[Not[FileExistsQ[fn]], Return[fn]];
  ]
]

(* Convert the arguments to string and write it into the specified file. *)
MkFile[filename_, arguments__] := Module[{f},
  (* The BinaryFormat is needed for the BinaryWrite in WrString. *)
  f = OpenWrite[filename, BinaryFormat->True];
  If[f === $Failed, Error["MkFile: failed to open ", filename, " for writing"]];
  BinaryWrite[f, MkString[arguments]];
  Close[f];
]

(* Find all unique occurrences of pat in ex. *)
CaseUnion[ex_, pat_] := Cases[ex, pat, {0, Infinity}] // Union
CaseUnion[pat_] := CaseUnion[#, pat]&

(* Return two maps: one to rename all indices in an expression to
 * FORM syntax, and the other to rename them back. *)
RenameIndicesToFromFORM[ex_] := Module[{original, renamed, i, idx, n},
  original = CaseUnion[ex, _lor | _adj | _fun | _spn | _flv | _flvt];
  renamed = original // Map[Replace[{
       i : (idx_)[n_?Negative] :> ToExpression[ToString[idx] <> "m" <> ToString[-n]],
       i : (idx_)[n_] :> ToExpression[ToString[idx] <> "p" <> ToString[n]]
       }]];
  {MapThread[Rule, {original, renamed}], MapThread[Rule, {renamed, original}]}
]

(* Convert an expression to FORM syntax. Note that indices should be
 * renamed before this. *)
ToFORM[expression_] := ToString[expression, InputForm] // StringReplace[{"\"" -> "", " " -> "", "[" -> "(", "]" -> ")"}]

(* Save a list of expressions to a FORM file, execute a FORM program
 * that transforms them, read the results back. *)
RunThroughFORM[exprlist_List, formcode__] := Module[{tmpfile, tmplogfile, resultfiles, toform, fromform, results},
  tmpfile = MkTemp["tmp", ".frm"];
  tmplogfile = tmpfile // StringReplace[".frm" ~~ EndOfString -> ".log"];
  resultfiles = Table[MkString[tmpfile, ".", i, ".m"], {i, Length[exprlist]}];
  {toform, fromform} = RenameIndicesToFromFORM[exprlist];
  MkFile[tmpfile,
    "#include alltheFORMcode.frm\n",
    Table[
      {"local EXPR", i, " = (", exprlist[[i]]/.toform//ToFORM, ");\n"},
      {i, Length[exprlist]}
    ],
    "#call input\n",
    formcode,
    "#call output\n",
    Table[
      {"#call saveto(EXPR", i, ",", tmpfile, ".", i, ".m)\n"},
      {i, Length[exprlist]}
    ],
    ".end\n"
  ];
  (*Print[tmpfile];*)
  Run[MkString["form -l ", tmpfile]];
  results = Table[
    Get[MkString[tmpfile, ".", i, ".m"]], 
    {i, Length[exprlist]}
  ] /. fromform;
  DeleteFile[{tmpfile, tmplogfile}];
  DeleteFile[resultfiles];
  results
]
RunThroughFORM[expr_, formcode__] := RunThroughFORM[{expr}, formcode] // First

(* Just a test: with no code in the middle, all expressions should
 * just be returned unchanged. *)
RunThroughFORM[amplitudes, ""] - amplitudes//Expand

amplitudes2 = RunThroughFORM[amplitudes,
  "#call applydeltas\n",
  "#call colorsum\n",
  "#call flavorsum\n",
  "#call diracsum\n",
  "#call contractmomenta\n",
  "id mt1^2 = mt2;\n"
]
Print["Non-zero amplitudes: ", amplitudes2 // Count[Except[0]]];

(* To double-check that we've got only what we expected to get. *)

amplitudes2 // CaseUnion[_Symbol]

amplitudes2 // CaseUnion[(head_)[___] :> head]

(*
 * # Investigate the amplitudes that became zero
 *)
If[False,

  amplitudes[[3]]

  amplitudes2[[3]]

  (* Answer: look at the diagrams. *)

  (* But we can also investigate further. *)

  Factors[ex_Times] := List @@ ex
  Factors[ex_] := {ex}

  colorfactor = amplitudes[[3]]//Factor//Factors//Select[FreeQ[_adj|_fun]/*Not]//Apply[Times]

  RunThroughFORM[colorfactor,
    "#call applydeltas\n",
    "#call colorsum\n"
  ]

  colorfactors = amplitudes // Map[
    Factor/*Factors/*Select[FreeQ[_adj|_fun]/*Not]/*Apply[Times]
  ] //
    RunThroughFORM[#, "#call applydeltas\n", "#call colorsum\n" ]&

]; (* End investigation. *)

(*
 * # Integral families (IBP topologies)
 *)

DropLeadingSign[ex_ /; (FactorTermsList[ex] // First // Negative)] := -ex
DropLeadingSign[ex_] := ex

NormalizeDens[ex_] := ex /. den[p_,x___] :> den[DropLeadingSign[p], x]

loopmomenta = diagrams // CaseUnion[l1|l2|l3|l4|l5];
externalmomenta = diagrams // CaseUnion[q|q1|q2|q3|q4|q5|p1|p2|p3|p4|p5];
Print["External momenta: ", externalmomenta];
Print["Loop momenta: ", loopmomenta, " (expected: ", NLOOPS, ")"];

denominatorsets = amplitudes2 // DeleteCases[1] // NormalizeDens // Map[
  CaseUnion[_den] /* Select[FreeQ[Alternatives@@loopmomenta]/*Not]
];
Print["Unique denominator sets: ", denominatorsets // DeleteCases[{}] // Union // Length];

denominatorsets // Map[Print];

(* Return a list of momenta maps, such that applying them to
 * the list of feynman integral families makes symmetries and
 * subtopology relations explicit. So, families that are symmetric
 * will have identical sets of denominators after the maps are
 * applied. Families that are symmetric to a subtopology of a
 * bigger family will have a subsets of the denominators.
 *
 * The families are defined by their set of den[]s.
 *
 * You can use UniqueSupersetMapping[] to figure out the
 * topmost supertopologies after this.
 *
 * Uses Feynson (https://github.com/magv/feynson) underneath.
 *)
SymmetryMaps[denominatorsetlist_List, loopmom_List, extmom_List] :=
Module[{densets, uniqdensets, densetindices, uniqdensetmaps},
  densets = denominatorsetlist // NormalizeDens // Map[
    CaseUnion[_den] /* Union /* Select[FreeQ[Alternatives@@loopmom] /* Not]
  ];
  densets = densets /.
    den[p_] :> p^2 /.
    den[p_, m_] :> p^2-m;
  RunThrough["./feynson/feynson symmetrize -j1 -q -", {densets, loopmom, {}}] //
    Map[Map[Apply[Rule]]]
];

momentamaps = SymmetryMaps[denominatorsets, loopmomenta, externalmomenta];
Print["Found ", momentamaps // DeleteCases[{}] // Length, " momenta mappings"];

momentamaps//Map[Print];

symmetrizeddenominatorsets = MapThread[ReplaceAll, {denominatorsets, momentamaps}] // NormalizeDens;

(* Return the first index of the given list where `f[element]` is true. *)
ElementIndex[l_List, f_, default_] := FirstPosition[l, _?f, {default}, {1}, Heads->False] // First

(* Among a list of sets, find such a sublist such that all other
 * sets are subsets of these ones. Return the list, and a list of
 * indices indicating which set belongs to which superset.
 *
 * Example:
 *
 *   {{3},{1,2,3},{2,3,1},{2},{1,4,3},{4}}//UniqueSupersetMapping
 *   > { {{1,2,3}, {1,4,3}}, {1,1,1,1,2,2} }
 *)
UniqueSupersetMapping[sets_List, subsetq_:SubsetQ] := Module[{supersets, idx, set, IdxOf},
  supersets = {};
  IdxOf[set_] := IdxOf[set] = (
    idx = ElementIndex[supersets, subsetq[#, set]&, None];
    If[idx === None,
      supersets = Append[supersets, set // Sort];
      Length[supersets]
      ,
      idx
    ]
  );
  sets // SortBy[Length /* Minus] // Map[IdxOf];
  {supersets, sets // Map[IdxOf] }
]

{denominatorsupersets, supersetindices} = UniqueSupersetMapping[symmetrizeddenominatorsets];
Print["Total integral families: ", denominatorsupersets//Length];
denominatorsupersets // Map[Print];

(* An IBP basis has these components:
 * - "id": its number (or name);
 * - "denominators": the list of its denominators (den[] expression);
 * - "loopmom": the list of loop momenta names;
 * - "extmom": the list of external momenta names;
 * - "sprules": the list of external invariant substitutions (i.e. {sp[q,q]->sqrq});
 * - "numrules": the map from scalar products involving loop momenta to denominator indices;
 * - "denrules": the map from den[] expressions to denominator indices.
 *
 * CompleteIBPBasis[] takes the first 5, and completes the last two, returning an Association.
 *)
CompleteIBPBasis[id_, denominators_List, loopmom_List, externalmom_List, sprules_List] :=
Module[{L, M, p, l, k, dens, nums, vars, c, mx, candidates, newnum, newc, newmx, Complete, rels},
  L = loopmom // Apply[Alternatives];
  M = Join[loopmom, externalmom] // Apply[Alternatives];
  dens = denominators // NormalizeDens // Sort;
  nums = dens /.
    den[p_] :> p^2 /.
    den[p_, m2_, ___] :> p^2 - m2 //
    Expand;
  nums = nums /.
    (l:M) (k:M) :> Sort[sp[l, k]] /.
    (l:M)^2 :> sp[l, l] /.
    sprules;
  vars = Tuples[{loopmom, Join[loopmom, externalmom]}] //
    Map[Sort /* Apply[sp]] //
    Union;
  Print["Independent scalar products: ", Length[vars], ", ", vars];
  {c, mx} = CoefficientArrays[nums, vars] // Normal;
  (* nums == c + mx.vars *)
  If[MatrixRank[mx] < Length[mx],
    Error["CompleteIBPBasis: denominators ", denominators, " are already linearly dependent"]
  ];
  If[Length[mx] =!= Length[vars],
    (*Error["CompleteIBPBasis: denominators ", denominators, " are incomplete"]*)
    candidates =
       Subsets[Join[loopmom, externalmom], {1, Infinity}] //
       Map[Apply[Plus]] //
       Select[FreeQ[L] /* Not];
    Print["Candidate irreducible denominators: ", candidates];
    While[Length[mx] < Length[vars] && candidates =!= {},
      newnum =
        Expand[candidates[[1]]^2] /.
        (l:M) (k:M) :> Sort[sp[l, k]] /.
        (l:M)^2 :> sp[l, l] /.
        sprules;
      {newc, newmx} = CoefficientArrays[newnum, vars] // Normal;
      If[MatrixRank[Append[mx, newmx]] === Length[mx] + 1,
        Print["Adding denominator: ", den[candidates[[1]], 0, irr]];
        AppendTo[mx, newmx];
        AppendTo[dens, den[candidates[[1]], 0, irr]];
        AppendTo[c, newc];
        ,
        candidates = candidates[[2 ;;]];
      ];
    ];
  ];
  <|
    "id" -> id,
    "denominators" -> dens,
    "loopmom" -> loopmom,
    "externalmom" -> externalmom,
    "sprules" -> sprules,
    "denmap" -> (
      MapIndexed[#1 -> DEN @@ #2 &, dens] //
      DeleteCases[den[_, _, irr] -> _] //
      ReplaceAll[(den[p_, x___] -> y_) :> {den[p, x] -> y, den[-p, x] -> y}] //
      Flatten
    ),
    "nummap" -> (
      Inverse[mx].(Map[1/DEN[#] &, Range[Length[mx]]] - c) //
      Collect[#, _DEN, Factor]& //
      MapThread[Rule, {vars, #}]& //
      Join[#, # /. sp[a_, b_] :> sp[b, a]]&
    )
  |>
]
CompleteIBPBasis[1, denominatorsupersets[[1]], loopmomenta, externalmomenta, {sp[q,q]->sqrq}]

(* Test auto-completion of the basis *)
CompleteIBPBasis[1, {den[l1]}, loopmomenta, externalmomenta, {sp[q,q]->sqrq}];

bases = denominatorsupersets //
  MapIndexed[CompleteIBPBasis[First[#2], #1, loopmomenta, externalmomenta, {sp[q,q]->sqrq}]&];

(*
 * # Convert amplitudes to B notation
 *
 * B notation: B[basis id, den1 power, den2 power, ..., denN power]
 *)

FORMCodeToB[bases_List] := Module[{basis},
  {
    "#procedure toBIDandDEN\n",
    "  ",
    Table[
      {
        "if (match(only, BID^", basis["id"], "));\n",
        basis["denmap"] //
          Map[{
            "    id ", #[[1]] // ToFORM, " = ",
            #[[2]] /. basis["sprules"] /. DEN[n_] :> MkString["DEN", n] // ToFORM,
            ";\n"
          }&] // Union,
        basis["nummap"] //
          Map[{
            "    id ", #[[1]] /. sp->(Dot/*Sort) // ToFORM, " = ",
            #[[2]] /. basis["sprules"] /. DEN[n_] :> MkString["DEN", n] // ToFORM,
            ";\n"
          }&] // Union,
        basis["sprules"] //
          Map[{
            "    id ", #[[1]] /. sp->(Dot/*Sort) // ToFORM, " = ",
            #[[2]] /. DEN[n_] :> MkString["DEN", n] // ToFORM,
            ";\n"
          }&] // Union
      },
      {basis, bases}
    ] //
      Riffle[#, "  else"]&,
    "  else;\n",
    "    exit \"ERROR: toBden: got a term without a proper BID^n factor.\";\n",
    "  endif;\n",
    "#endprocedure\n",
    "#call toB(toBIDandDEN,", Length[bases[[1, "denominators"]]], ")\n"
  }
]
FORMCodeToB[bases]//MkString//Print;

amplitudesB =
  MapThread[ReplaceAll, {amplitudes2, momentamaps}] //
  # * BID^supersetindices & //
  RunThroughFORM[#, FORMCodeToB[bases]] &;

amplitudesB // CaseUnion[_Symbol]

amplitudesB // CaseUnion[(head_)[___] :> head]

integrallist = amplitudesB // CaseUnion[_B];
Print["Need to know ", integrallist // Length, " different integrals"];

(*
 * # Filter out zero integrals early.
 *)

(* Subsector of a family is a subset of its denominators.
 * Sector Id is a numerical notation for this subset: if
 * out of all D only D_{k} are present in a subsector,
 * then its Sector Id is Sum[2^(k-1)]. *)
IndicesToSectorId[idx_List] := Table[If[idx[[i]] === 0, 0, 2^(i-1)], {i, Length[idx]}] // Apply[Plus]
SectorIdToIndices[sector_Integer, ndens_Integer] := IntegerDigits[sector, 2, ndens] // Reverse

IndicesToR[idx_List] := idx // Cases[n_ /; n > 0 :> n] // Apply[Plus]
IndicesToDots[idx_List] := idx // Cases[n_ /; n>1 :> n-1] // Apply[Plus]
IndicesToT[idx_List] := idx // Count[n_ /; n > 0]
IndicesToS[idx_List] := idx // Cases[n_ /; n < 0 :> -n] // Apply[Plus]

(* List zero sectors of a given basis in B notation. *)
ZeroSectors[bases_List] := bases // Map[ZeroSectors] // Apply[Join]
ZeroSectors[basis_] :=
  RunThrough["./feynson/feynson zero-sectors -sq -j1 -", {
      basis["denominators"] /.
        den[p_] :> p^2 /.
        den[p_, m_] :> p^2-m /.
        den[p_, m_, irr] :> p^2-m,
      basis["denominators"] /.
        den[_, _, cut] -> 1 /.
        den[___] -> 0,
      basis["loopmom"],
      basis["sprules"] /.
        Rule->List /.
        sp -> Times
    }] //
    Map[B[basis["id"], Sequence@@SectorIdToIndices[#, Length[basis["denominators"]]]]&]

zerosectors = ZeroSectors[bases]

ZeroSectorPattern[basis_] :=
  ZeroSectors[basis] //
  Map[Replace[B[bid_, idx__] :> B[bid, {idx} /. 1 -> _ /. 0 -> _?NonPositive // Apply[Sequence]]]] //
  Apply[Alternatives]

zerosectorpattern = ZeroSectorPattern[bases];
integrallist2 = integrallist // DeleteCases[zerosectorpattern];
Print["Out of all ", integrallist//Length, " there are only ", integrallist2//Length, " non-zero ints"];

FORMCodeZeroSectors[zerosectors_List] :=
  zerosectors // Map[Replace[
    B[bid_, idx__] :> {
      "id B(", bid, MapIndexed[ReplaceAll[#1, {
        0 -> {", tmpx", #2, "?neg0_"},
        1 -> {", tmpx", #2, "?"}
      }]&, {idx}], ") = 0;\n"
    }
  ]]

FORMCodeZeroSectors[zerosectors] // MkString // Print;

amplitudesB2 =
  MapThread[ReplaceAll, {amplitudes2, momentamaps}] //
  # * BID^supersetindices & //
  RunThroughFORM[#,
    FORMCodeToB[bases],
    FORMCodeZeroSectors[zerosectors]
  ] &;
Print["Dropping zero integrals early is ", ByteCount[amplitudesB]/ByteCount[amplitudesB2]//N, " times better"];

(*
 * # Solving IBP relations with Kira
 *
 * https://gitlab.com/kira-pyred/kira
 * https://kira.hepforge.org/
 *
 * https://arxiv.org/abs/1705.05610
 * https://arxiv.org/abs/2008.06494
 *)

(* To define a job for Kira we need multiple files (in a separate directory):
 * - "config/kinematics.yaml" with the definition of kinematics;
 * - "config/integralfamilies.yaml" with the definition of integral families;
 * - files with lists of integrals we care about (one file per integral family);
 * - a jobs file (e.g. "jobs.yaml") tying it all together.
 *)

Quiet[CreateDirectory["kira-files/config"], {CreateDirectory::filex}];

(* Kira sorts bases by name instead of adhering to the order
 * of definition. We shall make sure that both the numerical
 * and the lexicographic orders match, which will prevent Kira
 * from messing it up.
 *)
KiraBasisName[basisid_] := MkString["b", IntegerDigits[basisid, 10, 5]]

KiraBasisName[1]

(* Create Kira integral list files in the given directory.
 * Kira insists on separate files for each integral family,
 * so multiple files with names "b00xxx.integrals" will be created,
 * one for each basis.
 *)
MkKiraIntegralList[dirname_, blist_] := Module[{basisids, basisid, idlist},
  basisids = blist // CaseUnion[B[basisid_, ___] :> basisid];
  Do[
    MkFile[dirname <> "/" <> KiraBasisName[basisid] <> ".integrals",
      blist //
        CaseUnion[B[basisid, idx__] :> {idx}] //
        Map[{KiraBasisName[basisid], "[", Riffle[#, ","], "]\n"}&]
    ];
    ,
    {basisid, basisids}];
]

(*
MkKiraIntegralList["kira-files/", integrallist2];
*)

MkKiraKinematicsYaml[filename_, extmom_List, sprules_List] :=
  MkFile[filename,
    "kinematics:\n",
    " incoming_momenta: [", extmom // Riffle[#, ", "]&, "]\n",
    " kinematic_invariants:\n",
    "  - [mt2, 2]\n",
    sprules[[;;,2]] // Cases[_Symbol] // Union // Map[{"  - [", # , ", 2]\n"}&],
    sprules[[;;,2]] // Cases[_Symbol^2] // Union // Map[{"  - [", # , ", 1]\n"}&],
    " scalarproduct_rules:\n",
    sprules //
      ReplaceAll[sp -> (sp /* Sort)] //
      Union //
      Map[Replace[{
        (sp[p_] -> v_) :> {"  - [[", p//InputForm, ",", p//InputForm, "], ", v//InputForm, "]\n"},
        (sp[p1_, p2_] -> v_) :> {"  - [[", p1//InputForm, ",", p2//InputForm, "], ", v//InputForm, "]\n"}
      }]],
    "# symbol_to_replace_by_one: sqrq"
  ];

(*
MkKiraKinematicsYaml["kira-files/config/kinematics.yaml", bases[[1, "externalmom"]], bases[[1, "sprules"]]];
*)

MkKiraIntegralFamiliesYaml[filename_, bases_List] :=
  MkKiraIntegralFamiliesYaml[filename, bases,
    (#["denominators"] // Map[Replace[den[_, _, irr] -> 0, _den -> 1]] // IndicesToSectorId)&]
MkKiraIntegralFamiliesYaml[filename_, bases_List, topsectors_] :=
Module[{loopmom, extmom, dens, basis},
  MkFile[filename,
    "integralfamilies:\n",
    Table[
      loopmom = basis["loopmom"];
      extmom = basis["externalmom"];
      dens = basis["denominators"];
      {
        "  - name: \"", KiraBasisName[basis["id"]], "\"\n",
        "    loop_momenta: [", Riffle[loopmom, ", "], "]\n",
        "    top_level_sectors: [", dens // Map[Replace[den[_, _, irr] -> 0, _den -> 1]] // IndicesToSectorId, "]\n",
        "    propagators:\n",
        dens // Map[Replace[{
          den[p_] | den[p_, 0, ___] :> {"      - [\"", CForm[p], "\", 0]\n"},
          den[p_, m_, ___] :> {"      - [\"", CForm[p], "\", \"", CForm[m /. sp[q] -> qq], "\"]\n"},
          d_ :> Error["MkKiraConfig: bad denominator form: ", d]
        }]],
        If[FreeQ[dens, cut]//Not,
          {"    cut_propagators: [",
          Riffle[Range[Length[dens]] // Select[MatchQ[dens[[#]], den[_, _, cut]]&], ", "],
          "]\n"
          }
          ,
          {}
        ]
      }
      ,
      {basis, bases}]
  ];
]

MkKiraJobsYaml[filename_, bids_List, topsectors_, mode_String] := Module[{bid, sector}, 
  MkFile[filename,
    "jobs:\n",
    " - reduce_sectors:\n",
    "    reduce:\n",
    Table[
        {"     - {topologies: [", KiraBasisName[bid], "], sectors: [", sector["id"], "], r: ", sector["r"], ", s: ", sector["s"], "}\n"}
        ,
        {bid, bids},
        {sector, topsectors[bid]}],
    "    select_integrals:\n",
    "     select_mandatory_list:\n",
    Table[
        {"      - [", KiraBasisName[bid], ", \"", KiraBasisName[bid], ".integrals\"]\n"}
        ,
        {bid, bids}],
    "#     select_mandatory_recursively:\n",
    Table[
        {
        "#      - {topologies: [", KiraBasisName[bid],
            "], sectors: [", sector["id"],
            "], r: ", sector["r"],
            ", s: ", sector["s"],
            ", d: ", sector["d"], "}\n"},
        {bid, bids},
        {sector, topsectors[bid]}],
    "    integral_ordering: 8\n",
    "    run_symmetries: true\n",
    "    run_initiate: true\n",
    "    run_triangular: true\n",
    "    run_back_substitution: true\n",
    " - kira2math:\n",
    "    target:\n",
    Table[
      {"     - [", KiraBasisName[bid], ", \"", KiraBasisName[bid], ".integrals\"]\n"},
      {bid, bids}],
    Table[
        {"#     - {topologies: [", KiraBasisName[bid], "], sectors: [", sector["id"], "], r: ", sector["r"], ", s: ", sector["s"], ", d: ", sector["d"], "}\n"},
        {bid, bids},
        {sector, topsectors[bid]}],
    "    reconstruct_mass: false\n"
  ];
]

KiraTopSectors[idxlist_List] :=
Module[{tops, sector2r, sector2s, sector2d, s2sectors, int, sector, r, s, d, sectors, done, i, ss},
  tops = idxlist // Map[IndicesToS] // Max[#, 1]&;
  sector2r = <||>;
  sector2s = <||>;
  sector2d = <||>;
  s2sectors = Association @@ Table[s -> {}, {s, 0, tops}];
  Do[
      sector = IndicesToSectorId[int];
      r = IndicesToR[int];
      s = IndicesToS[int];
      d = IndicesToDots[int];
      AppendTo[s2sectors[s], sector];
      sector2r[sector] = Max[r, sector2r[sector] /. _Missing -> 0];
      sector2d[sector] = Max[d, sector2d[sector] /. _Missing -> 0];
      (* Note: s=0 makes Kira produce false masters. It's not
       * clear if we should only fix s=0 case, or if we need
       * to add +1 to all s. Currently we're doing the former.
       *)
      sector2s[sector] = Max[s, sector2s[sector] /. _Missing -> 1];
      ,
      {int, idxlist}
  ];
  Print["* Sectors by numerator power sum (s)"];
  sectors = {};
  done = {};
  For[s = tops, s >= 0, s--,
      s2sectors[s] = s2sectors[s] // Union // Reverse;
      Do[
          If[MemberQ[done, sector], Continue[]];
          i = FirstPosition[done, ss_ /; (BitAnd[ss, sector] === sector)];
          If[MatchQ[i, _Missing],
              AppendTo[done, sector];
              AppendTo[sectors, sector];
              Print["Unique sector: ", sector, ", nprops=", DigitCount[sector, 2, 1], ", r=", sector2r[sector], ", s=", sector2s[sector], ", d=", sector2d[sector]];
              ,
              i = i[[1]];
              Print["Subsector of ", done[[i]], ": ", sector, ", nprops=", DigitCount[sector, 2, 1], ", r=", sector2r[sector], ", s=", sector2s[sector], ", d=", sector2d[sector]];
              sector2r[done[[i]]] = Max[sector2r[sector], sector2r[done[[i]]]];
              sector2d[done[[i]]] = Max[sector2d[sector], sector2d[done[[i]]]];
              sector2s[done[[i]]] = Max[sector2s[sector], sector2s[done[[i]]]];
              ];
          ,
          {sector, s2sectors[s]}
      ];
  ];
  (* We need to make sure each sector has more integrals than
   * masters, otherwise Kira will have nothing to work with, and
   * we'll miss masters in the IBP table.
   *)
  Do[
      sector2s[sector] = Max[sector2s[sector], 1];
      sector2d[sector] = Max[sector2d[sector], 1];
      sector2r[sector] = Max[sector2r[sector], DigitCount[sector, 2, 1]];
      ,
      {sector, sectors}];
  Print["Final sectors:"];
  Do[
    Print["- ", sector, " ", IntegerDigits[sector, 2, Length[First[idxlist]]]//Reverse, ", nprops=", DigitCount[sector, 2, 1], ", r=", sector2r[sector], ", s=", sector2s[sector], ", d=", sector2d[sector]];
    ,
    {sector, sectors}];
  Table[
    <|"id" -> sector, "r" -> sector2r[sector], "s" -> sector2s[sector], "d" -> sector2d[sector]|>
    ,
    {sector, sectors}]
]

MkKiraConfig[dirname_, bases_List, blist_] :=
Module[{bid, bids, bid2topsectors, idxlist},
  If[Not[FileExistsQ[dirname]], CreateDirectory[dirname]];
  If[Not[FileExistsQ[dirname <> "/config"]], CreateDirectory[dirname <> "/config"]];
  bids = blist // CaseUnion[B[bid_, ___] :> bid];
  bid2topsectors = Table[
    idxlist = blist // CaseUnion[B[bid, idx__] :> {idx}];
    bid -> (idxlist // KiraTopSectors // Sort)
    ,
    {bid, bids}] // Association;
  MkKiraKinematicsYaml[dirname <> "/config/kinematics.yaml",
    bases[[1,"externalmom"]], bases[[1,"sprules"]]];
  MkKiraIntegralFamiliesYaml[dirname <> "/config/integralfamilies.yaml", bases, bid2topsectors];
  MkKiraJobsYaml[dirname <> "/jobs.yaml", bids, bid2topsectors, "all"];
  MkKiraIntegralList[dirname, blist];
]

Run["rm -rf kira-files"];
MkKiraConfig["kira-files", bases, integrallist2]

Run[MkString[
  "cd kira-files && ",
  "rm -rf results sectormappings tmp && ",
  "env FERMATPATH=$PWD/../fermat/fer64 ",
  "../kira/kira jobs.yaml"
]];

Print["Master integrals:"];
Run["cat kira-files/results/b*/masters.final"];

KiraApplyResults[ex_, confdir_String, bases_List] := Module[{exx, basisids, basisid, filename, map},
  exx = ex;
  basisids = ex // CaseUnion[B[basisid_, __] :> basisid];
  Do[
    filename = MkString[confdir, "/results/", KiraBasisName[basisid], "/kira_", KiraBasisName[basisid], ".integrals.m"];
    map = Get[filename];
    If[map === $Failed, Error["Failed to load: ", filename]];
    map = map /. ToExpression[KiraBasisName[basisid]][idx__] :> B[basisid, idx];
    exx = exx /. map;
    map = None;
    ,
    {basisid, basisids}];
  exx
]

amplitudesIBP = KiraApplyResults[amplitudesB2, "kira-files", bases]

masters = amplitudesIBP // CaseUnion[_B];
Print["Previously we had ", amplitudesB2 // CaseUnion[_B] // Length, " integrals"];
Print["Now we have ", masters // Length, " integrals"];

fullamplitude = amplitudesIBP // Apply[Plus] // Collect[#, _B, Factor]&;

fullamplitude

fullamplitude /. (Nt|SumQt|SumQt2) :> 0 /. mt2->0 

(*
At NLOOPS=1 this would be the known leading order R ratio:

fullamplitude / (fullamplitude /. Nc->1 /. SumQf2 -> 1) /. mt2->0 /. SumQt2 -> 0

*)

(*
 * # Numerical evaluation of masters with pySecDec
 *
 * Note these environment variables:
 * export SECDEC_CONTRIB=$PWD/pySecDec-1.4.5/
 * export PYTHONPATH=$PWD/pySecDec-1.4.5/pySecDec-1.4.5:$PYTHONPATH
 *)

Quiet[CreateDirectory["secdec-files"], {CreateDirectory::filex}];

IntegralName[integral_B] := integral //
  ToString //
  StringReplace[" " -> ""] //
  StringReplace["," -> "_"] //
  StringReplace["[" -> ""] //
  StringReplace["]" -> ""]
basedir = "secdec-files";
Do[
  name = IntegralName[integral];
  Print["* Working on ", basedir, "/", name, ".*"];
  basisid = integral[[1]];
  indices = integral[[2;;]] // Apply[List];
  basis = bases[[basisid]];
  MkFile[basedir <> "/" <> name <> ".generate.py",
    "#!/usr/bin/env python3\n",
    "import pySecDec as psd\n",
    "loopint = psd.loop_integral.LoopIntegralFromPropagators(\n",
    "  loop_momenta = ['", loopmomenta // Riffle[#, "','"]&, "'],\n",
    "  external_momenta = ['", externalmomenta // Riffle[#, "','"]&, "'],\n",
    "  propagators = [\n",
    basis["denominators"] /. {
      den[p_] :> {"    '(", p//CForm, ")^2'"},
      den[p_,m_,___] :> {"    '(", p//CForm, ")^2-", m//CForm, "'"}
    } // Riffle[#, ",\n"]&,
    "\n",
    "  ],\n",
    "  powerlist = [", indices // Riffle[#, ","]&, "],\n",
    "  replacement_rules = [('q*q', 'sqrq')]\n",
    ")\n",
    "psd.loop_integral.loop_package(\n",
    "    name = '", name, "',\n",
    "    loop_integral = loopint,\n",
    "    real_parameters = ['sqrq', 'mt2'],\n",
    "    additional_prefactor = '1',\n",
    "    requested_order = 2,\n",
    "    form_work_space = '50M',\n",
    "    contour_deformation=True\n",
    ")\n"
  ];
  MkFile[basedir <> "/" <> name <> ".integrate.py",
    "#!/usr/bin/env python3\n",
    "import sys\n",
    "import pySecDec as psd\n",
    "parameters = [float(parameter) for parameter in sys.argv[1:]]\n",
    "sys.stderr.write(f'Parameters: {parameters}\\n')\n",
    "lib = psd.integral_interface.IntegralLibrary('", name, "/", name, "_pylink.so')\n",
    "lib.use_Vegas(epsrel=1e-4, epsabs=1e-07, maxeval=1000000)\n",
    "int_wo_prefactor, prefactor, int_with_prefactor = lib(real_parameters=parameters)\n",
    "print('\"', int_wo_prefactor, '\"')\n"
  ];
  ,
  {integral, masters}];

Do[
  name = IntegralName[integral];
  Run[MkString["rm -rf secdec-files/", name, "/"]];
  Run[MkString[
    "cd secdec-files && env SECDEC_CONTRIB=$PWD/../pySecDec-1.4.5/",
    " PYTHONPATH=$PWD/../pySecDec-1.4.5/pySecDec-1.4.5:$PYTHONPATH",
    " python3 ./", name, ".generate.py"
  ]];
  Run[MkString[
    "env SECDEC_CONTRIB=$PWD/pySecDec-1.4.5/",
    " make -j3 -C secdec-files/", name, "/"
  ]];
  ,
  {integral, masters}];

FromPySecDec[stringresult_String] :=
  stringresult //
  StringReplace["+ O(eps^" ~~ n:(_ ...) ~~ ")" ~~ _ ... ~~ EndOfString :> "+ {{1,0},{1,0}}*O[xxx]^" <> n] //
  StringReplace["+ O(eps^" ~~ n:(_ ...) ~~ ")" ~~ _ ... ~~ EndOfString -> "O[ep]^\n"] //
  StringTrim[#, (" " | "+")...]& //
  StringReplace["(" -> "{"] //
  StringReplace[")" -> "}"] //
  StringReplace[" +/- " -> ", "] //
  StringReplace[" + " -> ", "] //
  StringReplace["eps" -> "xxx" ] //
  StringReplace["e" -> "*10^" ] //
  StringReplace["xxx" -> "ep" ] //
  "X[" <> # <> "]"& //
  ToExpression //
  {
    #[[;;,1]] // Map[Apply[#1 + I * #2&]] // Apply[Plus],
    #[[;;,2]] // Map[Apply[#1 + I * #2&]] // Apply[Plus]
  }&

Prefactor[B[basisid_, idx__]] := Module[{n, nu},
  n = {idx} // Apply[Plus];
  (
    (-1)^n
    Gamma[n-NLOOPS*d/2]
    1/Product[Gamma[nu], {nu, {idx} // DeleteCases[0]}]
    (*(I Pi^(d/2) / (2 Pi)^d)^NLOOPS*)
    (I)^NLOOPS
  )
]

fullamplitudefmt = fullamplitude /.
  d -> 4-2ep //
  Collect[#,
    gs|Nf|Nt|SumQf|SumQt|SumQf2|SumQt2|Tf|Nc|Na|Cf|Ca|d33|d44|Xi,
    Collect[#, _B, Factor]& /* Coef
  ]& //
  Expand

Put[fullamplitudefmt, "fullamplitude.m"];

(* sqrq, mt2 *)
parameters = {sqrq -> 1, mt2 -> 1/5};
valuemap = <||>;
errormap = <||>;
Do[
  name = IntegralName[integral];
  Print["* Integrating ", name];
  stringresult = RunThrough[MkString[
    "cd secdec-files && env PYTHONPATH=$PWD/../pySecDec-1.4.5/pySecDec-1.4.5:$PYTHONPATH",
    " python3 ./", name, ".integrate.py ", N[{sqrq, mt2} /. parameters] // Riffle[#, " "]&
  ], ""];
  {value, error} = stringresult // FromPySecDec;
  prefactor = Prefactor[masters[[1]]] /. d->4-2ep;
  valuemap[integral] = value * prefactor;
  errormap[integral] = error * prefactor;
  ,
  {integral, masters}];

Put[Table[i -> {valuemap[i], errormap[i]}, {i, valuemap//Keys}], "valuemap.m"];

fullvalue = fullamplitude /.
  integral_B :> integral*Prefactor[integral] /.
  d -> 4-2ep /.
  parameters //
  Collect[#, gs|Nf|Nt|SumQf|SumQt|SumQf2|SumQt2|Tf|Nc|Na|Cf|Ca|d33|d44|Xi, Coef]& //
  ReplaceAll[Coef[x_] :> Coef[x/.valuemap, x/.errormap]] //
  Expand

Put[fullvalue, "fullvalue.m"];
