"""
This script does the exact same as 'NerveSegmenetation2D.ipynb'
"""
import numpy as np
import matplotlib.pyplot as plt
from skimage.io import imread
from scipy.ndimage.interpolation import map_coordinates
from slgbuilder import GraphObject,MaxflowBuilder

def draw_segmentations(data, helper):
    """Draw all segmentations for objects in the helper on top of the data."""
    
    # Create figure.
    plt.figure(figsize=(10, 10))
    plt.imshow(data, cmap='gray')
    plt.xlim([0, data.shape[1]-1])
    plt.ylim([data.shape[0]-1, 0])

    # Draw segmentation lines.
    for i, obj in enumerate(helper.objects):

        # Get segmentation.
        segment = helper.get_labels(obj)

        # Create line.
        line = np.count_nonzero(segment, axis=0)

        # Get actual points.
        point_indices = tuple(np.asarray([line - 1, np.arange(len(line))]))
        points = obj.sample_points[point_indices]
        # Close line.
        points = np.append(points, points[:1], axis=0)

        # Plot points.
        plt.plot(points[..., 1], points[..., 0])

    plt.show()
    
def unfold_image(img, center, max_dists=None, r_min=1, r_max=20, angles=30, steps=15):
    """ Unfolds around a point in an image with radial sampling """
    # Sampling angles and radii.
    angles = np.linspace(0, 2*np.pi, angles, endpoint=False)
    distances = np.linspace(r_min, r_max, steps, endpoint=True)
    
    if max_dists is not None:
        max_dists.append(np.max(distances))
    
    # Get angles.
    angles_cos = np.cos(angles)
    angles_sin = np.sin(angles)
    
    # Calculate points positions.
    x_pos = center[0] + np.outer(angles_cos, distances)
    y_pos = center[1] + np.outer(angles_sin, distances)
    
    # Create list of sampling points.
    sampling_points = np.array([x_pos, y_pos]).transpose()
    sampling_shape = sampling_points.shape
    sampling_points_flat = sampling_points.reshape((-1, 2))
    
    # Sample from image.
    samples = map_coordinates(img, sampling_points_flat.transpose(), mode='nearest')
    samples = samples.reshape(sampling_shape[:2])
        
    return samples, sampling_points
#%%

in_dir = 'data/'
data = imread(in_dir+'nerves2D.png').astype(np.int32)

# Get centers
path = in_dir+'nerveCenters.png'
centers = imread(path)[:,:,0]

centers_small = np.transpose(np.where(centers==1))
centers_large = np.transpose(np.where(centers==2))


# Show image with centers.
plt.imshow(data, cmap='gray')
plt.scatter(centers_small[..., 1], centers_small[..., 0], color='red', s=6)
plt.scatter(centers_large[..., 1], centers_large[..., 0], color='blue', s=6)
plt.show()

print(f'Number of objects: {len(centers_small)+len(centers_large)}')

#%% Example of radial unfolding



samples, sample_points = unfold_image(data, centers_small[4],r_max=20,angles=30,steps=30)

plt.figure(figsize=(13, 5))
ax = plt.subplot(1, 3, 1, title='Sample positions in data')
ax.imshow(data, cmap='gray')
ax.scatter(sample_points[..., 1], sample_points[..., 0], s=2, color='red')
ax = plt.subplot(1, 3, 2, title='Sample positions and intensities')
ax.scatter(sample_points[..., 1], sample_points[..., 0], c=samples, cmap='gray')
ax = plt.subplot(1, 3, 3, title='Unfolded image')
ax.imshow(samples, cmap='gray')
plt.show()

#%% Detect layers in object
# Create gradient-based objects.
diff_samples = np.diff(samples, axis=0)
outer_nerve = GraphObject(255 - diff_samples)
inner_nerve = GraphObject(diff_samples)

f,ax = plt.subplots(1,2,figsize=(10,5))
ax[0].imshow(outer_nerve.data, cmap='gray')
ax[0].set_title('Outer never data')
ax[1].imshow(inner_nerve.data, cmap='gray')
ax[1].set_title('Inner never data')
plt.show()

#%% Segment one sample

helper = MaxflowBuilder()
helper.add_objects([outer_nerve, inner_nerve])
helper.add_layered_boundary_cost()
helper.add_layered_smoothness(delta=2)
helper.add_layered_containment(outer_nerve, inner_nerve, min_margin=3, max_margin=6)

flow = helper.solve()
print('Maximum flow/minimum energy:', flow)

