"""
Module that hold techniques to project the brain volume onto 2D images for visualisation purposes
"""
from functools import lru_cache

import pandas as pd
import numpy as np
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt
import matplotlib.colors

from iblutil.numerical import ismember
from iblutil.util import Bunch, get_logger
from iblutil.io.hashfile import md5
import one.remote.aws as aws

from ibllib.atlas.atlas import AllenAtlas, BrainRegions


_logger = get_logger(__file__)


@lru_cache(maxsize=1, typed=False)
def circles(N=5, atlas=None, display='flat'):
    """
    :param N: number of circles
    :param atlas: brain atlas at 25 m
    :param display: "flat" or "pyramid"
    :return: 2D map of indices, ap_coordinate, ml_coordinate
    """
    atlas = atlas if atlas else AllenAtlas()

    sz = np.array([])
    level = np.array([])
    for k in np.arange(N):
        nlast = 2000  # 25 um for 5mm diameter
        n = int((k + 1) * nlast / N)
        r = .4 * (k + 1) / N
        theta = (np.linspace(0, 2 * np.pi, n) + np.pi / 2)
        sz = np.r_[sz, r * np.exp(1j * theta)]
        level = np.r_[level, theta * 0 + k]

    atlas.compute_surface()
    iy, ix = np.where(~np.isnan(atlas.top))
    centroid = np.array([np.mean(iy), np.mean(ix)])
    xlim = np.array([np.min(ix), np.max(ix)])
    ylim = np.array([np.min(iy), np.max(iy)])

    s = Bunch(
        x=np.real(sz) * np.diff(xlim) + centroid[1],
        y=np.imag(sz) * np.diff(ylim) + centroid[0],
        level=level,
        distance=level * 0,
    )

    # compute the overall linear distance for each circle
    d0 = 0
    for lev in np.unique(s['level']):
        ind = s['level'] == lev
        diff = np.abs(np.diff(s['x'][ind] + 1j * s['y'][ind]))
        s['distance'][ind] = np.cumsum(np.r_[0, diff]) + d0
        d0 = s['distance'][ind][-1]

    fcn = interp1d(s['distance'], s['x'] + 1j * s['y'], fill_value='extrap')
    d = np.arange(0, np.ceil(s['distance'][-1]))

    s_ = Bunch({
        'x': np.real(fcn(d)),
        'y': np.imag(fcn(d)),
        'level': interp1d(s['distance'], level, kind='nearest')(d),
        'distance': d
    })

    if display == 'flat':
        ih = np.arange(atlas.bc.nz)
        iw = np.arange(s_['distance'].size)
        image_map = np.zeros((ih.size, iw.size), dtype=np.int32)
        iw, ih = np.meshgrid(iw, ih)
        # i2d = np.ravel_multi_index((ih[:], iw[:]), image_map.shape)
        iml, _ = np.meshgrid(np.round(s_.x).astype(np.int32), np.arange(atlas.bc.nz))
        iap, idv = np.meshgrid(np.round(s_.y).astype(np.int32), np.arange(atlas.bc.nz))
        i3d = atlas._lookup_inds(np.c_[iml.flat, iap.flat, idv.flat])
        i3d = np.reshape(i3d, [atlas.bc.nz, s_['x'].size])
        image_map[ih, iw] = i3d

    elif display == 'pyramid':
        for i in np.flipud(np.arange(N)):
            ind = s_['level'] == i
            dtot = s_['distance'][ind]
            dtot = dtot - np.mean(dtot)
            if i == N - 1:
                ipx = np.arange(np.floor(dtot[0]), np.ceil(dtot[-1]) + 1)
                nh = atlas.bc.nz * N
                X0 = int(ipx[-1])
                image_map = np.zeros((nh, ipx.size), dtype=np.int32)

            iw = np.arange(np.sum(ind))
            iw = np.int32(iw - np.mean(iw) + X0)
            ih = atlas.bc.nz * i + np.arange(atlas.bc.nz)

            iw, ih = np.meshgrid(iw, ih)
            iml, _ = np.meshgrid(np.round(s_.x[ind]).astype(np.int32), np.arange(atlas.bc.nz))
            iap, idv = np.meshgrid(np.round(s_.y[ind]).astype(np.int32), np.arange(atlas.bc.nz))
            i3d = atlas._lookup_inds(np.c_[iml.flat, iap.flat, idv.flat])
            i3d = np.reshape(i3d, [atlas.bc.nz, s_['x'][ind].size])
            image_map[ih, iw] = i3d
    x, y = (atlas.bc.i2x(s.x), atlas.bc.i2y(s.y))
    return image_map, x, y
    # if display == 'flat':
    #     fig, ax = plt.subplots(2, 1, figsize=(16, 5))
    # elif display == 'pyramid':
    #     fig, ax = plt.subplots(1, 2, figsize=(14, 12))
    # ax[0].imshow(ba._label2rgb(ba.label.flat[image_map]), origin='upper')
    # ax[1].imshow(ba.top)
    # ax[1].plot(centroid[1], centroid[0], '*')
    # ax[1].plot(s.x, s.y)


