text: # -*- coding: utf-8 -*-
  from sys import exit
  
  from sympy.printing.mathematica import MCodePrinter
  from sympy.printing.precedence import precedence
  from sympy import Indexed, Symbol, zeros
  
  from Definitions import mSymbol
  
  from Logging import loggingCritical
  
  class MathematicaExport():
      def __init__(self, model):
          self._Name = model._Name
          self.string = ""
  
          # BetaFunc definition
          self.betaFactor = model.betaFactor
  
          self.translation = {'GaugeCouplings': 'Gauge Couplings',
                              'Yukawas': 'Yukawa Couplings',
                              'QuarticTerms': 'Quartic Couplings',
                              'TrilinearTerms' : 'Trilinear Couplings',
                              'ScalarMasses': 'Scalar Mass Couplings',
                              'FermionMasses': 'Fermion Mass Couplings',
                              'Vevs': 'Vacuum-expectation Values'}
  
          self.cListNames = {k:v.replace(' ', '') for k,v in self.translation.items()}
          self.cListNames['Vevs'] = 'Vevs'
  
          self.allCouplings = {}
          self.couplingStructure = {}
  
          self.couplingStructure = {mathematica(model.allCouplings[k][1]): v for k,v in model.couplingStructure.items()}
  
          self.yukLikeCouplings = {}
  
          self.conjugatedCouplings = {}
  
          self.inconsistentRGset = (model.NonZeroCouplingRGEs != {} or model.NonZeroDiagRGEs != {})
  
          anomalous = (model.scalarAnomalous != {} or model.fermionAnomalous != {})
  
          self.preamble(model, anomalous)
  
          self.RGEs(model, anomalous)
  
          if anomalous:
              self.anomalous(model)
  
          if model.runSettings['MathematicaSolver'] is True and model.symbolicGen is False:
              self.RGsolver(model)
  
          self.replacements()
  
      def write(self, fileName):
          try:
              self.file = open(fileName, 'w')
          except:
              exit('ERROR while opening the Mathematica output file')
  
          self.file.write(self.string)
          self.file.close()
  
      def preamble(self, model, anomalous):
          self.string += f"""(* ::Package:: *)\n
  (* ::Text:: *)
  (*This file was automatically generated by PyR@TE 3.*)
  (*Model : {model._Name}*)
  (*Author : {model._Author}*)
  (**)"""
          if model.runSettings['MathematicaSolver'] is True:
              self.string += f"""
  (*Section 1 contains the renormalization group equations""" + (" and the anomalous dimensions of the fields." if anomalous else ".") + """*)
  (*Section 2 is a proposition of a minimal working code to solve the RGEs and plot the running couplings.*)
  (*Section 3 should be completed by the user since it contains the initial conditions for the running couplings.*)\n"""
  
      def RGEs(self, model, anomalous):
          sec = '1.' if not anomalous else '1.A)'
          s = f"(* ::Subsection:: *)\n(*{sec} RGEs*)\n\n"
  
          if self.inconsistentRGset:
              s += """(* ::Text:: *)
  (*< WARNING : > The RGE set  is inconsistent. Please refer to the .tex output to solve this.*)\n\n"""
  
          exponent = mathematica(model.betaExponent(Symbol('i')))
          s += f"""(* \\[Beta]-function definition *)
  
  \\[Beta][x_] := {mathematica(1/self.betaFactor)}*Sum[If[With[{{j=i}}, ValueQ[\\[Beta][x, j]]] && (!ValueQ[loops[x]] || i <= loops[x]),  1/(4 Pi)^({exponent})*\[Beta][x,i], 0], {{i,{max(model.loopDic.values())}}}]*Log[10];\n"""
  
          substitutedCouplings = [str(k) for subDic in model.substitutions.values() for k in subDic]
  
          for cType in model.toCalculate:
              if 'Anomalous' in cType:
                  continue
  
              cDic = {}
              for k,v in model.allCouplings.items():
                  if v[0] == cType and k not in substitutedCouplings:
                      # The conjugated couplings are removed, and must be replaced by Conjugate[ ... ]
                      if not k[-2:] == '^*' and not k[-4:] == '^{*}' and not k[-4:] == 'star':
                          cDic[k] = mathematica(v[1])
                      else:
                          candidates = [el for el in model.allCouplings if el in k and el != k]
                          if len(candidates) == 1:
                              self.conjugatedCouplings[k] = candidates[0]
                          else:
                              lengths = [len(el) for el in candidates]
                              i, maxLen = lengths.index(max(lengths)), max(lengths)
                              lengths.remove(maxLen)
  
                              if maxLen not in lengths:
                                  self.conjugatedCouplings[k] = candidates[i]
                              else:
                                  loggingCritical(f"Warning in Mathematica export: could not determine the conjugate quantity of {k} automatically." +
                                                  "\n -> The user will have to modify the output Mathematica file manually.")
  
              s += f"\n\n(* {self.translation[cType]} *)\n\n{self.cListNames[cType]} = "
              s += '{' + ', '.join([c for c in cDic.values()]) + '};\n'
  
              if model.allRGEs[cType] == {}:
                  continue
  
              if cType == 'Vevs':
                  if model.gaugeFixing is None:
                      s += "\n(* For vevs the gauge must be fixed. Let's use for instance the Landau gauge : *)\n"
                      s += "xiGauge = 0;\n"
                  else:
                      s += "\n(* Gauge fixing *)\n"
                      s += f"xiGauge = {mathematica(model.gaugeFixing)};\n"
  
  
              self.allCouplings[cType] = []
              for nLoop, RGEdic in model.couplingRGEs[cType].items():
  
                  RGEdic = {v:RGEdic[k] for k,v in cDic.items() if k in RGEdic}
                  for k, RGE in RGEdic.items():
                      s += '\n\\[Beta][' + k + f', {nLoop+1}] = '
                      betaFunc = mathematica(RGE, k in self.couplingStructure)
  
                      # Yukawa singlet workaround : take the trace of the beta-function if needed
                      if k in self.couplingStructure and self.couplingStructure[k] == (1,1) and 'Dot[' in betaFunc:
                          s += 'Tr[' + betaFunc + '];'
                      else:
                          s += betaFunc + ';'
  
                      if k not in self.allCouplings[cType]:
                          self.allCouplings[cType].append(k)
  
                  s += '\n'
  
          self.string += s
  
      def anomalous(self, model):
  
          s = f"(* ::Subsection:: *)\n(*1.B) Anomalous dimensions*)\n\n"
  
          maxLoops = max(model.loopDic['FermionAnomalous'], model.loopDic['ScalarAnomalous'])
  
          s += f"""(* \\[Gamma] definition *)
  
  \\[Gamma][x_, y_] := Sum[If[With[{{j=i}}, ValueQ[\\[Gamma][x, y, j]]], 1/(4 Pi)^(2i)*\[Gamma][x,y,i], 0], {{i,{maxLoops}}}];\n\n"""
  
          s += "(* Gauge fixing ? *)\n\n"
  
          if model.gaugeFixing is not None:
              s += f"xi = {mathematica(model.gaugeFixing)};\n\n"
          else:
              s += "xi = \\[Xi];\n\n"
  
          if model.fermionAnomalous != {}:
              s += "\n(* Fermions *)\n"
  
              for k,v in model.fermionAnomalous.items():
                  if model.saveSettings['FermionAnomalous'] == 'All':
                      if all([el[k] == 0 for el in model.couplingRGEs['FermionAnomalous'].values()]):
                          continue
  
                  fields = Printer.extractAnomalousFields(k)
                  gamma = lambda n: f'\\[Gamma][{fields}, {n+1}]'
  
                  for n, RGEdic in model.couplingRGEs['FermionAnomalous'].items():
                      RGE = RGEdic[k]
  
                      s += '\n' + gamma(n) + ' = '
                      betaFunc = mathematica(RGE).replace('xiGauge', 'xi')
  
                      # Yukawa singlet workaround : take the trace of the beta-function if needed
                      if k in self.couplingStructure and self.couplingStructure[k] == (1,1) and 'Dot[' in betaFunc:
                          s += 'Tr[' + betaFunc + '];'
                      else:
                          s += betaFunc + ';'
  
                  s += '\n'
              s += '\n'
  
          if model.scalarAnomalous != {}:
              s += "\n(* Scalars *)\n"
  
              for k,v in model.scalarAnomalous.items():
                  if model.saveSettings['ScalarAnomalous'] == 'All':
                      if all([el[k] == 0 for el in model.couplingRGEs['ScalarAnomalous'].values()]):
                          continue
  
                  fields = Printer.extractAnomalousFields(k)
                  gamma = lambda n: f'\\[Gamma][{fields}, {n+1}]'
  
                  for n, RGEdic in model.couplingRGEs['ScalarAnomalous'].items():
                      RGE = RGEdic[k]
  
                      s += '\n' + gamma(n) + ' = '
                      betaFunc = mathematica(RGE).replace('xiGauge', 'xi')
  
                      s += betaFunc + ';'
  
                  s += '\n'
              s += '\n'
          s += '\n'
  
          self.string += s
  
      def RGsolver(self, model):
          #########################
          ## Initial conditions ###
          # This is printed later #
          #########################
  
          sInit = """
  (* ::Subsection:: *)
  (*3. Actual solving*)
  
  
  (* Some useful quantities. Energy and mass are in units of Log10[ GeV ] *)
  
  MZ = 91.1876                 // Log10; (* Z boson mass *)
  MP = 1.22*10^19              // Log10; (* Planck mass *)
  Mp = 1.22*10^19 / Sqrt[8 Pi] // Log10; (* Reduced Planck mass *)
  
  MT = 173.0 // Log10; (* Pole top mass *)
  
  vev = 246.22 // Log10; (* SM Higgs vacuum-expectation value *)
  
  
  (*** Initial conditions ; Energy range ***)
  
  initialScale = MZ;
  range = {0, MP};
  """
          for cType, cList in self.allCouplings.items():
              sInit += f"\n(* {self.translation[cType]} *)\n\n"
  
              for i, c in enumerate(cList):
                  if not c in self.couplingStructure or self.couplingStructure[c] == (1,1):
                      sInit += f"init[{c}] = 0;\n"
                  else:
                      value = mathematica(zeros(*self.couplingStructure[c]))
                      value = value.replace('}, {', '},\n' + ' '*(len(c)+10) + '{')
                      sInit += f"init[{c}] = {value};\n"
                      if i < len(cList)-1:
                          sInit += '\n'
  
                      self.yukLikeCouplings[c] = self.couplingStructure[c]
  
          sInit += """
  
  (* Definition of the running scheme *)
  
  """
  
          for cType in self.allCouplings:
              n = model.loopDic[cType]
              sInit += f"(loops[#] = {n})& /@ {self.cListNames[cType]};\n"
  
          sInit += """
  
  (* Now call the functions defined in section 2 *)
  
  solutions = RGsolve[initialScale, range];
  solutions = DiscardZeroCouplings[range, solutions];
  
  RGplot[range, solutions]"""
  
  
          #######################
          # RG-solving function #
          #######################
  
          s = """(* ::Subsection:: *)
  (*2. Solving and plotting functions*)
  
  (* ::Subsubsection::Closed:: *)
  (*RG-solving functions*)
  
  (* Solving function *)
  
  RGsolve[initialScale_, range_] := Block[{allCouplings, couplingsToFunction, differentialEquations, initialConditions, allEquations, solutions, yukawaLike},
  \t(* Define and solve the system of coupled differential equations *)
  \tallCouplings = Join[""" + ', '.join([self.cListNames[el] for el in self.allCouplings]) + """];
  \tcouplingsToFunction = # -> #[t] & /@ allCouplings;
  
  \tdifferentialEquations = (#'[t] == (\[Beta][#] /. couplingsToFunction) & /@ allCouplings);
  \tinitialConditions = (#[initialScale] == init[#] & /@ allCouplings);
  \tallEquations = Join[differentialEquations, initialConditions];
  
  \tsolutions = NDSolve[allEquations, allCouplings /. couplingsToFunction, {t, range[[1]], range[[2]]}][[1]];"""
  
          # This part is for spliting the individual components of the matrix couplings
          if self.yukLikeCouplings != {}:
              s += """\n
  \t(* Now split matrix couplings into their individual components *)
  \tyukawaLike = {""" + ', '.join(['{'+k+', '+str(v).replace('(', '{').replace(')', '}')+'}' for k,v in self.yukLikeCouplings.items()]) + """};
  \tsolutions = Block[{coupling, function, shape},
  \t\t\t\t\tcoupling = #[[1,0]];
  \t\t\t\t\tfunction = #[[2]];
  \t\t\t\t\tIf[MemberQ[Flatten@yukawaLike, coupling],
  \t\t\t\t\t\tshape = yukawaLike[[Position[yukawaLike, coupling][[1,1]], 2]];
  \t\t\t\t\t\tSequence @@ Flatten @ Table[coupling[i,j] -> FunctionInterpolation[If[NumericQ[t], function[[i,j]]], {t, range[[1]], range[[2]]}][t], {i, shape[[1]]}, {j, shape[[2]]}],
  \t\t\t\t\t\tcoupling -> function]]& /@ solutions;
  """
  
          s += """
  \tReturn[solutions];\n];
  
  (* The following function removes from the list the couplings which seem to identically vanish over the energy range *)
  DiscardZeroCouplings[range_, solutions_] := Block[{N, testPoints},
  	N = 20;
  	testPoints =  Range[range[[1]], range[[2]], (range[[2]] - range[[1]])/N];
  	Return @ Select[solutions, (DeleteDuplicates[#[[2, 0]] /@ testPoints] != {0}) &];
  ];"""
  
          self.string += s
  
          #################
          # Plot function #
          #################
  
          s = """
  (* ::Subsubsection::Closed:: *)
  (*Plot function*)
  
  (* Plotting function handling both real and complex quantities *)
  
  Complexplot[funcs_, {min_, max_}, legends_, plotRange_] := Block[{cplxPos, newFuncs, newLegends},
  	cplxPos = If[Head[Chop[funcs[[#]] /. t -> 1/2(max-min)]] === Complex,#, Nothing]& /@ Range[Length[funcs]];
  	newFuncs = If[MemberQ[cplxPos, #], Sequence @@ {Re@funcs[[#]], Im@funcs[[#]]}, funcs[[#]]]& /@ Range[Length[funcs]];
  	newLegends = If[MemberQ[cplxPos, #], Sequence @@ {Re@legends[[#]], Im@legends[[#]]}, legends[[#]]]& /@ Range[Length[funcs]];
  
  	Return @ Plot[Evaluate @ newFuncs, {t, min, max}, PlotLegends->newLegends, PlotRange->plotRange]
   ]
  
  
  RGplot[range_, solutions_] := Block[{selectType, toPlot},
  \tselectType[type_] := Select[solutions, MemberQ[type, If[Head@#[[1]] =!= Symbol, Head@#[[1]], #[[1]]]]&];"""
  
          for cType in self.allCouplings:
              s += '\n\n\ttoPlot = selectType[' + self.cListNames[cType] + '];\n'
              s += '\tIf[toPlot =!= {},\n'
              s += '\t\tPrint["### ' + self.translation[cType] + ' ###"];\n'
              s += '\t\tPrint @ Complexplot[toPlot[[All, 2]], range, toPlot[[All, 1]], All];\n'
              s += '\t];'
  
          s += "\n];"
          self.string += s
  
          self.string += sInit
  
  
  
      def replacements(self):
          for k,v in self.conjugatedCouplings.items():
              self.string = self.string.replace(k, 'Conjugate[' + v + ']')
  
  
  class Printer(MCodePrinter):
  
      def __init__(self, mat=None):
          MCodePrinter.__init__(self)
          self.mat = mat
  
      def _print_Symbol(self, expr):
          tex = super(MCodePrinter, self)._print_Symbol(expr)
          return tex.replace('{', '').replace('}', '')
  
      def _print_Trace(self, expr):
          return 'Tr[' + mathematica(expr.args[0]) + ']'
  
      def _print_transpose(self, expr):
          return 'Transpose[' + mathematica(expr.args[0]) + ']'
  
      def _print_Pow(self, expr):
          arg, p = expr.args
          if not isinstance(arg, mSymbol) or arg.is_commutative:
              return mathematica(arg)+'^'+mathematica(p)
  
          return '.'.join(p*[mathematica(arg)])
  
      def _print_Mul(self, expr):
          PREC = precedence(expr)
          c, nc = expr.args_cnc()
  
          if c:
              res = super(MCodePrinter, self)._print_Mul(expr.func(*c))
              if nc:
                  res += '*'
          else:
              res = ''
  
          if nc:
              res += '.'.join(self.parenthesize(a, PREC) for a in nc)
  
          return res
  
      def _print_adjoint(self, expr):
          return 'ConjugateTranspose[' + mathematica(expr.args[0]) + ']'
  
      def extractAnomalousFields(fields):
          newFields = []
  
          for f in fields:
              if not isinstance(f, Indexed):
                  newFields.append(f)
              else:
                  newFields.append(Indexed(str(f.base), *[i+1 for i in f.indices]))
  
          return ', '.join([str(el) for el in newFields])
  
  def mathematica(expr, mat=None):
      ret = Printer(mat).doprint(expr)
      ret = ret.replace('_', '').replace('\\', '')
  
      return ret
  