#
# Read("~/Workspace/groupsSB/epi/B2/handle.gi");
#

type:="B";
rank:=2;
nr_pos_roots:=4;

Read("~/Workspace/groupsSB/epi/group.gi");

#sList(rels,r->Value(r,[vars[22]],[vars[2]+vars[12]]);

#
# c_1, c_2 not 0 => c_3=c_1/2
#
# a test that the coeffs are right
#a:=root_group(1,xvars[1]);
#b:=root_group(2,xvars[2]);
#c:=root_group(3,xvars[3]);
#d:=root_group(4,xvars[4]);
#test1:=evaluate_U(a*b*c*d,[[[xvars[1]],[One(APR)]],[[xvars[2]],[One(APR)]],[[xvars[3]],[-One(APR)/2]],[[xvars[4]],[-2*One(APR)/3]]]);
#test2:=evaluate_U(a*b*c*d,[[[xvars[1]],[One(APR)*2]],[[xvars[2]],[One(APR)*2]],[[xvars[3]],[-One(APR)*2]],[[xvars[4]],[-16*One(APR)/3]]]);
#Set(Concatenation(test2-test1^2));
handleReg:=function()
	local ua,ub,uab,uabb,u,e,nn,vals,rels,i,v;
	ua:=root_group(1,1);
	ub:=root_group(2,1);
	uab:=root_group(3,-1/2);
	uabb:=root_group(4,-2/3);
	u:=ua*ub*uab*uabb;
	e:=u-id_mat;

	nn:=List([1..2*nr_pos_roots+rank],i->xvars[i]);

	vals:=[];
 	Append(vals,[[[xvars[6]],[Zero(APR)]]]);
# 	Append(vals,[[[xvars[4]],[Zero(APR)]]]);
  	Append(vals,[[[xvars[5]],[Zero(APR)]]]);
  	Append(vals,[[[xvars[7]],[Zero(APR)]]]);
  	Append(vals,[[[xvars[8]],[Zero(APR)]]]);
  	Append(vals,[[[xvars[9]],[2*xvars[10]]]]);
 	Append(vals,[[[xvars[10]],[Zero(APR)]]]);
  	Append(vals,[[[xvars[2]],[xvars[1]]]]);
  	Append(vals,[[[xvars[3]],[Zero(APR)]]]);
#  	Append(vals,[[[xvars[3]],[-2*xvars[1]/3]]]);
#  	##
	#Append(vals,[[[xvars[1]],[Zero(APR)]]]);
	#Append(vals,[[[xvars[3]],[Zero(APR)]]]);
	#Append(vals,[[[xvars[7]],[Zero(APR)]]]);

	#
	# atentie la ordinea transpunerilor
	#
	rels:=nn*TransposedMat(u)-nn;
	rels:=evaluate_rels(rels,vals);
	nn:=evaluate_rels(nn,vals);

	return [u,nn,rels,vals];
end;

#
# Done
#
handleRegModule:=function()
	local ua,ub,uab,uabb,u,e,nn,vals,rels,i,v;
	ua:=root_group(1,1);
	ub:=root_group(2,1);
	uab:=root_group(3,-1/2);
	uabb:=root_group(4,-2/3);
	u:=ua*ub*uab*uabb;

	nn:=One(APR)*[0,0,0,0,xvars[1],xvars[2],xvars[3],0,0,0];
	#
	# atentie la ordinea transpunerilor
	#
	nn:=nn*TransposedMat(u);

	rels:=[];
	Append(rels,[nn[1]-nn[2]]);
	# Append(rels,[nn[3]-2*nn[1]/3]);

	vals:=[];
	Append(vals,[[[xvars[3]],[-6*xvars[1]/5+8*xvars[2]/5]]]); # because I want the projection on U3 to be trivial
	Append(vals,[[[xvars[2]],[3*xvars[1]/4]]]); # because I want the projections on U1 and U2 to be equal
