** Polynomial Regression model with order (dd-1): eta(x,theta)=theta1+theta2*x+...+theta_dd*x**(dd-1);
** find the Phi_p optimal design using OWEA algorithm;

proc iml;

** each element of step denotes the number of support points, we divided the design space equally; 
** if there is more than one element in step, it means we first search in a rough grid space;
** then serach in a finer grid, so on and so forth;  
step={10000};
pp=0;
dd=6;
** matrix for interested parameters;
WB=i(dd);
theta=j(1,dd,1);
** exist design, last column refers to the number of replications; 
** if 0, means one stage design;  
exist_design={0 0};
*exist_design={0 10, 1 10, 2 10, 3 10};

** sample size to be added in the next stage;
sn=80; 


** design space;
bound={-1 1}; 

** initial design points;
x0=(1:dd)`/5-1; 


** the maximum number of iteration when we compute the optimal weights for given support;
max_iter=50; 


w0=1/nrow(x0)*j(nrow(x0),1,1);

** information matrix for a single design point t;
** you need to change this function for a different model;
start infor1(t,theta) global (dd);
    ft=j(dd,1,1);
	ft[1,1]=1;
	do i=2 to dd;
	     
		ft[i,1]=t**(i-1);
	end;
    infor=ft*ft`*(1-t**2);
    return(infor);
finish;


** information matrix for a given design;
** the last column of d is the weight for the design point;
start infor(d,theta);
    n=nrow(d);
    m=ncol(d);
    dim=nrow(theta)*ncol(theta);
    infor_d=j(dim,dim,0);
    do i=1 to n;
    infor_d=infor_d+d[i,m]*infor1(d[i,1:m-1],theta);
    end;
    return(infor_d);
finish;

** verify equivalence theorem;
start verify_equiv(pp,x_temp,opt_infor,theta,exist_design,sn,WB);

	infor0=infor(exist_design,theta)/sn;

	temp=-100000;
	nx=nrow(x_temp);
	np=ncol(x_temp);
	result=j(1,np+1,-10000);
	coeff=1;

	temp_inv=wb*ginv(opt_infor+infor0);
	if pp=0 then do;
	    part=temp_inv`*inv(temp_inv*wb`)*temp_inv;
	    diff1=trace(opt_infor*part);
	end;
	if pp>0 then do;
		var_cov=temp_inv*wb`;
   		part=temp_inv`*((var_cov)**(pp-1))*temp_inv;
		diff1=trace(opt_infor*part);
		coeff=(1/nrow(wb))**(1/pp)*(trace((var_cov)**pp))**(1/pp-1);
	end;

	do i=1 to nx;
		x=x_temp[i,]; 
		inforx=infor1(x,theta);
		diff=(trace(inforx*part)-diff1)*coeff;
		if diff>temp then do;
			temp=diff;
			result[1,1:np]=x; 
			result[1,np+1]=temp;
		end;
	end;
	*print result;
	return(result);

finish;

* weight1 is to find the critical points;
* x is a give design points set;  
* wo is the initial weight vector of x; 

