%% K-modes


% Data: 
% Each subject is a row that contains entry 1, 2, 3 (for the chosen option)
% in each cell.


% NOTE: 
% In order to get exactly reproducible results, this code needs to be run
% with 6 workers in the parallel computing toolbox. The reason is that due
% to the construction of the toolbox, each worker's RNG needs to be
% initialized separately. 


% IMPORTANT: Before running this script, need to run Stata script makeDataForIncentivizedChoiceClustering.do

% TO DO: 
% Check correctness of the script: 
%   1. In simulated data in which all subjects follow a prespecified rule, 
%      no endogenous clusters should emerge. 
%           Yes, checked
%   2. In simulated data in which all subjects follow the same rule that 
%      differs from the prespecified rule, the algorithm should detect that
%           Yes, checked. 



clear
clc

%% Set parameters

nRestarts = 10000;
maxIter = 20; 

kSelection = [1, 2, 3, 5, 10];
maxK = numel(kSelection);

rerunAlgo = 1;

%% Run clustering algorithm 

if rerunAlgo == 1

    T = readtable('../dataAnalysis/data/dataKmodesExtendedRules.csv');
    subjChoicesOrig = T{:,:};

    R = readtable('data/rulePredictions3optExpandedReady.csv');
    ruleChoices = R{:,:};

    load('data/ruleLabelsExpanded');


    [nPrespec, nProfiles] = size(ruleChoices);
    [nSubj, ~] = size(subjChoicesOrig);
    nPoints = nPrespec + nSubj;
    choices = cat(1, ruleChoices, subjChoicesOrig);


    % Exclude perfectly rule-following subjects as cluster centers  

    rulePerfect = zeros(nSubj+nPrespec, 1);

    for i = nPrespec + 1 : nPoints
        for j = 1 : nPrespec
            rulePerfect(i, j) = sum(choices(i, :) == ruleChoices(j,:));
        end
    end

    rulePerfect = max(rulePerfect, [], 2);
    rulePerfect = +(rulePerfect == 17);

    
    numNonfitSubj = sum(1 - rulePerfect(nPrespec+1:nPoints));

    mapNonfit = zeros(nPoints, numNonfitSubj);
    j=1;
    for i = nPrespec + 1:nPoints
        if rulePerfect(i) == 0
            mapNonfit(i,j) = 1;
            j=j+1;
        end
    end

    optimalKcentroids = zeros(maxK,nProfiles,maxK);

    for i = 1 : maxK 

        k = kSelection(i);
       
        % For k <= 1: Comprehensive search
            
        if k <= 1
            subsets = nchoosek(1:nSubj, k);
            nRuns = nchoosek(nSubj, k);
            centroidsSet = zeros(k, nProfiles, nRuns);
            totalCost = zeros(nRuns, 1);
            iterations = zeros(nRuns, 1);
                            
            parfor r = 1:nRuns
                rng(r);     % RNG seed must be set inside parfor-loop, so each worker receives it's own (distinct) seed. If the RNG is only set in the client, the workers' RNG will not be initialized
                [centroids, costAssignment, iter] = kModesFunction(k, maxIter, choices, nPrespec, mapNonfit, numNonfitSubj, subsets(r,:));

                totalCost(r) = costAssignment; 
                iterations(r) = iter;
                centroidsSet(:,:,r) = centroids(nPrespec+1:nPrespec+k,:);
                
                disp(r)

            end
            
            [~, bestClustering] = min(totalCost);

            optimalKcentroids([1:k],:,i) = centroidsSet(:,:,bestClustering);
                        
        end
        
        % For k >= 2: Loop over the algorithm for random restarts
        if k >= 2
            centroidsSet = zeros(k, nProfiles, nRestarts);
            totalCost = zeros(nRestarts, 1);
            iterations = zeros(nRestarts, 1);

            % Use parfor here!
            parfor r = 1 : nRestarts 
                rng(r);     % RNG seed must be set inside parfor-loop, so each worker receives it's own (distinct) seed. If the RNG is only set in the client, the workers' RNG will not be initialized
                [centroids, costAssignment, iter] = kModesFunction(k, maxIter, choices, nPrespec, mapNonfit, numNonfitSubj, []);

                totalCost(r) = costAssignment; 
                iterations(r) = iter;
                centroidsSet(:,:,r) = centroids(nPrespec+1:nPrespec+k,:);

            end

            [~, bestClustering] = min(totalCost);

            optimalKcentroids([1:k],:,i) = centroidsSet(:,:,bestClustering);
        end
        
        disp('Current k:')
        disp(k)

    end     % end loop over k

end


