# Randomly move surface blobs keeping them within a masked region.
# This is to test whether colocalization in 3D light microscopy is due to random chance.

def random_move_blobs(session, surface, mask_volume):
    if surface.scene_position != mask_volume.scene_position:
        from chimerax.core.errors import UserError
        raise UserError('Surface and mask volume must have the same scene position')
    vertices = surface.vertices
    vxyz_min, vxyz_max = mask_volume.xyz_bounds()	# Volume coords
    from chimerax.surface.dust import Blob_Masker
    bm = Blob_Masker(vertices, surface.masked_triangles)
    blobs = bm.blob_list()
    from chimerax.geometry import point_bounds
    for vi, ti in blobs:
        vb = vertices[vi]
        b = point_bounds(vb)
        shift_min, shift_max = vxyz_min - b.xyz_min, vxyz_max - b.xyz_max
        shift = _random_shift(vb, shift_min, shift_max, mask_volume)
        if shift is None:
            session.logger.info('Could not find a shift for blob with vertex vi[0]')
        vertices[vi] += shift

    surface.set_geometry(vertices, surface.normals, surface.triangles, triangle_mask = surface.triangle_mask)
    session.logger.info(f'Random shifted {len(blobs)} surface blobs')
    
def _random_shift(vertices, shift_min, shift_max, mask_volume, max_tries = 1000):
    from random import uniform
    for t in range(max_tries):
        shift = [uniform(shift_min[a], shift_max[a]) for a in (0,1,2)]
        vv = mask_volume.interpolated_values(vertices + shift)
        if (vv != 0).all():
            # All vertices in mask region.
            return shift
    return None

def coloc_random(session, surface, near_surface, distance, mask_volume, iterations = 1):
    vsave = surface.vertices.copy()
    tmsave = None if surface.triangle_mask is None else surface.triangle_mask.copy()
    surface.auto_remask_triangles = None
    from chimerax.surface.hidefarblobs import surface_hide_far_blobs
    shown, hidden = surface_hide_far_blobs(session, surface, near_surface, distance, symmetric = False)
    session.logger.info(f'Initial blobs {shown} near, {hidden} far')
    shown_counts = []
    surface.triangle_mask = tmsave
    for i in range(iterations):
        random_move_blobs(session, surface, mask_volume)
        shown, hidden = surface_hide_far_blobs(session, surface, near_surface, distance, symmetric = False)
        shown_counts.append(shown)
        surface.set_geometry(vsave, surface.normals, surface.triangles, triangle_mask = tmsave)
    from numpy import mean, std
    m,sd = mean(shown_counts), std(shown_counts)
    session.logger.info(f'{iterations} iterations, near count mean {m}, stdev {sd}')

def register_command(logger):
    from chimerax.core.commands import CmdDesc, register, SurfaceArg, FloatArg, IntArg
    from chimerax.map import MapArg
    desc = CmdDesc(
        required = [('surface', SurfaceArg)],
        keyword = [('mask_volume', MapArg)],
        required_arguments = ['mask_volume'],
        synopsis = 'Move surface blobs by random amounts within mask volume'
    )
    register('moveblobs', desc, random_move_blobs, logger=logger)

    desc = CmdDesc(
        required = [('surface', SurfaceArg)],
        keyword = [('near_surface', SurfaceArg),
                   ('distance', FloatArg),
                   ('mask_volume', MapArg),
                   ('iterations', IntArg)],
        required_arguments = ['near_surface', 'distance', 'mask_volume'],
        synopsis = 'Move surface blobs by random amounts within mask volume'
    )
    register('colocrand', desc, coloc_random, logger=logger)

register_command(session.logger)