#############################################################################
##
#W  utilsfrgrp.gi              automgrp package                Yevgen Muntyan
##                                                             Dmytro Savchuk
##
#Y  Copyright (C) 2003 - 2018 Yevgen Muntyan, Dmytro Savchuk
##

#############################################################################
##
##  AG_InverseLessThanForLetters(<w1>, <w2>)
##
##  Compares w1 and w2 according to lexicografic ordering given by
##  x1 < x1^-1 < x2 < x2^-1 < ...
##
BindGlobal("AG_InverseLessThanForLetters",
function(w1, w2)
  local i, er1, er2;

  if Length(w1) <> Length(w2) then
    return Length(w1) < Length(w2);
  fi;

  er1 := LetterRepAssocWord(w1);
  er2 := LetterRepAssocWord(w2);
  for i in [1..Length(er1)] do
    if AbsInt(er1[i]) <> AbsInt(er2[i]) then
      return AbsInt(er1[i]) < AbsInt(er2[i]);
    fi;
    if er1[i] <> er2[i] then
      return er1[i] > er2[i];
    fi;
  od;

  return false;
end);

#############################################################################
##
##  AG_ReducedListOfWordsByNielsen(<words_list>)
##  AG_ReducedListOfWordsByNielsenBack(<words_list>)
##  AG_ReducedListOfWordsByNielsen(<words_list>, <string>)
##  AG_ReducedListOfWordsByNielsenBack(<words_list>, <string>)
##
InstallMethod(AG_ReducedListOfWordsByNielsen, [IsAssocWordCollection],
function(words)
  return AG_ReducedListOfWordsByNielsen(words, \<);
end);

InstallMethod(AG_ReducedListOfWordsByNielsenBack, [IsAssocWordCollection],
function(words)
  return AG_ReducedListOfWordsByNielsen(words, \<);
end);

InstallMethod(AG_ReducedListOfWordsByNielsen, [IsAssocWordCollection, IsString],
function(words, string)
  return AG_ReducedListOfWordsByNielsen(words, AG_InverseLessThanForLetters);
end);

InstallMethod(AG_ReducedListOfWordsByNielsenBack, [IsAssocWordCollection, IsString],
function(words, string)
  return AG_ReducedListOfWordsByNielsen(words, AG_InverseLessThanForLetters);
end);

#############################################################################
##
#M  AG_ReducedListOfWordsByNielsen(<words_list>, <lt>)
##
InstallMethod( AG_ReducedListOfWordsByNielsen,
                    [IsAssocWordCollection, IsFunction],
function(words_list, lt)
  local result, transform, did_something, n, i, j, try_again, tmp;

  n := Length(words_list);
  result := ShallowCopy(words_list);
  transform := ShallowCopy(FreeGeneratorsOfFpGroup(FreeGroup(n)));
  did_something := false;
  try_again := true;

  for i in [1..n] do
    if not IsAssocWord(result[i]) then
      Print("error in AG_ReducedListOfWordsByNielsen(IsAssocWordCollection, IsFunction):\n");
      Print("  ", i, "-th element of list is not an associative word\n");
      return fail;
    fi;
  od;

  while try_again do
    try_again := false;

    for i in [1..n] do
    for j in [1..n] do

      if i = j then
        if lt(result[i]^-1, result[i]) then
          result[i] := result[i]^-1;
          transform[i] := transform[i]^-1;
          did_something := true;
          try_again := true;
        fi;
        continue;
      fi;

      if i > j and lt(result[i], result[j]) then
        tmp := result[i];
        result[i] := result[j];
        result[j] := tmp;
        tmp := transform[i];
        transform[i] := transform[j];
        transform[j] := tmp;
        did_something := true;
        try_again := true;
      fi;

      if lt(result[i]*result[j], result[i]) then
        result[i] := result[i]*result[j];
        transform[i] := transform[i]*transform[j];
        did_something := true;
        try_again := true;
      fi;
      if lt(result[i]*result[j], result[j]) then
        result[j] := result[i]*result[j];
        transform[j] := transform[i]*transform[j];
        did_something := true;
        try_again := true;
      fi;

      if lt(result[i]^-1*result[j], result[i]) then
        result[i] := result[i]^-1*result[j];
        transform[i] := transform[i]^-1*transform[j];
        did_something := true;
        try_again := true;
      fi;
      if lt(result[i]^-1*result[j], result[j]) then
        result[j] := result[i]^-1*result[j];
        transform[j] := transform[i]^-1*transform[j];
        did_something := true;
        try_again := true;
      fi;

      if lt(result[i]*result[j]^-1, result[i]) then
        result[i] := result[i]*result[j]^-1;
        transform[i] := transform[i]*transform[j]^-1;
        did_something := true;
        try_again := true;
      fi;
      if lt(result[i]*result[j]^-1, result[j]) then
        result[j] := result[i]*result[j]^-1;
        transform[j] := transform[i]*transform[j]^-1;
        did_something := true;
        try_again := true;
      fi;

      if lt(result[i]^-1*result[j]^-1, result[i]) then
        result[i] := result[i]^-1*result[j]^-1;
        transform[i] := transform[i]^-1*transform[j]^-1;
        did_something := true;
        try_again := true;
      fi;
      if lt(result[i]^-1*result[j]^-1, result[j]) then
        result[j] := result[i]^-1*result[j]^-1;
        transform[j] := transform[i]^-1*transform[j]^-1;
        did_something := true;
        try_again := true;
      fi;
    od;
    od;
  od;

  return [result, transform, did_something];
end);

#############################################################################
##
#M  AG_ReducedListOfWordsByNielsenBack(<words_list>, <lt>)
##
InstallMethod(AG_ReducedListOfWordsByNielsenBack,
                   [IsAssocWordCollection, IsFunction],
function(words_list, lt)
  local result, transform, did_something, n, i, j, try_again, tmp;

  n := Length(words_list);
  result := ShallowCopy(words_list);
  transform := ShallowCopy(FreeGeneratorsOfFpGroup(FreeGroup(n)));
  did_something := false;
  try_again := true;

  for i in [1..n] do
    if not IsAssocWord(result[i]) then
      Print("error in AG_ReducedListOfWordsByNielsenBack(IsAssocWordCollection, IsFunction):\n");
      Print("  ", i, "-th element of list is not an associative word\n");
      return fail;
    fi;
  od;

  while try_again do
    try_again := false;

    for i in [1..n] do
    for j in [1..n] do

      if i = j then
        if lt(result[i]^-1, result[i]) then
          result[i] := result[i]^-1;
          transform[i] := transform[i]^-1;
          did_something := true;
          try_again := true;
        fi;
        continue;
      fi;

      if i > j and lt(result[i], result[j]) then
        tmp := result[i];
        result[i] := result[j];
        result[j] := tmp;
        tmp := transform[i];
        transform[i] := transform[j];
        transform[j] := tmp;
        did_something := true;
        try_again := true;
      fi;

      if lt(result[i]^-1*result[j]^-1, result[j]) then
        result[j] := result[i]^-1*result[j]^-1;
        transform[j] := transform[i]^-1*transform[j]^-1;
        did_something := true;
        try_again := true;
      fi;
      if lt(result[i]^-1*result[j]^-1, result[i]) then
        result[i] := result[i]^-1*result[j]^-1;
        transform[i] := transform[i]^-1*transform[j]^-1;
        did_something := true;
        try_again := true;
      fi;

      if lt(result[i]*result[j]^-1, result[j]) then
        result[j] := result[i]*result[j]^-1;
        transform[j] := transform[i]*transform[j]^-1;
        did_something := true;
        try_again := true;
      fi;
      if lt(result[i]*result[j]^-1, result[i]) then
        result[i] := result[i]*result[j]^-1;
        transform[i] := transform[i]*transform[j]^-1;
        did_something := true;
        try_again := true;
      fi;

      if lt(result[i]^-1*result[j], result[i]) then
        result[i] := result[i]^-1*result[j];
        transform[i] := transform[i]^-1*transform[j];
        did_something := true;
        try_again := true;
      fi;
      if lt(result[i]^-1*result[j], result[j]) then
        result[j] := result[i]^-1*result[j];
        transform[j] := transform[i]^-1*transform[j];
        did_something := true;
        try_again := true;
      fi;

      if lt(result[i]*result[j], result[j]) then
        result[j] := result[i]*result[j];
        transform[j] := transform[i]*transform[j];
        did_something := true;
        try_again := true;
      fi;
      if lt(result[i]*result[j], result[i]) then
        result[i] := result[i]*result[j];
        transform[i] := transform[i]*transform[j];
        did_something := true;
        try_again := true;
      fi;

    od;
    od;

  od;

  return [result, transform, did_something];
end);

