#
# Place markers on cell protrusions located by subtracting actin and membrane labeled channels,
# smoothing the difference and placing a marker on each connected blob enclosing at least
# a specified volume.
#
# To record movie:
#
# movie record super 3 ; vseries play #0 showMark true ; wait 250 ; movie encode ~/Desktop/cell15.mp4 quality high
#
def mark_protrusions(
        n = None,	# Number of maps in series to process, none means all. 
        actin_id = 1,	# Model id for actin channel series, subids 1 to n
        memb_id = 0,	# Model id for membrane channel series, subids 1 to n
        volume = 5,	# Minimum enclosed volume for protrusion in difference map, cubic microns.
        level = 85,	# Contour level for smoothed difference map
        sdev = 0.5,	# Smoothing of difference map, standard deviation, microns
        fit_level = 100, # Contour level used for fitting.
        cell_rgba = (.8, .8, .4, 1),		# Cell color
        protrusion_rgba = (1, .3, .2, 1),	# Marker color
        link_rgba = (.2,.5,.8,1),		# Linker color
        link_radius = 1,			# microns
        diff_id = 1000,				# Temporary model number for difference map
        smooth_id = 1001			# Temporary model number for smoothed difference.
):

    if n is None:
        from chimera import openModels
        n = len(openModels.list(id = actin_id))

    from VolumePath import Marker_Set, Link
    mset = Marker_Set('protrusions')

    mlist = []
    from chimera import runCommand as run
    for i in range(n):
        aid = '#%d.%d' % (actin_id, i+1)
        mid = '#%d.%d' % (memb_id, i+1)
        # Fitting is optional and is to correct misalignment between the two channels
        if fit_level is not None:
            run('volume %s level %.5g' % (aid, fit_level))
            run('fitmap %s %s rotate false' % (aid, mid))
        # TODO: Unsigned 16-bit maps subtract to give unsigned result, but signed is needed.
        #       For now have to make sure input maps have signed values.
        run('vop subtract %s %s minRMS true model #%d' % (mid, aid, diff_id))
        run('vop gaussian #%d sdev %.5g model #%d' % (diff_id, sdev, smooth_id))
        run('volume #%d level %.5g' % (smooth_id, level))
        # Use hide dust to find surface blobs with enclosed volume large enough.
        run('sop hideDust #%d metric volume size %.5g update false' % (smooth_id, volume))

        cc = cell_center_marker((actin_id,i+1), cell_rgba, mset)
        plist = make_protrusion_markers(smooth_id, volume, protrusion_rgba, i, mset)
        for m in plist:
            Link(m, cc, link_rgba, link_radius)
        mlist.append((i,cc,plist))

        run('close #%d,%d' % (diff_id, smooth_id))

    # Print protrusion statistics
    for i, cc, plist in mlist:
        print protrusion_stats(i, cc, plist)

    return mset

def make_protrusion_markers(smooth_id, volume, protrusion_rgba, frame, mset):
    # Mark protrusions
    from chimera import openModels
    ms = openModels.list(id = smooth_id)[0]
    markers = []
    for p in ms.surfacePieces:
        b = p.blobs
        tval = b.triangle_values('volume')
        for vi, ti in b.blob_list():
            if tval[ti[0]] < volume:
                continue
#            va, ta = b.varray[vi,:], b.tarray[ti,:]
            ta = b.tarray[ti,:].copy()
            center, radius = surface_centroid(b.varray,ta)
            m = mset.place_marker(center, protrusion_rgba, radius)
            m.set_note('%d' % frame)	# Label marker with time series frame number.
            markers.append(m)
    return markers

def cell_center_marker(model_id, cell_rgba, mset):
    id,subid = model_id
    from chimera import openModels
    m = openModels.list(id = id, subid = subid)[0]
    p = m.surface_piece_list[0]
    va, ta = p.geometry
    center, radius = surface_centroid(va,ta)
    c = mset.place_marker(center, cell_rgba, radius)
    c.set_note('%d' % (subid-1))	# Label marker with time series frame number.
    return c

def surface_centroid(va, ta):
    from _surface import vertex_areas, enclosed_volume
    areas = vertex_areas(va, ta)
    n = len(areas)
    center = (areas.reshape((n,1))*va).sum(axis = 0)/areas.sum()       # Area weighted centroid
    volume, hole_count = enclosed_volume(va, ta)
    from math import pow, pi
    radius = pow(volume/(4*pi/3), 1.0/3)	# Marker volume = protrusion volume.
    return center, radius

def print_stats(mset = None):
    
    from VolumePath import selected_markers
    mlist = selected_markers() if mset is None else mset.markers()

    mt = {}
    for m in mlist:
        t = int(m.note()) 
        mt.setdefault(t,[]).append(m)

    for t in sorted(mt.keys()):
        plist = mt[t]
        plist.sort(key = lambda m: m.radius(), reverse = True)
        print protrusion_stats(t, plist[0], plist[1:])
        
def protrusion_stats(t, cell_center, protrusions, np = 5):

    psort = list(protrusions)
    psort.sort(key = lambda m: m.radius(), reverse = True)
    radii = [m.radius() for m in psort[:np]]
    radii += [None for i in range(np-len(psort))]

    angles = []
    for i,p1 in enumerate(psort[:np]):
        for p2 in psort[i+1:np]:
            angles.append(angle(p1.xyz(), cell_center.xyz(), p2.xyz()))
    a = sum(angles) / len(angles) if angles else None

    unk = '%6s' % '?'
    line = ''.join(['%3d' % t, unk if a is None else '%6.0f' % a] +
                   [(unk if r is None else ('%6.2f'%r)) for r in radii])
    return line

def angle(p1, p2, p3):
    import Matrix
    return Matrix.vector_angle([a-b for a,b in zip(p1,p2)],
                               [a-b for a,b in zip(p3,p2)])

mark_protrusions()
#print_stats()