start weight1(pp,x,theta,exist_design,sn,w0,WB);
	infor0=infor(exist_design,theta)/sn;
	n=nrow(x); *n is the number of design points; 
	last_infor=infor1(x[n,],theta);
	p=ncol(theta)*nrow(theta);
	dim_wb=nrow(wb);
	do i=1 to n-1;
		temp_infor=infor1(x[i,],theta);
		dinfor_m=dinfor_m||(temp_infor-last_infor);
		infor_m=infor_m||temp_infor;
	end;
	infor_m=infor_m||last_infor;

	weight=w0[1:n-1]; 
	
	diff=1; repli=1; indic=1; delta=1;

    do while((diff>0.0000000000000001)&(indic>0.5)&(repli<40));

        last_infor=infor_m[,p*(n-1)+1:p*n];
        infor=(1-sum(weight))*last_infor;
        do i=1 to n-1;
            infor=infor+weight[i]*infor_m[,p*(i-1)+1:p*i];
        end;

        d1w=j(n-1,1); * d1w is the first derivative of omega;
        d2w=i(n-1);   * d2w is the second derivative of omega;
        temp_inv=ginv(infor+infor0);
        /*tempxx=eigval(infor+infor0); print tempxx;*/
        var_cov=wb*temp_inv*wb`;
		inv_var=(var_cov)**(pp-1);
        part1=wb*temp_inv*dinfor_m;
		part2=-wb*temp_inv*dinfor_m[,1:p]*temp_inv*wb`;
		if n>2 then do;
	        do i=2 to n-1;
		        part2=part2||(-wb*temp_inv*dinfor_m[,p*(i-1)+1:p*i]*temp_inv*wb`);
			end;
		end;
		var_cov_stack=i(dim_wb);
		if pp>2 then do;
			do i=1 to pp-2;
				var_cov_stack=var_cov_stack||((var_cov)**i);
			end;
		end;

        do i=1 to n-1;
            d1w[i]=trace(inv_var*part2[,dim_wb*(i-1)+1:dim_wb*i]);
            do j=1 to i;
				temp1=inv_var*part1[,p*(i-1)+1:p*i]*temp_inv;
                temp2=part1[,p*(j-1)+1:p*j];
				if pp=0 then do; 
	                d2w[i,j]=trace(2*temp1*temp2`-inv_var*part2[,dim_wb*(i-1)+1:dim_wb*i]*inv_var*part2[,dim_wb*(j-1)+1:dim_wb*j]);
				end;
				else if pp=1 then do;
					d2w[i,j]=2*trace(temp1*temp2`);
				end;
				else do;
	                d2w[i,j]=2*trace(2*temp1*temp2`);
					if pp-2=2*floor((pp-2)/2) then do;
						d2w[i,j]=d2w[i,j]+trace(var_cov_stack[,dim_wb*floor((pp-2)/2)+1:dim_wb*(floor((pp-2)/2)+1)]*part2[,dim_wb*(i-1)+1:dim_wb*i]*var_cov_stack[,dim_wb*(floor((pp-2)/2))+1:dim_wb*(floor((pp-2)/2)+1)]); 
						if floor((pp-2)/2)>0 then do;
							do count=0 to floor((pp-2)/2)-1;
								d2w[i,j]=d2w[i,j]+2*trace(var_cov_stack[,dim_wb*count+1:dim_wb*(count+1)]*part2[,dim_wb*(i-1)+1:dim_wb*i]*var_cov_stack[,dim_wb*(pp-2-count)+1:dim_wb*(pp-1-count)]);
							end;
						end;
					end; 
					else do;
						do count=0 to floor((pp-2)/2);
							d2w[i,j]=d2w[i,j]+2*trace(var_cov_stack[,dim_wb*count+1:dim_wb*(count+1)]*part2[,dim_wb*(i-1)+1:dim_wb*i]*var_cov_stack[,dim_wb*(pp-2-count)+1:dim_wb*(pp-1-count)]);
						end;
					end;
				end;
				if d2w[i,j]>10e15 then d2w[i,j]=10e15;
				if d2w[i,j]<-10e15 then d2w[i,j]=-10e15;
                d2w[j,i]=d2w[i,j];
            end;
        end;
		scale=max(abs(d2w));
        new_weight=weight-delta*ginv(d2w/scale)*d1w/scale;
        if ((min(new_weight)<0)|(sum(new_weight)>1)) then do;
            if delta>0.00001 then delta=delta/2;
            else indic=0;
            end;
        else do;
            diff=sum((new_weight-weight)##2);
            repli=repli+1;
        weight=new_weight;
        end;
        *print diff delta weight;
end;
    *print d1 repli delta diff;
temp=1-sum(new_weight);
weight=new_weight//temp;
return(weight);
finish;

** derive the critical point under boundary;
start weight2(pp,x,theta,exist_design,sn,w0,WB);
    new_x=x;
    n1=nrow(new_x);
    if n1>1 then do;
	    opt_weight=j(nrow(x),1,0);
	    weight=weight1(pp,x,theta,exist_design,sn,w0,WB);
	    do while((nrow(new_x)>1)&(min(weight)<0.000001));
	        if weight[weight[>:<]]<0.000001 then d_point=weight[>:<];
	        index=j(1,nrow(new_x));
	        do i=1 to nrow(new_x); index[i]=i; end;
	        new_x=new_x[remove(index,d_point),]; *print weight d_point new_xy;
	        new_w0=weight[remove(index,d_point),];
	        if nrow(new_x)>1 then weight=weight1(pp,new_x,theta,exist_design,sn,new_w0,WB);
	    end;
	end;
	else weight=1;
    *print weight;
    n1=nrow(new_x);
    if n1=1 then weight=1;
    opt_weight=weight;
    design=new_x||opt_weight;
    opt_infor=infor(design,theta);
    return(design);
finish;

** return optimal design for given a set of design points, then verifying whether it is. If not, add potential points, search again;
** design is the set of design points;
** x0 is the initial design points;

start opt_fixed(pp,design_points,x0,w0,theta,exist_design,sn,WB,max_iter);
    n=nrow(x0);
    temp=weight2(pp,x0,theta,exist_design,sn,w0,WB);
    opt_infor=infor(temp,theta);
    temp1=verify_equiv(pp,design_points,opt_infor,theta,exist_design,sn,WB);
    iter=1;
    do while ((temp1[,ncol(temp1)]>0.00001)&(iter<max_iter));
        newx=temp[,1:ncol(temp)-1]//temp1[,1:ncol(temp1)-1];
        n=nrow(newx);
        w0=temp[,ncol(temp)]//0;
        temp=weight2(pp,newx,theta,exist_design,sn,w0,WB);
        opt_infor=infor(temp,theta);
        temp1=verify_equiv(pp,design_points,opt_infor,theta,exist_design,sn,WB);
        iter=iter+1;
    end;
	temp=temp//temp1;
    return(temp);
finish;

** generate fixed design points. These points are uniformaly distribued on the intervals;
start generate_design_points(step,bound);

	p=nrow(bound);
	x_step=j(p,1);

	do i=1 to p;
		x_step[i]=(bound[i,2]-bound[i,1])/step;
	end;
	do i=1 to p;
		temp1=j(step+1,1);
		do j=1 to step+1; 
			temp1[j,1]=bound[i,1]+x_step[i]*(j-1);
		end;
		if i=1 then temp2=temp1@j((step+1)**(p-1),1);
		else if i=p then temp2=j((step+1)**(p-1),1)@temp1;
		else temp2=j((step+1)**(i-1),1)@temp1@j((step+1)**(p-i),1);
		x_temp=x_temp||temp2;
	end;
	return(x_temp);
finish;

* this is a new version, supposed to be faster;
start select_design_points(step1,step2,bound,temp);

	p=nrow(bound);
	n_temp=nrow(temp);
	x_step1=j(p,1);
	x_step2=j(p,1);

	do i=1 to p;
		x_step1[i]=(bound[i,2]-bound[i,1])/step1;
		x_step2[i]=(bound[i,2]-bound[i,1])/step2;
	end;

	do i=1 to n_temp;
		do j=1 to p;
			temp_bound1=max(temp[i,j]-x_step1[j],bound[j,1]);
			temp_bound2=min(temp[i,j]+x_step1[j],bound[j,2]);
			temp_p=temp_bound1;
			temp_pp=0;
			do while(temp_p<=temp_bound2);
				temp_pp=temp_pp//temp_p;
				temp_p=temp_p+x_step2[j];
			end;
			temp_pp=temp_pp[2:nrow(temp_pp)];
			if j=1 then temp_points=temp_pp;
			else do;
				n1=nrow(temp_points);
				n2=nrow(temp_pp);
				temp_points=temp_points@j(n2,1,1);
				temp_points=temp_points||(j(n1,1,1)@temp_pp);
			end;
		end;
		select_points=select_points//temp_points;
	end;
	
	return(select_points);
finish;

start phi_value(pp,design,theta,exist_design,sn,wb);
    opt_infor=infor(design,theta)*sn+infor(exist_design,theta);
	if pp=0 then value=log(det(wb*ginv(opt_infor)*wb`));
	else value=(trace((wb*ginv(opt_infor)*wb`)**pp)/nrow(wb))**(1/pp);
    return(value);
finish;

dp=nrow(bound);

design_points=generate_design_points(step[1],bound); 
temp=opt_fixed(pp,design_points,x0,w0,theta,exist_design,sn,WB,max_iter);
temp_design=temp[1:nrow(temp)-1,];
value=phi_value(pp,temp_design,theta,exist_design,sn,wb);
*print temp_design value;

if nrow(step)>1 then do;
	do i=2 to nrow(step);

		x0=temp_design[,1:dp];
		w0=temp_design[,dp+1];
		select_points=select_design_points(step[i-1],step[i],bound,x0);
		temp=opt_fixed(pp,select_points,x0,w0,theta,exist_design,sn,WB,max_iter);
		temp_design=temp[1:nrow(temp)-1,];
		*print temp_design;

		value=phi_value(pp,temp_design,theta,exist_design,sn,wb);
	end;
end;
		verify_equiv=temp[nrow(temp),];
		*print verify_equiv;
		print temp_design value;

quit;