#############################################################################
##
##  AG_ComputeMihailovaSystemPairs(<pairs_list>)
##
InstallGlobalFunction(AG_ComputeMihailovaSystemPairs,
function(pairs)
  local result, i, nie, m, n, w, tmp,
        did_smth, npairs, transform,
        generate_full_group, nielsen_mihaylov, nielsen_low, rank,
        number_of_letters;

  if not IsDenseList(pairs) then
    Print("error in AG_ComputeMihailovaSystemPairs:  \n");
    Print("  argument is not an IsDenseList\n");
    return fail;
  fi;
  if not IsList(pairs[1]) then
    Print("error in AG_ComputeMihailovaSystemPairs:  \n");
    Print("  first element of list is not an IsList\n");
    return fail;
  fi;
  if Length(pairs[1]) <> 2 then
    Print("error in AG_ComputeMihailovaSystemPairs:  \n");
    Print("  can work only with pairs\n");
    return fail;
  fi;
  if not IsAssocWord(pairs[1][1]) then
    Print("error in AG_ComputeMihailovaSystemPairs:  \n");
    Print("  <arg>[1][1] is not IsAssocWord\n");
    return fail;
  fi;

  #############################################################################
  ##
  ##  generate_full_group
  ##
  generate_full_group := function(list, rank)
    local nie, i;
    nie := AG_ReducedListOfWordsByNielsen(list)[1];
    if Length(Difference(nie, [One(nie[1])])) <> rank then
      return false;
    fi;
    for i in [1..Length(nie)] do
      if Length(nie[i]) > 1 then return false; fi;
    od;
    return true;
  end;

  #############################################################################
  ##
  ##  rank
  ##
  rank := function(words)
    return Length(Difference(AG_ReducedListOfWordsByNielsen(words)[1],
                             [One(words[1])]));
  end;

  #############################################################################
  ##
  ##  number_of_letters
  ##
  number_of_letters := function(list)
    local letters, i, j;
    letters := [];
    for i in [1..Length(list)] do
      for j in [1..NumberSyllables(list[i])] do
        AddSet(letters, GeneratorSyllable(list[i], j));
      od;
    od;
    return Length(letters);
  end;

  #############################################################################
  ##
  ##  nielsen_mihaylov
  ##
  nielsen_mihaylov := function(words_list, m, n)
    local result, transform, did_something, try_again, nie, i, j, tf, pair,
          good_tf, good_pair, tmp;

    result := StructuralCopy(words_list);
    transform := ShallowCopy(FreeGeneratorsOfFpGroup(FreeGroup(m+n)));
    did_something := false;
    try_again := true;

    while try_again do
      try_again := false;

      nie := nielsen_low(result, m, n, \<);
      if nie[3] then
        did_something := true;
        try_again := true;
        result := nie[1];
        transform := AG_CalculateWords(nie[2], transform);
      fi;

      nie := AG_ReducedListOfWordsByNielsen(result{[m+1..m+n]});
      if nie[3] then
        did_something := true;
        try_again := true;
        result := Concatenation(result{[1..m]}, nie[1]);
        transform := Concatenation( transform{[1..m]},
                                    AG_CalculateWords(nie[2],
                                    transform{[m+1..m+n]}));
      fi;

      if rank(result{[m+1..m+n]}) = n then
        if List(result{[m+1..m+n]}, w -> Length(w)) = List([1..n], i -> 1) then
          ## ok
          try_again := false;
        else
          ##  try to minimize sum of lengths
          good_pair := false;
          for pair in ListX([m+1..m+n], [1..m], function(i,j) return [i,j]; end) do
            good_tf := false;
            for tf in [  [1,1,2,1],[2,1,1,1],[1,-1,2,1],[2,-1,1,1],
                        [1,1,2,-1],[2,1,1,-1],[1,-1,2,-1],[2,-1,1,-1]  ] do
              tmp := StructuralCopy(result);
              tmp[pair[1]] := tmp[pair[tf[1]]]^tf[2] * tmp[pair[tf[3]]]^tf[4];
              if  rank(tmp{[m+1..m+n]}) = n and
                  number_of_letters(tmp{[m+1..m+n]}) =
                    number_of_letters(result{[m+1..m+n]}) and
                  Sum(List(tmp{[m+1..m+n]}, w -> Length(w))) <
                    Sum(List(result{[m+1..m+n]}, w -> Length(w)))
              then
                good_tf := true;
                break;
              fi;
            od;
            if good_tf then
              good_pair := true;
              break;
            fi;
          od;
          if not good_pair then
            ##  give up
            return [result, transform, did_something];
          else
            result[pair[1]] :=  result[pair[tf[1]]]^tf[2] *
                                result[pair[tf[3]]]^tf[4];
            transform[pair[1]] := transform[pair[tf[1]]]^tf[2] *
                                  transform[pair[tf[3]]]^tf[4];
            try_again := true;
            did_something := true;
          fi;
        fi;
      else
        ##  try to make rank bigger
        for i in [1..m] do
          good_tf := false;
          pair := [m+1, i];
          for tf in [  [1,1,2,1],[2,1,1,1],[1,-1,2,1],[2,-1,1,1],
                      [1,1,2,-1],[2,1,1,-1],[1,-1,2,-1],[2,-1,1,-1]  ] do
            tmp := StructuralCopy(result);
            tmp[pair[1]] := tmp[pair[tf[1]]]^tf[2] * tmp[pair[tf[3]]]^tf[4];
            if   rank(tmp{[m+1..m+n]}) > rank(result{[m+1..m+n]}) and
                number_of_letters(tmp{[m+1..m+n]}) >=
                  number_of_letters(result{[m+1..m+n]})
            then
              good_tf := true;
              break;
            fi;
          od;
          if good_tf then
            good_pair := true;
            break;
          fi;
        od;
        if not good_pair then
          ##  give up
          return [result, transform, did_something];
        else
          result[pair[1]] :=  result[pair[tf[1]]]^tf[2] *
                              result[pair[tf[3]]]^tf[4];
          transform[pair[1]] := transform[pair[tf[1]]]^tf[2] *
                                transform[pair[tf[3]]]^tf[4];
          try_again := true;
          did_something := true;
        fi;
      fi;

    od;

    return [result, transform, did_something];
  end;

  #############################################################################
  ##
  ##  nielsen_low
  ##
  nielsen_low := function(words_list, m, n, lt)
    local result, transform, did_something, i, j, try_again, tmp, nie;

    result := ShallowCopy(words_list);
    transform := ShallowCopy(FreeGeneratorsOfFpGroup(FreeGroup(m+n)));
    did_something := false;
    try_again := true;

    while try_again do
      try_again := false;

      nie := AG_ReducedListOfWordsByNielsen(result{[1..m]});
      if nie[3] then
        result := Concatenation(AG_CalculateWords(nie[2], result{[1..m]}),
                                result{[m+1..m+n]});
        transform := Concatenation( AG_CalculateWords(nie[2], transform{[1..m]}),
                                    transform{[m+1..m+n]} );
        did_something := true;
        try_again := true;
      fi;

      for i in [1..m] do
        for j in [m+1..m+n] do
          if lt(result[i]^result[j], result[i]) then
            result[i] := result[i]^result[j];
            transform[i] := transform[i]^transform[j];
            did_something := true;
            try_again := true;
          fi;
          if lt(result[i]^(result[j]^-1), result[i]) then
            result[i] := result[i]^(result[j]^-1);
            transform[i] := transform[i]^(transform[j]^-1);
            did_something := true;
            try_again := true;
          fi;
          if lt((result[i]^-1)^result[j], result[i]) then
            result[i] := (result[i]^-1)^result[j];
            transform[i] := (transform[i]^-1)^transform[j];
            did_something := true;
            try_again := true;
          fi;
          if lt((result[i]^-1)^(result[j]^-1), result[i]) then
            result[i] := (result[i]^-1)^(result[j]^-1);
            transform[i] := (transform[i]^-1)^(transform[j]^-1);
            did_something := true;
            try_again := true;
          fi;
        od;
      od;
    od;

    return [result, transform, did_something];
  end;

  #############################################################################
  ##
  ##  MihailovaSystem body
  ##
  n := Length(FreeGeneratorsOfWholeGroup(Group(pairs[1][1])));
  m := Length(pairs) - n;
  npairs := StructuralCopy(pairs);
  transform := StructuralCopy(FreeGeneratorsOfFpGroup(FreeGroup(n+m)));
  did_smth := false;
  if not generate_full_group(List(pairs, p -> p[1]), n)
      or not generate_full_group(List(pairs, p -> p[2]), n)
  then
    Print("error in AG_ComputeMihailovaSystemPairs:  \n");
    Print("  projections do not generate full free group\n");
    return fail;
  fi;

  ##  if rank equals number of pairs then just make one coordinate nicer
  if m = 0 then
    nie := AG_ReducedListOfWordsByNielsen(List(npairs, p -> p[1]), "r");
    if nie[3] then
      tmp := List(npairs, p -> []);
      for i in [1..m+n] do
        tmp[i][1] := AG_CalculateWord(nie[2][i], List(npairs, p -> p[1]));
        tmp[i][2] := AG_CalculateWord(nie[2][i], List(npairs, p -> p[2]));
      od;
      npairs := StructuralCopy(tmp);
      transform := StructuralCopy(nie[2]);
      did_smth := true;
    fi;
    return [npairs, transform, did_smth];
  fi;

  ##  else try to do as much as possible

  ##  1. Apply Nielsen to first coordinate
  nie := AG_ReducedListOfWordsByNielsen(List(npairs, p -> p[1]), "r");
  if nie[3] then
    tmp := StructuralCopy(npairs);
    for i in [1..m+n] do
      tmp[i][1] := AG_CalculateWord(nie[2][i], List(npairs, p -> p[1]));
      tmp[i][2] := AG_CalculateWord(nie[2][i], List(npairs, p -> p[2]));
    od;
    npairs := StructuralCopy(tmp);
    transform := StructuralCopy(nie[2]);
    did_smth := true;
  fi;

  ##  2. Now apply nielsen_mihaylov to the second coordinate
  nie := nielsen_mihaylov(List(npairs, p -> p[2]), m, n);
  if nie[3] then
    tmp := StructuralCopy(npairs);
    for i in [1..m+n] do
      tmp[i][1] := AG_CalculateWord(nie[2][i], List(npairs, p -> p[1]));
      tmp[i][2] := AG_CalculateWord(nie[2][i], List(npairs, p -> p[2]));
    od;
    npairs := StructuralCopy(tmp);
    tmp := StructuralCopy(transform);
    for i in [1..m+n] do
      tmp[i] := AG_CalculateWord(nie[2][i], transform);
    od;
    transform := StructuralCopy(tmp);
    did_smth := true;
  fi;

  ##  3. Try to get nice generators on first coordinate
  nie := AG_ReducedListOfWordsByNielsenBack(List(npairs{[m+1..m+n]}, p -> p[1]), "r");
  if nie[3] then
    tmp := StructuralCopy(npairs);
    for i in [1..n] do
      tmp[m+i][1] := AG_CalculateWord(nie[2][i], List(npairs{[m+1..m+n]}, p -> p[1]));
      tmp[m+i][2] := AG_CalculateWord(nie[2][i], List(npairs{[m+1..m+n]}, p -> p[2]));
    od;
    npairs := StructuralCopy(tmp);
    tmp := StructuralCopy(transform);
    for i in [1..n] do
      tmp[m+i] := AG_CalculateWord(nie[2][i], transform{[m+1..m+n]});
    od;
    transform := StructuralCopy(tmp);
    did_smth := true;
  fi;

  return [npairs, transform, did_smth];
end);

#############################################################################
##
#M  AG_ReducedByNielsen(<words_list>)
##
InstallMethod(AG_ReducedByNielsen,
              "for [IsList and IsAssocWordCollection]",
              [IsList and IsAssocWordCollection],
function(words)
  if AG_Globals.use_inv_order_in_apply_nielsen then
    return AG_ReducedListOfWordsByNielsen(words, "back")[1];
  else
    return AG_ReducedListOfWordsByNielsen(words)[1];
  fi;
end);

