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

Class(PDTT_Base, TaggedNonTerminal, rec(
    abbrevs := [ n -> Checked(IsPosIntSym(n), [n]) ],
    dims := self >> [self.params[1], self.params[1]],
    isReal := True,
    print := (self,i,is) >> Print(self.__name__, "(", PrintCS(self.params), ")",
	When(self.transposed, Print(".transpose()")),
        When(self.tags<>[], Print(".withTags(", self.tags, ")"))),

    normalizedArithCost := self >> let(n := self.params[1],
	floor(2.5 * n * log(n) / log(2.0)) )
));

#
# Discrete Trigonometric Transforms (DTTs) via PRF
#
Declare(PDST3, PDCT3);

Class(PDST4, PDTT_Base, rec(
    transpose := self >> Copy(self),
    terminate := self >> Mat(DST_IVunscaled(EvalScalar(self.params[1]))),
));

Class(PDCT4, PDTT_Base, rec(
    transpose := self >> Copy(self),
    terminate := self >> Mat(DCT_IVunscaled(EvalScalar(self.params[1]))),
));

Class(PDCT2, PDTT_Base, rec(
    terminate := self >> Mat(DCT_IIunscaled(EvalScalar(self.params[1]))),
    transpose := self >> PDCT3(self.params[1])
));

Class(PDST2, PDTT_Base, rec(
    terminate := self >> Mat(DST_IIunscaled(EvalScalar(self.params[1]))),
    transpose := self >> PDST3(self.params[1])
));

Class(PDCT3, PDTT_Base, rec(
    terminate := self >> Mat(DCT_IIIunscaled(EvalScalar(self.params[1]))),
    transpose := self >> PDCT2(self.params[1])
));

Class(PDST3, PDTT_Base, rec(
    terminate := self >> Mat(DST_IIIunscaled(EvalScalar(self.params[1]))),
    transpose := self >> PDST2(self.params[1])
));

RulesFor(PDST4, rec(
    PDST4_Base2 := DST4_Base, #BaseRule(PDST4, 2),
    PDST4_CT := rec(
	isApplicable     := P -> P[1] > 2, #not IsPrime(P),
	forTransposition := false,
	allChildren  := P -> let(N := P[1], Map2(DivisorPairs(2*N), (m,n) -> Cond(
		IsEvenInt(m) and IsEvenInt(n), [ PRDFT3(m,-1).transpose(), PRDFT3(n) ],
		#IsEvenInt(m) and IsEvenInt(n), [ PDHT3(m).transpose(), PDHT3(n) ],
		IsEvenInt(m) and IsOddInt(n),  [ PRDFT3(m,-1).transpose(), PRDFT3(n), PDST4(m/2) ],
		IsOddInt(m)  and IsEvenInt(n), [ PRDFT3(m,-1).transpose(), PRDFT3(n), PDST4(n/2) ],
		IsOddInt(m)  and IsOddInt(n),  Error("This can't happen")))),

	rule := (P,C) -> let(N := P[1], m := Rows(C[1]), n := Cols(C[2]),
	    nf:=Int(n/2), nc:=Int((n+1)/2), mf:=Int(m/2), mc:=Int((m+1)/2), j:=Ind(nf), i:=Ind(mf),
	    # mult twid by -E(4) for RDFT, by 1/2 for DHT
	    t := ScaledTwid(2*N, mf, 1, 1/2, 1/2, j, -E(4)),
	    T := When(IsOddInt(m), Diag(diagDirsum(t,fConst(1,1))), Diag(t)),
	    Et := Tensor(I(mf),Mat([[0,-1],[-1,0]])),
	    Ett := Tensor(I(mf),Mat([[0,-1],[1,0]])),
	    SUM(
		ISum(j,
		    Scat(BH(N, 2*N-1, m, j, n)) * C[1] * #Ett *
		    RC(T) *
		    #Et * # DHT only
		    RC(Gath(H(mc*nc, mc, j, nc)))
		),
		When(IsEvenInt(n), [], Scat(H(N,m/2,nf,n))*C[3]*Gath(H(2*mc*nc, m/2, 2*nf, 2*nc)))
	    ) *
	    SUM(
		ISum(i, RC(Scat(H(mc*nc, nc, i*nc, 1))) * C[2] * Gath(BH(N, 2*N-1, n, i, m))),
                # output is imaginary, but we output into real slot, because subseq transform
		# is jIR2, which we implement as IR2 * (-j)
		When(IsEvenInt(m), [], Scat(H(2*mc*nc, nc, 2*mf*nc, 2))*C[3]*Gath(H(N,n/2,mf,m)))
	    ))),

    PDST4_CT_SPL := rec(
	isApplicable     := P -> P[1] > 2 and ForAny(DivisorPairs(2*P[1]), d->IsEvenInt(d[1]) and IsEvenInt(d[2])),
	forTransposition := false,

	allChildren  := P -> let(N := P[1],
	    Map2(Filtered(DivisorPairs(2*N), d -> IsEvenInt(d[1]) and IsEvenInt(d[2])),
		(m,n) -> [ PRDFT3(m,-1).transpose(), PRDFT3(n) ])),

	rule := (P,C) -> let(N := P[1], m := Rows(C[1]), n := Cols(C[2]),
            TT := Diag(fPrecompute(diagMul(fConst(N/2, -E(4)),
                        fCompose(dOmega(8 * N, 1),
                            diagTensor(dLin(N/m, 2, 1, TInt), dLin(m/2, 2, 1, TInt)))))),

	    Prm(Refl(N, 2*N-1, N, L(m*n, m))) *
	    Tensor(I(n/2), C[1]) *
            RC(TT) *
            RC(L(m*n/4, n/2)) *
	    Tensor(I(m/2), C[2]) *
	    Prm(Refl(N, 2*N-1, N, L(m*n, m)))
	)
    )
));

RulesFor(PDCT4, rec(
    PDCT4_Base2 := rec(isApplicable:=P->P[1]=2, rule:=(P,C)->DCT4(2).terminate()),
    PDCT4_Base4 := rec(isApplicable:=P->P[1]=4,
	               allChildren := P -> [[ DCT4(4) ]],
		       rule:=(P,C)->C[1]),

    PDCT4_CT := rec(
	isApplicable     := P -> P[1] > 2, #not IsPrime(P),
	forTransposition := false,
	allChildren  := P -> let(N := P[1], Map2(DivisorPairs(2*N), (m,n) -> Cond(
		IsEvenInt(m) and IsEvenInt(n), [ PRDFT3(m,-1).transpose(), PRDFT3(n) ],
		#IsEvenInt(m) and IsEvenInt(n), [ PDHT3(m).transpose(), PDHT3(n) ],
		IsEvenInt(m) and IsOddInt(n),  [ PRDFT3(m,-1).transpose(), PRDFT3(n), PDCT4(m/2) ],
		IsOddInt(m)  and IsEvenInt(n), [ PRDFT3(m,-1).transpose(), PRDFT3(n), PDCT4(n/2) ],
		IsOddInt(m)  and IsOddInt(n),  Error("This can't happen")))),

	rule := (P,C) -> let(N := P[1], m := Rows(C[1]), n := Cols(C[2]),
	    nf:=Int(n/2), nc:=Int((n+1)/2), mf:=Int(m/2), mc:=Int((m+1)/2), j:=Ind(nf), i:=Ind(mf),
	    t := fPrecompute(Twid(2*N, mf, 1, 1/2, 1/2, j)),
	    T := When(IsOddInt(m), Diag(diagDirsum(t,fConst(1,1))), Diag(t)),
	    Et := Tensor(I(mf),Mat([[0,1],[1,0]])),
	    SUM(
		ISum(j,
		    Scat(BH(N, 2*N-1, m, j, n)) * Diag(BHN(m)) * C[1] *
		    # mult twid by 1/2 for DHT
		    RC(T) *
		    #Et * # DHT only
		    RC(Gath(H(mc*nc, mc, j, nc)))
		),
		When(IsEvenInt(n), [], Scat(H(N,m/2,nf,n))*C[3]*Gath(H(2*mc*nc, m/2, 2*nf, 2*nc)))
	    ) *
	    SUM(
		ISum(i, RC(Scat(H(mc*nc, nc, i*nc, 1))) * C[2] * Diag(BHN(n)) * Gath(BH(N, 2*N-1, n, i, m))),
		When(IsEvenInt(m), [], Scat(H(2*mc*nc, nc, 2*mf*nc, 2))*C[3]*Gath(H(N,n/2,mf,m)))
	    ))),
));

