%Generate data for bagging and/or boosting
%make some fake data in dim dimensions.  You'll need to tailor the
%distributions if adjusting dim from the value 3.
nr=100; % number of red, blue, and green points
nb=100;
ng=100;
nt0=nr+nb+ng;
dim=3;

%training data
sd=1; % a value for the standard deviation of the normals
xr=mvnrnd([-1,0,0],sd*eye(dim),nr); % generate data, using different means to make the groups somewhat distinct
xb=mvnrnd([1,0.85,3],sd*eye(dim),nb);
xg=mvnrnd([1,-0.85,-2],sd*eye(dim),ng);
data0=[xr;xb;xg];
data0=[data0,[ones(nr,1);2*ones(nb,1);3*ones(ng,1)]]; %for y values red is 1, blue is 2, green is 3

%test dataset
xr=mvnrnd([-1,0,0],sd*eye(dim),nr); % generate data, using different means to make the groups somewhat distinct
xb=mvnrnd([1,0.85,3],sd*eye(dim),nb);
xg=mvnrnd([1,-0.85,-2],sd*eye(dim),ng);
testdata=[xr;xb;xg];
testdata=[testdata,[ones(nr,1);2*ones(nb,1);3*ones(ng,1)]]; %for y values red is 1, blue is 2, green is 3

%

%bagging for regression trees.  Can do random forest by adjusting fdim
%below

fdim=ceil(sqrt(dim)); %how many dimensions to use for random forest option; next line of code determines if this is used or not by commenting
%fdim=dim; %regular bagging; comment out if you want random forest
m=4; %define minimum number of points for a leaf of the tree.  Trees here should be deep, so this should probably be small
boots=100; %the number of bootstrap samples for bagging

mses=zeros(boots,1); %various storage for results like out of bag counts and out of bag estimates
testmses=mses;
noob=zeros(nt0,1);
yhatsum=noob;
testyhatsum=noob;