#############################################################################
##
#M  AG_ReducedByNielsen(<autom_list>)
##
InstallMethod(AG_ReducedByNielsen,
              "for [IsList and IsAutomCollection]",
              [IsList and IsAutomCollection],
function(automs)
  local words;
  if IsEmpty(automs) then
    return [];
  fi;
  words := AG_ReducedByNielsen(List(automs, a -> a!.word));
  return List(words, w -> Autom(w, automs[1]));
end);

#E


# Apply a function to a value, given variable names for both function and value
CheckEval := function(fun, val)
	local f, x;
	if IsBoundGlobal(fun) and IsBoundGlobal(val) then
		f := ValueGlobal(fun);
		x := ValueGlobal(val);
		return f(x);
	fi;
end;

bloop := -1;
CheckEval("AbsInt", "bloop");
# 1

# Sum of integer variables
GlobalIntegers := function()
	local s, x;
	s := 0;
	for name in SortedList(NamesGVars()) do
		if IsBoundGlobal(name) then
			x := ValueGlobal(name);
			if IsInt(x) then
				Print(name, " ", x, "\n");
				s := s + x;
			fi;
		fi;
	od;
	return s;
end;


##  Copyright (c) 2018-2021, Carnegie Mellon University
##  See LICENSE for details

# Extension of RulesRC/Sums/Diag to SIMT operators.

RewriteRules(RulesRC, rec(
    RC_SIMTISum := Rule([RC, @(1, SIMTISum)], e -> SIMTISum(@(1).val.simt_dim, @(1).val.var, @(1).val.domain, RC(@(1).val.child(1)))),
     RC_SIMTSUM := Rule([RC, @(1, SIMTSUM)], e -> SIMTSUM(@(1).val.simt_dim, List(@(1).val.children(), RC)))
    ));

RewriteRules(RulesSums, rec(
    SIMTSUMAssoc := ARule( SIMTSUM, [ @(1,SIMTSUM) ],
                            function(e) # Substitute rngs from inner SIMTSUM in simt_dim
                                local pos;

                                pos := Position(e.children(), @(1).val);
                                e.simt_dim := e.simt_dim.update_simt_range(pos, @(1).val.simt_dim);
                                return @(1).val.children();
                            end ),

    PullInRightSIMTSums := ARule( Compose,
            [ @(1, [Prm, Scat, ScatAcc, TCast, PushR, PushLR, Conj, ConjL, ConjR, ConjLR, FormatPrm]),
              @(2, [SIMTISum, SIMTSUM]) ],
              e -> [ CopyFields(@(2).val, rec(
                        _children :=  List(@(2).val._children, c -> @(1).val * c),
                        dimensions := [Rows(@(1).val), Cols(@(2).val)] )) ]),

    PullInLeftSIMTSums := ARule( Compose,
           [ @(1, [SIMTISum, SIMTSUM]),
             @(2, [Prm, Gath, TCast, PushL, PushLR, Conj, ConjL, ConjR, ConjLR, FormatPrm]) ],
             e -> [ CopyFields(@(1).val, rec(
                    _children := List(@(1).val._children, c -> c * @(2).val),
                    dimensions := [Rows(@(1).val), Cols(@(2).val)] )) ]),
    ));

RewriteRules(RulesDiag, rec(
    DiagSIMTISumLeft := ARule( Compose,
           [ @(1, SIMTISum, canReorder), @(2, RightPull) ],
      e -> [ SIMTISum(@1.val.simt_dim, @1.val.var, @1.val.domain, @1.val.child(1) * @2.val).attrs(@(1).val) ]),

    DiagSIMTISumRight := ARule( Compose,
             [ @(1, LeftPull), @(2, SIMTISum, canReorder) ],
        e -> [ SIMTISum(@2.val.simt_dim, @2.val.var, @2.val.domain, @1.val * @2.val.child(1)).attrs(@(2).val) ]),

     DiagSIMTSUMPullInRight := ARule( Compose,
           [ @(1, LeftPull), @(2, SIMTSUM) ],
      e -> [ ApplyFunc(ObjId(@2.val), [@2.val.simt_dim]::List(@2.val.children(), c -> @1.val * c)) ]),

     DiagSIMTSUMPullInLeft  := ARule( Compose,
           [ @(1, SIMTSUM), @(2, RightPull) ],
      e -> [ ApplyFunc(ObjId(@1.val), [@1.val.simt_dim]::List(@1.val.children(), c -> c * @2.val)) ]),

    ));

RewriteRules(RulesDiagStandalone, rec(
    # Gath * SIMTIDirSum
    CommuteGathSIMTIDirSum := ARule( Compose,
            [ [@(1, Gath), [@(3,fTensor), ..., [fId,@(2)]]],
              [@(0, [PushL, PushLR]), @(4, SIMTIDirSum, e -> (@(2).val mod Rows(e.child(1))) = 0)]
            ],
            e -> let(
                    ch := DropLast(@(3).val.children(), 1),
                    ratio := @(2).val / Rows(@(4).val.child(1)),
                    f := When(ratio=1, fTensor(ch), fTensor(Concatenation(ch, [fId(ratio)]))),
                    simt_dim := @(4).val.simt_dim,
                    j := @(4).val.var,
                    jj := Ind(f.domain()),

                    [ ObjId(@(0).val)(SIMTIDirSum(simt_dim, jj,      #SubstVars(@(4).val.child(1), rec((j.id) := f.at(jj)))),
                        Data(j, f.at(jj), SubstTopDown(@(4).val.child(1), @(5,fBase,e->e.params[2]=j),
                                                e -> fCompose(f, fBase(jj)))))),
                        @(1).val ]
                )
    ),

    # SIMTIDirSum * Scat
    CommuteSIMTIDirSumScat := ARule( Compose,
            [ [@(0,[PushR, PushLR]), @(4, SIMTIDirSum)],
              [@(1, Scat), [@(3,fTensor), ..., [fId,@(2).cond(e -> (e mod Cols(@(4).val.child(1))) = 0)]]]
            ],

            e -> let(
                    ch := DropLast(@(3).val.children(), 1),
                    ratio := @(2).val / Cols(@(4).val.child(1)),
                    f := When(ratio=1, fTensor(ch), fTensor(Concatenation(ch, [fId(ratio)]))),
                    simt_dim := @(4).val.simt_dim,
                    j := @(4).val.var,
                    jj := Ind(f.domain()),

                   [ @(1).val,
                     ObjId(@(0).val)(SIMTIDirSum(simt_dim, jj,      #SubstVars(@(4).val.child(1), rec((j.id) := f.at(jj)))),
                           Data(j, f.at(jj), SubstTopDown(@(4).val.child(1), @(5,fBase,e->e.params[2]=j),
                                               e -> fCompose(f, fBase(jj)))))) ]
                )
    ),

));


# Copyright (c) 2018-2021, Carnegie Mellon University
# See LICENSE for details

Class(TInterval, Typ, rec(
     __call__ := (self, size) >>
        WithBases(self, rec(
        size := Checked((IsInt(size) and size >= 0) or IsSymbolic(size), size),
        operations := TypOps)),

    hash := (self, val, size) >> (Sum(val, x->x.t.hash(x.v, size)) mod size)+1,

    equals := (self, o) >> IsBound(o.__bases__) and self.__bases__ = o.__bases__
                           and self.size = o.size,

    rChildren := self >> [self.size],
    rSetChild := rSetChildFields("size"),

    print := self >> When(IsBound(self.size), Print(self.name, "(", self.size, ")"), Print(self.name)),
    isIntervalT := true
));

#F IntervalFunc - base class for interval generating functions
Class(IntervalFunc, Function, Sym, rec(
    isSPL := false,
    range := self >> TInterval(self.size),
    domain := self >> TInterval(self.size),
#    tensor_op := (v,rows,cols,fl,gl) ->
#        Lambda(v, fl.at(idiv(v, cols)) * gl.at(imod(v, cols))),
#
    sums := self >> self,
    rSetChild := meth(self, n, newChild)
        self.params[n] := newChild;
    end,

    print := meth(self, i, is)
        local params;
    params := let(p:=self.params, When(IsList(p), p, [p]));
        Print(self.name, "(",
          DoForAllButLast(params, x -> Print(x, ", ")),
          Last(params), ")");
    end,

    def := arg -> rec(),

    toset := self >> Filtered([0..self.size-1], i ->self.elOf(i).ev() = V(1)),

    fromDef := meth(arg)
        local result, self, params, h, lkup;
    self := arg[1];
    params := arg{[2..Length(arg)]};
    params := self.canonizeParams(params);
        self.checkParams(params);
    params := When(IsList(params), params, [params]);

    h := self.hash;
    if h<>false then
        lkup := h.objLookup(self, params);
        if lkup[1] <> false then return lkup[1]; fi;
    fi;

    result := SPL(WithBases(self,  Inherit(
        rec(params := params),
        ApplyFunc(self.def,params))));

    if h<>false then return h.objAdd(result, lkup[2]);
    else return result;
    fi;
    end,
    __call__ := ~.fromDef
));

#F ivII - interval function
#F
#F ivII(<N>) - diagonal of N 1's
#F ivII(<N>, <to>) - diagonal with 1's from 0..<to>, and 0's after
#F ivII(<N>, <from>, <to>) - diagonal with 1's in <from>..<to>, 0's elsewhere
#F
Class(ivII, IntervalFunc, rec(
    abbrevs := [ (N)         -> [N,0,N],
             (N,to)      -> [N,0,to],
             (N,from,to) -> [N,from,to] ],
    def := (N, from, to) -> Checked(IsPosInt0(N),IsPosInt0(from),IsPosInt0(to),
                                rec(size:=N)),
    elOf := (self, i) >> let(from:=self.params[2], to:=self.params[3], cond(leq(from, i, to-1), 1, 0))

));

Class(ivElOf, FuncClassOper, rec(
    updateDomain  := self >> TVoid,
    updateRange := self >> TInt,
    eval := self >> self.child(1).elOf(self.child(2))
));

Class(ivUnion, FuncClassOper, rec(
    updateDomain := self >> self.child(1).domain(),
    updateRange  := self >> self.child(1).range(),
    elOf := (self, i) >> When(ForAny(self.children(), j->j.elOf(i).ev() = V(1)), V(1), V(0)),
    toset := self >> Filtered([0..self.child(1).size-1], i ->self.elOf(i).ev() = V(1)),
));