NewRulesFor(PDCT4, rec(
    PDCT4_CT_SPL := rec(
        libApplicable := t -> eq(imod(t.params[1], 2), 0),
	applicable := t -> IsSymbolic(t.params[1]) or (t.params[1] > 2 and (t.params[1] mod 4) = 0),
        extraLeftTags := [],
	forTransposition := false,

        freedoms := t -> [ divisorsIntNonTriv(t.params[1]/2) ],

	child  := (self, t, fr) >> let(N := t.params[1], f := fr[1],
	    [ spiral.sym.ASP(-1).RDFT3(2*f).withTags(self.extraLeftTags).transpose(),
	      spiral.sym.ASP.RDFT3(div(N,f)) ]),

	apply := (t,C,Nonterms) -> let(N := t.params[1], m := Rows(C[1]), n := Cols(C[2]),
            mh := Rows(C[1])/2, nh := Cols(C[2])/2,
            D := RC(Diag(fPrecompute(fCompose(dOmega(8 * N, 1),
                            diagTensor(dLin(N/m, 2, 1, TInt), dLin(m/2, 2, 1, TInt)))))),

            Grp(Scat(Refl1(nh, m)) *
                Tensor(I(nh), Diag(BHN(m)) * C[1]) *
                D *
                Tensor(I(nh), Tr(2, mh))) *
            #RC(Tr(mh, nh)) *
            Grp(Tr(mh, n) *
                Tensor(I(mh), C[2] * Diag(BHN(n))) *
                Gath(Refl1(mh, n)))
	)
    )
));

RulesFor(PDCT4, rec(
    # Variant better suited for vectorization
    PDCT4_CT_SPL_Vec := rec(
	isApplicable     := P -> P[1] > 2 and (P[1] mod 2) = 0,
	forTransposition := false,

	allChildren  := P -> let(N := P[1],
	    Map2(Filtered(DivisorPairs(2*N), d -> IsEvenInt(d[1]) and IsEvenInt(d[2])),
		(m,n) -> [ PRDFT3(m,-1).transpose(), PRDFT3(n) ])),

	rule := (P,C) -> let(N := P[1], m := Rows(C[1]), n := Cols(C[2]),
            TT := Diag(fPrecompute(fCompose(dOmega(8 * N, 1),
                        diagTensor(dLin(N/m, 2, 1, TInt), dLin(m/2, 2, 1, TInt))))),

            Prm(condIJ(N, n/2)) *
            Prm(L(N, m)) *
	    Tensor(I(n/2),
                   condM(m,m/2) * Diag(BHN(m)) * C[1]) *
            RC(TT) *
            Prm(L(N, n/2)) *
	    Tensor(I(m/2),
                   L(n, 2) * C[2] * Diag(BHN(n)) * condK(n, 2)) *
            Prm(L(N, m/2)) *
            Prm(condIJ(N, m/2))
	)
    )

));

testPDCT4 := function(n)
    local opts, r, s;
    opts := CopyFields(SpiralDefaults, rec(
            breakdownRules := rec(
                PDCT4 := [PDCT4_Base2, PDCT4_CT_SPL_Vec],
                PRDFT1 := [PRDFT1_Base1, PRDFT1_Base2, PRDFT1_CT],
                PRDFT3 := [PRDFT3_Base1, PRDFT3_Base2, PRDFT3_CT],
                DFT := [DFT_Base, DFT_CT]
            )));
    r := RandomRuleTree(PDCT4(n), opts);
    s := SumsRuleTree(r, opts);
    return [opts,r,s];
end;

NewRulesFor(PDCT2, rec(
    PDCT2_Base2 := rec(
	forTransposition := true,
        applicable := t -> t.params[1] = 2,
        apply := (t, C, Nonterms) -> DCT2(2).terminate()
    ),

    PDCT2_Base4 := rec(
	forTransposition := true,
        applicable := t -> t.params[1] = 4,
        apply := (t, C, Nonterms) ->
            LIJ(4) *
            DirectSum(
                Diag(FList(TReal, [ 1, 0.70710678118654757 ])) * F(2),
                Rot(cospi(13/8), sinpi(13/8)) * J(2)) *
            (Tensor(I(2), F(2))) ^ LIJ(4)
    ),

    PDCT2_CT_SPL := rec(
	forTransposition := true,
        libApplicable := t -> eq(imod(t.params[1], 2), 0),
        extraLeftTags := [],

	applicable := t -> let(N:=t.params[1],
            IsSymbolic(N) or
            (N > 2 and IsEvenInt(N) and (N mod 4) = 0)),

	freedoms := t -> [ divisorsIntNonTriv(t.params[1]/2) ],

        # N/2 = k*m
        child := (self, t, fr) >> let(N:=t.params[1], k:=fr[1], m:=N/2/k,
            [ PDCT2(2*k).withTags(self.extraLeftTags),
              spiral.sym.ASP(-1).RDFT3(2*k).transpose().withTags(self.extraLeftTags),
              spiral.sym.ASP(+1).URDFT(2*m) ]),

	apply := (t,C,Nonterms) -> let(N:=t.params[1], k:=Rows(C[1])/2, m:=Cols(C[3])/2,
            D := RC(Diag(fPrecompute(fCompose(dOmega(4 * N, 1),
                        diagTensor(dLin(m-1, 1, 1, TInt), dLin(k, 2, 1, TInt)))))),
            Grp(
                Scat(Refl0_u(m, 2*k)) *
                DirectSum(
                    C[1] * KK(k, 2),
                    Tensor(I(m-1), Diag(BHN(2*k)) * C[2]) * D
                )) *
            Grp(
                RC(Tr(k, m)) *
                Tensor(I(k), C[3]) *
                Gath(Refl1(k, 2*m))
            )
	)
    )
));

NewRulesFor(PDST2, rec(
    PDST2_Base2 := rec(
        applicable := t -> t.params[1] = 2,
        apply := (t, C, Nonterms) -> DST2(2).terminate()
    ),

    PDST2_CT_SPL := rec(
	forTransposition := true,
	applicable := t -> let(N:=t.params[1], N > 2 and IsEvenInt(N) and (N mod 4) = 0),
	freedoms := t -> [ DivisorsIntNonTriv(t.params[1]/2) ],

        # N/2 = k*m
        child := (t, fr) -> let(N:=t.params, k:=fr[1], m:=N/2/k,
                [ DST2(2*k), PRDFT3(2*k,-1).transpose(), URDFT(2*m)  ]),

	apply := (t,C,Nonterms) -> let(N:=t.params[1], k:=Rows(C[1])/2, m:=Cols(C[3])/2,
            TT := Diag(fPrecompute(fCompose(dOmega(4 * N, 1),
                        # ie multiply all twiddles by -E(4)
                        diagAdd(diagDirsum(fConst(k, 0), fConst(k*(m-1), 3*N)),
                                diagTensor(dLin(m, 1, 0, TInt), dLin(k, 2, 1, TInt)))))),

            J(N) *
            Kp(N, 2*k) *
	    DirectSum(
                J(2*k) * C[1] * Diag(BHN(2*k)) * K(2*k, 2),
                Tensor(I(m-1), J(2*k) * M(2*k, k) * C[2])
            ) *
            RC(TT) *
            RC(L(N/2, m)) *
	    Tensor(I(k), C[3] * Diag(BHN(2*m))) *
	    Prm(Refl(N, 2*N-1, N, L(2*N, 2*k)))
	)
    )

));

RulesFor(DCT5, rec(
  DCT5_Rader := rec(
	forTransposition := false,
	isApplicable     := P -> P > 2 and IsPrime(2*P-1),
	allChildren      := P -> [[ IPRDFT(P-1,-1), PRDFT(P-1,-1) ]],

	diag := N -> Sublist(
	    DFT_Rader.raderDiag(2*N-1,1,PrimitiveRootMod(2*N-1)),
	    [1..Int((N-1)/2)]*2),

	# 3rd col with 0's is for PRDFT, not necessary for (non-packed) RDFT
	raderMid := (self, N) >> let(Fsize := 2*N-1,
	             DirectSum(Mat([[1, 1, 0], [1, -1/(Fsize-1), 0], [0,0,0]]),
			       RC(Diag(FData(self.diag(N)))))),

	# NOTE, special case for even P-1
	rule := (self,P,C) >> let(N := P,
	    #RealRR(N).transpose() *

	    Gath(H(2*P-1, P, 0, 1)) *
	    RR(2*P-1).transpose() *
	    DirectSum(I(1), VStack(I(P-1), I(P-1))) *

	    DirectSum(I(1), C[1]) *
#		C[1]*Diag(1,1, Replicate(2*Int((P-2)/2), 2),
#		               When(IsEvenInt(P-1), [1,1],[]))) *
	    self.raderMid(N) *
	    DirectSum(I(1), C[2]) *

            # replace by RealRR
	    Gath(H(2*P-1, P, 0, 1)) *
	    RR(2*P-1) * DirectSum(I(1), VStack(I(P-1), J(P-1)))
    ))
));

#q:=ExpandSPL(DCT5(6))[1];
#s:=SPLRuleTree(q);
#me := MatSPL(s);
#them := remat(MatSPL(q.node));

#RDFT-11,  DCT5(6) 15(6/9 rots) + 19(12/7 irdft) + 17(12/5 rdft) = 30/21 = 51
#           DST5(5)
#           10(blk)
# fftw 60/50

# RDFT-13
#  DCT5(7)  18(8/10=4/8+0/1+2/1 rots) + 18(14/4 irdft) + 18(14/4 rdft) = 36/18=54
#  DST5(6)
#  12(blk)
#  fftw 76/34=110

