% Code to fit the stationary probability distribution to the voting data.
% We then maximise the path likelihood to fit the timescale.

clear all

% Group sizes
N_vals = [12,13,16,18,19,20,22,25,28,31,33];
N_vals = [31]; % Change here to fit the parameters to the desired group size.

for Nit = 1:length(N_vals)

    N = N_vals(Nit);

    % read data from .csv files
    data = readtable(".\"+N+"_1.csv");
    %data = readtable(".\"+N+"_2.csv"); % For experiment 19_2
    
    k = 1:120;
    
    % define the headings to gather the data from the dataframe
    lab_evotes = "my_voting_"+k+"_player_election_vote";
    lab_groupA = "my_voting_"+k+"_group_A_votes";
    lab_groupB = "my_voting_"+k+"_group_B_votes";
    lab_groupAbs = "my_voting_"+k+"_group_abstain_votes";
    
    % read votes from each participant, number of A votes, B votes and
    % abstentions in the group at each round of the game
    elec_vote = data{:,lab_evotes};
    avotes = data{:,lab_groupA};
    bvotes = data{:,lab_groupB};
    absvotes = data{:,lab_groupAbs};
    
    avotes = reshape(avotes(~isnan(avotes)),[N,120]);
    bvotes = reshape(bvotes(~isnan(bvotes)),[N,120]);
    absvotes = reshape(absvotes(~isnan(absvotes)),[N,120]);
    
    evote = reshape(elec_vote(~isnan(elec_vote)),[N,120]);
    
    % define alignment (polarisation)
    align = avotes(1,:) - bvotes(1,:);
    
    % computes the empirical SPD (used to fit the parameters of the
    % theoretical SPD)
    pdf_vals = histcounts([align -align],'BinEdges',-(N+0.5):(N+0.5));
    pdf_x = -N:N;

    % define function to be minimised, taking logs helps with convergence
    % the function is the squared error between the empirical SPD and the
    % theoretical SPD given a set of parameters.
    errorfun = @(cf) log(log(fit_spd(pdf_x,pdf_vals,N,cf)));
    
    % define options for the minimisation algorithm
    optionsswa = optimset('MaxFunEvals',100000,'TolFun',1e-500,'Display','None','MaxTime',Inf,'ObjectiveLimit',-Inf,'UseParallel',0);
    
    tic

    % Upper bounds to optimise the convergence of the optimisation
    % algorithm. These have been tested extensively.
    % 12: [.01 2 .01 10 10 20 10 10 .1 inf]
    % 13: [1 .2 2 2 2 4 2 10 .01 inf]
    % 16: [.1 2 .01 2 2 10 10 10 .1 inf]
    % 18: [1 .2 2 2 2 4 2 10 .01 inf]
    % 19_1: [.01 .2 2 2 2 10 10 15 .01 inf]
    % 19_2: [.01 .2 2 4 2 10 10 12 .01 inf]
    % 20: [.01 .2 2 4 2 10 10 16 .01 inf]
    % 22: [.01 .2 2 4 2 10 10 12 .01 inf]
    % 25: [.01 2 .01 10 10 20 10 10 .1 inf]
    % 28: [.01 2 .01 10 12 20 12 10 .1 inf]
    % 31: [.01 2 .01 10 12 20 12 10 .1 inf]
    % 33: [.01 .2 2 4 2 10 10 12 .01 inf]

    coef = zeros(20,10);
    fval = zeros(20,1);
    % to find the optimal parameters, initialise the search 20 times and
    % get the parameter set with the minimum error
    parfor k=1:20
        [coef(k,:),fval(k),exitflag,output] = particleswarm(errorfun, 10, [0.00342 0.0019 0.0090 0 0.03 3.07 1.24 3.69 0], [.01 2 .01 10 12 20 12 10 .1 inf], optionsswa); %[1 .2 2 1 2 4 7 7 .1 inf]
    end
    % find the minimum error
    [m,i] = min(fval);
    % get the parameters that minimise error
    coef = coef(i,:)
    
    % fit timescale using time series data (minimises the negative
    % loglikelihood of the data)
    error_ts = @(ts) log(-loglike(avotes,bvotes,evote,N,coef,ts));
    
    [ts,fvals,exitflag,output] = particleswarm(error_ts, 1, 0, 1, optionsswa);
    
    % multiply parameters by the timescale
    coef = coef(1:9)*ts
    
    % save parameters
    % writematrix(coef,['./fitted_rates/rates_',num2str(N),'_1.txt']);
    toc
end



