% --- Data quantities ---
enum Patient; % Patients for which samples have been taken
set of int: Horizon = -5..2; % (Relative) sampling times

% --- Data input ---
float: DLMO_threshold = 13.043;
array[Patient] of opt float: dlmo;
array[Patient, Horizon] of opt float: sample;

% consistency
array[Patient] of opt float: consistency;

% grouping parameters
int: nr_groups;
float: success_rate;

array[_] of float: known_consistency = [deopt(consistency[pt]) | pt in Patient where occurs(consistency[pt])];
test has_dlmo(Patient: pt) = occurs(dlmo[pt]);

% --- Derived data ---

% For each patient: after which sample could we stop and have found dlmo? (round must then include `critical_test[PT]` and its successor to catch it).
array[Patient] of var Horizon: critical_sample = [
	if has_dlmo(pt) then
		min([i | i in -5..1 where sample[pt, i] <= DLMO_threshold /\ sample[pt, i+1] >= DLMO_threshold])
	else 
		min(Horizon)
	endif
| pt in Patient];

% --- Decisions ---
enum Group = G(1..nr_groups);
array[Group] of var min(known_consistency)..max(known_consistency): split;
array[Patient] of var Group: group;

array[Group] of var int: group_count :: output = [grp: count(pt in Patient)(group[pt] = grp) | grp in Group];

% First round:
array[Group] of var Horizon: first_sample; % When to start sampling on the first round 
array[Group] of var 2..length(Horizon): num_samples; % How many samples to take on the first round 

% --- Decision Consequences ---

% For each patient: will their DLMO be detected on the first round?
array[Patient] of var bool: first_detected = [
	has_dlmo(pt) /\ (critical_sample[pt] >= first_sample[group[pt]] /\ critical_sample[pt] + 1 <= first_sample[group[pt]] + (num_samples[group[pt]] - 1))
	| pt in Patient
];

% --- Constraint ---
% Start first group with any positive
constraint split[G(1)] = min(known_consistency);
constraint forall(grp in Group where grp != min(Group)) (split[grp] > split[enum_prev(grp)]);
% Ensure consistency of patient is between split
constraint forall(pt in Patient where occurs(consistency[pt])) (deopt(consistency[pt]) >= split[group[pt]]);
constraint forall(pt in Patient where occurs(consistency[pt])) (group[pt] == max(Group) \/ deopt(consistency[pt]) < split[enum_next(group[pt])]);

% The DLMO should be found for patients with success_rate
constraint forall(grp in Group)(
	sum(pt in Patient where has_dlmo(pt) /\ group[pt] == grp) (first_detected[pt]) >= (success_rate * group_count[grp])
);

% --- Objective ---
solve minimize total_tests;
% Pareto Front objectives
var int: total_tests ::output =	sum(grp in Group)(count(pt in Patient) (group[pt] == grp) * num_samples[grp]);

% var int: total_not_detected ::output = length([pt | pt in Patient where has_dlmo(pt)]) - count(first_detected);

% % --- Output ---
%
% output ["First Round:\n"]
% ++ [format(2, i) ++ "   " | i in Horizon] ++ ["\n |"]
% ++ [if i in fix(first_sample)..fix(first_sample + (num_samples-1) - 1) then "====|" else "----|" endif | i in -5..1]
% ++ ["\n\nFirst try (%): \(count(first_detected) / length([pt | pt in Patient where has_dlmo(pt)]) * 100)\n"]
% ++ ["\n\nSecond Round:\n"]
% ++ [format(2, i) ++ "   " | i in Horizon] ++ ["\n |"]
% ++ [if i in fix(first_sample-(num_samples-1))..fix(first_sample-1) \/ i in fix(first_sample + (num_samples - 1))..fix(first_sample + (num_samples - 1) + (num_samples - 1) - 1) then "====|" else "----|" endif | i in -5..1]
% ++ ["\n\nTotal number of tests = \(total_tests);"]
%
%
% ++ ["\n\n\nfirst_sample = \(first_sample);\n"]
% ++ ["num_samples = \(num_samples);\n"]