# DFT-13
# fftw 176+68=244
# us = 54*2 (dct5) + 24(blk)*2 + dst5


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

# input: imaginary, output: co1 (?)
_j := N -> Tensor(I(approx.CeilingRat(N/2)),J(2)*Diag(-1,1));

RealRR_Out := (N, root) -> let(
    rrind := RR(N, 1, root).tolist(),
    outind := List(rrind{[1..(N+1)/2]}, x->When(x.v >= (N+1)/2, N-x.v, x.v)),
    #same as Refl((N+1)/2, N, (N+1)/2, RR(N,1,root)).tolist(),

    # Reflective RR
    Scat(fTensor(FList((N+1)/2, outind), fId(2))) *
    # Conjugate complex elements that will be reflected
    Diag(ConcatList([0..(N-1)/2], x->When(rrind[x+1].v >= (N+1)/2, [1,-1], [1,1])))
);

RulesFor(PRDFT, rec(

   # PRDFT_PD : Projection of complex DFT partial diagonalization rule to real DFT
   # See also : PRDFT_Rader
   #
   PRDFT_PD := rec(
	forTransposition := true,
	maxSize          := 13,
	isApplicable     := (self, P) >> P[1] > 2 and P[1] <= self.maxSize and IsPrime(P[1]),

	rule := (self,P,C) >> let(N:=P[1], n:=N-1, k:=P[2], root:=PrimitiveRootMod(N),
	    BB(RealRR_Out(N, root) *
	    DirectSum(Mat([[1],[0]]), L(n, n/2)) *
	    Mat(MatSPL(DFT_PD.core(N, k, root, true) * DFT_PD.A(N))) *
	    DirectSum(I(1), Tensor(F(2), I(n/2))) *
	    DirectSum(I(1), OS(n, -1)) *
	    Gath(RR(N, 1, root)))
	)
   ),

   # PRDFT_Rader : Projection of complex DFT Rader rule to real DFT
   #
   #   Note: PRDFT of types 2--4 of odd size (which includes all primes > 2)
   #         can be converted without arithmetic cost to PRDFT1,
   #         which means that this rule enables implementation of a prime
   #         size PRDFT of any type.
   PRDFT_Rader := rec(
	isApplicable := (self, P) >> P[1] > 3 and IsPrime(P[1]),

	diag := (N, k, root) -> let(fulldiag := DFT_Rader.raderDiag(N,k,root),
	    last := Int((N-1)/2),
	    Concatenation(2*fulldiag{[1..last-1]}, [fulldiag[last]])),

	# 3rd col with 0's is for padding used by PRDFT
	raderMid := (self, N, k, root) >> let(n:=N-1,
	    DirectSum(Mat([[1, 1, 0], [1, -1/n, 0], [0,0,0]]),
		      RCDiag(RCData(FData(self.diag(N, k, root)))))),

	allChildren := P -> let(n:=P[1]-1,
	    [[ PRDFT(n/2).transpose(), PRDFT3(n/2).transpose(), PRDFT(n, -1) ]]),

	rule := (self,P,C) >> let(N:=P[1], n:=N-1, k:=P[2], root:=PrimitiveRootMod(N),
	    RealRR_Out(N, root) *
	    DirectSum(Mat([[1],[0]]), ##
		      L(n, n/2) *
		      DirectSum(C[1], C[2] * _j(n/2)) *
		      RC(L_or_OS(n/2+1, 2))) *
	    self.raderMid(N, k, root) *
	    DirectSum(I(1), C[3]) *
	    Gath(RR(N, 1, root))
	)
   )
));

RulesFor(IPRDFT, rec(

   # IPRDFT_PD : Projection of complex DFT partial diagonalization rule to inverse real DFT
   # See also : IPRDFT_Rader
   #
   IPRDFT_PD := rec(
	forTransposition := true,
	maxSize          := 13,
	isApplicable     := (self, P) >> P[1] > 2 and P[1] <= self.maxSize and IsPrime(P[1]),
	rule := (self,P,C) >> let(N:=P[1], n:=N-1, k:=P[2], root:=PrimitiveRootMod(N),
	    BB(Scat(RR(N, 1, root)) *
	    DirectSum(I(1), OS(n, -1)) *
	    DirectSum(I(1), Tensor(F(2), I(n/2))) *
	    Mat(MatSPL(TransposedSPL(DFT_PD.core(N, -k, root, true) * DFT_PD.A(N)) * DirectSum(I(1), 2*I(n)))) *
	    DirectSum(Mat([[1,0]]), L(n, 2)) *
	    RealRR_Out(N, root).transpose())
	)
   ),

   # IPRDFT_Rader : Projection of complex DFT Rader rule to inverse real DFT
   #
   #   Note: IPRDFT of types 2--4 of odd size (which includes all primes > 2)
   #         can be converted without arithmetic cost to IPRDFT1,
   #         which means that this rule enables implementation of a prime
   #         size IPRDFT of any type.
   IPRDFT_Rader := CopyFields(PRDFT_Rader, rec(
	# 3rd col with 0's is for padding, we also scale everything but 1st elt by 2,
	# as IPRDFT requires.
	raderMid := (self, N, k, root) >> let(n:=N-1,
	    DirectSum(Mat([[1, 2, 0], [1, -2/n, 0], [0,0,0]]),
		      RCDiag(RCData(FConj(FData(2*self.diag(N, -k, root))))))),

	allChildren := P -> let(n:=P[1]-1,
	    [[ PRDFT(n/2), PRDFT3(n/2), PRDFT(n,-1).transpose() ]]),

	rule := (self,P,C) >> let(N:=P[1], n:=N-1, k:=P[2], root:=PrimitiveRootMod(N),
	    Scat(RR(N, 1, root)) *
	    DirectSum(I(1), C[3]) *
	    self.raderMid(N, k, root) *
	    DirectSum(Mat([[1,0]]), ##
		      RC(L_or_OS(n/2+1, 2).transpose()) *
		      DirectSum(C[1], _j(n/2).transpose() * C[2]) *
		      L(n, 2)) *
	    RealRR_Out(N, root).transpose()
	)
   ))
));


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

_vlist := (l,t)->List(Flat(l), x->t.value(x));

Class(fNearestNeighbor, FuncClass, rec(
#    def := (N,n,base,stride) -> Checked(IntFloat(base + (n-1)*stride) < N, rec(N:=N, n:=n)),
    def := (N,n,base,stride) -> rec(N:=N, n:=n),
    domain := self >> self.params[2],
    range := self >> self.params[1],
    lambda := self >> let(h := var.fresh_t("h", TInt), Lambda(h, floor(self.params[3]+h*self.params[4]+0.5))),
    # NOTE. FF: that is a temporary solution to make this function work inside Lambda. need to talk to YSV how to do it right.
    t := TFunc
));

Class(fZeroPadMiddle, FuncClass, rec(
    def := (N,n) -> Checked(IsEvenInt(n), rec(N:=N, n:=n)),
    domain := self >> self.params[2],
    range := self >> self.params[1],
    lambda := self >> let(N:=self.params[1], n:=self.params[2], fStack(fAdd(N, n/2, 0), fAdd(N, n/2, N-n/2)))
));

Class(Downsample, TaggedNonTerminal, rec(
    abbrevs := [ (dsfunc) -> [dsfunc] ],
    dims      := self >> [self.params[1].domain(), self.params[1].range()],
    terminate := self >> Gath(self.params[1]),
    isReal    := True,
    normalizedArithCost := self >> 0,
    TType := T_Complex(T_Real(64)),
    doNotMeasure := true,
    HashId := self >> let(h := [ ObjId(self.params[1]), self.params[1].range() ],
        When(IsBound(self.tags), Concatenation(h, self.tags), h)),

));

Class(Upsample, TaggedNonTerminal, rec(
    abbrevs := [ (usfunc) -> [usfunc] ],
    dims      := self >> [self.params[1].range(), self.params[1].domain()],
    terminate := self >> Scat(self.params[1]),
    isReal    := True,
    normalizedArithCost := self >> 0,
    TType := T_Complex(T_Real(64))
));

NewRulesFor(Upsample, rec(
    Upsample_base := rec(
        applicable := (self, nt) >> not nt.hasTags(),
        apply := (nt, C, cnt) -> Mat(MatSPL(Scat(nt.params[1])))
    )
));

NewRulesFor(Downsample, rec(
    Downsample_base := rec(
        applicable := (self, nt) >> not nt.hasTags(),
        apply := (nt, C, cnt) -> NeedInterleavedComplex(ScatGath(fId(nt.params[1].domain()), nt.params[1]))
    )
));


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

# NOTE: MISSING (YV)
Class(VScat_red);

Class(BaseVScat, BaseMat, SumsBase, rec(
    isVScat := true,
    sums := self >> self,
    dims := self >> self.dimensions,
    isReal := self >> true,
    needInterleavedLeft := False,
    needInterleavedRight := False,
    area := self >> self.transpose().area(),
    toAMat := self >> TransposedAMat(self.transpose().toAMat()),
    conjTranspose := self >> self.transpose(),
));

