import os
import subprocess
import numpy as np
import time
from .utils import ProgressBar
from astropy import log
import warnings

__all__ = ['ytCube']

class ytCube(object):
    """ Light wrapper of a yt object with ability to translate yt<->wcs
    coordinates """

    def __init__(self, cube, dataset, spectral_factor=1.0):
        self.cube = cube
        self.wcs = cube.wcs
        self.dataset = dataset
        self.spectral_factor = spectral_factor


    def world2yt(self, world_coord, first_index=0):
        """
        Convert a position in world coordinates to the coordinates used by a
        yt dataset that has been generated using the ``to_yt`` method.

        Parameters
        ----------
        world_coord: `astropy.wcs.WCS.wcs_world2pix`-valid input
            The world coordinates
        first_index: 0 or 1
            The first index of the data.  In python and yt, this should be
            zero, but for the FITS coordinates, use 1
        """
        yt_coord = self.wcs.wcs_world2pix([world_coord], first_index)[0]
        yt_coord[2] = (yt_coord[2] - 0.5)*self.spectral_factor+0.5
        return yt_coord

    def yt2world(self, yt_coord, first_index=0):
        """
        Convert a position in yt's coordinates to world coordinates from a
        yt dataset that has been generated using the ``to_yt`` method.

        Parameters
        ----------
        world_coord: `astropy.wcs.WCS.wcs_pix2world`-valid input
            The yt pixel coordinates to convert back to world coordinates
        first_index: 0 or 1
            The first index of the data.  In python and yt, this should be
            zero, but for the FITS coordinates, use 1
        """
        yt_coord = np.array(yt_coord) # stripping off units
        yt_coord[2] = (yt_coord[2] - 0.5)/self.spectral_factor+0.5
        world_coord = self.wcs.wcs_pix2world([yt_coord], first_index)[0]
        return world_coord


    def quick_render_movie(self, outdir, size=256, nframes=30,
                           camera_angle=(0,0,1), north_vector=(0,0,1),
                           rot_vector=(1,0,0),
                           colormap='doom',
                           cmap_range='auto',
                           transfer_function='auto',
                           start_index=0,
                           image_prefix="",
                           output_filename='out.mp4',
                           log_scale=False,
                           rescale=True):
        """
        Create a movie rotating the cube 360 degrees from
        PP -> PV -> PP -> PV -> PP

        Parameters
        ----------
        outdir: str
            The output directory in which the individual image frames and the
            resulting output mp4 file should be stored
        size: int
            The size of the individual output frame in pixels (i.e., size=256
            will result in a 256x256 image)
        nframes: int
            The number of frames in the resulting movie
        camera_angle: 3-tuple
            The initial angle of the camera
        north_vector: 3-tuple
            The vector of 'north' in the data cube.  Default is coincident with
            the spectral axis
        rot_vector: 3-tuple
            The vector around which the camera will be rotated
        colormap: str
            A valid colormap.  See `yt.show_colormaps`
        transfer_function: 'auto' or `yt.visualization.volume_rendering.TransferFunction`
            Either 'auto' to use the colormap specified, or a valid
            TransferFunction instance
        log_scale: bool
            Should the colormap be log scaled?
        rescale: bool
            If True, the images will be rescaled to have a common 95th
            percentile brightness, which can help reduce flickering from having
            a single bright pixel in some projections
        start_index : int
            The number of the first image to save
        image_prefix : str
            A string to prepend to the image name for each image that is output
        output_filename : str
            The movie file name to output.  The suffix may affect the file type
            created.  Defaults to 'out.mp4'.  Will be placed in ``outdir``

        Returns
        -------


        """
        try:
            import yt
        except ImportError:
            raise ImportError("yt could not be imported.  Cube renderings are not possible.")

        scale = np.max(self.cube.shape)

        if not os.path.exists(outdir):
            os.makedirs(outdir)
        elif not os.path.isdir(outdir):
            raise OSError("Output directory {0} exists and is not a directory.".format(outdir))

        if cmap_range == 'auto':
            upper = self.cube.max().value
            lower = self.cube.std().value * 3
            cmap_range = [lower,upper]

        if transfer_function == 'auto':
            tfh = self.auto_transfer_function(cmap_range, log=log_scale)
            tfh.tf.map_to_colormap(cmap_range[0], cmap_range[1], colormap=colormap)
            tf = tfh.tf
        else:
            tf = transfer_function

        center = self.dataset.domain_center
        cam = self.dataset.h.camera(center, camera_angle, scale, size, tf,
                                    north_vector=north_vector, fields='flux')

        im  = cam.snapshot()
        images = [im]

        pb = ProgressBar(nframes)
        for ii,im in enumerate(cam.rotation(2 * np.pi, nframes,
                                            rot_vector=rot_vector)):
            images.append(im)
            im.write_png(os.path.join(outdir,"%s%04i.png" % (image_prefix,
                                                             ii+start_index)),
                         rescale=False)
            pb.update(ii+1)
        log.info("Rendering complete in {0}s".format(time.time() - pb._start_time))

        if rescale:
            _rescale_images(images, os.path.join(outdir, image_prefix))

        pipe = _make_movie(outdir, prefix=image_prefix,
                           filename=output_filename)

        return images

    def auto_transfer_function(self, cmap_range, log=False, colormap='doom',
                               **kwargs):

        from yt.visualization.volume_rendering.transfer_function_helper import TransferFunctionHelper
        tfh = TransferFunctionHelper(self.dataset)
        tfh.set_field('flux')
        tfh.set_bounds(bounds=cmap_range)
        tfh.set_log(log)
        tfh.build_transfer_function()
        return tfh

    def quick_isocontour(self, level='3 sigma', title='', description='',
                         color_map='hot', color_log=False,
                         export_to='sketchfab', filename=None,
                         **kwargs):
        """
        Export isocontours to sketchfab

        Requires that you have an account on https://sketchfab.com and are
        logged in

        Parameters
        ----------
        level: str or float
            The level of the isocontours to create.  Can be specified as
            n-sigma with strings like '3.3 sigma' or '2 sigma' (there must be a
            space between the number and the word)
        title: str
            A title for the uploaded figure
        description: str
            A short description for the uploaded figure
        color_map: str
            Any valid colormap.  See `yt.show_colormaps`
        color_log: bool
            Whether the colormap should be log scaled.  With the default
            parameters, this has no effect.
        export_to: 'sketchfab', 'obj', 'ply'
            You can export to sketchfab, to a .obj file (and accompanying .mtl
            file), or a .ply file.  The latter two require ``filename``
            specification
        filename: None or str
            Optional - prefix for output filenames if ``export_to`` is 'obj',
            or the full filename when ``export_to`` is 'ply'.  Ignored for
            'sketchfab'
        kwargs: dict
            Keyword arguments are passed to the appropriate yt function

        Returns
        -------
        The result of the `yt.surface.export_sketchfab` function
        """
        if isinstance(level, str):
            sigma = self.cube.std().value
            level = float(level.split()[0]) * sigma

        self.dataset.periodicity = (True,True,True)
        surface = self.dataset.surface(self.dataset.all_data(),
                                       "flux",
                                       level)
        if export_to == 'sketchfab':
            if filename is not None:
                warnings.warn("sketchfab export does not expect a filename entry")
            return surface.export_sketchfab(title=title,
                                            description=description,
                                            color_map=color_map,
                                            color_log=color_log, **kwargs)
        elif export_to == 'obj':
            if filename is None:
                raise ValueError("If export_to is not 'sketchfab',"
                                 " a filename must be specified")

            surface.export_obj(filename, color_field='ones',
                               color_map=color_map, color_log=color_log,
                               **kwargs)
        elif export_to == 'ply':
            if filename is None:
                raise ValueError("If export_to is not 'sketchfab',"
                                 " a filename must be specified")

            surface.export_ply(filename, color_field='ones',
                               color_map=color_map, color_log=color_log,
                               **kwargs)
        else:
            raise ValueError("export_to must be one of sketchfab,obj,ply")