%% Attraction to individual subjects

bestDist = 100*ones(nSubj,1);
assignment = zeros(nSubj, 1);
subjAttraction = zeros(nSubj, 1);
clusterClosenessMean = 100*ones(nSubj,1);
clusterClosenessMin = 100*ones(nSubj,1);
clusterClosenessMax = 100*ones(nSubj,1);

attractionMode = zeros(nSubj, 17);

for i = 1:nSubj
        
    choicesExt = cat(1, ruleChoices, subjChoicesOrig(i,:));
    
    for subj = nPrespec+1:nPoints
        subjChoices = repmat(choices(subj,:),[nPrespec + 1, 1]);
        dist = subjChoices ~= choicesExt; 
        dist = sum(dist,2);
        [bestDist(subj), assignment(subj)] = min(dist); 
    end
     
    subjAttraction(i) = sum(assignment == nPrespec + 1);
    
    attractionMode(i,:) = mode(choices(assignment == nPrespec + 1, :));
    
    assignment(nPrespec+i) = 0;
    if numel(bestDist(assignment == nPrespec + 1))>0
        clusterClosenessMean(i) = mean(bestDist(assignment == nPrespec + 1));
        clusterClosenessMin(i) = min(bestDist(assignment == nPrespec + 1));
        clusterClosenessMax(i) = max(bestDist(assignment == nPrespec + 1));
    end
end

% Choices of subjects that attract the largest number: 
maxAttr = max(subjAttraction);
disp('Attraction to individual subjects');
disp(sum(subjAttraction == maxAttr))
disp(attractionMode(subjAttraction == maxAttr, :));


%% Table: 
% For each k, report the fraction of subjects assigned to each endogenous
% cluster, and the nature of that cluster (i.e. the choice vector)


fracEndo = zeros(maxK, kSelection(maxK));

for i = 1:maxK
    k = kSelection(i);
    centroidsPlusRules = [choices(1:nPrespec, :); optimalKcentroids([1:k],:,i)];
        
    % 1. Find best assignment

    points = 1:nPoints;
    assignment = points;
    bestDist = zeros(nPoints, 1);

    for subj = nPrespec+1:nPoints
        subjChoices = repmat(choices(subj,:),[nPrespec + k, 1]);
        dist = subjChoices ~= centroidsPlusRules; 
        dist = sum(dist,2);
        [bestDist(subj), assignment(subj)] = min(dist); 
    end
    
    for j = 1:kSelection(maxK)
        fracEndo(i,j) = sum(assignment == nPrespec+j)/nSubj;
    end
end

fracEndo = fracEndo';
fracEndoTotal = sum(fracEndo);



labels = 'ABC';
tableData = '';
for i = 1 : maxK
    [fracEndo(:,i), sortOrder] = sort(fracEndo(:,i), 'descend');
    centroids = optimalKcentroids(1:kSelection(i), :, i);
    centroids = centroids(sortOrder(1:kSelection(i)), :);
    tableData_i = labels(centroids);
    tableData_i = cat(2, tableData_i);
    tableData = cat(1, tableData, tableData_i);
end

[tableRows, tableCols] = size(tableData);


insertString = '';
for i=1:tableRows
    insertString = [insertString; '&'];
end

for j = 1:tableCols-1
    [~, tableColsNew] = size(tableData);
    jInv = tableCols - j;
    tableData = [tableData(:, 1:jInv), insertString, tableData(:,jInv+1:tableColsNew)];
end


insertString = '';
for i=1:tableRows
    insertString = [insertString; '\\'];
end

tableData = [tableData, insertString];

fracEndoTotal = fracEndoTotal*100;
fracEndo = fracEndo*100;

myFile = fopen('../writeup/tables/clusterResults.tex', 'w');
%fprintf(myFile, '\\begin{tabular}{lcccccccccccccccccc} \\toprule \n');
counter = 0;
for i = 1:maxK
    fprintf(myFile, '\\\\ \\multicolumn{18}{c}{Endogenous clusters, $k=%4.0f$} \\\\  \n', kSelection(i));
    fprintf(myFile, '\\cmidrule(r){2-18} \n');
    fprintf(myFile, '\\%% subjects \\\\ \n');
    for j = 1:kSelection(i)
        counter = counter + 1;
        fprintf(myFile, '\\quad %s \n', [num2str(fracEndo(j, i), '%4.2f') '&' tableData(counter,:)]);
    end
    fprintf(myFile, 'Total: %4.2f \\\\ \n', fracEndoTotal(i));
    
end
%fprintf(myFile, '\\bottomrule \\end{tabular} \n');
fclose(myFile);