Class(BaseSVScat, BaseMat, SumsBase, rec(
    isVScat := true,
    isSVScat := true,
    #-----------------------------------------------------------------------
    abbrevs        := BaseSVGath.abbrevs,
    new            := BaseSVGath.new,
    print          := BaseSVGath.print,
    getConstantRem := BaseSVGath.getConstantRem,
    #-----------------------------------------------------------------------
    dims := self >> Error(ObjId(self), ".dims() is undefined"),
    #-----------------------------------------------------------------------
    sums := self >> self,
    isReal := self >> true,
    needInterleavedLeft := True,
    needInterleavedRight := False,
    #-----------------------------------------------------------------------
    conjTranspose := self >> self.transpose(),
    area := self >> self.transpose().area(),
    toAMat := self >> TransposedAMat(self.transpose().toAMat())
));

IsVScat := x -> IsRec(x) and IsBound(x.isVScat) and x.isVScat;
IsSVScat := x -> IsRec(x) and IsBound(x.isSVScat) and x.isSVScat;

# ==========================================================================
# VScat(<func>, <v>) - vector scatter (write) matrix, assumes aligned output
# ==========================================================================
Class(VScat, BaseVScat, rec(
    #-----------------------------------------------------------------------
    rChildren := self >> [self.func],
    rSetChild := rSetChildFields("func"),
    from_rChildren := (self, rch) >> ObjId(self)(rch[1], self.v),
    #-----------------------------------------------------------------------
    new := (self, func, v) >> SPL(WithBases(self,
        rec(func := func, v := v))).setDims(),
    #-----------------------------------------------------------------------
    dims := self >> [self.v * self.func.range(), self.v * self.func.domain()],
    #-----------------------------------------------------------------------
    transpose := self >> VGath(self.func, self.v),
    #-----------------------------------------------------------------------
    print := (self,i,is) >> Print(self.__name__, "(", self.func, ", ", self.v, ")", self.printA()),
 ));

# ==========================================================================
# VScat_u(<func>, <v>) -- same as VScat but assumes that output is unaligned
# ==========================================================================
Class(VScat_u, VScat, rec(
    dims := self >> [ self.func.range(), self.v * self.func.domain()],
    transpose := self >> VGath_u(self.func, self.v)
));

# ==========================================================================
# VScat_zero(N, n, <v>) - vector scatter (write) matrix
# ==========================================================================
Class(VScat_zero, BaseVScat, rec(
    #-----------------------------------------------------------------------
    rChildren := self >> [],
    rSetChild := rSetChildFields(),
    from_rChildren := (self, rch) >> ObjId(self)(self.N, self.n, self.v),
    #-----------------------------------------------------------------------
    new := (self, N, n, v) >> SPL(WithBases(self,
        rec(n := n, N := N, v := v))).setDims(),
    #-----------------------------------------------------------------------
    dims := self >> [self.v * self.N, self.v * self.n],
    #-----------------------------------------------------------------------
    transpose := self >> VGath_zero(self.N, self.n, self.v),
    #-----------------------------------------------------------------------
    print := (self,i,is) >> Print(self.__name__, "(", self.N, ", ",
        self.n, ", ", self.v,")", self.printA()),
 ));

# ==========================================================================
# VScat_sv(<func>, <v>, <sv>) - vector scatter (write) matrix on subvectors
# ==========================================================================
Class(VScat_sv, BaseSVScat, rec(
    #-----------------------------------------------------------------------
    rChildren := self >> [self.func],     # self >> [self.func, self.v, self.sv],
    rSetChild := rSetChildFields("func"), # rSetChildFields("func", "v", "sv"),
    from_rChildren := (self, rch) >> ObjId(self)(rch[1], self.v, self.sv, self.rem),
    #-----------------------------------------------------------------------
    dims := self >> [self.sv * self.func.range(),
                     _roundup(self.sv * self.func.domain(), self.v)],
    #-----------------------------------------------------------------------
    transpose := self >> VGath_sv(self.func, self.v, self.sv, self.rem),
));

# ==========================================================================
# RCVScat_sv(<func>, <v>, <sv>) - vector scatter (write) matrix on subvectors
# ==========================================================================
Class(RCVScat_sv, BaseSVScat, rec(
    #-----------------------------------------------------------------------
    rChildren := self >> [self.func],     # self >> [self.func, self.v, self.sv],
    rSetChild := rSetChildFields("func"), # rSetChildFields("func", "v", "sv"),
    from_rChildren := (self, rch) >> ObjId(self)(rch[1], self.v, self.sv, self.rem),
    #-----------------------------------------------------------------------
    dims := self >> let(
	n := self.func.domain() * self.sv,  nv := _roundup(n, self.v),
	N := self.func.range() * self.sv,
	[2*N, 2*nv]),
    #-----------------------------------------------------------------------
    transpose := self >> RCVGath_sv(self.func, self.v, self.sv, self.rem),
));

# ==========================================================================
# VStretchScat(<func>, part, <v>) - vector scatter (write) matrix on subvectors
# ==========================================================================
Class(VStretchScat, BaseSVScat, rec(
    #-----------------------------------------------------------------------
    rChildren := self >> [self.func],
    rSetChild := rSetChildFields("func"),
    from_rChildren := (self, rch) >> ObjId(self)(rch[1], self.part, self.v),
    #-----------------------------------------------------------------------
    dims := self >> let(
	n  := self.func.domain(), N := self.func.range(),
	nv := self.part * _roundup(n / self.part, self.v),
        [N, nv]),
    #-----------------------------------------------------------------------
    abbrevs := [],
    new := (self, func, part, v) >> SPL(WithBases(self,
	rec(func := func, part := part, v := v))).setDims(),
    #-----------------------------------------------------------------------
    transpose := self >> VStretchGath(self.func, self.part, self.v),
    #-----------------------------------------------------------------------
    print := (self,i,is) >> Print(self.__name__, "(", self.func, ", ",
        self.part, ", ", self.v, ")", self.printA()),
));

# ==========================================================================
# vRCStretchScat(<func>, part, <v>) - vector scatter (write) matrix on subvectors
# ==========================================================================
Class(vRCStretchScat, BaseSVScat, rec(
    #-----------------------------------------------------------------------
    rChildren := self >> [self.func],
    rSetChild := rSetChildFields("func"),
    from_rChildren := (self, rch) >> ObjId(self)(rch[1], self.part, self.v),
    #-----------------------------------------------------------------------
    dims := self >> let(
	n  := self.func.domain(), N := self.func.range(),
	nv := self.part * _roundup(n / self.part, self.v/2),
        2*[N, nv]),
    #-----------------------------------------------------------------------
    abbrevs := [],
    new := (self, func, part, v) >> SPL(WithBases(self,
        rec(func := func, part := part, v := v))).setDims(),
    #-----------------------------------------------------------------------
    transpose := self >> vRCStretchGath(self.func, self.part, self.v),
    #-----------------------------------------------------------------------
    print := (self,i,is) >> Print(self.__name__, "(", self.func, ", ",
        self.part, ", ", self.v, ")", self.printA()),
));

# ==========================================================================
# RCVStretchScat(<func>, part, <v>) - vector scatter (write) matrix on subvectors
# ==========================================================================
Class(RCVStretchScat, BaseSVScat, rec(
    #-----------------------------------------------------------------------
    rChildren := self >> [self.func],
    rSetChild := rSetChildFields("func"),
    from_rChildren := (self, rch) >> ObjId(self)(rch[1], self.part, self.v),
    #-----------------------------------------------------------------------
    dims := self >> let(
	n  := self.func.domain(), N := self.func.range(),
	nv := self.part * _roundup(n / self.part, self.v),
        2*[N, nv]),
    #-----------------------------------------------------------------------
    abbrevs := [],
    new := (self, func, part, v) >> SPL(WithBases(self, rec(
        func := func, part := part, v := v))).setDims(),
    #-----------------------------------------------------------------------
    transpose := self >> RCVStretchGath(self.func, self.part, self.v),
    #-----------------------------------------------------------------------
    print := (self,i,is) >> Print(self.__name__, "(", self.func, ", ",
        self.part, ", ", self.v, ")", self.printA()),
 ));

# ==========================================================================
# VScat_pc(N, n, ofs, v) - vector scatter (write) matrix writing unaligned
#                         partial contiguous vectors
# ==========================================================================
Class(VScat_pc, BaseSVScat, rec(
    #-----------------------------------------------------------------------
    rChildren := self >> [self.N, self.n, self.ofs, self.v],
    rSetChild := rSetChildFields("N", "n", "ofs", "v"),
    #-----------------------------------------------------------------------
    dims := self >> let(nv := _roundup(self.n, self.v), [self.N, nv]),
    #-----------------------------------------------------------------------
    abbrevs := [],
    new := (self, N, n, ofs, v) >> SPL(WithBases(self, rec(
        N := N, n:= n, ofs := ofs, v := v))).setDims(),
    #-----------------------------------------------------------------------
    transpose := self >> VGath_pc(self.N, self.n, self.ofs, self.v),
    #-----------------------------------------------------------------------
    print := (self,i,is) >> Print(self.__name__, "(", self.N, ", ", self.n, ", ",
        self.ofs, ", ", self.v,")", self.printA()),
));