Class(ivIntersect, FuncClassOper, rec(
    updateDomain := self >> self.child(1).domain(),
    updateRange  := self >> self.child(1).range(),
    elOf := (self, i) >> When(ForAll(self.children(), j->j.elOf(i).ev() = V(1)), V(1), V(0)),
    toset := self >> Filtered([0..self.child(1).size-1], i ->self.elOf(i).ev() = V(1)),
));

Class(ivInvert, FuncClassOper, rec(
    updateDomain := self >> self.child(1).domain(),
    updateRange  := self >> self.child(1).range(),
    elOf := (self, i) >> When(self.child(1).elOf(i).ev() = V(1), V(0), V(1)),
    toset := self >> Filtered([0..self.child(1).size-1], i ->self.elOf(i).ev() = V(0)),
));

Class(ivTensor, fTensorBase, rec(
    print := FuncClassOper.print,
    updateDomain := self >> TInterval(Product(List(self.children(), i->i.size))),
    updateRange := self >> TInterval(Product(List(self.children(), i->i.size))),
    split_op := self >> let(
        c := self.children(), v := Ind(self.domain().size),
        c1dom := c[1].domain().size, c2dom := c[2].domain().size,
        Checked(Length(c)=2,
            [ Cond(c1dom=1, 0, c2dom=1, v, idiv(v, c2dom)),
            Cond(c1dom=1, v, c2dom=1, 0, imod(v, c2dom)),
            v ])),
    combine_op := (self, split, F, G) >> F.elOf(split[1]) * G.elOf(split[2]),
    elOf := (self, i) >>
        Cond(self.numChildren() > 2, self.rightBinary().elOf(i),
         let(split := self.split_op(),
            Lambda(split[3],
             self.combine_op(split, self.child(1), self.child(2))).at(i))),
    toset := self >> Filtered([0..self.domain().size-1], i ->self.elOf(i).ev() = V(1)),
));


# Copyright (c) 2018-2021, Carnegie Mellon University
# See LICENSE for details

@1 := @(1);
@2 := @(2);
zero   := @.cond(e -> e=0 or  (IsValue(e) and (e.v=0  or e.v=0.0)));
one    := @.cond(e -> e=1 or  (IsValue(e) and (e.v=1  or e.v=1.0)));
negone := @.cond(e -> e=-1 or (IsValue(e) and (e.v=-1 or e.v=-1.0)));

IS := objid -> When(IsList(objid),
    o -> ObjId(o) in objid,
    o -> Same(ObjId(o), objid)
);

IS_A := (o, objid) -> Same(ObjId(o), objid);

CHILDREN_ARE := (objid) -> (o -> ForAll(o.children(), IS(objid)));

NOT_ALL_CHILDREN_ARE := (objid) -> (o -> not ForAll(o.children(), IS(objid)));

NOT_ALL_ARGS_ARE := (objid) -> (o -> not ForAll(o.args, IS(objid)));

countRealAdds := c -> Sum(Collect(c, @(1, [add,sub], x->IsRealT(x.t))), op -> Length(op.args)-1);
countRealMuls := c -> Sum(Collect(c, @(1, mul, x->IsRealT(x.t))), op -> Length(op.args)-1);

countVectMuls := c -> Sum(Collect(c, @(1, mul, x->IsVecT(x.t))), op -> Length(op.args)-1);
countVectAdds := c -> Sum(Collect(c, @(1, [add,sub], x->IsVecT(x.t))), op -> Length(op.args)-1);


# Copyright (c) 2018-2021, Carnegie Mellon University
# See LICENSE for details

Import(paradigms.common, paradigms.vector, paradigms.vector.sigmaspl);

_bDivisorPairs := N -> Concatenation(DivisorPairs(N), [[N, 1]]);

VBfly1 := VBase(F(2), 2);
VBfly3 := VBase(Mat([[1,E(4)],[1,-E(4)]]), 2);

Class(v_lo2, Exp);
Class(v_hi2, Exp);
Class(v_rev2, Exp);
Class(v_revhi2, Exp);
Class(v_neg01, Exp);
Class(v_neg03, Exp);
Class(v_neg23, Exp);
Class(v_neg2, Exp);
Class(v_mul1j, Exp);

VBaseHash := HashTableSPL();

VectorCodegen.VBase := (self, o, y, x, opts) >> let(
    lkup := HashLookup(VBaseHash, o),
    Cond(lkup <> false,
         lkup.code(o, y, x, opts),
         Error("VBase implementation of <o> not found in VBaseHash")));

# NOTE: increase vlen to 4 and RC the matrices
HashAdd(VBaseHash, VBase(L(4,2), 2),
    rec(
        spl := VBase(L(4,2), 2),
        code := (o, y, x, opts) -> chain(
            assign(vvref(y,0,2), v_lo2(vvref(x,0,2), vvref(x,1,2))),
            assign(vvref(y,1,2), v_hi2(vvref(x,0,2), vvref(x,1,2))))
    ));

HashAdd(VBaseHash, VBfly1,
    rec(
        spl := VBfly1,
        code := (o, y, x, opts) -> let(
            hi := TempVar(TVect(x.t.t, 2)),
            lo := TempVar(TVect(x.t.t, 2)),
            chain(
                assign(lo, v_lo2(vvref(x,0,2), vvref(x,0,2))),
                assign(hi, v_hi2(vvref(x,0,2), vvref(x,0,2))),
                assign(hi, v_neg03(hi)),
                assign(vvref(y, 0, 2), hi+lo)))
    ));

HashAdd(VBaseHash, VBfly3,
    rec(
        spl := VBfly3,
        code := (o, y, x, opts) -> let(
            hi := TempVar(TVect(x.t.t, 2)),
            lo := TempVar(TVect(x.t.t, 2)),
            chain(
                assign(lo, v_lo2(vvref(x,0,2), vvref(x,0,2))),
                assign(hi, v_revhi2(vvref(x,0,2), vvref(x,0,2))),
                assign(hi, v_neg03(hi)),
                assign(vvref(y, 0, 2), hi+lo)))
    ));

HashAdd(VBaseHash, VBase(J(2), 2),
    rec(
        spl := VBase(J(2), 2),
        code := (o, y, x, opts) ->
            assign(vvref(y, 0, 2), v_rev2(vvref(x,0,2)))
    ));

HashAdd(VBaseHash, VBase(I(2), 2),
    rec(
        spl := VBase(J(2), 2),
        code := (o, y, x, opts) ->
            assign(vvref(y, 0, 2), vvref(x,0,2))
    ));

# NOTE: Pull out twiddles, for 2-way nothing else needs to be done! (1-way cplx)
RulesFor(PSkewDFT3, rec(
   PSkewDFT3_toSkewDFT := rec(
       isApplicable := P -> IsEvenInt(P[1]),
       allChildren := P -> [[ SkewDFT(P[1]/2, P[2]) ]],
       rule := (P, C) -> let(n := P[1],
           M(n, n/2) *
           L(n,2) *
           Tensor(I(n/2), Mat([[1, E(4)], [1, -E(4)]])) *
           RC(C[1]) *
           L(n,n/2))
   )
));

RulesFor(BSkewDFT3, rec(
   BSkewDFT3_Base2 := rec(
       isApplicable := P -> P[1]=2,
       rule := (P, C) -> let(
       D := Dat1d(TReal, 2),
       re := nth(D, 0), im := nth(D, 1),
       Data(D, fPrecompute(FList(TReal, [cospi(2*P[2]), sinpi(2*P[2])])),
       F(2)*Mat([[1,re],
             [0,E(4)*im]])))
   ),

   BSkewDFT3_Base4 := rec(
       isApplicable := P -> P[1]=4,
       allChildren := P -> let(n:=P[1], a:=P[2], aa:=[a/2, (1-a)/2],
       [[ BSkewDFT3(n/2, aa[1]), BSkewDFT3(n/2, aa[2]), BRDFT3(4,a) ]]),
       rule := (P, C) ->
           K(P[1],P[1]/2) * DirectSum(C[1], C[2]) * C[3]
   ),

   BSkewDFT3_Fact := rec(
       switch := false,
       isApplicable := P -> P[1] = 4,
       allChildren := P -> let(n:=P[1], a:=P[2], aa:=[a/2, (1-a)/2],
       [[ BSkewDFT3(n/2, aa[1]), BSkewDFT3(n/2, aa[2]) ]]),
       rule := (P, C) -> K(P[1],P[1]/2) * DirectSum(C) * bruun2(P[1],P[2])
   ),

   BSkewDFT3_Decomp := rec(
       isApplicable := P -> P[1] > 4,
       allChildren := P -> let(N:=P[1], a:=P[2],
       Map2(DivisorPairs(N/2), (m,n) -> let(i := Ind(n),
           [BSkewDFT3(2*m, fr(2*n,i,2*a)), BRDFT3(2*n,a), InfoNt(i)]))),

       rule := (P, C, Nonterms) -> let(N := P[1], m := Cols(C[1])/2,
       K(N,2*m) * IterDirectSum(Nonterms[3].params[1], C[1]) * Tensor(C[2], I(m)))
   )
));

RulesFor(BSkewDFT4, rec(
   BSkewDFT4_Base2 := rec(
       isApplicable := P -> P[1]=2,
       rule := (P, C) -> let(w:=E(2*Denominator(P[2]))^Numerator(P[2]), cc:=Global.Conjugate,
       Mat([[w, w^3],
        [-cc(w), -cc(w)^3]]))
   ),

   BSkewDFT4_Base4 := rec(
       isApplicable := P -> P[1]=4,
       allChildren := P -> let(n:=P[1], a:=P[2], aa:=[a/2, (1-a)/2],
       [[ BSkewDFT4(n/2, aa[1]), BSkewDFT4(n/2, aa[2]), BRDFT3(4,a) ]]),
       rule := (P, C) ->
           K(P[1],P[1]/2) * DirectSum(C[1], C[2]) * C[3]
   ),

   BSkewDFT4_Decomp := rec(
       isApplicable := P -> P[1] > 4,
       allChildren := P -> let(N:=P[1], a:=P[2],
       Map2(DivisorPairs(N/2), (m,n) -> Concatenation(
           List([0..n-1], i->BSkewDFT4(2*m, fr(2*n,i,2*a).ev())),
           [BRDFT3(2*n,a)]))),
       rule := (P, C) -> let(N := P[1], m := Cols(C[1])/2,
       K(N,2*m) * DirectSum(DropLast(C,1)) * Tensor(Last(C), I(m)))
   )
));

