"""Segmentation algorithm used to identify the different structures
that are formed in the encounter. This file can be called from the
command line to make an illustrative plot of the algorithm.
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches

import utils


def segmentEncounter(data, plot=False, mode='all'):
    """Segment the encounter into tail, bridge, orbitting and
    stolen particles as described in the report.

    Parameters:
        data: A data instance as saved by the simulation to a pickle file
        plot: If true the segmentation will be plotted and shown. Highly
            useful for debugging.
        mode (string): If mode is 'all' all parts of the encounter will be
            identified. If mode is 'bridge' only the bridge will be
            identified. This is useful when there may be no tail.

    Returns:
        masks (tupple): tupple of array corresponding to the masks of the 
            (bridge, stolen, orbitting, tail) particles. One can then use
            e.g. data['r_vec'][bridgeMask].
        shape (tupple): tupple of (distances, angles) as measured from the
            center of mass and with respect to the x axis. They define the
            shape of the tail
        length (float): total length of the tail.
    """
    nRings = 100 # number of rings to use when segmenting the data

    # Localize the central masses
    r_vec = data['r_vec']
    centers = r_vec[data['type'][:,0]=='center']
    rCenters_vec = centers[1] - centers[0]
    rCenters = np.linalg.norm(rCenters_vec)
    rCenters_unit = rCenters_vec/np.linalg.norm(rCenters_vec)
    # Take particles to be on the tail a priori and
    # remove them as they are found in other structures
    particlesLeft = np.arange(0, len(r_vec))


    if plot:
        colour = '#c40f4c'
        f, axs = plt.subplots(1, 3,  figsize=(9, 4), sharey=False)
        f.subplots_adjust(hspace=0, wspace=0)
        axs[0].scatter(r_vec[:,0], r_vec[:,1], c=colour, alpha=0.1, s=0.1)
        axs[0].axis('equal')
        utils.plotCenterMasses(axs[0], data)
        axs[0].axis('off')

    # Step 1: project points to see if they are part of the bridge
    parallelProjection = np.dot(r_vec - centers[0], rCenters_unit)
    perpendicularProjection = np.linalg.norm(r_vec - centers[0][np.newaxis] 
        - parallelProjection[:,np.newaxis] * rCenters_unit[np.newaxis], axis=-1)
    bridgeMask = np.logical_and(np.logical_and(0.3*rCenters < parallelProjection,
        parallelProjection < .7*rCenters), perpendicularProjection < 2)

    # Remove the bridge
    notInBridge = np.logical_not(bridgeMask)
    r_vec = r_vec[notInBridge]
    particlesLeft = particlesLeft[notInBridge]

    if mode == 'bridge':
        return (bridgeMask, None, None, None), None, None

    # Step 2: select stolen particles by checking distance to centers
    stolenMask = (np.linalg.norm(r_vec - centers[0][np.newaxis], axis=-1) 
        > np.linalg.norm(r_vec - centers[1][np.newaxis], axis=-1))
    # Remove the stolen part
    notStolen = np.logical_not(stolenMask)
    r_vec = r_vec[notStolen]
    particlesLeft, stolenMask = particlesLeft[notStolen], particlesLeft[stolenMask]

    # Step 3: segment data into concentric rings (spherical shells really)
    r_vec = r_vec - centers[0]
    r = np.linalg.norm(r_vec, axis=-1)
    edges = np.linspace(0, 30, nRings) # nRings concentric spheres 
    indices = np.digitize(r, edges) # Classify particles into shells

    if plot:
        axs[1].scatter(r_vec[:,0], r_vec[:,1], c=colour, alpha=.1, s=.1)
        axs[1].axis('equal')
        axs[1].scatter(0, 0, s=100, marker="*", c='black', alpha=.7)
        axs[1].axis('off')
    
    # Step 4: find start of tail
    start = None
    for i in range(1, nRings+1):
        rMean = np.mean(r[indices==i])
        rMean_vec = np.mean(r_vec[indices==i], axis=0)
        parameter = np.linalg.norm(rMean_vec)/rMean

        if plot:
            circ = patches.Circle((0,0), edges[i-1], linewidth=0.5,edgecolor='black',facecolor='none', alpha=.7)
            axs[1].add_patch(circ)
            txtxy = edges[i-1] * np.array([np.sin(i/13), np.cos(i/13)])
            axs[1].annotate("{:.2f}".format(parameter), xy=txtxy, backgroundcolor='#ffffff55')
        
        if start is None and parameter>.8 : 
            start = i #Here starts the tail
            startDirection = rMean_vec/np.linalg.norm(rMean_vec)
            if not plot: break;

    if start is None: #abort if nothing found
        raise Exception('Could not identify tail')

    # Step 5: remove all circles before start
    inInnerRings = indices < start
    # Remove particles on the opposite direction to startDirection. 
    # in the now innermost 5 rings. Likely traces of the bridge.
    oppositeDirection = np.dot(r_vec, startDirection) < 0
    in5InnermostRings = indices <= start+5
    orbitting = np.logical_or(inInnerRings, 
        np.logical_and(oppositeDirection, in5InnermostRings))
    orbittingMask = particlesLeft[orbitting]
    r_vec = r_vec[np.logical_not(orbitting)]
    tailMask = particlesLeft[np.logical_not(orbitting)]

    if plot:
        axs[2].scatter(r_vec[:,0], r_vec[:,1], c=colour, alpha=0.1, s=0.1)
        axs[2].axis('equal')
        axs[2].scatter(0, 0, s=100, marker="*", c='black', alpha=.7)
        axs[2].axis('off')

    # Step 6: measure tail length and shape
    r = np.linalg.norm(r_vec, axis=-1)
    indices = np.digitize(r, edges)
    # Make list of barycenters
    points = [list(np.mean(r_vec[indices==i], axis=0))
        for i in range(1, nRings) if len(r_vec[indices==i])!=0]
    points = np.array(points)
    # Calculate total length
    lengths = np.sqrt(np.sum(np.diff(points, axis=0)**2, axis=1))
    length = np.sum(lengths)
    # Shape (for 2D only)
    angles = np.arctan2(points[:,1], points[:,0])
    distances = np.linalg.norm(points, axis=-1)
    shape = (distances, angles)

    if plot:
        axs[2].plot(points[:,0], points[:,1], c='black', linewidth=0.5, marker='+')

    if plot:
        plt.show()

    return (bridgeMask, stolenMask, orbittingMask, tailMask), shape, length


if __name__ == "__main__":
    data = utils.loadData('200mass', 10400)
    segmentEncounter(data, plot=True)