# ==========================================================================
# IxVScat_pc(k, N, n, ofs, v) - vector gather (read)
#   matrix reading unaligned
#   partial REAL contiguous vectors
# ==========================================================================
Class(IxVScat_pc, BaseSVScat, rec(
    #-----------------------------------------------------------------------
    rChildren := self >> [self.k, self.N, self.n, self.ofs, self.v],
    rSetChild := rSetChildFields("k", "N", "n", "ofs", "v"),

    from_rChildren := (self, rch) >> ObjId(self)(self.k, self.N, self.n, self.ofs, self.v),
    #-----------------------------------------------------------------------
    dims := self >> let(nv := _roundup(self.n, self.v), [self.k*self.N, self.k*nv]),
    #-----------------------------------------------------------------------
    abbrevs := [],
    new := (self, k, N, n, ofs, v) >> SPL(WithBases(self, rec(
        k := k, N := N, n:= n, ofs := ofs, v := v))).setDims(),
    #-----------------------------------------------------------------------
    transpose := self >> IxVGath_pc(self.k, self.N, self.n, self.ofs, self.v),
    #-----------------------------------------------------------------------
    toloop := (self, bksize) >> self.transpose().toloop(bksize).transpose(),
    #-----------------------------------------------------------------------
    print := (self,i,is) >> Print(self.__name__, "(", self.k, ", ", self.N, ", ",
        self.n, ", ", self.ofs, ", ", self.v,")", self.printA()),
));

# ==========================================================================
# IxRCVScat_pc(k, N, n, ofs, v) - vector gather (read)
#   matrix reading unaligned
#   partial REAL contiguous vectors
# ==========================================================================
Class(IxRCVScat_pc, BaseSVScat, rec(
    #-----------------------------------------------------------------------
    rChildren := self >> [],         # self >> [self.ofs],  # SEEMS NOT TO WORK?!
    rSetChild := rSetChildFields(),  #  rSetChildFields("ofs"),
    from_rChildren := (self, rch) >> ObjId(self)(self.k, self.N, self.n, self.ofs, self.v),
    #-----------------------------------------------------------------------
    dims := self >> let(nv := _roundup(self.n, self.v), 2 * [self.k*self.N, self.k*nv]),
    #-----------------------------------------------------------------------
    abbrevs := [],
    new := (self, k, N, n, ofs, v) >> SPL(WithBases(self, rec(
        k := k, N := N, n := n, ofs := ofs, v := v))).setDims(),
    #-----------------------------------------------------------------------
    transpose := self >> IxRCVGath_pc(self.k, self.N, self.n, self.ofs, self.v),
    #-----------------------------------------------------------------------
    print := (self,i,is) >> Print(self.__name__, "(", self.k, ", ", self.N, ", ",
        self.n, ", ", self.ofs, ", ", self.v,")", self.printA()),
));

# ==========================================================================
# Accumulative scatter versions
# NOTE: to be replaced by a Sigma-SPL lowering step, CodeSumsAcc, and
#        a specialized CodegenAcc
#
# NOTE(!!!): .transpose().transpose() on below classes will lead to non-accumulative
#             versions, and thus invalid code!
Class(VScatAcc,    VScat,    rec());
Class(VScatAcc_u,  VScat_u,  rec());
Class(VScat_svAcc, VScat_sv, rec());
Class(VScat_pcAcc, VScat_pc, rec());

# ==========================================================================
# Methods for vectorization  of ScatGath constructs
# NOTE: This is a hack

# BB needed, as basic block size is not known due to varying function domain
ScatGath.toloopRCVec := (self, _bksize, v) >> let(
    bksize := When(_bksize = 1, v, _bksize),
    dom    := RulesMergedStrengthReduce(_divideFunc(self.gfunc.domain(), bksize)),
    i      := Ind(dom), #ok := Error("caught"),
    sfunc  := RulesFuncSimp(self.sfunc),
    ISum(i, dom,
        BB(
            Compose(When(ObjId(sfunc)=fId, [], [RCVScat_sv(sfunc, v, 1)]) ::
		    [VScat(fTensor(fBase(i), fId(2*bksize/v)), v)]) *
            RCVGath_sv(fCompose(self.gfunc, fTensor(fBase(i), fId(bksize))).setDomain(bksize), v, 1)
        )
    )
);

ScatGath.toloopVec := (self, _bksize, v) >> let(
    bksize := When(_bksize = 1, v, _bksize),
    dom    := RulesMergedStrengthReduce(_divideFunc(self.gfunc.domain(), bksize)),
    i      := Ind(dom),
    sfunc  := RulesFuncSimp(self.sfunc),
    ISum(i, dom,
        BB(
            Compose(When(ObjId(sfunc)=fId, [], [VScat_sv(sfunc, v, 1)]) ::
		    [ VScat(fTensor(fBase(i), fId(bksize/v)), v)]) *
            VGath_sv(fCompose(self.gfunc, fTensor(fBase(i), fId(bksize))).setDomain(bksize), v, 1)
        )
    )
);


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

NewRulesFor(TRDiag, rec(
    # TRDiag is a tall matrix.
    # TRDiag.transpose() is a wide matrix
    #
    TRDiag_Vec := rec(
        forTransposition := true,
	requiredFirstTag := [AVecReg, AVecRegCx],
        applicable := t -> not AnySyms(t.params[1], t.params[2]),

        apply := (t, C, Nonterms) -> let(
            v := t.firstTag().v,
            N := t.params[1],
            n := t.params[2],
            nv := _roundup(n, v),

            VScat_zero(N/v, nv/v, v) *
	    VDiag(fStretch(t.params[3], nv, n), v) *
	    VGath_pc(n, n, 0, v)
        )
    ),

    TRDiag_VecN := rec(
        forTransposition := true,
	requiredFirstTag := [AVecReg, AVecRegCx],
        applicable := t -> true, #AnySyms(t.params[1], t.params[2]),

        apply := (t, C, Nonterms) -> let(
            v := t.firstTag().v,
            N := t.params[1],
            n := t.params[2],
            ndiv := idiv(n, v),
	    nmod := imod(n, v),
	    nmodrem := imod(v-nmod, v),
	    j := Ind(ndiv),

	    # * NOTE: redundant zero-padding operations when transposing the below
	    #   VGath_zero/VScat_zero type operations fix this
	    # * Virtual pad seems to be a hack
	    # * O.toloop() should be implemented, VO is missing
	    # * Middle summand is super ugly, use ScatGath?
	    #
	    SUM(
		VirtualPad(N, n,
		    ISum(j, j.range,
			VScat(HH(N/v, ndiv, 0, [1]), v) *
			VDiag(fCompose(t.params[3], fAdd(n, ndiv*v, j*v)), v) *
			VGath(HH(n/v, ndiv, 0, [1]), v))),

		# Below could be simplified, because the block completely disappears
		# when nmod==0
		Scat(HH(N, nmod+nmodrem, ndiv*v, [1])) *
		VStack(
		    Diag(fCompose(t.params[3], fAdd(n, nmod, ndiv*v))).toloop(1),
		    O(nmodrem, nmod) # XXX NOTE: should have O.toloop() here
		) *
		Gath(HH(n, nmod, ndiv*v, [1])),

                VirtualPad(N, n,
		    VScat_zero(idiv(N,v), idiv(N,v)-ndiv, 0, v)) # XXX NOTE: VScat_zero needs to take 4th parameter (here it is 2nd) which is the scatter offset. This code won't work until VScat_zero is patched.
	    )
        )
    )
 ));


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

Class(Function, rec(
   isFunction := true,

   range := self >> self._range,
   tolist := self >> self.lambda().tolist(),

   lambda := self >> Error("Not implemented"),

   free   := self >> self.lambda().free(),
   domain := self >> self.lambda().domain(),
   at := (self, pos) >> self.lambda().at(pos),

   from_rChildren := (self, rch) >> ApplyFunc(ObjId(self), rch),

   _domain := false,
   _range := false,
   setRange := meth(self, range)
       self._range := range;
       return self;
   end,
   setDomain := meth(self, domain)
       self._domain := spiral.code.RulesStrengthReduce(domain);
       return self;
   end,

   # NOTE: below is a hack that already causes problems if we compute w/ symbolic higher order functions
   # 2 methods below allow to use constant functions (ie domain=1) as expressions

   # NOTE: don't remember where the hack above is used, so I removed it, now need to test what happens
   eval := self >> self, #When(self.domain() = 1 and self.rank() = 0, self.at(0), self),
   ev := self >> self,   #Checked(self.domain() = 1, self.at(0).ev()),
));

IsFunction := x -> IsRec(x) and IsBound(x.isFunction) and x.isFunction;

Declare(Lambda, LambdaOps, FList, FData);