RulesFor(BRDFT1, rec(
   BRDFT1_Base2 := rec(
       isApplicable := P -> P[1]=2,
       rule := (P, C) -> F(2)),

   BRDFT1_Base3 := rec(
       isApplicable := P -> P[1]=3,
       rule := (P, C) -> BRDFT1(P[1], P[2]).terminate()),

   BRDFT1_Decomp := rec(
       # NOTE: add support for odd sizes
       isApplicable := P -> P[1] > 2 and IsEvenInt(P[1]),
       allChildren := P -> let(N:=P[1],
       Map2(_bDivisorPairs(N/2), (m,n) -> let(i := Ind(n-1),
           When(n=1,
           [BRDFT1(m), BRDFT3(m), BRDFT1(2*n)],
           [BRDFT1(m), BRDFT3(m), BRDFT1(2*n), InfoNt(i), BRDFT3(2*m, (1+i)/(2*n))])))),

       rule := (P, C, Nonterms) -> let(N := P[1], m := Cols(C[1]), n := N/m/2,
       RC(Scat(IP(N/2, OS(n, -1)))) * RC(L(N/2, m)) * # RC(K'(N/2,m)) *
       DirectSum(RC(L(m, m/2)) * DirectSum(C[1], C[2]),
                 When(n=1, [], IterDirectSum(Nonterms[4].params[1], C[5]))) *
       Tensor(C[3], I(m)))
   )
));

# URDFT_Decomp:
# MatSPL(Kp(16,4))^-1 * MatSPL(DFT(16))*MatSPL(Tensor(URDFT(8), I(2)))^-
#
NewRulesFor(DFT, rec(
    DFT_Bruun_Decomp := rec(
        switch := false,
        # NOTE: add support for odd sizes
        # only apply if there are no tSPL tags.
        applicable := nt -> nt.params[1] > 2 and IsEvenInt(nt.params[1]) and not nt.hasTags(),

        children := nt -> let(
            N := nt.params[1],
            Map2(_bDivisorPairs(N/2), (m,n) -> let(
                i := Ind(n-1),
                When(n=1,
                    [DFT(m), BSkewDFT3(m), BRDFT1(2*n)],
                    [DFT(m), BSkewDFT3(m), BRDFT1(2*n), InfoNt(i), BSkewDFT3(2*m, (1+i)/(2*n))]
                )
            ))
        ),

        apply := (nt, C, nonterms) -> let(
            N := nt.params[1],
            m := Cols(C[1]),
            n := N/m/2,

            Scat(IP(N, OS(n, -1)))
            * L(N, 2*m)
#            * RC(K'(N/2,m))
            * DirectSum(
                L(2*m, m)
                * DirectSum(C[1], C[2]),
                When(n=1,
                    [],
                    IterDirectSum(Nonterms[4].params[1], C[5])
                )
            )
            * Tensor(C[3], I(m))
        )

#D       isApplicable := (self, P) >> P[1] > 2 and IsEvenInt(P[1]) and PHasNoTags(self.nonTerminal, P),
#D       allChildren := P -> let(N:=P[1],
#D       Map2(_bDivisorPairs(N/2), (m,n) -> let(i := Ind(n-1),
#D           When(n=1,
#D           [DFT(m), BSkewDFT3(m), BRDFT1(2*n)],
#D           [DFT(m), BSkewDFT3(m), BRDFT1(2*n), InfoNt(i), BSkewDFT3(2*m, (1+i)/(2*n))])))),
#D
#D       rule := (P, C, Nonterms) -> let(N := P[1], m := Cols(C[1]), n := N/m/2,
#D       Scat(IP(N, OS(n, -1))) * L(N, 2*m) * # RC(K'(N/2,m)) *
#D       DirectSum(L(2*m, m) * DirectSum(C[1], C[2]),
#D                 When(n=1, [], IterDirectSum(Nonterms[4].params[1], C[5]))) *
#D       Tensor(C[3], I(m)))
    )
));

Class(VWrapRC, VWrapBase, rec(
    __call__ := self >> self,
    wrap := (self, r, t, opts) >> @_Base(RC(t), r),
    twrap := (self, t, opts) >> RC(t)
));

NewRulesFor(DFT, rec(
   DFT_URDFT_Decomp := rec(
       switch := false,

       # NOTE: add support for odd sizes
       applicable := (self, t) >> let(n:=t.params[1],
          n > 4 and IsEvenInt(n) and not IsPrime(n/2)) and t.getTags() = [],

       maxRadix := 128,

       children := (self, t) >> let(N:=t.params[1], maxrad := self.maxRadix, rot := t.params[2],
               Map2(Filtered(DivisorPairs(N/2), x -> x[1] <= maxrad/2),
                   (m,n) -> [ DFT(2*m, rot), DFT(m, rot).addWrap(VWrapRC()),
                              GT(URDFT1(2*n, rot), GTVec, GTVec, [m]).withTags(t.getTags()) ])),

       apply := (t, C, Nonterms) -> let(
           N := t.params[1], m := Cols(C[2]), n := N/m/2, rot := t.params[2],
           j := Ind(n-1),

           Scat(Refl0_u(n, 2*m)) *
       DirectSum(
               BB(C[1]),
               IterDirectSum(j, BB(
                   #condM(2*m, m) *
                   L(2*m, 2) *
                   Tensor(I(m), Mat([[1, E(4)], [1, -E(4)]])) *
                   RC(C[2]) *
                   RC(DirectSum(I(1), Diag(fPrecompute(
                                   fCompose(Twid(N, m, rot, 0, 0, j+1), fAdd(m, m-1, 1)))))) *
                   L(2*m, m))
                   # can replace by jj * b * (rdft(a) x I2)
                   # b = b' * (I x L^4_2)
               )
           ) *
       C[3] #Tensor(C[3], I(m))
       )
   ),

   DFT_Base4_VecCx := rec(
       switch := false,
       requiredFirstTag := AVecRegCx,
       applicable := (self, t) >> t.params[1]=4 and t.firstTag().v = 2,
       apply := (t, C, Nonterms) ->
           VTensor(F(2), 2) *
           VDiag(Tw1(4,2,t.params[2]), 2) *
           VBase(L(4,2), 2) *
           VTensor(F(2), 2)
   ),

   DFT_URDFT_Decomp_VecCx := rec(
       switch := false,
       requiredFirstTag := AVecRegCx,

       # NOTE: add support for odd sizes
       applicable := (self, t) >> let(n:=t.params[1], k:=t.params[2],
           k = 1 and n > 4 and IsEvenInt(n)), # and t.getTags() = []),

       maxRadix := 32,

       children := (self, t) >> let(N:=t.params[1], maxrad := self.maxRadix, tags := t.getTags(),
               Map2(Filtered(DivisorPairs(N/2), x -> x[1] <= maxrad/2),
                   (m,n) -> let(j := Ind(n-1),
                       [ DFT(2*m).withTags(tags),
                         rDFT(2*m, (j+1)/(2*n), 1).withTags(tags),
                         GT(URDFT1(2*n), GTVec, GTVec, [m]).withTags(tags),
                         InfoNt(j) ]))),

       apply := (t, C, Nonterms) -> let(
           N := t.params[1], m := Cols(C[2])/2,
           j := Nonterms[4].params[1],
           v := t.getTags()[1].v,

           condKp(N, 2*m) *
       DirectSum(C[1],
               IDirSum(j, K(2*m, 2) * Tensor(I(m), VBfly3) * C[2])) *
       C[3]
       )
   )
));

RulesFor(PkRDFT1, rec(
    PkRDFT1_Base2 := rec(
    isApplicable := P -> P[1]=2,
    rule := (P, C) -> F(2)
    ),

    PkRDFT1_Bruun_Decomp := CopyFields(BRDFT1_Decomp, rec(
       allChildren := P -> let(N:=P[1],
       Map2(_bDivisorPairs(N/2), (m,n) -> let(i := Ind(n-1),
           When(n=1,
           [PkRDFT1(m), BSkewPRDFT(m), BRDFT1(2*n), InfoNt(i)],
           [PkRDFT1(m), BSkewPRDFT(m), BRDFT1(2*n), InfoNt(i), BSkewPRDFT(2*m, (1+i)/(2*n))])))),
       )),
));

RulesFor(BRDFT3, rec(
   BRDFT3_Base2 := rec(
       isApplicable := P -> P[1]=2,
       rule := (P, C) -> I(2)),

   BRDFT3_Base3 := rec(
       isApplicable := P -> P[1]=3,
       rule := (P, C) -> ApplyFunc(BRDFT3, P).terminate()),

   BRDFT3_Base4 := rec(
       isApplicable := P -> P[1]=4,
       rule := (P, C) -> let(
       a := P[2], D := Dat1d(TReal, 2), q := nth(D, 0), q2 := nth(D, 1),
       Data(D, fPrecompute(FList(TReal, [2*cospi(a), 4*cospi(a)^2 - 1])),
           L(4,2) * VStack(
           F(2)*Mat([[1, 0,-1, 0],
                 [0, 0, 0,-q]]),
           F(2)*Mat([[0, 1, 0,q2],
                 [0, 0, q, 0]]))))
   ),

   BRDFT3_Fact := rec(
       switch := false,
       isApplicable := P -> P[1] > 2,
       allChildren := P -> let(n:=P[1], a:=P[2], [[ BRDFT3(n/2, a/2), BRDFT3(n/2, (1-a)/2) ]]),
       rule := (P, C) -> K(P[1],P[1]/2) * DirectSum(C) * bruun2(P[1],P[2])
   ),

   BRDFT3_Trig := rec(
       switch := false,
       isApplicable := P -> P[1] > 2 and (P[1] mod 2) = 0,
       allChildren := P -> let(n:=P[1], a:=P[2], [[ PolyDTT(SkewDTT(DST3(n/2), 2*a)) ]]),
       rule := (P, C) -> let(n := P[1], nn := n/2-1,
       b := Diag(BHD(n/2, 1, -1)) *
            VStack(
           HStack(DirectSum(I(1), -I(nn-1)),  O(nn,1), -J(nn),   O(nn,1)),
           RowVec(Replicate(nn,0), -1, Replicate(nn+1, 0)),
           HStack(O(nn,1), I(nn),  O(nn,1), J(nn)),
           RowVec(Replicate(nn+1,0), 1, Replicate(nn, 0))),
       M(P[1],P[1]/2) * Tensor(C[1], I(2)) * b)
   ),

   BRDFT3_Decomp := rec(
       isApplicable := P -> P[1] > 4,
       allChildren := P -> let(N:=P[1], a:=P[2], rot := P[3],
       Map2(DivisorPairs(N/2), (m,n) -> let(i := Ind(n),
           [BRDFT3(2*m, fr(2*n,i,2*a), rot), BRDFT3(2*n,a,rot), InfoNt(i)]))),

       rule := (P, C, Nonterms) -> let(N := P[1], m := Cols(C[1])/2,
       RC(K(N/2,m)) * IterDirectSum(Nonterms[3].params[1], C[1]) * Tensor(C[2], I(m)))
   )
));

