"""

sph
===

routines for plotting smoothed quantities

"""

import matplotlib
import numpy as np
import pylab as p

from .. import config, sph, units as _units


def sideon_image(sim, *args, **kwargs):
    """

    Rotate the simulation so that the disc of the passed halo is
    side-on, then make an SPH image by passing the parameters into
    the function image

    For a description of keyword arguments see :func:`~pynbody.plot.sph.image`.

    """

    from ..analysis import angmom

    with angmom.sideon(sim):
        return image(sim, *args, **kwargs)


def faceon_image(sim, *args, **kwargs):
    """

    Rotate the simulation so that the disc of the passed halo is
    face-on, then make an SPH image by passing the parameters into
    the function image

    For a description of keyword arguments see :func:`~pynbody.plot.sph.image`.

    """

    from ..analysis import angmom

    with angmom.faceon(sim):
        return image(sim, *args, **kwargs)


def velocity_image(sim, width="10 kpc", vector_color='black', edgecolor='black', quiverkey_bg_color=None,
                   vector_resolution=40, scale=None, mode='quiver', key_x=0.3, key_y=0.9,
                   key_color='white', key_length="100 km s**-1", quiverkey=True, density=1.0,
                   vector_qty='vel', **kwargs):
    """

    Make an SPH image of the given simulation with velocity vectors overlaid on top.

    For a description of additional keyword arguments see :func:`~pynbody.plot.sph.image`,
    or see the `tutorial <http://pynbody.github.io/pynbody/tutorials/pictures.html#velocity-vectors>`_.

    **Keyword arguments:**

    *vector_color* (black): The color for the velocity vectors

    *edgecolor* (black): edge color used for the lines - using a color
     other than black for the *vector_color* and a black *edgecolor*
     can result in poor readability in pdfs

    *vector_resolution* (40): How many vectors in each dimension (default is 40x40)

    *quiverkey_bg_color* (none): The color for the legend (scale) background

    *scale* (None): The length of a vector that would result in a displayed length of the
    figure width/height.

    *mode* ('quiver'): make a 'quiver' or 'stream' plot

    *key_x* (0.3): Display x (width) position for the vector key (quiver mode only)

    *key_y* (0.9): Display y (height) position for the vector key (quiver mode only)

    *key_color* (white): Color for the vector key (quiver mode only)

    *key_length* (100 km/s): Velocity to use for the vector key (quiver mode only)

    *density* (1.0): Density of stream lines (stream mode only)

    *quiverkey* (True): Whether or not to inset the key

    *vector_qty* ('vel'): The name of the vector field to plot
    """

    subplot = kwargs.get('subplot', False)
    av_z = kwargs.get('av_z',None)
    if subplot:
        p = subplot
    else:
        import matplotlib.pylab as p

    vx_name, vy_name, _ = sim._array_name_ND_to_1D(vector_qty)

    vx = image(sim, qty=vx_name, width=width, log=False,
               resolution=vector_resolution, noplot=True,av_z=av_z)
    vy = image(sim, qty=vy_name, width=width, log=False,
               resolution=vector_resolution, noplot=True,av_z=av_z)
    key_unit = _units.Unit(key_length)

    if isinstance(width, str) or issubclass(width.__class__, _units.UnitBase):
        if isinstance(width, str):
            width = _units.Unit(width)
        width = width.in_units(sim['pos'].units, **sim.conversion_context())

    width = float(width)

    pixel_size = width / float(vector_resolution)
    X, Y = np.meshgrid(np.linspace(-width / 2 + pixel_size/2, width / 2 - pixel_size/2, vector_resolution),
                       np.linspace(-width / 2 + pixel_size/2, width / 2 - pixel_size/2, vector_resolution))

    im = image(sim, width=width, **kwargs)

    if mode == 'quiver':
        if scale is None:
            Q = p.quiver(X, Y, vx, vy, color=vector_color, edgecolor=edgecolor)
        else:
            Q = p.quiver(X, Y, vx, vy, scale=_units.Unit(scale).in_units(
                sim['vel'].units), color=vector_color, edgecolor=edgecolor)
        if quiverkey:
        	qk = p.quiverkey(Q, key_x, key_y, key_unit.in_units(sim['vel'].units, **sim.conversion_context()),
                    r"$\mathbf{" + key_unit.latex() + "}$", labelcolor=key_color, color=key_color, fontproperties={'size': 24})
        if  quiverkey_bg_color is not None:
            qk.text.set_backgroundcolor(quiverkey_bg_color)
    elif mode == 'stream' :
        Q = p.streamplot(X, Y, vx, vy, color=vector_color, density=density)

	# RS - if a axis object is passed in, use the right limit call
    if subplot:
        p.set_xlim(-width/2, width/2)
        p.set_ylim(-width/2, width/2)
    else:
        p.xlim(-width/2, width/2)
        p.ylim(-width/2, width/2)

    return im