#F Lambda(<vars>, <expr>)  - symbolic representation of a function mapping vars -> expr
#F
#F <vars> must be a list of variables or a single variable
#F <expr> is a symbolic expression using variables from <vars>
#F
Class(Lambda, Function, rec(
   rChildren := self >> [self.vars, self.expr, self._domain, self._range],
   rSetChild := rSetChildFields("vars", "expr", "_domain", "_range"),
   isLambda := true,
   lambda := self >> self,

   from_rChildren := (self, rch) >> ObjId(self)(rch[1], rch[2]).setDomain(rch[3]).setRange(rch[4]),

   domain := self >> When(self._domain <> false,
       self._domain,
       let(v := Last(self.vars), # ignore rank-associated vars
           When(IsBound(v.range) and not IsMeth(v.range), v.range, v.t))),

   # FF: need to be able to override range
   range := self >> When(self._range <> false, self._range, self.expr.t),

   tolist := self >> Checked(Length(self.vars)=1, let(dom := _unwrapV(self.domain()),
       Cond(IsType(dom), Error("Can't convert <self> to a list, ",
	       "because its .domain() is not a bounded interval"),
	    IsSymbolic(dom), Error("Can't convert <self> to a list, ",
	       "because its .domain() interval size is not known"),
	    List([0 .. dom-1], r -> self.at(r).eval())))),

   free := self >> Difference(self.expr.free(), self.vars),

   # FF: need to be able to override range
   print := self >> Print(self.name, "(", self.vars, ", ", self.expr, ")",
        When(self._range <> false, Print(".setRange(", self._range, ")"), ""),
        When(self._domain <> false, Print(".setDomain(", self._domain, ")"), "")
   ),

   # at(args), strict version, must have Length(args) = Length(vars)
   at := meth(arg)
       local self, args, i, bind;
       self := arg[1];
       args := Drop(arg, 1);
       # support variable-length call through lists; required by OL
       if IsList(args[1]) then args := args[1]; fi;

       Constraint(Length(args) = Length(self.vars));
       bind := tab();
       for i in [1..Length(self.vars)] do
           if IsBound(args[i]) then
               bind.(self.vars[i].id) := toExpArg(args[i]);
           fi;
       od;
       return SubstVars(Copy(self.expr), bind);
    end,

    # relaxed_at(lvars, v), relaxed version of at() for multi-ranks, allows lvars to be of any length
    # equivalent to at(Concatenation(lvars, [v])), if Length(lvars)=Length(self.vars)-1
    # v is the "standard" variable, everything in lvars is a loop index
    relaxed_at := meth(self, lvars, v)
       local i, bind, la, lv, args, vars;
       args := Concatenation(lvars, [v]);

       # for multi-rank functions loop indices can stay unbound
       vars := self.vars;
       [la, lv] := [Length(args), Length(vars)];
       if   la > lv then args := Drop(args, la - lv);
       elif lv > la then vars := Drop(vars, lv - la);
       fi;

       bind := tab();
       for i in [1..Length(vars)] do
           bind.(vars[i].id) := toExpArg(args[i]);
       od;
       return SubstVars(Copy(self.expr), bind);
    end,

#    computeType := self >> ApplyFunc(TFunc, Concatenation(
#           List(self.vars, x ->
#               When(IsBound(x.range) and not IsMeth(x.range), x.range, x.t)),
#           [self.expr.t])),

    computeType := self >> ApplyFunc(TFunc,
           List(self.vars, x ->
                Cond(IsBound(x.range) and not IsMeth(x.range), x.range,
                     Collect(self.expr, x)=[], TDummy,
                     x.t)
           ) :: [self.expr.t]),

    normalize := self >> let(
	t := self.computeType(),
	first_non_dummy := PositionProperty(DropLast(t.params,2), x->x<>TDummy),
	k := Cond(first_non_dummy=false, Length(self.vars), first_non_dummy),
	Lambda(self.vars{ [k .. Length(self.vars)]},
	       self.expr)),

    forCmp := meth(self)
        local map, i, t;
        map := tab();
        for i in [1..Length(self.vars)] do
            t := self.vars[i].t;
            map.(self.vars[i].id) := param(When(IsType(t), t, InferType(t)), "v"::StringInt(i));
        od;
        return SubstParamsCustom([self.vars, self.expr], map, [var]);
    end,

    __call__ := meth(self, vars, expr)
        local usage, res, rank, loopvars;
        usage := Concatenation("Usage: Lambda(<argList>, <expr>)\n",
                            "Usage: Lambda(<arg>, <expr>)");
            if not IsList(vars) then vars := [vars]; fi;
        if not ForAll(vars, x -> IsVar(x))
            then Error("all elements in <argList> must be variables"); fi;
        if Length(vars) <> Length(Set(vars))
            then Error("<argList> should not contain any duplicates"); fi;

        expr := toExpArg(expr);
            rank := _rank(expr);
            # if <expr> is ranked, add extra arguments to lambda
            if rank > 0 and Length(vars)=1 then
                loopvars := List([1..rank], i->var.fresh_t("q", TInt));
                vars := Concatenation(loopvars, vars);
                expr := _downRankFull(expr, loopvars);
            fi;
        res := WithBases(self, rec(expr := expr, vars := vars, operations := LambdaOps));
        res.t := res.computeType();
	# NOTE: the below line should not be necessary, not sure what the purpose was
        # if IsExp(Last(res.vars).range) then res:=res.setDomain(Last(res.vars).range); fi;
        return res;
    end,

    downRank := meth(self, loopid, ind)
        local v, vars, exp;
        if Length(self.vars) <= loopid then
            return self;
        else
            v    := self.vars[ Length(self.vars) - loopid ];
            vars := ListWithout(self.vars, Length(self.vars)-loopid);
            exp  := SubstVars(Copy(self.expr), tab((v.id) := ind));
            return ObjId(self)(vars, exp);
        fi;
    end,
));

IsLambda := obj -> IsRec(obj) and IsBound(obj.isLambda) and obj.isLambda = true;

CompatLambdas := (x, y) -> Length(x.vars) = Length(y.vars);

LambdaCompose := function(l1, l2)
    if ObjId(l1)<>Lambda then l1 := l1.lambda(); fi;
    if ObjId(l2)<>Lambda then l2 := l2.lambda(); fi;
    Constraint(Length(l1.vars)=1); Constraint(Length(l2.vars)=1);
    return let(v:=l2.vars[1].clone(), Lambda(v, l1.at(l2.at(v))));
end;

Class(LambdaOps, rec(
    Print := x->x.print(),
    \+ := (l1, l2) -> let(ll1:=l1.lambda(), ll2:=l2.lambda(), v:=l1.vars[1].clone(),
    Lambda(v, ll1.at(v) + ll2.at(v))),
    \- := (l1, l2) -> let(ll1:=l1.lambda(), ll2:=l2.lambda(), v:=l1.vars[1].clone(),
    Lambda(v, ll1.at(v) - ll2.at(v))),
    \* := (l1, l2) -> let(ll1:=l1.lambda(), ll2:=l2.lambda(), v:=l1.vars[1].clone(),
    Lambda(v, ll1.at(v) * ll2.at(v))),
    \/ := (l1, l2) -> let(ll1:=l1.lambda(), ll2:=l2.lambda(), v:=l1.vars[1].clone(),
    Lambda(v, ll1.at(v) / ll2.at(v))),
    \^ := (l1, l2) -> let(ll1:=l1.lambda(), ll2:=l2.lambda(), v:=l1.vars[1].clone(),
    Lambda(v, ll1.at(v) ^ ll2.at(v))),
    \= := (l1, l2) -> IsLambda(l1) and IsLambda(l2) and l1.lambda().forCmp() = l2.lambda().forCmp(),
    \< := (l1, l2) -> When( IsLambda(l1) and IsLambda(l2),
                              l1.lambda().forCmp() < l2.lambda().forCmp(),
                              BagAddr(l1) < BagAddr(l2) ),
));

#F FDataOfs(<datavar>, <len>, <ofs>)
#F
Class(FDataOfs, Function, rec(
    __call__ := (self, datavar, len, ofs) >> WithBases(self, rec(
    var := datavar,
    operations := PrintOps,
    ofs := toExpArg(ofs),
    len := Checked(IsPosIntSym(len), len)
    )),

# <-Daniele's changes
#    rChildren := self >> [ self.var, self.len, self.ofs],
#    rSetChild := rSetChildFields("var", "len", "ofs"),

    rChildren := self >> [ self.var, self.len, self.ofs, self._domain, self._range],
    rSetChild := rSetChildFields("var", "len", "ofs", "_domain", "_range"),
    from_rChildren := (self, rch) >> ObjId(self)(rch[1], rch[2], rch[3]).setDomain(rch[4]).setRange(rch[5]),
# ->

    domain := self >> self.len,
    print := self >> Print(self.name,"(",self.var,", ",self.len,", ",self.ofs,")"),

    at := (self, n) >> When(IsInt(n) and IsValue(self.ofs) and IsBound(self.var.value),
        self.var.value.v[n + self.ofs.v + 1],
        nth(self.var, n + self.ofs)),

    tolist := self >> List([0..EvalScalar(self.len-1)], i -> nth(self.var, self.ofs+i)),
    lambda := self >> let(x := Ind(self.domain()), Lambda(x, nth(self.var, self.ofs+x))),

    domain := self >> self.len,
    range := self >> When(self._range=false, self.var.t.t, self._range),
    inline := true,
    free := self >> self.ofs.free()
));