NewRulesFor(rDFT, rec(
    rDFT_Base4 := rec(
       requiredFirstTag := ANoTag,
       applicable := t -> t.params[1]=4,
       apply := (t, C, Nonterms) -> let(
       a := t.params[2],
           Diag(1,1,1,-1) *
           Tensor(F(2), I(2)) *
           DirectSum(I(2), RCDiag(fPrecompute(FList(TReal, [cospi(a), sinpi(a)])))) *
           L(4,2))
   ),

   rDFT_Base4_Vec := rec(
       requiredFirstTag := AVecRegCx,
       applicable := t -> t.params[1]=4 and t.firstTag().v = 2,
       apply := (t, C, Nonterms) -> let(
       a :=t.params[2], D := Dat1d(TReal, 2),
           VTensor(F(2), 2) *
           DirectSum(VBase(I(2), 2), VBase(J(2), 2)*VDiag(FList(TReal, [1,-1]), 2)) *
           VRCDiag(fPrecompute(VData(FList(TReal, [1, cospi(a), 0, -sinpi(a)]), 2)), 2)
       )  # for some reason its not 1 c 0 s ...
   ),

   rDFT_Decomp := rec(
       applicable := t -> t.params[1] > 4,
       children := t -> let(N := t.params[1], a := t.params[2], tags := t.getTags(),
       Map2(DivisorPairs(N/2), (m,n) -> let(i := Ind(n),
           [ rDFT(2*m, fr(2*n,i,2*a)).withTags(tags),
                 GT(rDFT(2*n, a), GTVec, GTVec, [m]).withTags(tags),
                 InfoNt(i)]))),

       apply := (t, C, Nonterms) -> let(N := t.params[1], m := Cols(C[1])/2,
       RC(K(N/2,m)) * IterDirectSum(Nonterms[3].params[1], C[1]) * C[2]),
   )
));
#s:=SPLRuleTree(ExpandSPL(rDFT(4).withTags([AVecRegCx(SSE_4x32f)]), SpiralDefaults)[2]);

# packed fmt
RulesFor(BSkewPRDFT, rec(
   BSkewPRDFT_Decomp := CopyFields(BRDFT3_Decomp, rec(
       allChildren := P -> let(N:=P[1], a:=P[2],
       Map2(DivisorPairs(N/2), (m,n) -> let(i := Ind(n),
           [BSkewPRDFT(2*m, fr(2*n,i,2*a)), BRDFT3(2*n,a), InfoNt(i) ]))),
   )),

   BSkewPRDFT_Base2 := rec(
       isApplicable := P -> P[1]=2,
       rule := (P, C) -> I(2)
   ),

   BSkewPRDFT_Base4 := rec(
       isApplicable := P -> P[1]=4,
       rule := (P, C) -> let(
       a := P[2], D := Dat1d(TReal, 6), d := Diag(1,-1),
       Data(D, fPrecompute(FList(TReal, [cospi(a), cospi(2*a), cospi(3*a),
                                             sinpi(a), sinpi(2*a), sinpi(3*a) ])),
           L(4,2) * VStack(
           F(2)*Mat([[1,        0,   nth(D,1),          0 ],
                 [0,  nth(D,0),         0,  nth(D, 2) ]]),
         d*F(2)*Mat([[0,        0,   nth(D,4),          0 ],
                 [0,  nth(D,3),         0,  nth(D, 5) ]]))))
   )
));

# algebraic fmt
RulesFor(BSkewRDFT, rec(
   BSkewRDFT_Decomp := rec(
       isApplicable := P -> P[1] > 4,
       allChildren := P -> let(N:=P[1], a:=P[2],
       Map2(DivisorPairs(N/2), (m,n) -> Concatenation(
           List([0..n-1], i->BSkewRDFT(2*m, fr(2*n,i,2*a).ev())),
           [BRDFT3(2*n,a)]))),
       rule := (P, C) -> let(N := P[1], m := Cols(C[1])/2,
       K(N,2*m) * DirectSum(DropLast(C,1)) * Tensor(Last(C), I(m)))
   )
));

RulesFor(DCT4, rec(
   DCT4_BSkew_Decomp := rec(
       switch := false,
       isApplicable := P -> P > 2,
       allChildren := P -> let(N:=P, a:=1/2,
       Map2(DivisorPairs(N), (m,n) -> Concatenation(
           List([0..n-1], i->SkewDTT(DCT4(m), fr(n,i,a).ev())),
           [BRDFT3(2*n,a/2)]))),

       rule := (P, C) -> let(N := P, m := Cols(C[1])/2, n:=Cols(Last(C))/2,
       i := Ind(m),
       K(N,2*m) *
       DirectSum(DropLast(C,1)) *
       K(N, 2*n) *
       ISum(i,
           Scat(H(N,2*n,(2*n)*i,1)) * Last(C) *
           DirectSum(I(n+1), -I(n-1)) * OS(2*n,-1) * Gath(BH(N,2*N-1,2*n,i,2*m))))
   )
));

# right := (m,n) -> let(N:=2*m*n, i:=Ind(m),
#     ISum(i,
#   Scat(H(N,2*n,(2*n)*i,1)) * BRDFT3(2*n,1/4) *
#   DirectSum(I(n), -I(n)) * Gath(BH1(2*N,2*n,i,2*m))));

# RC.toAMat := self >> AMatMat(RCMatCyc(MatSPL(self.child(1))));

# bb:=(n,a)->MatSPL(BRDFT3(n,a));;
# ts:=TensorProductMat;;
# ii:=IdentityMat;;
# ms:=MatSPL;
# rms := x->RCMatCyc(MatSPL(x));

# BRDFT3(4,a) -> 6a+3m
# BRDFT3(4,1/4) -> 6a+2m

# costs:
#         4, 6+2   = 8
# PRDFT3, 8, 22+10 = 32
#        16, 66+30 = 96
#
#           4, 8(1/4) or 9(other)
# BRDFT3, 8, 34
#           16, 104


# Copyright (c) 2018-2021, Carnegie Mellon University
# See LICENSE for details