for boot=1:boots %begin bagging loop

    nt=nt0; %create bootstrap sample
    sample=randi(nt,nt,1);
    data=data0(sample,:);

    xminmax=[min(data0(:,1))-.1,max(data0(:,1))+.1]; %bounding box and other items for full dataset (root of tree)
    for i=2:dim
        xminmax=[xminmax,[min(data0(:,i))-.1,max(data0(:,i))+.1]];
    end
    yhat=mean(data(:,end));
    err=(data(:,end)-yhat)'*(data(:,end)-yhat);
    regions=[xminmax,err,0,nt,yhat,true]; %track current regions by bounding box, error before split, index of parent region, number of points, yhat, and flag on whether split or not yet

    tosplit=regions(:,end)==1; %find regions still to be split
    ntosplit=sum(tosplit);

    while ntosplit>0 %split those regions and continue until no more regions can be split
        splitters=find(tosplit);
        for k=1:ntosplit
            region=regions(splitters(k),:);
            err=region(2*dim+1);
            xminmax=region(1,1:2*dim);
            box=reshape(xminmax,[2,dim])';
            inregion=(box(1,1)<data(:,1))&(data(:,1)<box(1,2));
            for i=2:dim
                inregion=inregion&(box(i,1)<data(:,i))&(data(:,i)<box(i,2));
            end

            subdata=data(inregion,:);
            nt=sum(inregion);

            dimuse=randsample(1:dim,fdim); %for random forest, only use some dimensions
            subdata=[subdata(:,dimuse), subdata(:,end)];

            errs1=9*nt*ones(nt,dim);
            errs2=errs1;
            n1s=errs1;
            n2s=errs1;
            yhat1s=errs1;
            yhat2s=errs1;
            for d=1:fdim
                sortsubdata=sortrows(subdata,d);
                for i=m:nt-m
                    s=sortsubdata(i,d);
                    inr1=subdata(:,d)<=s;
                    n1s(i,d)=sum(inr1);
                    inr2=subdata(:,d)>s;
                    n2s(i,d)=sum(inr2);
                    yhat1=mean(subdata(inr1,end));
                    yhat1s(i,d)=yhat1;
                    yhat2=mean(subdata(inr2,end));
                    yhat2s(i,d)=yhat2;
                    errs1(i,d)=(subdata(inr1,end)-yhat1)'*(subdata(inr1,end)-yhat1);
                    errs2(i,d)=(subdata(inr2,end)-yhat2)'*(subdata(inr2,end)-yhat2);
                end
            end
            errors=errs1+errs2;
            [mins,locs]=min(errors);
            [err0,d0]=min(mins);
            sortsubdata=sortrows(subdata,d0);
            ind0=locs(d0);
            s0=sortsubdata(ind0,d0);
            subset=unique(sortsubdata(:,d0));
            [nsub,~]=size(subset);
            where=find(subset==s0);

            d0t=dimuse(d0); %translate back into the true dimension for random forest
            if where<nsub
                s0=(s0+subset(where+1))/2;
            else
                s0=box(d0t,2);
            end
            if err0<err
                box1=box;
                box2=box;
                box1(d0t,2)=s0;
                box2(d0t,1)=s0;
                xminmax1=reshape(box1',[1,2*dim]);
                xminmax2=reshape(box2',[1,2*dim]);
                bb=[xminmax1,errs1(ind0,d0),splitters(k),n1s(ind0,d0),yhat1s(ind0,d0),2*m<=n1s(ind0,d0);xminmax2,errs2(ind0,d0),splitters(k),n2s(ind0,d0),yhat2s(ind0,d0),2*m<=n2s(ind0,d0)];
                regions=[regions;bb];
            end
            regions(splitters(k),end)=false;
        end
        tosplit=regions(:,end)==1;
        ntosplit=sum(tosplit);
    end

    %find leaves
    [nreg,~]=size(regions);
    branches=unique(regions(:,2*dim+2));
    branches=branches(2:end); %remove the extra "0" region, which doesn't acutally exist but is listed as the parent of the full dataset region
    [nbranch,~]=size(branches);
    nleaves=nreg-nbranch;
    leaves=setdiff((1:nreg)',branches);

    leafxminmax=regions(leaves,1:2*dim); %leaf regions

    oob=~ismember(data0,data,'rows'); %find oob data and estimate those
    noob=noob+oob;
    oobcount=sum(oob);
    oobdata=data0(oob,:);
    leafind=oobdata(:,1);
    for i=1:oobcount
        inleaf=(oobdata(i,1)>leafxminmax(:,1))&(oobdata(i,1)<=leafxminmax(:,2));
        for d=2:dim
            inleaf=inleaf&(oobdata(i,d)>leafxminmax(:,2*d-1))&(oobdata(i,d)<=leafxminmax(:,2*d));
        end
        leafind(i)=leaves(inleaf);
    end
    yhatsum(oob)=yhatsum(oob)+regions(leafind,2*dim+4);

    %get overall oob error given the bootstraps done so far
    everoob=noob>0;
    testy=data0(everoob,end);
    testyhat=yhatsum(everoob)./noob(everoob);
    resid=testy-testyhat;
    rss=resid'*resid;
    mses(boot)=rss/sum(everoob);

    %also get test error, which we can compare with oob
    leafind=testdata(:,1);
    usepoint=true(nt0,1);
    for i=1:nt0
        inleaf=(testdata(i,1)>leafxminmax(:,1))&(testdata(i,1)<=leafxminmax(:,2));
        for d=2:dim
            inleaf=inleaf&(testdata(i,d)>leafxminmax(:,2*d-1))&(testdata(i,d)<=leafxminmax(:,2*d));
        end
        if sum(inleaf)==0
            usepoint(i)=false;
        else
            leafind(i)=leaves(inleaf);
        end
    end
    yhattest=regions(leafind(usepoint),2*dim+4);
    testyhatsum(usepoint)=testyhatsum(usepoint)+yhattest;
    resid=testyhatsum(usepoint)/boot-testdata(usepoint,end);
    rss=resid'*resid;
    testmses(boot)=rss/sum(usepoint);

end
plot(1:boots,mses,1:boots,testmses)

%

%boosting for regression trees

m=4; %define minimum number of points for a leaf of the tree.  for boosting this probably won't be used, as parameter J below sets max number of leaves and should be smallish
boosts=100; %the number of boosts to perform.  Afterward you can go back and pick an optimal value via cross validation or the test mse in this example
J=10; %max number of leaves
nu=.1; %learning rate, should be <=1 and small is usually better

mses=zeros(boosts,1); %various storage for results 
testmses=mses;
yhatsum=zeros(nt0,1);
yhattest=yhatsum;

data=data0; %start with original dataset

for boost=1:boosts %begin boosting loop

    nt=nt0;
    xminmax=[min(data0(:,1))-.1,max(data0(:,1))+.1]; %bounding box and other items for full dataset (root of tree)
    for i=2:dim
        xminmax=[xminmax,[min(data0(:,i))-.1,max(data0(:,i))+.1]];
    end
    yhat=mean(data(:,end));
    err=(data(:,end)-yhat)'*(data(:,end)-yhat);
    regions=[xminmax,err,0,nt,yhat,true]; %track current regions by bounding box, error before split, index of parent region, number of points, yhat, and flag on whether split or not yet

    tosplit=regions(:,end)==1; %find regions still to be split
    ntosplit=sum(tosplit);
    nleaves=1;

    while (ntosplit>0)&&(nleaves<J/2) %only make trees with at most J leaves and ability to further split
        splitters=find(tosplit);
        for k=1:ntosplit
            region=regions(splitters(k),:);
            err=region(2*dim+1);
            xminmax=region(1,1:2*dim);
            box=reshape(xminmax,[2,dim])';
            inregion=(box(1,1)<data(:,1))&(data(:,1)<box(1,2));
            for i=2:dim
                inregion=inregion&(box(i,1)<data(:,i))&(data(:,i)<box(i,2));
            end

            subdata=data(inregion,:);
            nt=sum(inregion);

            errs1=9*nt*ones(nt,dim);
            errs2=errs1;
            n1s=errs1;
            n2s=errs1;
            yhat1s=errs1;
            yhat2s=errs1;
            for d=1:dim
                sortsubdata=sortrows(subdata,d);
                for i=m:nt-m
                    s=sortsubdata(i,d);
                    inr1=subdata(:,d)<=s;
                    n1s(i,d)=sum(inr1);
                    inr2=subdata(:,d)>s;
                    n2s(i,d)=sum(inr2);
                    yhat1=mean(subdata(inr1,end));
                    yhat1s(i,d)=yhat1;
                    yhat2=mean(subdata(inr2,end));
                    yhat2s(i,d)=yhat2;
                    errs1(i,d)=(subdata(inr1,end)-yhat1)'*(subdata(inr1,end)-yhat1);
                    errs2(i,d)=(subdata(inr2,end)-yhat2)'*(subdata(inr2,end)-yhat2);
                end
            end
            errors=errs1+errs2;
            [mins,locs]=min(errors);
            [err0,d0]=min(mins);
            sortsubdata=sortrows(subdata,d0);
            ind0=locs(d0);
            s0=sortsubdata(ind0,d0);
            subset=unique(sortsubdata(:,d0));
            [nsub,~]=size(subset);
            where=find(subset==s0);

            if where<nsub
                s0=(s0+subset(where+1))/2;
            else
                s0=box(d0,2);
            end
            if err0<err
                box1=box;
                box2=box;
                box1(d0,2)=s0;
                box2(d0,1)=s0;
                xminmax1=reshape(box1',[1,2*dim]);
                xminmax2=reshape(box2',[1,2*dim]);
                bb=[xminmax1,errs1(ind0,d0),splitters(k),n1s(ind0,d0),yhat1s(ind0,d0),2*m<=n1s(ind0,d0);xminmax2,errs2(ind0,d0),splitters(k),n2s(ind0,d0),yhat2s(ind0,d0),2*m<=n2s(ind0,d0)];
                regions=[regions;bb];
            end
            regions(splitters(k),end)=false;
        end
        tosplit=regions(:,end)==1;
        ntosplit=sum(tosplit);

        %find nleaves
        [nreg,~]=size(regions);
        branches=unique(regions(:,2*dim+2));
        branches=branches(2:end); %remove the extra "0" region, which doesn't acutally exist but is listed as the parent of the full dataset region
        [nbranch,~]=size(branches);
        nleaves=nreg-nbranch;

    end

    %find leaf regions
    leaves=setdiff((1:nreg)',branches);
    leafxminmax=regions(leaves,1:2*dim); %leaf region boxes

    
    leafind=data(:,1); %now predict all datapoints
    for i=1:nt0
        inleaf=(data(i,1)>leafxminmax(:,1))&(data(i,1)<=leafxminmax(:,2));
        for d=2:dim
            inleaf=inleaf&(data(i,d)>leafxminmax(:,2*d-1))&(data(i,d)<=leafxminmax(:,2*d));
        end
        leafind(i)=leaves(inleaf);
    end
    yhat=regions(leafind,2*dim+4);
    yhatsum=yhatsum+nu*yhat;
    resid=data0(:,end)-yhatsum;
    mses(boost)=resid'*resid/nt0;
    data(:,end)=resid; %next tree will train on current residuals

    %get test mse
    usepoint=true(nt0,1);
    for i=1:nt0
        inleaf=(testdata(i,1)>leafxminmax(:,1))&(testdata(i,1)<=leafxminmax(:,2));
        for d=2:dim
            inleaf=inleaf&(testdata(i,d)>leafxminmax(:,2*d-1))&(testdata(i,d)<=leafxminmax(:,2*d));
        end
        if sum(inleaf)==0
            usepoint(i)=false;
        else
            leafind(i)=leaves(inleaf);
        end
    end
    yhat=regions(leafind(usepoint),2*dim+4);
    yhattest(usepoint)=yhattest(usepoint)+nu*yhat;
    resid=testdata(usepoint,end)-yhattest(usepoint);
    testmses(boost)=resid'*resid/sum(usepoint);

end
plot(1:boosts,mses,1:boosts,testmses)
[mintestmse,optimalB]=min(testmses)

%

%These bits are for exploring HME via the simple example from class
%generate the data for our simple HME example
K=2; %number of expert nodes in the end
gamma=[0;-1;2]; %the gamma coefficients of the single split point
betas=[[5;4;-1],[1;2;3]]; %the beta coefficients for the expert nodes

prob=@(y,mn,var) exp(-(y-mn).^2/2/var)/sqrt(2*pi*var);

n=5000; %number of data points
xdata=rand(n,2)*10; %just randomly make up some x values for each point
xmat=[ones(n,1),xdata];
p=exp(xmat*gamma); %given the x values and gamma, start to find branch probabilities
p=p./(1+p); %actual probability for branch 1 (vs 0), from logistic formula
truedel=rand(n,1)<p; %now assign each point to one of the two branches probabilistically
branch=truedel+1; %just translate this into one of the two experts
ydata=zeros(n,1); %generate y values
for i=1:n
    ydata(i)=xmat(i,:)*betas(:,branch(i));
end
s2=.1; %variance of error in y values
ydata=ydata+randn(n,1)*sqrt(s2);

logl=truedel'*log(p)+(1-truedel)'*log(1-p)+sum(log(prob(ydata(branch==2),xmat(branch==2,:)*betas(:,2),s2)))+sum(log(prob(ydata(branch==1),xmat(branch==1,:)*betas(:,1),s2)));

%

%Use EM to get parameters of HME

gammahat=[-5;0;1]; %just split the data in half along x2 axis
%gammahat=gamma; %can also just start with the true gamma value to see
%what happens
betahatfull=(xmat'*xmat)\(xmat'*ydata); %linear regression on full dataset, for initial guess of beta values
betashat=[betahatfull,betahatfull];
%betashat=betas; %could start with true answer to see what happens


%do an E step starting with initial guesses
weights=exp(xmat*gammahat); %given the x values and gamma, start to find latent variable probabilities
weights=weights./(1+weights);
weights=[1-weights,weights];
%now include the probability of expert nodes to complete calculation
weights=[weights(:,1).*prob(ydata,xmat*betashat(:,1),s2),weights(:,2).*prob(ydata,xmat*betashat(:,2),s2)];
for i=1:n
    if sum(weights(i,:))==0
        weights(i,:)=[0.5, 0.5];
    else
        weights(i,:)=weights(i,:)/sum(weights(i,:));
    end
end

iters=50; %how many iterations to do
changes=zeros(iters,1); %track changes to the beta coefficients over iterations
oldb=zeros(6,iters);

for iter=1:iters

    %E step
    weights=exp(xmat*gammahat); %given the x values and gamma, start to find latent variable probabilities
    weights=weights./(1+weights);
    weights=[1-weights,weights];
    weights=[weights(:,1).*prob(ydata,xmat*betashat(:,1),s2),weights(:,2).*prob(ydata,xmat*betashat(:,2),s2)];
    for i=1:n
        if sum(weights(i,:))==0
           weights(i,:)=[0.5, 0.5];
        else
            weights(i,:)=weights(i,:)/sum(weights(i,:));
        end
    end

    %M step
    oldb(1:3,iter)=betashat(:,1);
    oldb(4:6,iter)=betashat(:,2);
    for k=1:K %compute best fit betas using current latent variables and weighted OLS
        wmat=diag(weights(:,k));
        tmp=(xmat'*wmat*xmat)\(xmat'*wmat*ydata);
        betashat(:,k)=tmp;
    end
    gammahat=glmfit(xdata,weights(:,2),'binomial'); %compute best fit gamma using logistic regression with output variable y_i=delta_i (prob of being at expert "1" instead of "0")

    %track change in beta estimates over time
    changes(iter)=norm(oldb(:,iter)-[betashat(:,1);betashat(:,2)]);

end