% Bandit a deux bras (TP6) (peut se generaliser au bandit a q bras)

clear;
clf;
A=input('Probabilite de gain associee au levier A:');
B=input('Probabilite de gain associee au levier B:');
n=input('Nombre d iterations:');
k=input('Nombre de realisations:');

TLC=[];
GAIN=[];
for i=1:k
   
% Initialisation des variables
NA=0;NB=0;TA=0;TB=0;
G=[];    %vecteur gain
ESTA=[];ESTB=[];    %vecteurs des estimations
d=2;
p=round(log(n)/log(d));
PA=[2:2:p];PB=[1:2:p];
IA=d.^PA;IB=d.^PB;

% On initialise le processus croissant avec une chance sur deux pour chaque levier
r=rand;
if r<0.5
   U=A;NA=1;
else 
   U=B;NB=1;
end

% Debut du programme
for i=1:n     %calcul de la sortie a l instant i
   r1=rand;
   if r1<1-U
      X=0;
   else
      X=1;
   end


% Estimateurs des moindres carres a l instant i
   if (U==A & X==1)
      TA=TA+(1-TA)/NA;
   else if (U==A & X==0)
         TA=TA-TA/NA;
      else TA=TA;
      end
   end
   if (U==B & X==1)
      TB=TB+(1-TB)/NB;
   else if (U==B & X==0)
         TB=TB-TB/NB;
      else TB=TB;
      end
   end

% Choix du levier a l instant i
hA=ismember(i,IA);    
hB=ismember(i,IB);
if hA==1      %on force le joueur a choisir le levier que l on veut
   U=A;NA=NA+1;
else if hB==1
      U=B;NB=NB+1;
      else if TA>=TB
   U=A;NA=NA+1;
else
   U=B;NB=NB+1;
end
end
end

% Stockage
G=[G X];
ESTA=[ESTA TA];
ESTB=[ESTB TB];

% Fin de la boucle for
end

TLC=[TLC sqrt(k)*(TA-A)];
GAIN=[GAIN sqrt(k)*(sum(G)-max(A,B))];
end

% Sorties graphiques
subplot(5,1,1)  %trace de courbes sur un meme graphique
plot( cumsum(G)./[1:length(G)],'b' )
title('Strategie optimale pour le bandit a deux bras')
ylabel('Gain moyen')
hold on
C=max(A,B);
plot( C*ones(n,1),'r--' )
grid

subplot(5,1,2)
plot( ESTA,'r' )
ylabel('Estimation de A')
hold on
plot( A*ones(n,1),'b--' )
grid

subplot(5,1,3)
plot( ESTB,'r' )
ylabel('Estimation de B')
xlabel('Nombre de realisations')
hold on
plot( B*ones(n,1),'b--' )
grid
hold off

% Theoreme central limite
subplot(5,1,4)
histo(TLC,sqrt(k),0,1)

subplot(5,1,5)
histo(GAIN,sqrt(k),0,1)