NEON.setRules(rec(
  binrules := [ rec(
          perm := rec(
              spl := L(8, 2),
              l := 1,
              N := 8,
              n := 2,
              r := 1 ),
          instr := [ rec(
                  instr := vpacklo_neon,
                  p := [  ] ), rec(
                  instr := vpackhi_neon,
                  p := [  ] ) ],
          v := 4,
          vperm := VPerm(L(8, 2), (y, x) -> chain(assign(vref(y, 0, 4), vpacklo_neon(vref(x, 0, 4), vref(x, 4, 4), [  ])), assign(vref(y, 4, 4), vpackhi_neon(vref(x, 0, 4), vref(x, 4, 4), [  ]))), 4, 2) ), rec(
          perm := rec(
              spl := L(8, 4),
              l := 1,
              N := 8,
              n := 4,
              r := 1 ),
          instr := [ rec(
                  instr := vunpacklo_neon,
                  p := [  ] ), rec(
                  instr := vunpackhi_neon,
                  p := [  ] ) ],
          v := 4,
          vperm := VPerm(L(8, 4), (y, x) -> chain(assign(vref(y, 0, 4), vunpacklo_neon(vref(x, 0, 4), vref(x, 4, 4), [  ])), assign(vref(y, 4, 4), vunpackhi_neon(vref(x, 0, 4), vref(x, 4, 4), [  ]))), 4, 2) ), rec(
          perm := rec(
              spl := Tensor(
                  L(4, 2),
                  I(2)
                ),
              l := 1,
              N := 4,
              n := 2,
              r := 2 ),
          instr := [ rec(
                  instr := vunpacklolo2_neon,
                  p := [  ] ), rec(
                  instr := vunpackhihi2_neon,
                  p := [  ] ) ],
          v := 4,
          vperm := VPerm(Tensor(
                L(4, 2),
                I(2)
              ), (y, x) -> chain(assign(vref(y, 0, 4), vunpacklolo2_neon(vref(x, 0, 4), vref(x, 4, 4), [  ])), assign(vref(y, 4, 4), vunpackhihi2_neon(vref(x, 0, 4), vref(x, 4, 4), [  ]))), 4, 2) ), rec(
          perm := rec(
              spl := Tensor(
                  I(2),
                  L(4, 2)
                ),
              l := 2,
              N := 4,
              n := 2,
              r := 1 ),
          instr := [ rec(
                  instr := vpacklo_neon,
                  p := [  ] ), rec(
                  instr := vpackhi_neon,
                  p := [  ] ) ],
          v := 4,
          vperm := VPerm(Tensor(
                L(4, 2),
                I(2)
              ), (y, x) -> chain(assign(vref(y, 0, 4), vunpacklolo2_neon(vref(x, 0, 4), vref(x, 4, 4), [  ])), assign(vref(y, 4, 4), vunpackhihi2_neon(vref(x, 0, 4), vref(x, 4, 4), [  ]))), 4, 2) *
            VPerm(Tensor(
                L(4, 2),
                I(2)
              ) *
              Tensor(
                I(2),
                L(4, 2)
              ), (y, x) -> chain(assign(vref(y, 0, 4), vpacklo_neon(vref(x, 0, 4), vref(x, 4, 4), [  ])), assign(vref(y, 4, 4), vpackhi_neon(vref(x, 0, 4), vref(x, 4, 4), [  ]))), 4, 4.2000000000000002) *
            VTensor(I(2), 4) ), rec(
          perm := rec(
              spl := L(8, 4),
              l := 1,
              N := 8,
              n := 4,
              r := 1 ),
          instr := [ rec(
                  instr := vtransposelo_neon,
                  p := [  ] ), rec(
                  instr := vtransposehi_neon,
                  p := [  ] ) ],
          v := 4,
          vperm := VPerm(Tensor(
                L(4, 2),
                I(2)
              ), (y, x) -> chain(assign(vref(y, 0, 4), vunpacklolo2_neon(vref(x, 0, 4), vref(x, 4, 4), [  ])), assign(vref(y, 4, 4), vunpackhihi2_neon(vref(x, 0, 4), vref(x, 4, 4), [  ]))), 4, 2) *
            VPerm(Tensor(
                L(4, 2),
                I(2)
              ) *
              L(8, 4), (y, x) -> chain(assign(vref(y, 0, 4), vtransposelo_neon(vref(x, 0, 4), vref(x, 4, 4), [  ])), assign(vref(y, 4, 4), vtransposehi_neon(vref(x, 0, 4), vref(x, 4, 4), [  ]))), 4, 4.2000000000000002) *
            VTensor(I(2), 4) ), rec(
          perm := rec(
              spl := Tensor(
                  I(2),
                  L(4, 2)
                ),
              l := 2,
              N := 4,
              n := 2,
              r := 1 ),
          instr := [ rec(
                  instr := vunpacklo_neon,
                  p := [  ] ), rec(
                  instr := vunpackhi_neon,
                  p := [  ] ) ],
          v := 4,
          vperm := VTensor(I(2), 4) *
            VPerm(Tensor(
                I(2),
                L(4, 2)
              ) *
              Tensor(
                L(4, 2),
                I(2)
              ), (y, x) -> chain(assign(vref(y, 0, 4), vunpacklo_neon(vref(x, 0, 4), vref(x, 4, 4), [  ])), assign(vref(y, 4, 4), vunpackhi_neon(vref(x, 0, 4), vref(x, 4, 4), [  ]))), 4, 4.2000000000000002) *
            VPerm(Tensor(
                L(4, 2),
                I(2)
              ), (y, x) -> chain(assign(vref(y, 0, 4), vunpacklolo2_neon(vref(x, 0, 4), vref(x, 4, 4), [  ])), assign(vref(y, 4, 4), vunpackhihi2_neon(vref(x, 0, 4), vref(x, 4, 4), [  ]))), 4, 2) ), rec(
          perm := rec(
              spl := L(8, 2),
              l := 1,
              N := 8,
              n := 2,
              r := 1 ),
          instr := [ rec(
                  instr := vtransposelo_neon,
                  p := [  ] ), rec(
                  instr := vtransposehi_neon,
                  p := [  ] ) ],
          v := 4,
          vperm := VTensor(I(2), 4) *
            VPerm(L(8, 2) *
              Tensor(
                L(4, 2),
                I(2)
              ), (y, x) -> chain(assign(vref(y, 0, 4), vtransposelo_neon(vref(x, 0, 4), vref(x, 4, 4), [  ])), assign(vref(y, 4, 4), vtransposehi_neon(vref(x, 0, 4), vref(x, 4, 4), [  ]))), 4, 4.2000000000000002) *
            VPerm(Tensor(
                L(4, 2),
                I(2)
              ), (y, x) -> chain(assign(vref(y, 0, 4), vunpacklolo2_neon(vref(x, 0, 4), vref(x, 4, 4), [  ])), assign(vref(y, 4, 4), vunpackhihi2_neon(vref(x, 0, 4), vref(x, 4, 4), [  ]))), 4, 2) ) ],
  unrules := [  ],
  x_I_vby2 := [ [ rec(
              instr := vunpacklolo2_neon,
              p := [  ] ), rec(
              instr := vunpacklohi2_neon,
              p := [  ] ) ], [ rec(
              instr := vunpackhilo2_neon,
              p := [  ] ), rec(
              instr := vunpackhihi2_neon,
              p := [  ] ) ] ] ));
NEON_HALF.setRules(rec(
  binrules := [ rec(
          perm := rec(
              spl := L(4, 2),
              l := 1,
              N := 4,
              n := 2,
              r := 1 ),
          instr := [ rec(
                  instr := vunpacklo_half,
                  p := [  ] ), rec(
                  instr := vunpackhi_half,
                  p := [  ] ) ],
          v := 2,
          vperm := VPerm(L(4, 2), (y, x) -> chain(assign(vref(y, 0, 2), vunpacklo_half(vref(x, 0, 2), vref(x, 2, 2), [  ])), assign(vref(y, 2, 2), vunpackhi_half(vref(x, 0, 2), vref(x, 2, 2), [  ]))), 2, 2) ) ],
  unrules := [  ],
  x_I_vby2 := [ [ rec(
              instr := vunpacklo_half,
              p := [  ] ), [  ] ], [ [  ], rec(
              instr := vunpackhi_half,
              p := [  ] ) ] ] ));


# Copyright (c) 2018-2021, Carnegie Mellon University
# See LICENSE for details

# ================================================================
# Sum of Monomials
# NOTE: fix sum_mons, mon_diags, mon_perms
# ================================================================
sum_mons := pat -> pat.target(SUMAcc).cond(e->ForAll(e.children(), IS(Mon)));
mon_vdiags := pat -> List(pat.val.children(), c -> c.vdiag);
mon_sdiags := pat -> List(pat.val.children(), c -> c.sdiag);
mon_perms := pat -> List(pat.val.children(), c -> c.perm);

compat_tensor_mons := (p,v,s) -> let(
    pp := Filtered(p, c->domain(c) > 1),
    vv := Filtered(v, c->domain(c) > 1),
    ss := Filtered(s, c->domain(c) > 1),
    fully_compat_tensor_chains(pp, vv, domain, domain) and
    fully_compat_tensor_chains(vv, ss, domain, domain));

merge_tensor_mons := function(pp, vv, ss, combine) #, fid, gid, hid)
   local i, j, k, np, nv, ns, p, v, s, res;
   np := Length(pp);
   nv := Length(vv);
   ns := Length(ss);
   res := []; i:=1; j:=1; k := 1;

   while i <= np or j <= nv or k <= ns do
       if (i <= np and j <= nv and k <= ns) and
           domain(pp[i]) = domain(vv[j]) and
           domain(vv[j]) = domain(ss[k])
       then
           Add(res, combine(pp[i], vv[j], ss[k]));
           i := i+1;
           j := j+1;
           k := k+1;
       else
           p := fId(1);
           v := II(1);
           s := II(1);

           if i<=np and domain(pp[i]) = 1 then
               p := pp[i]; i:=i+1; fi;
           if j<=nv and domain(vv[j]) = 1 then
               v := vv[j]; j:=j+1; fi;
           if k<=ns and domain(ss[k]) = 1 then
               s := ss[k]; k:=k+1; fi;

           Add(res, combine(p,v,s));
       fi;
   od;
   return res;
end;

#UseRuleSet(RulesSums);

# pulling in sum of monomials into ISum
ARule( Compose,
       [ @(1, ISum), sum_mons(@(2)) ],
  e -> [ ISum(@(1).val.var, @(1).val.domain, @(1).val.child(1) * @(2).val).attrs(@(1).val) ]);

# Gath * SUMAcc(Mon, Mon,...) -> SUMAcc(GMon, GMon, ...)
ARule( Compose,
       [ [Gath, @(1)], sum_mons(@(2)) ],
  e -> [ SUMAcc(
	  let(vdiags := mon_vdiags(@(2)),
	      sdiags := mon_sdiags(@(2)),
	      perms  := mon_perms(@(2)),
	      numdiags := Length(vdiags),

	      List([1..numdiags],
		  i -> GMon(fCompose(perms[i], @(1).val),
		            fCompose(vdiags[i], perms[i], @(1).val),
		            fCompose(sdiags[i], perms[i], @(1).val))))) ]
);

# low-priority suck in of Gath into SUMAcc
ARule( Compose,
       [ @(1, SUMAcc), @(2, Gath) ],
  e -> [ ApplyFunc(ObjId(@(1).val),
	           List(@(1).val.children(), c -> c * @(2).val)) ]);

# Combining GMon and Gath, SMon and Scat, Mon and Prm
ARule( Compose,
       [ @(1, GMon), @(2, Gath) ], # o 1-> 2->
  e -> [ GMon(fCompose(@(2).val.func, @(1).val.perm),
	      @(1).val.vdiag, @(1).val.sdiag) ] );

ARule( Compose,
       [ @(1, Scat), @(2, SMon) ], # o 1-> 2->
  e -> [ GMon(fCompose(@(1).val.func, @(2).val.perm),
	      @(2).val.vdiag, @(2).val.sdiag) ] );

ARule( Compose,
       [ @(1, Mon), @(2, Prm) ], # o 1-> 2->
  e -> [ GMon(fCompose(@(2).val.func, @(1).val.perm),
	      fCompose(@(1).val.vdiag, @(2).val.func.transpose()),
	      fCompose(@(1).val.sdiag, @(2).val.func.transpose())) ]);

#UseRuleSet(RulesFuncSimp);

# ================================================================
# Simplification of Intervals
# ================================================================

full_interval := (int, size) ->
    [int.target(II), size, zero, @.cond(e->e=size.val)];

left_interval := (int, size, endd) ->
    [int.target(II), size, zero, endd];

right_interval := (int, size,start) ->
    [int.target(II), size, start, @.cond(e->e=size.val)];

# full interval o index mapping func @(3) = interval with domain of @(3)
# II(n, 0, n) o f = II(domain(f), 0, domain(f))
ARule( fCompose,
      [ full_interval(@, @(1)), @(2) ],
 e -> [ II(pdomain(@(2))) ]
);

# partial interval o dirsum of perms
# II(n, 0, k) o fDirsum(perm_k, perm_?) -> II(n,0,k)
ARule( fCompose,
      [ left_interval(@(0), @(1), @(2)),
	[fDirsum, @(3).cond(e -> is_perm(e) and range(e)=@(2).val), ...] ],
 e -> [ @(0).val ]
);

# partial interval o dirsum of perms
# II(n, n-k, n) o fDirsum(perm_?, perm_k) -> II(n,n-k,n)
ARule( fCompose,
      [ right_interval(@(0), @(1), @(2)),
	[fDirsum, ..., @(4).cond(e -> is_perm(e) and range(e) = @(1).val-@(2).val)] ],
 e -> [ @(0).val ]
);

# fTensor(IIn, IIk) -> II(n*k)
ARule( fTensor, [full_interval(@(0), @(1)), full_interval(@(0), @(2))],
 e -> [ II(@(1).val * @(2).val) ]
);

