%set up network
nhidden = [500 500]; %for layers 1 through M-1
ninput = 18; %9 nodes for focal player's symbol (0 absent, 1 present), 9 nodes for oppoent's
units = [ninput nhidden 9]; %number of units for layers 0 through M; output nodes correspond to all possible actions
M = length(units)-1; %number of weight layers
W = cell(M,1); %weights at all layers
for m=1:M
    W{m} = randn(units(m+1),units(m)+1)/10; %weight matrix; (size of postsynaptic layer) x (size of presynaptic layer plus bias term)
end
e = .001; %learning rate
T = 1; %choice temperature
V = tttOptimal; %9D array indicates which states lead to X wins, O wins, or ties under optimal play
win = [1 2 3;4 5 6;7 8 9;1 4 7;2 5 8;3 6 9;1 5 9;3 5 7]; %all winning combinations

%play
n = 100000; %number of games
testinterval = 1;
score = zeros(n/testinterval,1); %score on test games
for g=1:n %loop games
    if mod(g,100)==0,disp(g),end
    Delta = cell(M,1); for m=1:M,Delta{m}=zeros(units(m+1),units(m)+1);end %initialize Delta to track total weight updates
    s = ones(9,1); %initial game state; 1=blank, 2=X, 3=O
    over = 0; %tracks when game has ended
    move = 2; %2 for X's turn, 3 for O's turn
    Xstate = zeros(1,9); %will track last state X moved from
    Ostate = zeros(1,9); %will track last state O moved from
    while ~over
        if move==2
            [v,~] = feedforward([s==2;s==3],W); %get network activation from X's perspective
            x = datasample(find(s==1),1,'weights',exp(v{M}(s==1)/T)); %softmax selection from available actions
            Xstate = s; %save state X moved from, for learning after O moves
            s(x) = 2; %update state
            if sum(Ostate)>0 %O has previously moved in this game; get learning update
                if max(sum(s(win)==2,2)) == 3 %X won
                    Oval = -1;
                    over = 1; %game will end after this learning update
                else
                    [v,~] = feedforward([s==3;s==2],W); %get network activation in new state from O's perspective
                    Oval = max(v{M}.*(s==1)); %maximum Q-value among available actions
                end
                [v,a] = feedforward([Ostate==3;Ostate==2],W); %get network activation from before O's previous move
%                 oldvalue = v{M}(o);
                feedback = v{M}; feedback(o) = Oval; %only learn about O's previous move (no prediction error on other moves)
                update = backprop([Ostate==3;Ostate==2],v,a,W,feedback); %get weight updates by backpropagation
                for m=1:M,Delta{m} = Delta{m} + update{m};end %add updates to running sum
