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


# Auxiliary Matrices
# ==================
#
# This file contains useful matrices to be used in rules
#

IsRealMat := mat -> ForAll(mat, r -> ForAll(r, IsRealNumber)); 

ReMat := mat -> MapMat(mat, e -> re(e));

# Permutations
# ------------

#F
#F GenPerm0(generator_func, size)
#F   returns a permutation generated by applying generator_func to
#F   list [0..size-1]
#F
GenPerm0 := function(size, generator_func) 
     Constraint(IsInt(size) and size > 0);
     return PermList(List([0..size-1], generator_func) + 1);
end;

#F
#F GenPerm1(generator_func, size)
#F   returns a permutation generated by applying generator_func to
#F   list [1..size]
#F
GenPerm1 := function(size, generator_func) 
     Constraint(IsInt(size) and size > 0);
     return PermList(List([1..size], generator_func));
end;

#F 
#F LexMap(r,s)
#F   returns list of lexicographical ordering pairs
#F
LexMap := function(r,s)
   local i,j, result;
   result:=[];
   for i in [0..r-1] do
      for j in [0..s-1] do
         Append(result, [[i,j]]);
      od;
   od;
   return result;
end;


#F
#F ToeplitzMat([..elements..])
#F    Returns a Toeplitz matrix given by element list
#F
ToeplitzMat := function(elements)
    local i, j, n, Ls, mat;
    Constraint(IsList(elements));
    Constraint(Length(elements) > 0);
    Constraint(IsOddInt(Length(elements)));
    n := (Length(elements)+1)/2;
    mat := [ ];
    for i in [1..n] do
        Ls := List([1..n], k->0);
	for j in [1..n] do
	    Ls[j] := elements[n+i-j];
	od;        
	Add(mat,Ls);
    od;
    return mat;
end;

# -----------------------------------------------------------------------------
# Matrices used in some breakdown rules
# NOTE: these should be moved to transforms/*
# -----------------------------------------------------------------------------

# perm2( <n> )
#   (x_0, .., x_n-1) --> (x_0, x_n-1, x_1, x_n-2, ..)
# For even <n> it is L(n,n/2) * IJ(n,n/2);
# for odd <n>    OddStride(n,(n+1)/2) * DirectSum(I((n+1)/2), J((n-1)/2))
#
perm2 := n -> When(IsEvenInt(n), M(n,n/2), LIJ(n));

# perm3( <n> )
#   (x_0, .., x_n-1) --> (x_0, x_n/2, x_1, x_n/2+1, ..),
# can be expressed by the previous two:
#   perm3(n) = perm2(n) * (I_n/2 dirsum J(n/2))
# for 2|n it is the stride permutation L^n_n/2
perm3 := n -> When(IsEvenInt(n), L(n, n/2), OS(n, (n+1)/2));

# perm4( <n> )
#   permutes a direct sum of 2x2 blocks
#   into star form, first block outermost etc.
#   for n even it is equal to perm2
perm4 := n -> When(IsEvenInt(n), LIJ(n), Z(n).transpose()*LIJ(n));

# perm5( <n> )
#   reverse the odd indices:
#   (x_0, x_1, .., x_2k) --> (x_0, x_2k-1, x_2, x_2k-3, .., x_1, x_2k)
perm5 := n -> LIJ(n) * L(n,2);

# perm6( <n> )
#   a permutation for rule DCT4_Iterative.
perm6 := SymFunc("perm6", function ( n )
  local k, p, i;

  k := LogInt(n, 2);
  p := PermSPL(perm4(n));
  for i in [1..k - 2] do
    p := 
      p *
      TensorProductPerm( 
        [ 2^(i - 1), n/2^(i - 1) ],
        [ (), 
          DirectSumPerm(
            [ n/2^i, n/2^i ],
            [ PermSPL(perm4(n/2^i)),
              PermSPL(perm4(n/2^i)) ^ PermSPL(J(n/2^i)) ]
          ) ]
      );
  od;

  return Perm(p, n);
end);


# perm7( <n> )
#   a permutation for converting an RDFT into a DCT2.
perm7 := SymFunc("perm7", function ( n )
  local k, L, i;

  k := (n-1)/2;
  L := [ ];
  for i in [0..k] do
    Add(L, (-1)^(i + k + 1)*i + k);
  od;
  for i in [k + 1..n - 1] do
    Add(L, (-1)^(i + k + 1)*(n - i) + k);
  od;

  return Perm(PermListList(L, [0..n-1]), n);
end);


# Block Matrices
# --------------