#F FData(<datavar>) -- symbolic function i -> datavar[i],
#F
#F domain = datavar.range
#F range = datavar.t
#F
#F
Class(FData, Function, rec(
   __call__ := arg >> let(
       self := arg[1],
       _val := Cond(Length(arg)=2, When(IsList(arg[2]), arg[2], [arg[2]]),
                               Drop(arg, 1)),
       val := When(Length(_val)=1 and IsLoc(_val[1]), _val[1], V(_val)),
       datavar := Cond(IsLoc(val), val,
                   Dat(val.t).setValue(val)),
       WithBases(self, rec(var := datavar, operations := PrintOps))),

   print := self >> Print(self.name, "(", self.var, ")"),
   rChildren := self >> [ self.var ],
   rSetChild := rSetChildFields("var"),

   at := (self, n) >> When(IsInt(n), self.var.value.v[n+1], self.lambda().at(n)),
   tolist := self >> When(IsBound(self.var.value), self.var.value.v, self.lambda().tolist()),
   lambda := self >> let(x := Ind(self.domain()), Lambda(x, nth(self.var, x))),

   domain := self >> self.var.t.size,
   range := self >> self.var.t.t,

   inline := true,
   free := self >> Set([]),
   part := (self, len, ofs) >> FDataOfs(self.var, len, ofs),
));

Class(RCData, Function, rec(
   __call__ := (self, func) >> WithBases(self, rec(
       func := Checked(IsFunction(func) or IsFuncExp(func), func),
       operations := PrintOps)),

   print := self >> Print(self.name, "(", self.func, ")"),
   rChildren := self >> [ self.func ],
   rSetChild := rSetChildFields("func"),

   at := (self, n) >> cond(n mod 2, im(self.func.at(idiv(n,2))), re(self.func.at(idiv(n,2)))),
   tolist := self >> ConcatList(self.func.tolist(), x->[re(x), im(x)]),

   domain := self >> let(d:=self.func.domain(),
       Cond(not IsType(d), d * 2,
            d=TInt,        TInt,
            Error("RCData: Invalid self.func.domain()"))),

   range := self >> self.func.range().realType(),
   # this .lambda corrently handles the case of multirank self.func
   lambda := self >> let(
       x := Ind(self.domain()),
       f := self.func.lambda(),
       jv := DropLast(f.vars, 1),
       Lambda(Concatenation(jv, [x]),
              cond(neq(imod(x, 2),0), im(f.relaxed_at(jv, idiv(x,2))), re(f.relaxed_at(jv, idiv(x,2)))))),

   inline := true,
   free := self >> self.func.free()
));

#F
#F CRData is the opposite of RCData: takes real function which assumed to be interleaved complex data
#F   and packs it into complex function:
#F   [re0, im0, re1, im1, ... ] -> [ (re0, im0), (re1, im1), ... ]
#F

Class(CRData, RCData, rec(

   at     := (self, n) >> cxpack(self.func.at(2*n), self.func.at(2*n+1)),
   tolist := self >> let( l := self.func.tolist(), List([1..Length(l)/2], i->cxpack(l[2*i-1], l[2*i]))),

   lambda := self >> let(
       x  := Ind(self.domain()),
       f  := self.func.lambda(),
       jv := DropLast(f.vars, 1),
       Lambda(Concatenation(jv, [x]), cxpack(f.relaxed_at(jv, 2*x), f.relaxed_at(jv, 2*x+1)))),

   domain := self >> let(d:=self.func.domain(),
       Cond(not IsType(d), div(d, 2),
            d=TInt,        TInt,
            Error("CRData: Invalid self.func.domain()"))),

   range := self >> self.func.range().complexType(),
));

Declare(FConj, FInv, FRConj);

Class(FConj, Function, rec(
   __call__ := (self, func) >> When(ObjId(func)=FConj, func.func, WithBases(self, rec(
       func := Checked(IsFunction(func) or IsFuncExp(func), func),
       operations := PrintOps))),

   print := self >> Print(self.name, "(", self.func, ")"),
   rChildren := self >> [ self.func ],
   rSetChild := rSetChildFields("func"),

   at := (self, n) >> conj(self.func.at(n)),
   tolist := self >> List(self.func.tolist(), conj),

   domain := self >> self.func.domain(),
   range := self >> self.func.range(),
   lambda := self >> let(x := Ind(self.domain()), Lambda(x, self.at(x))),

   inline := true,
   free := self >> self.func.free()
));

Class(FInv, Function, rec(
   __call__ := (self, func) >> When(ObjId(func)=FInv, func.func, WithBases(self, rec(
       func := Checked(IsFunction(func) or IsFuncExp(func), func),
       operations := PrintOps))),

   print := self >> Print(self.name, "(", self.func, ")"),
   rChildren := self >> [ self.func ],
   rSetChild := rSetChildFields("func"),

   at := (self, n) >> 1/(self.func.at(n)),
   tolist := self >> List(self.func.tolist(), x -> 1/x),

   domain := self >> self.func.domain(),
   range := self >> self.func.range(),
   lambda := self >> let(x := Ind(self.domain()), Lambda(x, self.at(x))),

   inline := true,
   free := self >> self.func.free()
));

# conjugation of N complex entries stored as 2*N reals r,i,r,i,...
Class(FRConj, Function, rec(
   __call__ := (self, func) >> When(ObjId(func)=FRConj, func.func, WithBases(self, rec(
       func := Checked(IsFunction(func) or IsFuncExp(func), func),
       operations := PrintOps))),

   print := self >> Print(self.name, "(", self.func, ")"),
   rChildren := self >> [ self.func ],
   rSetChild := rSetChildFields("func"),

   at := (self, n) >> cond(n mod 2, neg(self.func.at(n)), self.func.at(n)),
   tolist := self >> let(ll := self.func.tolist(),
       List([0..Length(ll)-1], i -> When(i mod 2 = 1, neg(ll[1+i]), ll[1+i]))),

   domain := self >> self.func.domain(),
   range := self >> self.func.range(),
   lambda := self >> let(x := Ind(self.domain()), Lambda(x, self.at(x))),

   inline := true,
   free := self >> self.func.free()
));

#F FList(<type>, <list>) -- symbolic function i -> list[i],
#F
#F Example: FList(TInt, [1,2,3])
#F See also: FData
Class(FList, Function, rec(
   __call__ := (self, t, list) >> Checked(IsList(list), WithBases(self, rec(
       list := list,
           t := t, operations := PrintOps))),

   print := self >> Print(self.name, "(", self.t, ", ", self.list, ")"),
   rChildren := self >> [self.t, self.list],
   rSetChild := rSetChildFields("t", "list"),

   domain := self >> Length(self.list),
   range  := self >> self.t,

   at := (self, n) >> When(IsInt(n), self.list[n+1], self.lambda().at(n)),

   _typ := self >> When(IsInt(self.t), TInt, self.t),
   tolist := self >> List(self.list, e->When(not IsSymbolic(e), Value(self._typ(), e), e)),
   tolistval := self >> Value(TArray(self._typ(), Length(self.list)), self.tolist()),
   lambda := self >> let(x := Ind(self.domain()), Lambda(x, nth(self.tolistval(), x))),

   todata := self >> FData(self.list),
   inline := false,
   free := self >> Union(List(self.list, FreeVars))
));

FF := x -> Cond(IsRec(x) and IsBound(x.isFunction), x,
                IsList(x), FList(x),
                Error("<x> must be a function or a list"));

#F FPerm(<perm>) -- function representing permutation
#F
#F Example: FPerm(Perm((1,2,3),4))
Class(FPerm, Function, rec(
   __call__ := (self, perm) >> WithBases(self, rec(
           perm := perm,
           operations := PrintOps)),

   print := self >> Print(self.name, "(", self.perm, ")"),
   rChildren := self >> [self.perm],
   rSetChild := rSetChildFields("perm"),

   domain := self >> self.perm.dims()[1],
   range  := self >> self.perm.dims()[1],

   at := (self, n) >> When(IsInt(n), self.list[n+1], self.lambda().at(n)),

   _typ := self >> TInt,
   tolist := self >> List(self.list, e->When(not IsSymbolic(e), Value(self._typ(), e), e)),
   tolistval := self >> Value(TArray(self._typ(), Length(self.list)), self.tolist()),
   lambda := self >> self.perm.sums().func.lambda(),

   inline := true,
   free := self >> FreeVars
));


local reps,tori,labels,I,II,ords,mizuno,blocks;

mizuno:=[[ 1, 2, 3, 4, 5, 6, fail,
           7, 8, 9, 10,
           11, 6, 12, 13,
           15, 14, 16, 11,
           17, 18, 19, 21,
           20, 16, 22, 23,
           24, 25, 21, 20,
           26, 27, 23, 28,
           25, 29, 30, 27,
           31, 28, 32, 33,
           30, 31, 34, 32,
           33, fail, 35, 34,
           fail, 36, 35, fail,
           36, fail, fail, fail,
           fail, fail, fail, fail ],
          []];