def volume(sim, qty='rho', width=None, resolution=200,
           color=(1.0,1.0,1.0),vmin=None,vmax=None,
           dynamic_range=4.0,log=True,
           create_figure=True):
    """Create a volume rendering of the given simulation using mayavi.

    **Keyword arguments:**

    *qty* (rho): The name of the array to interpolate

    *width* (None): The width of the cube to generate, centered on the origin

    *resolution* (200): The number of elements along each side of the cube

    *color* (white): The color of the volume rendering. The value of each voxel
       is used to set the opacity.

    *vmin* (None): The value for zero opacity (calculated using dynamic_range if None)

    *vmax* (None): The value for full opacity (calculated from the maximum
       value in the region if None)

    *dynamic_range*: The dynamic range to use if vmin and vmax are not specified

    *log* (True): log-scale the image before passing to mayavi

    *create_figure* (True): create a new mayavi figure before rendering
    """

    import mayavi
    from mayavi import mlab
    from tvtk.util.ctf import ColorTransferFunction, PiecewiseFunction



    if create_figure:
        fig = mlab.figure(size=(500,500),bgcolor=(0,0,0))

    grid_data = sph.to_3d_grid(sim,qty=qty,nx=resolution,
                               x2=None if width is None else width/2)



    if log:
        grid_data = np.log10(grid_data)
        if vmin is None:
            vmin = grid_data.max()-dynamic_range
        if vmax is None:
            vmax = grid_data.max()
    else:
        if vmin is None:
            vmin = np.min(grid_data)
        if vmax is None:
            vmax = np.max(grid_data)

    grid_data[grid_data<vmin]=vmin
    grid_data[grid_data>vmax]=vmax

    otf = PiecewiseFunction()
    otf.add_point(vmin,0.0)
    otf.add_point(vmax,1.0)

    sf = mayavi.tools.pipeline.scalar_field(grid_data)
    V = mlab.pipeline.volume(sf,color=color,vmin=vmin,vmax=vmax)




    V.trait_get('volume_mapper')['volume_mapper'].blend_mode = 'maximum_intensity'

    if color is None:
        ctf = ColorTransferFunction()
        ctf.add_rgb_point(vmin,107./255,124./255,132./255)
        ctf.add_rgb_point(vmin+(vmax-vmin)*0.8,200./255,178./255,164./255)
        ctf.add_rgb_point(vmin+(vmax-vmin)*0.9,1.0,210./255,149./255)
        ctf.add_rgb_point(vmax,1.0,222./255,141./255)

        V._volume_property.set_color(ctf)
        V._ctf = ctf
        V.update_ctf = True

    V._otf = otf
    V._volume_property.set_scalar_opacity(otf)


    return V

def contour(*args, **kwargs):
    """
    Make an SPH image of the given simulation and render it as contours.
    nlevels and levels are passed to pyplot's contour command.

    Other arguments are as for *image*.
    """

    import copy
    kwargs_image = copy.copy(kwargs)
    nlevels = kwargs_image.pop('nlevels',None)
    levels = kwargs_image.pop('levels',None)
    width = kwargs_image.get('width','10 kpc')
    kwargs_image['noplot']=True
    im = image(*args, **kwargs_image)
    res = im.shape

    units = kwargs_image.get('units',None)

    if isinstance(width, str) or issubclass(width.__class__, _units.UnitBase):
        if isinstance(width, str):
            width = _units.Unit(width)
        sim = args[0]
        width = width.in_units(sim['pos'].units, **sim.conversion_context())

    width = float(width)
    x,y = np.meshgrid(np.linspace(-width/2,width/2,res[0]),np.linspace(-width/2,width/2,res[0]))

    p.contour(x,y,im,nlevels=nlevels,levels=levels)