def swanson(filename="swanson2allen.npz"):
    # filename could be "swanson2allen_original.npz", or "swanson2allen.npz" for remapped indices to match
    # existing labels in the brain atlas
    OLD_MD5 = [
        'bb0554ecc704dd4b540151ab57f73822',  # version 2022-05-02 (remapped)
        '7722c1307cf9a6f291ad7632e5dcc88b',  # version 2022-05-09 (removed wolf pixels and 2 artefact regions)
    ]
    npz_file = AllenAtlas._get_cache_dir().joinpath(filename)
    if not npz_file.exists() or md5(npz_file) in OLD_MD5:
        npz_file.parent.mkdir(exist_ok=True, parents=True)
        _logger.info(f'downloading swanson image from {aws.S3_BUCKET_IBL} s3 bucket...')
        aws.s3_download_file(f'atlas/{npz_file.name}', npz_file)
    s2a = np.load(npz_file)['swanson2allen']  # inds contains regions ids
    return s2a


def plot_swanson(acronyms=None, values=None, ax=None, hemisphere=None, br=None,
                 orientation='landscape', annotate=False, empty_color='silver', **kwargs):
    """
    Displays the 2D image corresponding to the swanson flatmap.
    This case is different from the others in the sense that only a region maps to another regions, there
    is no correspondency from the spatial 3D coordinates.
    :param acronyms:
    :param values:
    :param hemisphere: hemisphere to display, options are 'left', 'right', 'both' or 'mirror'
    :param br: ibllib.atlas.BrainRegions object
    :param ax: matplotlib axis object to plot onto
    :param orientation: 'landscape' (default) or 'portrait'
    :param annotate: (False) if True, labels regions with acronyms
    :param empty_color: (grey) matplotlib color code or rgb_a int8 tuple defining the filling
     of brain regions not provided. Defaults to 'silver'
    :param kwargs: arguments for imshow
    :return:
    """
    mapping = 'Swanson'
    br = BrainRegions() if br is None else br
    s2a = swanson()
    # both hemishpere
    if hemisphere == 'both':
        _s2a = s2a + np.sum(br.id > 0)
        _s2a[s2a == 0] = 0
        _s2a[s2a == 1] = 1
        s2a = np.r_[s2a, np.flipud(_s2a)]
        mapping = 'Swanson-lr'
    elif hemisphere == 'mirror':
        s2a = np.r_[s2a, np.flipud(s2a)]
    if orientation == 'portrait':
        s2a = np.transpose(s2a)
    if acronyms is None:
        regions = br.mappings[mapping][s2a]
        im = br.rgba[regions]
        iswan = None
    else:
        user_aids = br.parse_acronyms_argument(acronyms)
        # if the user provided inputs are higher level than swanson propagate down
        swaids = br.id[np.unique(s2a)]
        maids = np.setdiff1d(user_aids, swaids)  # those are the indices not in Swanson
        for i, maid in enumerate(maids):
            if maid <= 1:
                continue
            childs_in_sw = np.intersect1d(br.descendants(maid)['id'][1:], swaids)
            if childs_in_sw.size > 0:
                user_aids = np.r_[user_aids, childs_in_sw]
                values = np.r_[values, values[i] + childs_in_sw * 0]
        # the user may have input non-unique regions
        df = pd.DataFrame(dict(aid=user_aids, value=values)).groupby('aid').mean()
        aids, vals = (df.index.values, df['value'].values)
        # apply mapping and perform another round of aggregation
        _, _, ibr = np.intersect1d(aids, br.id, return_indices=True)
        ibr = br.mappings['Swanson-lr'][ibr]
        df = pd.DataFrame(dict(ibr=ibr, value=vals)).groupby('ibr').mean()
        ibr, vals = (df.index.values, df['value'].values)
        # we now have the mapped regions and aggregated values, map values onto swanson map
        iswan, iv = ismember(s2a, ibr)
        im = np.zeros_like(s2a, dtype=np.float32)
        im[iswan] = vals[iv]
        im[~iswan] = np.nan
    if not ax:
        ax = plt.gca()
        ax.set_axis_off()  # unless provided we don't need scales here
    ax.imshow(im, **kwargs)
    # overlay the boundaries if value plot
    imb = np.zeros((*s2a.shape[:2], 4), dtype=np.uint8)
    # fill in the empty regions with the blank regions colours if necessary
    if iswan is not None:
        imb[~iswan] = (np.array(matplotlib.colors.to_rgba(empty_color)) * 255).astype('uint8')
    imb[s2a == 0] = 255
    # imb[s2a == 1] = np.array([167, 169, 172, 255])
    imb[s2a == 1] = np.array([0, 0, 0, 255])
    ax.imshow(imb)
    if annotate:
        annotate_swanson(ax=ax, orientation=orientation, br=br)

    # provides the mean to see the region on axis
    def format_coord(x, y):
        acronym = br.acronym[s2a[int(y), int(x)]]
        return f'x={x:1.4f}, y={x:1.4f}, {acronym}'

    ax.format_coord = format_coord
    return ax