segmentations = [helper.get_labels(o).astype(np.int32) for o in helper.objects]
segmentation_lines = [np.count_nonzero(s, axis=0) - 0.5 for s in segmentations]


f,ax = plt.subplots(1,3,figsize = (10,10))
ax[0].imshow(samples, cmap='gray')
ax[1].imshow(np.sum(segmentations,axis=0))
ax[2].imshow(samples, cmap='gray')
for line in segmentation_lines:
    ax[2].plot(line)
plt.show()

#%% Multiple object
# Lists for storing nerve objects.
nerve_samples = []
outer_nerves = []
inner_nerves = []

# For each center, create an inner and outer never.
for center in centers_small:
        samples, sample_points = unfold_image(data, center,r_max=35,angles=40,steps=30)
        nerve_samples.append(samples)
        
        # Create outer and inner nerve objects.
        diff_samples = np.diff(samples, axis=0)
        diff_sample_points = sample_points[:-1]
    

        outer_nerves.append(GraphObject(255 - diff_samples, diff_sample_points))
        inner_nerves.append(GraphObject(diff_samples, diff_sample_points))
for center in centers_large:
        samples, sample_points = unfold_image(data, center,r_max=60,angles=40,steps=30)
        nerve_samples.append(samples)
        
        # Create outer and inner nerve objects.
        diff_samples = np.diff(samples, axis=0)
        diff_sample_points = sample_points[:-1]
    

        outer_nerves.append(GraphObject(255 - diff_samples, diff_sample_points))
        inner_nerves.append(GraphObject(diff_samples, diff_sample_points))
    
helper = MaxflowBuilder()
helper.add_objects(outer_nerves + inner_nerves)
helper.add_layered_boundary_cost()
helper.add_layered_smoothness(delta=2)

for outer_nerve, inner_nerve in zip(outer_nerves, inner_nerves):
    helper.add_layered_containment(outer_nerve, inner_nerve, min_margin=3, max_margin=6)
    
flow = helper.solve()
print('Maximum flow/minimum energy:', flow)

# Get segmentations.
segmentations = []
for outer_nerve, inner_nerve in zip(outer_nerves, inner_nerves):
    segmentations.append(helper.get_labels(outer_nerve))
    segmentations.append(helper.get_labels(inner_nerve))

segmentation_lines = [np.count_nonzero(s, axis=0) - 0.5 for s in segmentations]

# Draw segmentations.
plt.figure(figsize=(15, 5))
for i, samples in enumerate(nerve_samples):
    ax = plt.subplot(3, len(nerve_samples) // 3 + 1, i + 1)
    ax.imshow(samples, cmap='gray')
    
    ax.plot(segmentation_lines[2*i])
    ax.plot(segmentation_lines[2*i + 1])

plt.show()
 
draw_segmentations(data, helper)


#%% Multi-object exclusion
from slgbuilder import QPBOBuilder

helper = QPBOBuilder()
helper.add_objects(outer_nerves + inner_nerves)
helper.add_layered_boundary_cost()
helper.add_layered_smoothness(delta=2)

for outer_nerve, inner_nerve in zip(outer_nerves, inner_nerves):
    helper.add_layered_containment(outer_nerve, inner_nerve, min_margin=3, max_margin=6)
    
twice_flow = helper.solve()
print('Two times maximum flow/minimum energy:', twice_flow)

if 2*flow == twice_flow:
    print('QPBO flow is exactly twice the Maxflow flow.')
else:
    print('Something is wrong...')

# Add exclusion constraints between all pairs of outer nerves.
for i in range(len(outer_nerves)):
    for j in range(i + 1, len(outer_nerves)):
        helper.add_layered_exclusion(outer_nerves[i], outer_nerves[j], margin=3)
        
twice_flow = helper.solve()
print('Two times maximum flow/minimum energy:', twice_flow)

draw_segmentations(data, helper)


#%% Region cost

# mu_inside = 90
# mu_ring = 70
# mu_outside = 90
# beta = 0.1

# for samples, outer_nerve, inner_nerve in zip(nerve_samples, outer_nerves, inner_nerves):
#     samples = samples[:-1]
    
#     inside_cost = np.abs(samples - mu_inside) * beta
#     ring_cost = np.abs(samples - mu_ring) * beta
#     outside_cost = np.abs(samples - mu_outside) * beta
    
#     helper.add_layered_region_cost(inner_nerve, ring_cost, inside_cost)
#     helper.add_layered_region_cost(outer_nerve, outside_cost, ring_cost)
    
# twice_flow = helper.solve()
# print('Two times maximum flow/minimum energy:', twice_flow)
# draw_segmentations(data, helper)