%                 disp('learn from O')
%                 disp(['current state: ' num2str(s') ', last X move: ' num2str(x)])
%                 disp(['previous O state: ' num2str(Ostate') ', last O move: ' num2str(o)])
%                 [v,~] = feedforward([Ostate==3;Ostate==2],W); newvalue = v{M}(o);
%                 disp(['old value: ' num2str(oldvalue) ', target: ' num2str(Oval) ', newvalue: ' num2str(newvalue)])
            end
            move = 3; %O's turn next
        else %move==3: O's turn
            [v,~] = feedforward([s==3;s==2],W); %get network activation from O's perspective
            o = datasample(find(s==1),1,'weights',exp(v{M}(s==1)/T)); %softmax selection from available actions
            Ostate = s; %save state O moved from, for learning after X moves
            s(o) = 3; %update state
            %now do learrning for X (who has necessarily already moved)
            if max(sum(s(win)==3,2)) == 3 %O won
                Xval = -1;
                over = 1; %game will end after this learning update
            else
                [v,~] = feedforward([s==2;s==3],W); %get network activation in new state from X's perspective
                Xval = max(v{M}.*(s==1)); %maximum Q-value among available actions
            end
            [v,a] = feedforward([Xstate==2;Xstate==3],W); %get network activation from before X's previous move
%                 oldvalue = v{M}(x);
            feedback = v{M}; feedback(x) = Xval; %only learn about X's previous move (no prediction error on other moves)
            update = backprop([Xstate==2;Xstate==3],v,a,W,feedback); %get weight updates by backpropagation
            for m=1:M,Delta{m} = Delta{m} + update{m};end %add updates to running sum
%                 disp('learn from X')
%                 disp(['current state: ' num2str(s') ', last O move: ' num2str(o)])
%                 disp(['previous X state: ' num2str(Xstate') ', last X move: ' num2str(x)])
%                 [v,~] = feedforward([Xstate==2;Xstate==3],W); newvalue = v{M}(x);
%                 disp(['old value: ' num2str(oldvalue) ', target: ' num2str(Xval) ', newvalue: ' num2str(newvalue)])
            move = 2; %X's turn next
        end
        if sum(s==1)==0 %no more moves available
            over = 1; %game over
        end
    end
    for m=1:M,W{m} = W{m} + e*Delta{m};end %update weights
    if mod(g,testinterval)==0 %time to test model's performance against optimal player
        score(g/testinterval) = 0; %count total number of moves before opponent can guarantee victory, over 2 games
        %first play as X
        s = ones(9,1);
        over = 0;
        while ~over
            [v,~] = feedforward([s==2;s==3],W); %get network activation from X's perspective
            x = datasample(find(s==1 & v{M}==max(v{M}(s==1))),1); %exploit: choose randomly from top-valued available actions
            s(x) = 2; %update state
            if V((3.^(0:8))*s-9840) == -1 %look up optimal value of new state (using vector unrolling of V) to see if O can force a win
                over = 1; %agent has lost
            else
                score(g/testinterval) = score(g/testinterval) + 1; %give credit for a non-losing move
                Ooptions = zeros(9,1); %all optimal moves for O
                for o=find(s==1) %check all open spaces
                    snew = s; snew(o) = 3; %candidate state for opponent to move to
                    if V((3.^(0:8))*snew-9840) == 0 %optimal moves are ones that lead to ties
                        Ooptions(o) = 1; %record o as an optimal move
                    end
                end
                o = datasample(find(Ooptions),1); %randomly select from optimal moves
                s(o) = 3; %update state
            end
            if sum(s==1)==1 %only one more move available; no reason to simulate it
                score(g/testinterval) = score(g/testinterval) + 1; 
                over = 1; %game over
            end
        end
        %now play as O
        s = ones(9,1); s(randi(9)) = 2; %random initial move by opponent
        over = 0;
        while ~over
            [v,~] = feedforward([s==3;s==2],W); %get network activation from O's perspective
            o = datasample(find(s==1 & v{M}==max(v{M}(s==1))),1); %exploit: choose randomly from top-valued available actions
            s(o) = 3; %update state
            if V((3.^(0:8))*s-9840) == 1 %look up optimal value of new state (using vector unrolling of V) to see if X can force a win
                over = 1; %agent has lost
            elseif sum(s==1)==1 %only one more move available (by opponent); no reason to simulate it
                score(g/testinterval) = score(g/testinterval) + 1; %give credit for a non-losing move
                over = 1; %game over
            else
                score(g/testinterval) = score(g/testinterval) + 1; %give credit for a non-losing move
                Xoptions = zeros(9,1); %all optimal moves for X
                for x=find(s==1) %check all open spaces
                    snew = s; snew(x) = 2; %candidate state for opponent to move to
                    if V((3.^(0:8))*snew-9840) == 0 %optimal moves are ones that lead to ties
                        Xoptions(x) = 1; %record x as an optimal move
                    end
                end
                x = datasample(find(Xoptions),1); %randomly select from optimal moves
                s(x) = 2; %update state
            end
        end
    end
end

plot(mean(reshape(score,100,length(score)/100)))