# blocks1( <n> )
#   returns an spl that represents a direct sum of 
#   DFT_2's in star form. If n is odd then it contains
#   a 1x1 block in the middle.
blocks1 := n -> When(IsEvenInt(n), 
    Tensor(I(Int(n/2)), F(2)) ^ perm4(n),
    DirectSum(I(1), Tensor(I(Int(n/2)), F(2))) ^ perm4(n));

blocks4 := (n, k) ->
    DirectSum(I(1), Tensor(I(Int(n/2)), F(2) * Diag([1, k*E(4)]))) ^ perm4(n);


# Sum Matrices
# ------------

# n x n matrix
# diagonal and upper diagonal 1's
sums1 := SymFunc("sums1", function ( n )
  local M, i;

  M := [ ];
  for i in [1..n] do
    Add(M, [i, i, 1]);
  od;
  for i in [1..n-1] do
    Add(M, [i, i+1, 1]);
  od;
  return Sparse(M);
end);

# n x 2n matrix
# [ identity, opposite identity ]
sums2 := SymFunc("sums2", n -> HStack(I(n), J(n)));


# n+1 x 3*n/2 matrix, e.g.,
#[ [   1,   0,   0,   0,   0,   0,   0,   0,   1,   0,   0,   0 ],
#  [   0,   1,   0,   0,   0,   0,   0,   0,   0,   1,   0,   0 ],
#  [   0,   0,   1,   0,   0,   0,   0,   0,   0,   0,   1,   0 ],
#  [   0,   0,   0,   1,   0,   0,   0,   0,   0,   0,   0,   1 ],
#  [   0,   0,   0,   0,   1,   0,   0,   0,   0,   0,   0,   0 ],
#  [   0,   0,   0,   0,   0,   1,   0,   0,   0,   0,   0,  -1 ],
#  [   0,   0,   0,   0,   0,   0,   1,   0,   0,   0,  -1,   0 ],
#  [   0,   0,   0,   0,   0,   0,   0,   1,   0,  -1,   0,   0 ],
#  [   1,   0,   0,   0,   0,   0,   0,   0,  -1,   0,   0,   0 ] ]
sums3 := SymFunc("sums3", function ( n )
  local M, i;

  M := [ ];
  for i in [1..n] do
    Add(M, [i, i, 1]);
  od;
  for i in [1..n/2] do
    Add(M, [i, n + i, 1]);
    Add(M, [n/2 + 1 + i, 3*n/2 - i + 1, -1]);
  od;
  Add(M, [n+1, 1, 1]);

  return Sparse(M);
end);

# n x 2n matrix
# [ identity, -opposite identity ]
sums4 := n -> HStack(I(n), -J(n));


# 2n x n matrix
# [ opposite lower diagonal 1;
#   identity
# ]

# n x 3n/2 matrix, e.g.,
#[ [  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0 ],
#  [  0,  1,  0,  0,  0,  0,  0,  0,  0,  0,  0,  1 ],
#  [  0,  0,  1,  0,  0,  0,  0,  0,  0,  0,  1,  0 ],
#  [  0,  0,  0,  1,  0,  0,  0,  0,  0,  1,  0,  0 ],
#  [  0,  0,  0,  0,  0,  0,  0,  0,  1,  0,  0,  0 ],
#  [  0,  0,  0,  0,  0,  1,  0,  0,  0,  1,  0,  0 ],
#  [  0,  0,  0,  0,  0,  0,  1,  0,  0,  0,  1,  0 ],
#  [  0,  0,  0,  0,  0,  0,  0,  1,  0,  0,  0,  1 ] ]
sums5 := SymFunc("sums5", function ( n )
  local M, i;

  M := [ ];
  for i in [2..n/2] do
    Add(M, [i, i, 1]);
    Add(M, [i, 3*n/2 - i + 2, 1]);
  od;
  Add(M, [n/2 + 1, n + 1, 1]);
  for i in [1..n/2 - 1] do
    Add(M, [n/2 + 1 + i, n/2 + 1 + i, 1]);
    Add(M, [n/2 + 1 + i, n + i + 1, 1]);
  od;

  return Sparse(M);
end);