def _rescale_images(images, prefix):
    """
    Save a sequence of images, at a common scaling
    Reduces flickering
    """

    cmax = max(np.percentile(i[:, :, :3].sum(axis=2), 99.5) for i in images)
    amax = max(np.percentile(i[:, :, 3], 95) for i in images)

    for i, image in enumerate(images):
        image = image.rescale(cmax=cmax, amax=amax).swapaxes(0,1)
        image.write_png("%s%04i.png" % (prefix, i), rescale=False)

def _make_movie(moviepath, prefix="", filename='out.mp4', overwrite=True):
    """
    Use ffmpeg to generate a movie from the image series
    """

    outpath = os.path.join(moviepath, filename)

    if os.path.exists(outpath) and overwrite:
        command = ['ffmpeg', '-y', '-r','5','-i',
                   os.path.join(moviepath,prefix+'%04d.png'),
                   '-r','30','-pix_fmt', 'yuv420p',
                   outpath]
    elif os.path.exists(outpath):
        log.info("File {0} exists - skipping".format(outpath))
    else:
        command = ['ffmpeg', '-r', '5', '-i',
                   os.path.join(moviepath,prefix+'%04d.png'),
                   '-r','30','-pix_fmt', 'yuv420p',
                   outpath]

    pipe = subprocess.Popen(command, stdout=subprocess.PIPE, close_fds=True)

    pipe.wait()

    return pipe