%% Set variables and create crease pattern

numTrain = 20;
numTest = 100;
numSamples = 300;
scp = 13;
DotRuleThresh = 0.85;

Colors = [];
hold on
for i = 1:1000
    h = plot(1:10,2:11);
    c = get(h,'Color');
    Colors = [Colors; c];
end
close()

% Setup crease pattern
CP = CreasePattern;

CP.intvertices = [    
    0.1589   -0.1662
    0.9744    0.0373
   -0.0161    0.9701
    0.8594    1.0604
    ];

CP.extvertices = [
    0.4631   -0.5981
    2.0702    0.6728
    0.5556    2.0472
   -0.9172    0.5667
    ];

CP.alledges = [
     1     2
     3     4
     1     3
     2     4
     1     5
     2     5
     2     6
     4     6
     3     7
     4     7
     1     8
     3     8
     1     4
    ];

CP = CP.fill_essentials();

% Setup force directions
R1 = [0.4434    0.0855   -0.8418   -0.2862   -0.0741];
R2 = [-0.7331   -0.0908   -0.2172   -0.5880    0.2476];

% Import training and text forces (data)
RedFS = importdata('RedFS.txt');
RedNS = importdata('RedNS.txt');
BlueFS = importdata('BlueFS.txt');
BlueNS = importdata('BlueNS.txt');


%% Plot train\test force sets in the null space

figure
hold on

Theta1 = abs(RedNS'*R1');
Theta2 = abs(RedNS'*R2');
h2 = scatter(Theta1(numTrain+1:end), Theta2(numTrain+1:end), 'o', 'filled', 'MarkerFaceColor', Colors(1,:));
h2.MarkerFaceAlpha = 0.5;
h2.MarkerFaceAlpha = 0.09;
h1 = scatter(Theta1(1:numTrain), Theta2(1:numTrain),80, 'D', 'filled', 'MarkerFaceColor', Colors(1,:));
h1.MarkerFaceAlpha = 1.0;

Theta1 = abs(BlueNS'*R1');
Theta2 = abs(BlueNS'*R2');
h2 = scatter(Theta1(numTrain+1:end), Theta2(numTrain+1:end), 'o', 'filled', 'MarkerFaceColor', Colors(2,:));
h2.MarkerFaceAlpha = 0.5;
h2.MarkerFaceAlpha = 0.09;
h1 = scatter(Theta1(1:numTrain), Theta2(1:numTrain),80, 'D', 'filled', 'MarkerFaceColor', Colors(2,:));
h1.MarkerFaceAlpha = 1.0;

%plot([0,1],[0,1],'k--');

axis equal
xlabel('F_1','FontSize',16)
ylabel('F_2','FontSize',16)
xlim([0 1]);
ylim([0 1]);



%% Check the folding results for forces on soft patterns (no stiffness on creases)
SThreshold = 0.95; %This is a threshold used to check the accuracy of the arrived configuration
FS = zeros(scp,1);
springs = ones(scp,1)*0.0;

rho = 0.5;
DotThresh = 0.975;

sampleModes = [];
sampleEnergies = [];
sampleAtt = [];

RedAttractors = []; RedFVs = [];
BlueAttractors = []; BlueFVs = [];

