# -*- coding: utf-8 -*-

"""Probe layout."""


#------------------------------------------------------------------------------
# Imports
#------------------------------------------------------------------------------

from collections import defaultdict

import numpy as np

from phylib.io.array import _flatten
from phylib.utils.geometry import _get_boxes, range_transform
from phylib.utils.color import selected_cluster_color


#------------------------------------------------------------------------------
# Layout
#------------------------------------------------------------------------------

def _iter_channel(positions):
    """Iterate through the channel positions."""
    size = 100
    margin = 5
    boxes = _get_boxes(positions, keep_aspect_ratio=False)
    xmin, ymin = boxes[:, :2].min(axis=0)
    xmax, ymax = boxes[:, 2:].max(axis=0)
    x = boxes[:, [0, 2]].mean(axis=1)
    y = - boxes[:, [1, 3]].mean(axis=1)
    positions = np.c_[x, y]
    tr = [margin, margin, size - margin, size - margin]
    positions_tr = range_transform([-1, -1, 1, 1], tr, positions)
    for x, y in positions_tr:
        yield x, y


def _disk(x, y, r, c, t=0):
    """Return a SVG disc."""
    return ('<circle cx="%.5f%%" cy="%.5f%%" '
            'r="%d" fill="%s" '
            'transform="translate(%d, 0)"'
            ' />') % (x, y, r, c, t)


def _rgba(rgb, a=1.):
    r, g, b = rgb
    return 'rgba(%d, %d, %d, %.3f)' % (255 * r, 255 * g, 255 * b, a)


def _iter_disks(positions, cluster_channels=None):
    """Iterate through the SVG dics for all channels and clusters.

    positions: Nx2 array
    cluster_channels: {cluster: channels}

    """

    color_masked = _rgba((.5, .5, .5), 1.)
    color_umasked = lambda clu: _rgba(selected_cluster_color(clu)[:3], 1)

    size_masked = 5
    size_unmasked = 7

    # n_channels = positions.shape[0]
    n_clusters = len(cluster_channels)

    cluster_channels = cluster_channels or {}
    channel_ids = set(_flatten(cluster_channels.values()))
    is_masked = {ch: False for ch in channel_ids}

    # List of clusters per channel.
    clusters_per_channel = defaultdict(lambda: [])
    for clu, channels in cluster_channels.items():
        for channel in channels:
            clusters_per_channel[channel].append(clu)

    # Enumerate the discs for each channel.
    for channel_id, (x, y) in enumerate(_iter_channel(positions)):
        masked = is_masked.get(channel_id, True)
        if masked:
            yield _disk(x, y, size_masked, color_masked)
            continue
        for clu in clusters_per_channel[channel_id]:
            # Translation.
            t = 10 * (clu - .5 * (n_clusters - 1))
            yield _disk(x, y, size_unmasked, color_umasked(clu), t=t)


def probe_layout(positions, cluster_channels):
    """Return a probe layout in SVG, showing selected clusters on relevant channels."""
    contents = '\n'.join(_iter_disks(positions, cluster_channels))
    return """
    <svg style="background: black; width:100%; height:100%;">
      {}
    </svg>
    """.format(contents)
