% Fix head orientation

function [flip_stack_x,theta_output] = FixHeadOrientation(optreg,is_moving,stack_number,transform_matrix,do_z)

if is_moving
    dataset = optreg.moving_dataset;
    dataset_directory = optreg.moving_dataset_directory;
    position = optreg.moving_position;
    strain_name = optreg.moving_strain_name;
else
    dataset = optreg.target_dataset;
    dataset_directory = optreg.target_dataset_directory;
    position = optreg.target_position;
    strain_name = optreg.target_strain_name;
end

use_h5 = optreg.use_h5;

cd(dataset_directory);

num_segmentation_variables = optreg.num_segmentation_variables;

img_height = optreg.img_height;
img_width = optreg.img_width;
num_slices = optreg.num_slices;

color_number = 0;
prefix = optreg.input_prefix;

filename = [prefix strain_name 'Ds' num2str(dataset) '_Pos' num2str(position) '_Co' num2str(color_number) '_St' num2str(stack_number)];
prob_array = zeros(img_height,img_width,num_slices,num_segmentation_variables);

if use_h5
    h5_input = h5read([filename '_Probabilities.h5'],'/exported_data');
%     h5_input = h5read([filename '_Probabilities.h5'],'/exported_data');

    for i = 1:num_segmentation_variables
        temp_stack = squeeze(h5_input(i,:,:,:));
        temp_stack2 = imwarp(temp_stack,affine3d([0 1 0 0; 1 0 0 0; 0 0 1 0; 0 0 0 1]),'OutputView',imref3d([img_height img_width num_slices]));
        size(prob_array)
        size(temp_stack)
        size(temp_stack2)
        prob_array(:,:,:,i) = temp_stack2;
    end
    clear h5_input;clear temp_stack; clear temp_stack2;
else
    size(prob_array)
    array_obj = Tiff([filename ' Probs.tiff']);
    size(array_obj)
    for q = 1:93
        prob_array(:,:,q,:) = array_obj.read;
        if (q~=93)
            array_obj.nextDirectory;
        end
    end
    close(array_obj)
end

seg_array = squeeze(prob_array(:,:,:,2)>.3);


num_points = nnz(seg_array)

points = zeros(num_points,3);

counter =0;

for q = 1:img_height
    for j = 1:img_width
        for k = 1:num_slices
            if seg_array(q,j,k) == 1;
                counter = counter + 1;
                points(counter,:) = [j q k];
            end
        end
    end
end


points_stretched = points * [1 0 0; 0 1 0; 0 0 3.077];

if num_points < 50000
    condensed_points = points_stretched;
    num_points_condensed = num_points;
else

num_points_condensed = 50000;

condensed_points = zeros(num_points_condensed,3);
for q = 1:num_points_condensed
    condensed_points(q,:) = points_stretched(randi([1 num_points],1),:);
end

end
isequal(condensed_points,points_stretched)

rotated_points = [condensed_points ones(num_points_condensed,1)]*(transform_matrix.');

initial_center = mean(rotated_points);

if initial_center(1)<(img_width/2)
    flip_stack_x = true;
else
    flip_stack_x = false;
end

if do_z

counter_1 =0;
counter_2 = 0;

for q = 1:num_points_condensed
    if flip_stack_x
        if rotated_points(q,1)<((img_width/2)+50)
        counter_1 = counter_1 +1;
        anterior_points(counter_1,:) = rotated_points(q,:);
        else
        counter_2 = counter_2 +1;
        posterior_points(counter_2,:) = rotated_points(q,:);
        end
    else
        if rotated_points(q,1)>=((img_width/2)-50)
        counter_1 = counter_1 +1;
        anterior_points(counter_1,:) = rotated_points(q,:);
        else
        counter_2 = counter_2 +1;
        posterior_points(counter_2,:) = rotated_points(q,:);
        end
    end
end

anterior_center = mean(anterior_points(:,2:3),1)
posterior_center = mean(posterior_points(:,2:3),1)

if flip_stack_x
    posterior_center(2) = 144 - posterior_center(2) + 144
    anterior_center(2) = 144 - anterior_center(2) + 144
end

first_vector = anterior_center - posterior_center
second_vector = [  -43.7551 -39.3530]

cos_theta = dot(first_vector,second_vector)/(norm(first_vector)*norm(second_vector))
trial_theta_1 = acos(cos_theta)
trial_theta_2 = -trial_theta_1

first_vector_shifted_1 = first_vector*[cos(trial_theta_1) -sin(trial_theta_1); sin(trial_theta_1) cos(trial_theta_1)]
first_vector_shifted_2 = first_vector*[cos(trial_theta_2) -sin(trial_theta_2); sin(trial_theta_2) cos(trial_theta_2)]

if dot(first_vector_shifted_1,second_vector)>dot(first_vector_shifted_2,second_vector)
    theta_output = trial_theta_1
else
    theta_output = trial_theta_2
end

else theta_output = 0;

end