for i = 1:numTrain
    disp([length(sampleEnergies)])
    [RHO,fvR] = CP.fn_fold_relax_fmincon_springs(RedFS(:,i) * rho, CP, CP.thresholds, springs); disp(abs(max(RHO))); RHO = RHO/norm(RHO);
    RedAttractors = [RedAttractors RHO];
    RedFVs = [RedFVs fvR];
    ee = exp(fvR);
    
    if isempty(sampleModes)
        sampleModes = [RHO];
        sampleModesCount = [1];
        sampleEnergies = [ee];
        sampleAtt = [1];
    else
        final_dot = abs(RHO'*sampleModes)>DotThresh;
        if ~any(final_dot)
            sampleModes = [sampleModes RHO];
            sampleModesCount = [sampleModesCount 1];
            sampleEnergies = [sampleEnergies ee];
            sampleAtt = [sampleAtt length(sampleEnergies)];
        else
            ai = find(final_dot);
            sampleModesCount(ai) = sampleModesCount(ai) + 1;
            sampleAtt = [sampleAtt ai(1)];
        end
    end
end
for i = 1:numTrain
    disp([length(sampleEnergies)])
    [RHO,fvB] = CP.fn_fold_relax_fmincon_springs(BlueFS(:,i) * rho, CP, CP.thresholds, springs); disp(abs(max(RHO))); RHO = RHO/norm(RHO);
    BlueAttractors = [BlueAttractors RHO];
    BlueFVs = [BlueFVs fvB];
    ee = exp(fvB);
    
    if isempty(sampleModes)
        sampleModes = [RHO];
        sampleModesCount = [1];
        sampleEnergies = [ee];
        sampleAtt = [1];
    else
        final_dot = abs(RHO'*sampleModes)>DotThresh;
        if ~any(final_dot)
            sampleModes = [sampleModes RHO];
            sampleModesCount = [sampleModesCount 1];
            sampleEnergies = [sampleEnergies ee];
            sampleAtt = [sampleAtt length(sampleEnergies)];
        else
            ai = find(final_dot);
            sampleModesCount(ai) = sampleModesCount(ai) + 1;
            sampleAtt = [sampleAtt ai(1)];
        end
    end
end
sampleModesCount = sampleModesCount/numSamples;

% Red and Blue reference states
RedMA = mean(RedAttractors')';    RedMA = RedMA/norm(RedMA);
BlueMA = mean(BlueAttractors')';  BlueMA = BlueMA/norm(BlueMA);

disp(['done!', length(sampleEnergies)])

%% Plot the training data classification result given no stiffness

figure
hold on

ANS = [RedNS(:,1:numTrain), BlueNS(:,1:numTrain), RedNS(:,numTrain+1:end), BlueNS(:,numTrain+1:end)];

Theta1 = ANS'*R1';
Theta2 = ANS'*R2';

for i = 1:max(sampleAtt)
    inds = find(sampleAtt == i) ;

    h1 = plot(Theta1(inds), Theta2(inds), 'o', 'MarkerSize',6);
    set(h1, 'markerfacecolor', get(h1, 'color'));   
end
axis equal

xlabel('F_1')
ylabel('F_2')
xlim([0,1])
ylim([0,1])


%% Train the sheet using the training forces


numTrain = 20;
numTest = 100;

ATrains = [];
ATests = [];
NAtts = [];
DotRules = [];

Att1 = [];
Att2 = [];
RL = [];
BL = [];
Same = [];


% Training reds and blues without feedback

num = numTrain+numTest;

rounds = 4000 * numTrain/20.;
Epochs = rounds / (numTrain*2);
LR = 1.e-4 * 20./numTrain;

Init_K2 = 0.02;


rho = 0.5;

CP.thresholds.fmin_tolfun = 1.e-4;
CP.thresholds.fmin_tolx = 1.e-4;

%K = 0;
%while abs(mean(K) - Init_K2) > 0.0005

K = ones(scp,1)*Init_K2; %Choose a uniform stiffness across the Origami actual creases
Ks = [K]; 

for r = 1:rounds

    %ChooseColor = ceil(rand()*2);
    %ChooseColor = rand() < pRed;
    %ChooseFNum = ceil(rand()*numTrain);

    ChooseColor = mod(r,2);
    ChooseFNum = ceil((1+mod(r-1,numTrain*2))/2);

    if ChooseColor == 1   % reds
        Teacher = RedFS(:,ChooseFNum);          
    else % blues
        Teacher = BlueFS(:,ChooseFNum);
    end 

    [fa,~] = CP.fn_fold_relax_fmincon_springs(Teacher * rho, CP, CP.thresholds, K);
    %fa = Teacher;
    RHO = fa; RHO = RHO/norm(RHO);


    if ChooseColor == 1   % reds
        DM = RedMA'*RHO;  % compute dot product with average red attractor 
    else % blues
        DM = BlueMA'*RHO;  % compute dot product with average blue attractor
    end

    if abs(DM) > DotRuleThresh  % check if folded state resembles mean force of right color
        K = K - LR * (abs(RHO)).^2; % make creases softer
    else
        K = K + LR * (abs(RHO)).^2; % make creases stiffer
    end

    K(K < 0) = 0.;
    Ks = [Ks K]; % Create array for storing the K values after each training round
    
    if max(K) == 0
        break
    end
    
    disp([r, RHO'*Teacher, DM, 1000*min(K),1000*max(K),1000*mean(K)])

end




%% Check the folding results for forces on stiff patterns

SThreshold = 0.95; %This is a threshold used to check the accuracy of the arrived configuration
FS = zeros(scp,1);

rho = 0.5;

Epochs = rounds / (numTrain*2);
GetRounds = [1, 1+linspace(numTrain*2,rounds,Epochs)];
%GetRounds = [1:200:4001];
%GetRounds = [200];

AccTrain = [];
AccTest = [];

RAccTrain = [];
RAccTest = [];
BAccTrain = [];
BAccTest = [];
Same = [];
RL = [];
BL = [];
Att1 = [];
Att2 = [];

NumAtt = [];

for TR = GetRounds
    epoch = (TR-1)/(numTrain*2);
    disp(epoch);

%TR = 200;
if TR == 1
    springs = Ks(:,TR)*0.;
else
    springs = Ks(:,TR);
end

sampleModes = [];
sampleModesCount = [];
sampleAtt1 = [];
sampleAtt2 = [];

RedAttractors = []; 
RedCounts1 = []; RedCounts2 = [];
BlueAttractors = [];
BlueCounts1 = []; BlueCounts2 = [];

TrainTest = [1:numTrain, length(RedNS)-numTest+1:length(RedNS)];

for i = TrainTest
    [RHO,~] = CP.fn_fold_relax_fmincon_springs(RedFS(:,i) * rho, CP, CP.thresholds, springs); RHO = RHO/norm(RHO);
    RedAttractors = [RedAttractors RHO];
    
    if isempty(sampleModes)
        sampleModes = [RHO];
        sampleModesCount = [1];
        if i <= numTrain
            sampleAtt1 = [sampleAtt1 1];
            RedCounts1 = [RedCounts1 1];
        else
            sampleAtt2 = [sampleAtt2 1];
            RedCounts2 = [RedCounts2 1];
        end
    else
        final_dot = abs(RHO'*sampleModes)>DotThresh;
        if ~any(final_dot)
            sampleModes = [sampleModes RHO];
            sampleModesCount = [sampleModesCount 1];
            if i <= numTrain
                sampleAtt1 = [sampleAtt1 length(sampleModesCount)];
                RedCounts1 = [RedCounts1 length(sampleModesCount)];
            else
                sampleAtt2 = [sampleAtt2 length(sampleModesCount)];
                RedCounts2 = [RedCounts2 length(sampleModesCount)];
            end
        else
            ai = find(final_dot);
            sampleModesCount(ai) = sampleModesCount(ai) + 1;
            if i <= numTrain
                sampleAtt1 = [sampleAtt1 ai(1)];
                RedCounts1 = [RedCounts1 ai(1)];
            else
                sampleAtt2 = [sampleAtt2 ai(1)];
                RedCounts2 = [RedCounts2 ai(1)];
            end
        end
    end
end

TrainTest = [1:numTrain, length(BlueNS)-numTest+1:length(BlueNS)];
for i = TrainTest
    [RHO,~] = CP.fn_fold_relax_fmincon_springs(BlueFS(:,i) * rho, CP, CP.thresholds, springs); RHO = RHO/norm(RHO);
    BlueAttractors = [BlueAttractors RHO];
    
    if isempty(sampleModes)
        sampleModes = [RHO];
        sampleModesCount = [1];
        if i <= numTrain
            sampleAtt1 = [sampleAtt1 1];
            BlueCounts1 = [BlueCounts1 1];
        else
            sampleAtt2 = [sampleAtt2 1];
            BlueCounts2 = [BlueCounts2 1];
        end
    else
        final_dot = abs(RHO'*sampleModes)>DotThresh;
        if ~any(final_dot)
            sampleModes = [sampleModes RHO];
            sampleModesCount = [sampleModesCount 1];
            if i <= numTrain
                sampleAtt1 = [sampleAtt1 length(sampleModesCount)];
                BlueCounts1 = [BlueCounts1 length(sampleModesCount)];
            else
                sampleAtt2 = [sampleAtt2 length(sampleModesCount)];
                BlueCounts2 = [BlueCounts2 length(sampleModesCount)];
            end
        else
            ai = find(final_dot);
            sampleModesCount(ai) = sampleModesCount(ai) + 1;
            if i <= numTrain
                sampleAtt1 = [sampleAtt1 ai(1)];
                BlueCounts1 = [BlueCounts1 ai(1)];
            else
                sampleAtt2 = [sampleAtt2 ai(1)];
                BlueCounts2 = [BlueCounts2 ai(1)];
            end
        end
    end
end
sampleModesCount = sampleModesCount/numSamples;

[RedMode, RFreq] = mode(RedCounts1);
[BlueMode, BFreq] = mode(BlueCounts1);
SameMode = (RedMode == BlueMode);
TrainingAccuracy = (1-SameMode) * ( sum(RedCounts1==RedMode) + sum(BlueCounts1==BlueMode) ) / (2. * numTrain);
TestAccuracy = (1-SameMode) * ( sum(RedCounts2==RedMode) + sum(BlueCounts2==BlueMode) ) / (2. * numTest);

if SameMode
    if RFreq >= BFreq
        RedMode = RedMode;
        if BFreq < numTrain
            BlueMode = mode(setdiff(BlueCounts1, RedMode));
        else
            BlueMode = BlueMode;
        end
    else
        BlueMode = BlueMode;
        if RFreq < numTrain
            RedMode = mode(setdiff(RedCounts1, BlueMode));
        else
            RedMode = RedMode;
        end
    end
    RTrainingAccuracy = ( sum(RedCounts1==RedMode) ) / (numTrain);
    BTrainingAccuracy = ( sum(BlueCounts1==BlueMode) ) / (numTrain);
    RTestAccuracy = ( sum(RedCounts2==RedMode)) / (numTest);
    BTestAccuracy = ( sum(BlueCounts2==BlueMode)) / (numTest);
else
    RTrainingAccuracy = ( sum(RedCounts1==RedMode) ) / (numTrain);
    BTrainingAccuracy = ( sum(BlueCounts1==BlueMode) ) / (numTrain);
    RTestAccuracy = ( sum(RedCounts2==RedMode) ) / (numTest);
    BTestAccuracy = ( sum(BlueCounts2==BlueMode) ) / (numTest);
end

AccTrain = [AccTrain TrainingAccuracy];
AccTest = [AccTest TestAccuracy];

RAccTrain = [RAccTrain RTrainingAccuracy];
RAccTest = [RAccTest RTestAccuracy];
BAccTrain = [BAccTrain BTrainingAccuracy];
BAccTest = [BAccTest BTestAccuracy];

disp([epoch, length(sampleModesCount), TrainingAccuracy, TestAccuracy])
disp([RTrainingAccuracy, RTestAccuracy, BTrainingAccuracy, BTestAccuracy])
NumAtt = [NumAtt length(sampleModesCount)];

Same = [Same SameMode];
RL = [RL RedMode];
BL = [BL BlueMode];
Att1 = [Att1; sampleAtt1];
Att2 = [Att2; sampleAtt2];


%disp(['done!'])

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Plot the blobs in the null space for stiff patterns

Order = 1:max(sampleAtt2);
Order = setdiff(Order, [RedMode, BlueMode]);
if SameMode
    Order = [RedMode, BlueMode, Order];
else
    Order = [RedMode, BlueMode, Order];
end

%Plot Training Date
ANS = [RedNS(:,1:numTrain), BlueNS(:,1:numTrain)];
Theta1 = abs(ANS'*R1');
Theta2 = abs(ANS'*R2');

ANS2 = [RedNS(:,length(RedNS)-numTest+1:length(RedNS)), BlueNS(:,length(BlueNS)-numTest+1:length(BlueNS))];
Theta12 = abs(ANS2'*R1');
Theta22 = abs(ANS2'*R2');

figure
hold on
for o = 1:length(Order)
    inds = find(sampleAtt2 == Order(o)) ;
    h2 = scatter(Theta12(inds), Theta22(inds), 'o', 'filled', 'MarkerFaceColor', Colors(o,:));
    h2.MarkerFaceAlpha = 0.2;
    
    inds = find(sampleAtt1 == Order(o)) ;
    h1 = scatter(Theta1(inds), Theta2(inds),80, 'D', 'filled', 'MarkerFaceColor', Colors(o,:));
    h1.MarkerFaceAlpha = 1.0;
end
axis equal

xlabel('F_1','FontSize',16)
ylabel('F_2','FontSize',16)
xlim([0 1]);
ylim([0 1]);

tl = strcat('Epoch',{' '}, num2str(epoch));
title(tl)
fn = strcat('Epoch', num2str(epoch), '.png');
saveas(gcf, fn);
close
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

end

ATrains = [ATrains; AccTrain];
ATests = [ATests; AccTest];
NAtts = [NAtts; NumAtt];



%% Train and Test errors
figure
hold on
f1 = AccTrain > -1;
X = 0:length(f1)-1;
plot(X(f1), AccTrain(f1), 'ok-')
plot(X(f1), AccTest(f1), 'ok--')

xlabel('Epoch');
ylabel('Accuracy');
legend('Training', 'Test');