@lru_cache(maxsize=None)
def _swanson_labels_positions():
    """
    This functions computes label positions to overlay on the Swanson flatmap
    :return: dictionary where keys are acronyms
    """
    NPIX_THRESH = 20000  # number of pixels above which region is labeled
    s2a = swanson()
    iw, ih = np.meshgrid(np.arange(s2a.shape[1]), np.arange(s2a.shape[0]))
    # compute the center of mass of all regions (fast enough to do on the fly)
    bc = np.maximum(1, np.bincount(s2a.flatten()))
    cmw = np.bincount(s2a.flatten(), weights=iw.flatten()) / bc
    cmh = np.bincount(s2a.flatten(), weights=ih.flatten()) / bc
    bc[0] = 1

    NWH, NWW = (200, 600)
    h, w = s2a.shape
    labels = {}
    for ilabel in np.where(bc > NPIX_THRESH)[0]:
        x, y = (cmw[ilabel], cmh[ilabel])
        # the polygon is convex and the label is outside. Dammit !!!
        if s2a[int(y), int(x)] != ilabel:
            # find the nearest point to the center of mass
            ih, iw = np.where(s2a == ilabel)
            iimin = np.argmin(np.abs((x - iw) + 1j * (y - ih)))
            # get the center of mass of a window around this point
            sh = np.arange(np.maximum(0, ih[iimin] - NWH), np.minimum(ih[iimin] + NWH, h))
            sw = np.arange(np.maximum(0, iw[iimin] - NWW), np.minimum(iw[iimin] + NWW, w))
            roi = s2a[sh][:, sw] == ilabel
            roi = roi / np.sum(roi)
            # ax.plot(x, y, 'k+')
            # ax.plot(iw[iimin], ih[iimin], '*k')
            x = sw[np.searchsorted(np.cumsum(np.sum(roi, axis=0)), .5) - 1]
            y = sh[np.searchsorted(np.cumsum(np.sum(roi, axis=1)), .5) - 1]
            # ax.plot(x, y, 'r+')
        labels[ilabel] = (x, y)
    return labels


def annotate_swanson(ax, acronyms=None, orientation='landscape', br=None, **kwargs):
    """
    Display annotations on the flatmap
    :param ax:
    :param acronyms: (None) list or np.array of acronyms or allen region ids. If None plot all.
    :param orientation:
    :param br: BrainRegions object
    :param kwargs: arguments for the annotate function
    :return:
    """
    br = br or BrainRegions()
    if acronyms is None:
        indices = np.arange(br.id.size)
    else:  # tech debt: here in fact we should remap and compute labels for hierarchical regions
        aids = br.parse_acronyms_argument(acronyms)
        _, indices, _ = np.intersect1d(br.id, br.remap(aids, 'Swanson-lr'), return_indices=True)
    labels = _swanson_labels_positions()
    for ilabel in labels:
        # do not display uwanted labels
        if ilabel not in indices:
            continue
        # rotate the labels if the dislay is in portrait mode
        xy = reversed(labels[ilabel]) if orientation == 'portrait' else labels[ilabel]
        ax.annotate(br.acronym[ilabel], xy=xy, ha='center', va='center', **kwargs)