def _units_imply_projection(sim, qty, units):
    try:
        sim[qty].units.ratio(units, **sim[qty].conversion_context())
        # if this fails, perhaps we're requesting a projected image?
        return False
    except _units.UnitsException:
        # if the following fails, there's no interpretation this routine
        # can cope with. The error will be allowed to propagate.
        sim[qty].units.ratio(
            units / (sim['x'].units), **sim[qty].conversion_context())
        return True


def image(sim, qty='rho', width="10 kpc", resolution=500, units=None, log=True,
          vmin=None, vmax=None, av_z=False, filename=None,
          z_camera=None, clear=True, cmap=None,
          title=None, qtytitle=None, show_cbar=True, subplot=False,
          noplot=False, ret_im=False, fill_nan=True, fill_val=0.0, linthresh=None,
          kernel_type='spline',
          **kwargs):
    """

    Make an SPH image of the given simulation.

    **Keyword arguments:**

    *qty* (rho): The name of the array to interpolate

    *width* (10 kpc): The overall width and height of the plot. If
     ``width`` is a float or an int, then it is assumed to be in units
     of ``sim['pos']``. It can also be passed in as a string
     indicating the units, i.e. '10 kpc', in which case it is
     converted to units of ``sim['pos']``.

    *resolution* (500): The number of pixels wide and tall

    *units* (None): The units of the output

    *av_z* (False): If True, the requested quantity is averaged down
            the line of sight (default False: image is generated in
            the thin plane z=0, unless output units imply an integral
            down the line of sight). If a string, the requested quantity
            is averaged down the line of sight weighted by the av_z
            array (e.g. use 'rho' for density-weighted quantity;
            the default results when av_z=True are volume-weighted).

    *z_camera* (None): If set, a perspective image is rendered. See
                :func:`pynbody.sph.image` for more details.

    *filename* (None): if set, the image will be saved in a file

    *clear* (True): whether to call clf() on the axes first

    *cmap* (None): user-supplied colormap instance

    *title* (None): plot title

    *qtytitle* (None): colorbar quantity title

    *show_cbar* (True): whether to plot the colorbar

    *subplot* (False): the user can supply a AxesSubPlot instance on
    which the image will be shown

    *noplot* (False): do not display the image, just return the image array

    *ret_im* (False): return the image instance returned by imshow

    *num_threads* (None) : if set, specify the number of threads for
    the multi-threaded routines; otherwise the pynbody.config default is used

    *fill_nan* (True): if any of the image values are NaN, replace with fill_val

    *fill_val* (0.0): the fill value to use when replacing NaNs

    *linthresh* (None): if the image has negative and positive values
     and a log scaling is requested, the part between `-linthresh` and
     `linthresh` is shown on a linear scale to avoid divergence at 0

    *kernel_type* ('spline'): SPH kernel to use for smoothing. Defaults to a
     cubic spline, but can also be set to 'wendlandC2'
    """

    if not noplot:
        import matplotlib.pylab as plt

    global config
    if not noplot:
        if subplot:
            p = subplot
        else:
            p = plt

    if qtytitle is None:
        qtytitle = qty

    if isinstance(units, str):
        units = _units.Unit(units)

    if isinstance(width, str) or issubclass(width.__class__, _units.UnitBase):
        if isinstance(width, str):
            width = _units.Unit(width)
        width = width.in_units(sim['pos'].units, **sim.conversion_context())

    width = float(width)

    if kernel_type == 'wendlandC2':
        kernel = sph.WendlandC2Kernel()
    elif kernel_type == 'spline':
        kernel = sph.Kernel()
    else:
        raise ValueError('Invalid kernel_type specified. Options are "spline" (default) and "wendlandC2".')

    perspective = z_camera is not None
    if perspective and not av_z:
        kernel = sph.Kernel2D(kernel)

    is_projected = False
    if units is not None:
        is_projected = _units_imply_projection(sim, qty, units)

    if is_projected:
        kernel = sph.Kernel2D(kernel)

    if av_z:
        if isinstance(kernel, sph.Kernel2D):
            raise _units.UnitsException(
                "Units already imply projected image; can't also average over line-of-sight!")
        else:
            kernel = sph.Kernel2D(kernel)
            if units is not None:
                aunits = units * sim['z'].units
            else:
                aunits = None

            if isinstance(av_z, str):
                if units is not None:
                    aunits = units * sim[av_z].units * sim['z'].units
                sim["__prod"] = sim[av_z] * sim[qty]
                qty = "__prod"

            else:
                av_z = "__one"
                sim["__one"] = np.ones_like(sim[qty])
                sim["__one"].units = "1"

            im = sph.render_image(sim, qty, width / 2, resolution, out_units=aunits, kernel=kernel,
                                  z_camera=z_camera, **kwargs)
            im2 = sph.render_image(sim, av_z, width / 2, resolution, kernel=kernel,
                                   z_camera=z_camera, **kwargs)

            top = sim.ancestor

            try:
                del top["__one"]
            except KeyError:
                pass

            try:
                del top["__prod"]
            except KeyError:
                pass

            im = im / im2

    else:
        im = sph.render_image(sim, qty, width / 2, resolution, out_units=units,
                              kernel=kernel,  z_camera=z_camera, **kwargs)

    if fill_nan:
        im[np.isnan(im)] = fill_val

    if not noplot:

        # set the log or linear normalizations
        if log:
            try:
                im[np.where(im == 0)] = abs(im[np.where(abs(im != 0))]).min()
            except ValueError:
                raise ValueError("Failed to make a sensible logarithmic image. This probably means there are no particles in the view.")

            # check if there are negative values -- if so, use the symmetric
            # log normalization
            if (vmin is None and (im < 0).any() ) or ((vmin is not None) and vmin<0):

                # need to set the linear regime around zero -- set to by
                # default start at 1/1000 of the log range
                if linthresh is None:
                    linthresh = np.nanmax(abs(im)) / 1000.
                norm = matplotlib.colors.SymLogNorm(
                    linthresh, vmin=vmin, vmax=vmax)
            else:
                norm = matplotlib.colors.LogNorm(vmin=vmin, vmax=vmax)

        else:
            norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax)

        #
        # do the actual plotting
        #
        if clear and not subplot:
            p.clf()

        if ret_im:
            return p.imshow(im[::-1, :].view(np.ndarray), extent=(-width / 2, width / 2, -width / 2, width / 2),
                            cmap=cmap, norm = norm)

        ims = p.imshow(im[::-1, :].view(np.ndarray), extent=(-width / 2, width / 2, -width / 2, width / 2),
                       cmap=cmap, norm = norm)

        u_st = sim['pos'].units.latex()
        if not subplot:
            plt.xlabel("$x/%s$" % u_st)
            plt.ylabel("$y/%s$" % u_st)
        else:
            p.set_xlabel("$x/%s$" % u_st)
            p.set_ylabel("$y/%s$" % u_st)

        if units is None:
            units = im.units



        if units.latex() == "":
            units=""
        else:
            units = "$"+units.latex()+"$"

        if show_cbar:
             plt.colorbar(ims).set_label(qtytitle+"/"+units)

        # colorbar doesn't work wtih subplot:  mappable is NoneType
        # elif show_cbar:
        #    import matplotlib.pyplot as mpl
        #    if qtytitle: mpl.colorbar().set_label(qtytitle)
        #    else:        mpl.colorbar().set_label(units)

        if title is not None:
            if not subplot:
                p.title(title)
            else:
                p.set_title(title)

        if filename is not None:
            p.savefig(filename)

        plt.draw()
        # plt.show() - removed by AP on 30/01/2013 - this should not be here as
        # for some systems you don't get back to the command prompt

    return im


def image_radial_profile(im, bins=100):

    xsize, ysize = np.shape(im)
    x = np.arange(-xsize / 2, xsize / 2)
    y = np.arange(-ysize / 2, ysize / 2)
    xs, ys = np.meshgrid(x, y)
    rs = np.sqrt(xs ** 2 + ys ** 2)
    hist, bin_edges = np.histogram(rs, bins=bins)
    inds = np.digitize(rs.flatten(), bin_edges)
    ave_vals = np.zeros(bin_edges.size)
    max_vals = np.zeros(bin_edges.size)
    min_vals = np.zeros(bin_edges.size)
    for i in np.arange(bin_edges.size):
        try:
            min_vals[i] = np.min(10 ** (im.flatten()[np.where(inds == i)]))
        except ValueError:
            min_vals[i] = float('nan')
        ave_vals[i] = np.mean(10 ** (im.flatten()[np.where(inds == i)]))
        try:
            max_vals[i] = np.max(10 ** (im.flatten()[np.where(inds == i)]))
        except ValueError:
            max_vals[i] = float('nan')

    return ave_vals, min_vals, max_vals, bin_edges