#	Append(vals,[[[xvars[3]],[13*xvars[1]/2-2*xvars[2]]]]);
#	Append(vals,[[[xvars[2]],[203*3*xvars[1]/(4*85)]]]);

	rels:=evaluate_rels(rels,vals);
	nn:=evaluate_rels(nn,vals);

	return [u,nn,rels,vals];
end;

#
# c_2=0 c_3=0
#
#
#
handle23:=function()
	local ua,ub,uab,uabb,u,e,nn,vals,rels,i,v;
	ua:=root_group(1,1);
	uabb:=root_group(4,1);
	u:=ua*uabb;
	id_mat:=u^0;
	e:=u-id_mat;

	nn:=List([1..2*nr_pos_roots+rank],i->xvars[i]);

	vals:=[];
	Append(vals,[[[xvars[7]],[Zero(APR)]]]);
	Append(vals,[[[xvars[8]],[Zero(APR)]]]);
	Append(vals,[[[xvars[5]],[Zero(APR)]]]);
	Append(vals,[[[xvars[10]],[Zero(APR)]]]); # for p>2
	Append(vals,[[[xvars[9]],[Zero(APR)]]]); # for p>2
	Append(vals,[[[xvars[2]],[-xvars[6]]]]);
	# ###
	# Append(vals,[[[xvars[1]],[Zero(APR)]]]);
	# Append(vals,[[[xvars[3]],[Zero(APR)]]]);
	# Append(vals,[[[xvars[4]],[Zero(APR)]]]);
	# Append(vals,[[[xvars[8]],[Zero(APR)]]]);
	# Append(vals,[[[xvars[9]],[Zero(APR)]]]);

	#
	# atentie la ordinea transpunerilor
	#
	rels:=nn*TransposedMat(u)-nn;
	rels:=evaluate_rels(rels,vals);
	nn:=evaluate_rels(nn,vals);

	return [u,nn,rels,vals];
end;

#
# c_2=0 c_4=0
#
handle24:=function()
	local ua,ub,uab,uabb,u,e,nn,vals,rels,i,v;
	ua:=root_group(1,1);
	uab:=root_group(3,1);
	u:=ua*uab;
	id_mat:=u^0;
	e:=u-id_mat;

	nn:=List([1..2*nr_pos_roots+rank],i->xvars[i]);

	vals:=[];
	Append(vals,[[[xvars[7]],[Zero(APR)]]]);
	Append(vals,[[[xvars[8]],[Zero(APR)]]]);
	Append(vals,[[[xvars[5]],[Zero(APR)]]]);
	Append(vals,[[[xvars[2]],[Zero(APR)]]]); # for p>2
	Append(vals,[[[xvars[9]],[Zero(APR)]]]);
	Append(vals,[[[xvars[10]],[-xvars[6]]]]); # for p>2
	# # ###
	# Append(vals,[[[xvars[1]],[Zero(APR)]]]);
	# Append(vals,[[[xvars[3]],[Zero(APR)]]]);
	# Append(vals,[[[xvars[4]],[Zero(APR)]]]);
	# Append(vals,[[[xvars[8]],[Zero(APR)]]]);
	# Append(vals,[[[xvars[9]],[Zero(APR)]]]);

	#
	# atentie la ordinea transpunerilor
	#
	rels:=nn*TransposedMat(u)-nn;
	rels:=evaluate_rels(rels,vals);
	nn:=evaluate_rels(nn,vals);

	return [u,nn,rels,vals];
end;