# fDirsum(IIn, IIk) -> II(n+k)
ARule( fDirsum, [full_interval(@(0), @(1)), full_interval(@(0), @(2))],
 e -> [ II(@(1).val + @(2).val) ]
);

# Interval Shifting, we support shift left (shift <= start), and right (shift >= end)
# II o Z = shifted II#
ARule( fCompose,
      [ [II, @(1), @(2), @(3)],
	[Z,  @(4), @(5).cond(e -> (e<=@(2).val) or (e >= @(3).val)) ]
      ],
 e -> [	II(@(1).val, (@(2).val-@(5).val) mod @(1).val,
           # this makes 0->size, because size mod size = 0
	   let(last:=(@(3).val-@(5).val) mod @(1).val,
	       When(last=0, @(1).val, last))) ]
);

# full_interval_diag * mat = mat
ARule( Compose,
      [ [Diag, [II, @(1), zero, @(2).cond(e->e=@(1).val)]], @(3) ],
 e -> [ @(3).val ]
);

# mat * full_interval_diag  = mat
ARule( Compose,
      [ @(3), [Diag, [II, @(1), zero, @(2).cond(e->e=@(1).val)]] ],
 e -> [ @(3).val ]
);

# ================================================================
# Sum of Monomials
# ================================================================
# *******************************************************************
Class( RulesGMonTensorPullOut, RuleSet );
# *******************************************************************

# GMon splitting
# GMon(fTensor, fTensor, fTensor) -> Tensor(GMon, GMon, ...)
#
tensor_ch := o -> When(IS_A(o, fTensor), o.children(), [o]);

Rule([GMon,
        @(1),
        @(2),
        @(3).cond(e ->
	    let(p := @(1).val, v := @(2).val, s := e,
	        ForAny([p,v,s], IS(fTensor)) and
	        compat_tensor_mons(tensor_ch(p), tensor_ch(v), tensor_ch(s))))],
 e -> let(
     cp := tensor_ch(@(1).val),
     cv := tensor_ch(@(2).val),
     cs := tensor_ch(@(3).val),
     Tensor( merge_tensor_mons(cp, cv, cs, GMon) ))
);

# ================================================================
# GMon -> SGMon
#  (delta * gath)
# ================================================================
# *******************************************************************
Class( RulesGMonToSGMon, RuleSet );
# *******************************************************************

Rule( [ GMon, @(1), @(2), full_interval(@(3),@(4)) ],
    e -> SGMon(fId(domain(@(1).val)), # scatter
	       @(1).val, # gather
	       @(2).val, # vdiag,
	       @(3).val));

## Diag(II) * Diag(..) * Gath(fId)
Rule( [ GMon, @(1,fId), @(2), @(3,II) ],
 e -> let(left := @(3).val.params[2],
          right:= @(3).val.params[3],
          size := @(3).val.params[1],
	  fadd := fAdd(size, right-left, left),
    SGMon(fadd, fadd,
          fCompose(@(2).val, fadd), II(right-left)))
);

## Diag(II) * Diag(..) * Gath(Z)
Rule( [ GMon,
	[Z, @n, @z],
	@(1),
	[II, @, @l, @(3).cond(r -> (@z.val <= @n.val-r) or
		                   (@z.val >= @n.val-@l.val)) ] ],
e -> let(size  := @n.val,
	 left  := @l.val,
         right := @(3).val,
	 shift := @z.val,
	 scat := fAdd(size, right-left, left),
	 gath := fAdd(size, right-left, (left + shift) mod size),
    SGMon(scat, gath,
	  fCompose(@(1).val, scat), II(domain(gath))))
);

## Diag(II o fBase) * Gath(fBase)
Rule([GMon,
	@(1,fBase),
	@(2),
	[fCompose, @(3,II), same_params(@(4),fBase,@(1)) ] ],

 e -> let(left := @(3).val.params[2], right := @(3).val.params[3],
          j    := @(1).val.params[2],
	  gath := @(1).val,
	  COND(leq(left, j, right-1),
	       SGMon(fId(1), gath, @(2).val, II(1)),
	       O(1, range(gath))))
);

## Diag(II o fBase) * Gath(fId(1))
Rule([skip,
	@(1,fId,e->range(e)=1),
	@(2),
	[fCompose, @(3,II), @(4,fBase) ] ],

 e -> let(left := @(3).val.params[2], right := @(3).val.params[3],
          j    := @(4).val.params[2],
	  COND(leq(left, j, right-1),
	       SGMon(fId(1), fId(1), @(2).val, II(1)),
	       O(1, 1)))
);

# NOTE:Make direct sum rules more general (... + p + ...) o fBase
#
## Diag(leftII o fBase) * Gath((p + ...) o fBase)
Rule([GMon,
        [fCompose,
	    [fDirsum, @(1).cond(e -> is_perm(e)), ...],
	    @(2,fBase) ],
	@(3),
	[fCompose, left_interval(@(10), @(20), @(4).cond(e->range(@(1).val)=e)),
	           same_params(@(5), fBase, @(2))] ],

 e -> let(right := @(4).val,
          size  := @(2).val.params[1],
          j     := @(2).val.params[2],
	  COND(leq(0, j, right-1),
	       SGMon(
		   fId(1),
		   fCompose(fAdd(size, right, 0), @(1).val, fBase(right, j)),
		   @(3).val,
		   II(1)),
	       O(1, size)))
);

## Diag(rightII o fBase) * Gath((... + p) o fBase)
Rule([GMon,
        [fCompose,
	    [fDirsum, ..., @(1).cond(e -> is_perm(e))],
	    @(2,fBase) ],
	@(3),
	[fCompose, right_interval(@(0), @(4), @(5).cond(e->range(@(1).val) = @(4).val - e)),
	           same_params(@(6), fBase, @(2)) ] ],

 e -> let(left := @(4).val,
          size  := @(2).val.params[1],
          j     := @(2).val.params[2],
	  COND(leq(left, j, size-1),
	       SGMon(
		   fId(1),
		   fCompose(fAdd(size, size-left, left), @(1).val, fBase(size-left, j-left)),
		   @(3).val,
		   II(1)),
	       O(1, size)))
);

# ================================================================
# Going back from Tensor(SGMon, ...) into SGMon(fTensor,...)
# ================================================================

Rule([@(0,Tensor), ..., @(1,COND), ...],
    e -> let( ch := e.children(),
	      left := Left(...),
	      right := Right(...),
	      leftch := ch{[1..left]},
	      rightch := ch{[right..Length(ch)]},
	      OP := ObjId(@(0).val),
	      COND(@(1).val.cond,
		  List(@(1).val.children(),
		       c -> OP(Concatenation(Copy(leftch), [c], Copy(rightch)))))));

# *******************************************************************
Class( RulesSGMonTensorPullIn, RuleSet );
# *******************************************************************

Rule( @(1,Tensor,CHILDREN_ARE(SGMon)),
 e -> let( ch:=e.children(),
     SGMon(
	 fTensor(List(ch, c->c.scat)),
	 fTensor(List(ch, c->c.gath)),
	 fTensor(List(ch, c->c.vdiag)),
	 fTensor(List(ch, c->c.sdiag)))));

# ================================================================
# Finalization
# ================================================================
Class( RulesFinalize, RuleSet );

#z:=SUMAcc(Gath(fId(4)), Scat(fAdd(4,3,0)) * Gath(fAdd(4,3,0)));

Elim := (elim_func, s) -> When(elim_func(s), [], s);

ARule(SUMAcc,
    [ [Gath, @R],
      [Compose,
	  [@(1,Scat), [@(2,fAdd), @, @, @k]],
	  [@(3,Gath), @S]] ],

    e -> let(n   := domain(@R.val), N   := range(@R.val),
	     nmk := domain(@S.val), nm  := nmk + @k.val,
	     m := n - nm,
    [ SUM( Elim(e->@k.val=0, Scat(fAdd(n, @k.val, 0)) * Gath(fCompose(@R.val, fAdd(n, @k.val, 0)))),
	   Scat(Copy(@(2).val)) *
	   SUMAcc(
	       Gath(fCompose(@R.val, fAdd(n, nmk, @k.val))),
	       Gath(@S.val)),
	   Elim(e->m=0, Scat(fAdd(n, m, nm)) * Gath(fCompose(@R.val, fAdd(n, m, nm))))
      ) ])
);

ARule(Compose,
    [ [Scat, @W],
      @(1, SUMAcc, e->ForAll(e.children(), c->ObjId(c)=Gath and domain(c.func)>1)) ],

  e -> [ let(n := domain(@W.val),
	     i := Ind(n),
	     fbase := fBase(n, i),
	  ISum(i, n,
	      Scat(fCompose(@W.val, fbase)) *
	      SUMAcc(List(@(1).val.children(), c -> Gath(fCompose(c.func, fbase)))))) ]
);

ARule(Compose,
    [ [Scat, @W], [Diag, @D], [Gath, @R] ],

  e -> [ let(n := domain(@R.val),
	     i := Ind(n),
	     fbase := fBase(n, i),
	  ISum(i, n,
	      Scat(fCompose(@W.val, fbase)) *
	      Blk1(@D.val.lambda().at(i)) *
	      Gath(fCompose(@R.val, fbase)))) ]
);

ARule(Compose, [@(1, O), @(2)], e->[ @(1).val ]);

Rule([Compose, [Scat, @(1)], @(2, O)], e->
    let(n := @(1).val.domain(),
	i := Ind(n),
	fbase := fBase(n,i),
	ISum(i, n,
	     Scat(fCompose(@(1).val, fbase)) *
	     Blk1(0) *
	     Gath(fCompose(fId(Cols(@(2).val)), fbase))))
);


multiply := function(a, b)
    return a*b;
end;


# Copyright (c) 2018-2021, Carnegie Mellon University
# See LICENSE for details

#   count number of vector instructions in a ruletree
#   used to build vector base cases with minimal instruction count

VCost := function(R, O)
   return SPLRuleTree(R).vcost();
end;

Tensor.vcost := self >> let(C := self.children(), When(IsIdentitySPL(C[1]), C[1].params*C[2].vcost(), 1E15));
VTensor.vcost := self >> 0;
Compose.vcost := self >> Sum(List(self.children(), i->i.vcost()));
BlockVPerm.vcost := self >> self.n * self.child(1).vcost();