reps:=[[[1,1],[2,1],[3,1],[4,1],[5,1],[6,1]],          # E_6
       [[2,1],[4,1],[5,1],[6,1],[8,1],[16,1]],         # E_6(a_1)
       [[1,1],[2,1],[10,1],[11,1],[12,1]],             # D_5
       [[8,1],[9,1],[12,1],[10,1],[11,1],[16,1]],      # A_5+A_1
       [[8,1],[9,1],[16,1],[12,1],[18,1]],             # D_5(a_1)
       [[1,1],[15,1],[16,1],[17,1],[6,1]],             # A_5
       [[14,1],[15,1],[16,1],[17,1],[12,1]],           # A_4+A_1
       [[2,1],[14,1],[16,1],[18,1]],                   # D_4
       [[14,1],[15,1],[17,1],[12,1]],                  # A_4
       [[14,1],[16,1],[22,1],[24,1]],                  # D_4(a_1)
       [[14,1],[22,1],[23,1],[24,1]],                  # A_3+A_1
       [[20,1],[21,1],[28,1],[23,1],[24,1]],           # 2A_2+A_1
       [[14,1],[22,1],[24,1]],                         # A_3
       [[26,1],[27,1],[28,1],[29,1]],                  # A_2+2A_1
       [[20,1],[21,1],[23,1],[24,1]],                  # 2A_2
       [[26,1],[27,1],[35,1]],                         # A_2+A_1
       [[26,1],[35,1]],                                # A_2
       [[37,1],[38,1],[40,1]],                         # 3A_1
       [[42,1],[43,1]],                                # 2A_1
       [[53,1]]                                        # A_1
       ];

labels:=["E_6",
         "E_6(a_1)",
         "D_5",
         "A_5+A_1",
         "D_5(a_1)",
         "A_5",
         "A_4+A_1",
         "D_4",
         "A_4",
         "D_4(a_1)",
         "A_3+A_1",
         "2A_2+A_1",
         "A_3",
         "A_2+2A_1",
         "2A_2",
         "A_2+A_1",
         "A_2",
         "3A_1",
         "2A_1",
         "A_1",
         ];

I:=[ [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ] ];

II:=[ [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ] ];

tori:=[ [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ], [  ] ];

#blocks:=[];
blocks:=[[[16, 4], [8, 1], [6, 1]], [[16, 2], [11, 2], [8, 2], [4, 2]], [[8, 9], [6, 1]], [[8, 8], [4, 2], [3, 2]], [[8, 5], [6, 1], [4, 8]], [[8, 8], [2, 6], [1, 2]], [[8, 4], [6, 2], [5, 2], [4, 4], [2, 4]], [[8, 1], [6, 9], [2, 8]], [[8, 2], [7, 4], [5, 2], [4, 2], [3, 4], [1, 4]], [[4, 18], [3, 2]], [[4, 16], [2, 6], [1, 2]], [[4, 16], [2, 6], [1, 2]], [[4, 16], [2, 4], [1, 6]], [[4, 10], [3, 2], [2, 16]], [[4, 16], [1, 14]], [[4, 8], [3, 6], [2, 10], [1, 8]], [[4, 2], [3, 18], [1, 16]], [[2, 38], [1, 2]], [[2, 32], [1, 14]], [[2, 22], [1, 34]]];

#ords:=[];
ords:=[ 16, 16, 8, 8, 8, 8, 8, 8, 8, 4, 4, 4, 4, 4, 4, 4, 4, 2, 2, 2 ];

return [reps,tori,labels,I,II,ords,blocks,mizuno];


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

#F PolyDTT( <dtt-nonterminal> )
#F Parameters:       <nonterminal for a (optionally skew) DCT or DST of type 1--8>
#F Definition:       Let <t> be an (otionally skew) DCT or DST of type 1--8. Then
#F                   Transform( "PolyDTT", <t> ) represents the (n x n)-matrix
#F                   obtained from dividing every row in <t> by its first
#F                   entry.
#F                   This transform is called the polynomial version of
#F                   the DCT or DST.
#F Note:             The transpose of PolyDTT(<t>) is
#F                   TPolyDTT(<t>^T).
#F Example:          PolyDTT(DCT2(8)).
#F
Class(PolyDTT, NonTerminal, rec(
    _short_print := true,
    abbrevs := [ nt -> Checked(IsNonTerminal(nt),
	    ObjId(nt) in
	    [ DCT1, DCT3, DCT5, DCT7, # PolyDTT(T)=T for these
	      DCT2, DCT4, DCT6, DCT8,
	      DST1, DST2, DST3, DST4,
	      DST5, DST6, DST7, DST8,
	      SkewDTT],
	    [ Copy(nt) ]
	)],
    dims := self >> self.params[1].dimensions,
    terminate := self >> Mat(PolynomialDTT(MatSPL(TerminateSPL(self.params[1])))),
    isReal := True
));

Class(fr, Exp, rec(
    ev := self >> let(
	m := self.args[1].ev(),
	i := self.args[2].ev(),
	r := self.args[3].ev(),
	Cond(i mod 2 = 0, (r + 2 * Int(i/2)) / m,
	     i mod 2 = 1, (2 - r + 2 * Int(i/2)) / m))));

#F Rules
#F -----
RulesFor(PolyDTT, rec(
    PolyDTT_ToNormal := rec(
	isApplicable := P -> ObjId(P[1]) in [DCT1, DCT3, DCT5, DCT7],
	allChildren := P -> [P[1]],
	rule := P -> P[1]
    ),

    #F PolyDTT_Base2: (base case for non-skew DTTs of size 2)
    #F
    #F   pDCT2_2 = F_2
    #F   pDCT4_2 = F_2 * [[1, -1], [0, sqrt(2)]]
    #F   pDCT6_2 = [ [ 1, 1 ], [ 1, -2 ] ]
    #F   pDCT8_2 = [ [ 1, cos(2*pi/5) ], [ 1, cos(4*pi/5) ] ]
    #F   pDST4_2 = F_2 * [[1, 1], [0, sqrt(2)]]
    #F
    #F   Pueschel/Moura: Discrete Cosine and Sine Transforms, in preparation
    #F
    PolyDTT_Base2 := rec (
	info           := "pDTT_2 -> F_2",
	isApplicable   := P -> Rows(P[1]) = 2 and ObjId(P[1]) <> SkewDTT,
	allChildren    := P -> [[ ]],
	rule := (P, C) -> let(nt := ObjId(P[1]), Cond(
		nt = DCT2, F(2),
		nt = DCT4, F(2) * Mat([[1, -1], [0,sqrt(2)]]),
		nt = DCT6, Mat([[1,1], [1,-2]]),
		nt = DCT8, Mat([[1,2*CosPi(2/5)], [1,2*CosPi(4/5)]]),
		nt = DST2, F(2),
		nt = DST3, F(2) * Diag([1, sqrt(2)]),
		nt = DST4, F(2) * Mat([[1, 1], [0, sqrt(2)]]),
		nt = DST5, Mat([[1,2*CosPi(2/5)], [1,2*CosPi(4/5)]]),
		nt = DST6, Mat([[1,2*CosPi(1/5)], [1,2*CosPi(3/5)]]),
		nt = DST7, Mat([[1,2*CosPi(1/5)], [1,2*CosPi(3/5)]]),
		nt = DST8, Mat([[1,2], [1,-1]]),
		Error("unrecognized <L.symbol>"))
	)
    ),

   PolyDTT_SkewBase2 := rec (
	info           := "pDTT_2 -> F_2",
	isApplicable   := P -> Rows(P[1]) = 2 and ObjId(P[1]) = SkewDTT and
	                       ObjId(P[1].params[1]) = DST3,
	allChildren    := P -> [[ ]],
	rule := (P, C) -> let(skewnt := ObjId(P[1].params[1]), r := P[1].params[2],
	    Cond(
		skewnt = DST3, F(2) * Diag(1, 2*CosPi(r/2)),
		Error("unrecognized <L.symbol>"))
	)
   ),

   PolyDTT_SkewDST3_CT := rec(
       isApplicable := P -> ObjId(P[1]) = SkewDTT and
                            ObjId(P[1].params[1]) = DST3 and
                            Rows(P[1]) > 2 and not IsPrime(Rows(P[1])),

       allChildren := P -> let(
	   N := Rows(P[1].params[1]), r := P[1].params[2],
	   List(DivisorPairs(N), d->
	       let(i := Ind(d[2]),
		   ri := fr(d[2], i, r),
		   [ PolyDTT(SkewDTT(DST3(d[1]), ri)),
		     PolyDTT(SkewDTT(DST3(d[2]), r)) ]))),

       rule := (P,C) -> let(
	   MN := Rows(P[1]),  N := Rows(C[1]),  M := Rows(C[2]),
	   r := P[1].params[2],
	   i := C[1].root.params.params[2].args[2],

	    K(MN, N) *
	    IterDirectSum(i, i.range, C[1]) *
	    Tensor(C[2], I(N)) *
	    B_DST3_U(MN, M)
       )
   )
));