#
# c_2=0
#
handle2:=function()
	local ua,ub,uab,uabb,u,e,nn,vals,rels,i,v;
	ua:=root_group(1,1);
	uab:=root_group(3,1);
	uabb:=root_group(4,1);
	u:=ua*uab*uabb;
	id_mat:=u^0;
	e:=u-id_mat;

	nn:=List([1..2*nr_pos_roots+rank],i->xvars[i]);

	vals:=[];
	Append(vals,[[[xvars[5]],[xvars[8]]]]);
	Append(vals,[[[xvars[7]],[-xvars[8]]]]);
	Append(vals,[[[xvars[2]],[-xvars[6]+xvars[9]]]]);
	Append(vals,[[[xvars[6]],[xvars[9]-xvars[10]]]]); # only in char 2
	###
	# Append(vals,[[[xvars[1]],[Zero(APR)]]]);
	# Append(vals,[[[xvars[3]],[Zero(APR)]]]);
	# Append(vals,[[[xvars[4]],[Zero(APR)]]]);
	# Append(vals,[[[xvars[8]],[Zero(APR)]]]);
	# Append(vals,[[[xvars[9]],[Zero(APR)]]]);

	#
	# atentie la ordinea transpunerilor
	#
	rels:=nn*TransposedMat(u)-nn;
	for i in [1..Length(rels)] do
		for v in vals do
			Print(rels[i],"\n");
			rels[i]:=One(APR)*Value(rels[i],v[1],v[2]);
			nn[i]:=One(APR)*Value(nn[i],v[1],v[2]);
		od;
	od;

	return [u,nn,rels,vals];
end;

#
# regular
# c_1, c_2 not, c_3 not 0
#
handle1:=function()

end;


Declare(CUFFTCall);

Class(CUFFTCall, BaseMat, SumsBase, rec(
    dims := self >> self.L.dims(),
    isReal := self >> self.L.isReal(),
    #-----------------------------------------------------------------------
    rChildren := self >> [self.L, self.codegen],
    rSetChild := rSetChildFields("L", "codegen"),
    #-----------------------------------------------------------------------
    new := (self, L, codegen) >> SPL(WithBases(self,
        rec(L   := L,
            codegen   := CopyFields(codegen, rec(
                plan_var := var.fresh_t("plan", TSym("cufftHandle")),
                size_var := var.fresh_t("size", TInt),
                data := (self, c) >> data(self.size_var, V(self.N), c),
                init := self >> call(rec(id:="cufftPlanMany"), addrof(self.plan_var), V(1), addrof(self.size_var), addrof(self.size_var), V(1), V(self.M),
    	            addrof(self.size_var), V((self.N * self.K) / (self.p1 * self.p2)), V(1), "CUFFT_Z2Z", V((self.N * self.K) / (self.p1 * self.p2))),
            )),
            dimensions     := L.dims())
    )),

    #-----------------------------------------------------------------------
    transpose := self >> CUFFTCall(self.L.transpose(), self.codegen),
    #-----------------------------------------------------------------------

    print := (self,i,is) >> Print(self.name, "(",
        self.L.print(i+is,is), ", <codegen>)"),
    #-----------------------------------------------------------------------
    toAMat := self >> self.L.toAMat(),
    #-----------------------------------------------------------------------
    isPermutation := False
));

RewriteRules(RulesRC, rec(
    RC_CUFFTCall := Rule([RC, @(1, CUFFTCall)], e -> CUFFTCall(RC(@(1).val.L), @(1).val.codegen)),
));

CudaCodegen.CUFFTCall := (self, o, y, x, opts) >> simt_block(
    call(rec(id := "cufftExecZ2Z", codegen := o.codegen), o.codegen.plan_var,
        tcast(TPtr(TSym("cufftDoubleComplex")), x), tcast(TPtr(TSym("cufftDoubleComplex")), y), When(o.codegen.k = 1, "CUFFT_INVERSE", "CUFFT_FORWARD")));


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

Class(RulesId, RuleSet);