# a factorization of a matrix like
#   [ [   1/2,     0,     0,     0,     0,     0,     0,     0 ],
#     [  -1/2,     1,     0,     0,     0,     0,     0,     0 ],
#     [   1/2,    -1,     1,     0,     0,     0,     0,     0 ],
#     [  -1/2,     1,    -1,     1,     0,     0,     0,     0 ],
#     [   1/2,    -1,     1,    -1,     1,     0,     0,     0 ],
#     [  -1/2,     1,    -1,     1,    -1,     1,     0,     0 ],
#     [   1/2,    -1,     1,    -1,     1,    -1,     1,     0 ],
#     [  -1/2,     1,    -1,     1,    -1,     1,    -1,     1 ] ]
# into 1 mult and n - 1 subs. If the first column is divided by 1/2,
# then we obtain the transpose - inverse of sums1.
sums6 := SymFunc("sums6", function ( n )
  local T, M, M1, i;

  # small cases
  if n = 1 then
    return Diag([1/2]);
  fi;
  if n = 2 then
    return Mat([[1, 0], [-1, 1]]) * Diag([1/2, 1]);
  fi;

  T    := List([1..n], i -> 1); 
  T[1] := 1/2; 
  T    := Diag(T);
  M1   := Mat([[1, 0], [-1, 1]]);
  T    := DirectSum(M1, I(n - 2)) * T;
  for i in [3..n - 1] do
    T := DirectSum(I(i - 2), M1, I(n - i)) * T;
  od;
  return DirectSum(I(n - 2), M1) * T;
end);


# sparse1( <n>, <r> )
#   returns an spl for the matrix
#     I_m  -Z_m
#           D_m
#   where m = n/2, Z_m has ones on the lower opposite diagonal,
#   and D_m = cos(r*pi) * diag(1, 2, ..., 2)
#   This matrix occurs in Steidl's rule for DCT2/3.
sparse1 := (n,r) -> CheckedD(IsInt(n), CanBeScalar(r) or IsScalar(r), n mod 2 = 0, 
    let(m := n/2,
	d := DirectSum(Diag(cospi(r)), Scale(2*cospi(r), I(n/2-1))),

	BlockMat(
	    [ [I(m), DirectSum(O(1), -J(m-1))],
	      [O(m),          d            ] ])));

# Copy Matrices
# -------------

# 2n x n+1 matrix, e.g., 
#[ [  1,  0,  0,  0,  0 ],
#  [  0,  1,  0,  0,  0 ],
#  [  0,  0,  1,  0,  0 ],
#  [  0,  0,  0,  1,  0 ],
#  [  0,  0,  0,  0,  1 ],
#  [  0,  0,  0,  1,  0 ],
#  [  0,  0,  1,  0,  0 ],
#  [  0,  1,  0,  0,  0 ] ]

copy1 := SymFunc("copy1", function ( n )
  local M, i;

  M := [ ];
  for i in [1..n + 1] do
    Add(M, [i, i, 1]);
  od;
  for i in [1..n - 1] do
    Add(M, [n + 1 + i, n - i + 1, 1]);
  od;

  return Sparse(M);
end);


# 2n x n matrix, e.g.,
#[ [   1,   0,   0,   0 ],
#  [   0,   1,   0,   0 ],
#  [   0,   0,   1,   0 ],
#  [   0,   0,   0,   1 ],
#  [   0,   0,   0,   0 ],
#  [   0,   0,   0,  -1 ],
#  [   0,   0,  -1,   0 ],
#  [   0,  -1,   0,   0 ] ]

copy2 := SymFunc("copy2", function ( n )
  local M, i;

  M := [ ];
  for i in [1..n] do
    Add(M, [i, i, 1]);
  od;
  for i in [1..n - 1] do
    Add(M, [n + 1 + i, n - i + 1, -1]); 
  od;

  return Sparse(M);
end);


# Monomial Matrices
# -----------------

# monomial1 ( <n> )
#   Let m = n/4. Then monomial1(<n>) returns an spl of type "sparse"
#   that represents
#   I_m+1  0   0    0
#    0     0   J_m  0
#    0     0   0   -J'_m
#    0   I_m-1 0    0
# where the prime indicates that the entry in the last row
# has opposite sign

monomial1 := SymFunc("monomial1", function ( n )
  local m, L, i;

  m := n/4;
  L := [ ];
  for i in [1..m+1] do
    Add(L, [i, i, 1]);
  od;
  for i in [1..m] do
    Add(L, [m+1+i, 3*m+1-i, 1]);
  od;
  for i in [1..m-1] do
    Add(L, [2*m+1+i, 4*m+1-i, -1]);
  od;

  # the exception
  Add(L, [3*m+1, 3*m+1, 1]);
  for i in [1..m-1] do
    Add(L, [3*m+1+i, m+1+i, 1]);
  od;

  return Sparse(L);
end);