RewriteRules(RulesId, rec(

 IdISumRight := ARule( Compose,
       [ @(1, Id), @(2, ISum, canReorder) ],
  e -> [ ISum(@2.val.var, @2.val.domain, @1.val * @2.val.child(1)).attrs(@(2).val) ]),

 IdISumLeft := ARule( Compose,
       [ @(1, ISum, canReorder), @(2, Id) ],
  e -> [ ISum(@1.val.var, @1.val.domain, @1.val.child(1) * @2.val).attrs(@(1).val) ]),

 IdIComposeRight := ARule( Compose,
       [ @(1, Id), @(2, ICompose, canReorder) ],
  e -> [ ICompose(@2.val.var, @2.val.domain, @1.val * @2.val.child(1)).attrs(@(2).val) ]),

 IdIComposeLeft := ARule( Compose,
       [ @(1, ICompose, canReorder), @(2, Id) ],
  e -> [ ICompose(@1.val.var, @1.val.domain, @1.val.child(1) * @2.val).attrs(@(1).val) ]),

 # Gath * Id
 CommuteGathId := ARule( Compose,
       [ @(1, Gath), @(2, Id) ], # o 1-> 2->
  e -> [ Id(fCompose(@2.val.element, @1.val.func)).attrs(@(2).val), @1.val ]),

 # Prm * Id
 CommutePrmId := ARule( Compose,
       [ @(1, Prm), @(2, Id) ], # o 1-> 2->
  e -> [ Id(fCompose(@2.val.element, @1.val.func)).attrs(@(2).val), @1.val ]),

 # Id * Scat
 CommuteIdScat := ARule( Compose,
       [ @(1, Id), @(2, Scat) ], # <-1 <-2 o
  e -> [ @2.val, Id(fCompose(@1.val.element, @2.val.func)).attrs(@(1).val) ]),

  #(A X B) o (C X D) -> (A o C) X (B o D)
  CommuteIdTensorfTensor := Rule([fCompose, [idTensor, @(1), @(2)],
      [@(5, fTensor).cond(e->Length(e.children()) = 2), @(3).cond(e->e.range() = @1.val.domain()), @(4).cond(e->e.range() = @2.val.domain())]],
      e -> idTensor(fCompose(@(1).val, @(3).val), fCompose(@(2).val, @(4).val))
  ),

 IdNoDiagPullinRight := ARule( Compose,
       [ @(1, Id), @(2, NoDiagPullinRight) ],
  e -> [ NoDiagPullinRight(@(1).val * @2.val.child(1)) ]),

 IdNoDiagPullinLeft := ARule( Compose,
       [ @(1, NoDiagPullinLeft), @(2, Id) ],
  e -> [ NoDiagPullinLeft(@1.val.child(1) * @(2).val) ]),

#(A X B) o ((C X D) X E) -> (A o (C X D)) X (B o E)
  CommuteIdTensorfTensor2 := Rule([fCompose, [idTensor, @(1), @(2)],
      [@(3, fTensor).cond(e->
           Length(e.children()) = 3 and
       fTensor(e.child(1), e.child(2)).range() = @1.val.domain() and
       e.child(3).range() = @2.val.domain()
       ),
       @(4), @(5), @(6)
      ]], e->

      idTensor(fCompose(@(1).val, fTensor(@(3).val.child(1), @(3).val.child(2))),
               fCompose(@(2).val, @(3).val.child(3)))
      ),

  # i o (j)n = (j)1->N, where i is an idId
  IdComposeConst := Rule([fCompose, @(1, idId), @(2, fBase)],
                    e -> idConst(1, @2.val.params[2])
  ),

  idTensorId1 := ARule(idTensor, [@(1), [@(2,[idId]), 1]], e -> [@(1).val]),
  idTensorId2 := ARule(idTensor, [[@(2,[idId,J]), 1], @(1)], e -> [@(1).val]),

  # This is a special-puprose modification of the IdComposeConst rule, above.
  # It's possible that this is not the best way to do this.  I'd be willing
  # to change it later (Peter).
  # i o ((j)n X (k)m ) = (j)1->N X (k)1->N
#  IdComposeTensConst := Rule([fCompose,
#     @(1, idId),
#     @(2, fTensor).cond(e->
#         Length(e.children()) = 2  and
#         ObjId(e.child(1)) = fBase and
#         ObjId(e.child(2)) = fBase)],
#      e -> @(2).val
#  ),

  # id(a x b) = id(c) where c is derived from merging a and b and
  # computing corresponding domains and ranges.
  FoldConstants := ARule(idTensor, [@(1, idConst), @(2, idConst)],
      e -> [idConst(@1.val.domain()*@2.val.domain(), @1.val.params[2]*@2.val.params[2])]
  )

));