"""
Skeleton objects that represent a few fundamental yt data types.



"""

#-----------------------------------------------------------------------------
# Copyright (c) 2013, yt Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------

import numpy as np
import abc
import json
import sys
import h5py as h5
import os
from uuid import uuid4
from yt.extern.six.moves import urllib
from yt.extern.six.moves import cPickle as pickle
from tempfile import TemporaryFile
from yt.config import ytcfg
from yt.funcs import \
    iterable, get_pbar, compare_dicts
from yt.extern.six import add_metaclass, string_types, b
from yt.utilities.exceptions import \
    YTHubRegisterError
from yt.utilities.logger import ytLogger as mylog
from yt.units.yt_array import \
    YTArray, \
    YTQuantity

if sys.version_info < (3, 0):
    from .poster.streaminghttp import register_openers
    from .poster.encode import multipart_encode
    register_openers()
else:
    # We don't yet have a solution for this, but it won't show up very often
    # anyway.
    pass


def _sanitize_list(flist):
    temp = []
    for item in flist:
        if isinstance(item, string_types):
            temp.append(b(item))
        elif isinstance(item, tuple) and \
                all(isinstance(i, string_types) for i in item):
            temp.append(tuple(_sanitize_list(list(item))))
        else:
            temp.append(item)
    return temp


def _serialize_to_h5(g, cdict):
    for item in cdict:
        if isinstance(cdict[item], (YTQuantity, YTArray)):
            g[item] = cdict[item].d
            g[item].attrs["units"] = str(cdict[item].units)
        elif isinstance(cdict[item], dict):
            _serialize_to_h5(g.create_group(item), cdict[item])
        elif cdict[item] is None:
            g[item] = "None"
        elif isinstance(cdict[item], list):
            g[item] = _sanitize_list(cdict[item])
        elif isinstance(cdict[item], tuple) and \
                all(isinstance(i, string_types) for i in cdict[item]):
            g[item] = tuple(_sanitize_list(cdict[item]))
        else:
            g[item] = cdict[item]


def _deserialize_from_h5(g, ds):
    result = {}
    for item in g:
        if item == "chunks":
            continue
        if "units" in g[item].attrs:
            if iterable(g[item]):
                result[item] = ds.arr(g[item][:], g[item].attrs["units"])
            else:
                result[item] = ds.quan(g[item][()],
                                       g[item].attrs["units"])
        elif isinstance(g[item], h5.Group):
            result[item] = _deserialize_from_h5(g[item], ds)
        elif g[item] == "None":
            result[item] = None
        else:
            try:
                result[item] = g[item][:]   # try array
            except ValueError:
                result[item] = g[item][()]  # fallback to scalar
    return result


class UploaderBar(object):
    pbar = None

    def __init__(self, my_name=""):
        self.my_name = my_name

    def __call__(self, name, prog, total):
        if self.pbar is None:
            self.pbar = get_pbar("Uploading %s " % self.my_name, total)
        self.pbar.update(prog)
        if prog == total:
            self.pbar.finish()


class ContainerClass(object):
    pass


@add_metaclass(abc.ABCMeta)
class MinimalRepresentation(object):

    def _update_attrs(self, obj, attr_list):
        for attr in attr_list:
            setattr(self, attr, getattr(obj, attr, None))
        if hasattr(obj, "ds"):
            self.output_hash = obj.ds._hash()
            self._ds_mrep = obj.ds._mrep
        if hasattr(obj, "data_source"):
            self.data_source_hash = obj.data_source._hash

    def __init__(self, obj):
        self._update_attrs(obj, self._attr_list)

    @abc.abstractmethod
    def _generate_post(self):
        pass

    @abc.abstractproperty
    def _attr_list(self):
        pass

    def _return_filtered_object(self, attrs):
        new_attrs = tuple(attr for attr in self._attr_list
                          if attr not in attrs)
        new_class = type('Filtered%s' % self.__class__.__name__,
                         (FilteredRepresentation,),
                         {'_attr_list': new_attrs})
        return new_class(self)

    @property
    def _attrs(self):
        return dict(((attr, getattr(self, attr)) for attr in self._attr_list))

    @classmethod
    def _from_metadata(cls, metadata):
        cc = ContainerClass()
        for a, v in metadata.values():
            setattr(cc, a, v)
        return cls(cc)

    def store(self, storage):
        if hasattr(self, "_ds_mrep"):
            self._ds_mrep.store(storage)
        metadata, (final_name, chunks) = self._generate_post()
        metadata['obj_type'] = self.type
        with h5.File(storage) as h5f:
            dset = str(uuid4())[:8]
            grp = h5f.create_group(dset)
            _serialize_to_h5(h5f[dset], metadata)
            if len(chunks) > 0:
                g = h5f[dset].create_group('chunks')
                g.attrs['final_name'] = final_name
                for fname, fdata in chunks:
                    if isinstance(fname, (tuple, list)):
                        fname = "*".join(fname)

                    if isinstance(fdata, (YTQuantity, YTArray)):
                        g.create_dataset(fname, data=fdata.d,
                                         compression="lzf")
                        g[fname].attrs["units"] = str(fdata.units)
                    else:
                        g.create_dataset(fname, data=fdata, compression="lzf")

    def restore(self, storage, ds):
        pass

    def upload(self):
        api_key = ytcfg.get("yt", "hub_api_key")
        url = ytcfg.get("yt", "hub_url")
        if api_key == '':
            raise YTHubRegisterError
        metadata, (final_name, chunks) = self._generate_post()
        if hasattr(self, "_ds_mrep"):
            self._ds_mrep.upload()
        for i in metadata:
            if isinstance(metadata[i], np.ndarray):
                metadata[i] = metadata[i].tolist()
            elif hasattr(metadata[i], 'dtype'):
                metadata[i] = np.asscalar(metadata[i])
        metadata['obj_type'] = self.type
        if len(chunks) == 0:
            chunk_info = {'chunks': []}
        else:
            chunk_info = {'final_name': final_name, 'chunks': []}
            for cn, cv in chunks:
                chunk_info['chunks'].append((cn, cv.size * cv.itemsize))
        metadata = json.dumps(metadata)
        chunk_info = json.dumps(chunk_info)
        datagen, headers = multipart_encode({'metadata': metadata,
                                             'chunk_info': chunk_info,
                                             'api_key': api_key})
        request = urllib.request.Request(url, datagen, headers)
        # Actually do the request, and get the response
        try:
            rv = urllib.request.urlopen(request).read()
        except urllib.error.HTTPError as ex:
            if ex.code == 401:
                mylog.error("You must create an API key before uploading.")
                mylog.error("https://data.yt-project.org/getting_started.html")
                return
            else:
                raise ex
        uploader_info = json.loads(rv)
        new_url = url + "/handler/%s" % uploader_info['handler_uuid']
        for i, (cn, cv) in enumerate(chunks):
            remaining = cv.size * cv.itemsize
            f = TemporaryFile()
            np.save(f, cv)
            f.seek(0)
            pbar = UploaderBar("%s, % 2i/% 2i" %
                               (self.type, i + 1, len(chunks)))
            datagen, headers = multipart_encode({'chunk_data': f}, cb=pbar)
            request = urllib.request.Request(new_url, datagen, headers)
            rv = urllib.request.urlopen(request).read()

        datagen, headers = multipart_encode({'status': 'FINAL'})
        request = urllib.request.Request(new_url, datagen, headers)
        rv = json.loads(urllib.request.urlopen(request).read())
        mylog.info("Upload succeeded!  View here: %s", rv['url'])
        return rv

    def load(self, storage):
        return pickle.load(open(storage, 'r'))

    def dump(self, storage):
        with open(storage, 'w') as fh:
            pickle.dump(self, fh)


class FilteredRepresentation(MinimalRepresentation):

    def _generate_post(self):
        raise RuntimeError


class MinimalDataset(MinimalRepresentation):
    _attr_list = ("dimensionality", "refine_by", "domain_dimensions",
                  "current_time", "domain_left_edge", "domain_right_edge",
                  "unique_identifier", "current_redshift", "output_hash",
                  "cosmological_simulation", "omega_matter", "omega_lambda",
                  "hubble_constant", "name")
    type = 'simulation_output'

    def __init__(self, obj):
        super(MinimalDataset, self).__init__(obj)
        self.output_hash = obj._hash()
        self.name = str(obj)

    def _generate_post(self):
        metadata = self._attrs
        chunks = []
        return (metadata, (None, chunks))


class MinimalMappableData(MinimalRepresentation):

    _attr_list = ("field_data", "field", "weight_field", "axis", "output_hash",
                  "vm_type")

    def _generate_post(self):
        nobj = self._return_filtered_object(("field_data",))
        metadata = nobj._attrs
        chunks = [(arr, self.field_data[arr]) for arr in self.field_data]
        return (metadata, ('field_data', chunks))

    def _read_chunks(self, g, ds):
        for fname in g.keys():
            if '*' in fname:
                arr = tuple(fname.split('*'))
            else:
                arr = fname
            try:
                self.field_data[arr] = ds.arr(g[fname][:],
                                              g[fname].attrs["units"])
            except KeyError:
                self.field_data[arr] = g[fname][:]


class MinimalProjectionData(MinimalMappableData):
    type = 'proj'
    vm_type = "Projection"
    _attr_list = ("field_data", "field", "weight_field", "axis", "output_hash",
                  "center", "method", "field_parameters",
                  "data_source_hash")

    def restore(self, storage, ds):
        if hasattr(self, "_ds_mrep"):
            self._ds_mrep.restore(storage, ds)
        metadata, (final_name, chunks) = self._generate_post()
        with h5.File(storage, 'r') as h5f:
            for dset in h5f:
                stored_metadata = _deserialize_from_h5(h5f[dset], ds)
                if compare_dicts(metadata, stored_metadata):
                    self._read_chunks(h5f[dset]["chunks"], ds)
                    return True
        return False


class MinimalSliceData(MinimalMappableData):
    type = 'slice'
    vm_type = "Slice"
    weight_field = "None"


class MinimalImageCollectionData(MinimalRepresentation):
    type = "image_collection"
    _attr_list = ("name", "output_hash", "images", "image_metadata")

    def _generate_post(self):
        nobj = self._return_filtered_object(("images",))
        metadata = nobj._attrs
        chunks = [(fn, d) for fn, d in self.images]
        return (metadata, ('images', chunks))

_hub_categories = ("News", "Documents", "Simulation Management",
                   "Data Management", "Analysis and Visualization",
                   "Paper Repositories", "Astrophysical Utilities",
                   "yt Scripts")


class MinimalProjectDescription(MinimalRepresentation):
    type = "project"
    _attr_list = ("title", "url", "description", "category", "image_url")

    def __init__(self, title, url, description,
                 category, image_url=""):
        assert(category in _hub_categories)
        self.title = title
        self.url = url
        self.description = description
        self.category = category
        self.image_url = image_url

    def _generate_post(self):
        metadata = self._attrs
        return (metadata, ("chunks", []))


class MinimalNotebook(MinimalRepresentation):
    type = "notebook"
    _attr_list = ("title",)

    def __init__(self, filename, title=None):
        # First we read in the data
        if not os.path.isfile(filename):
            raise IOError(filename)
        self.data = open(filename).read()
        if title is None:
            title = json.loads(self.data)['metadata']['name']
        self.title = title
        self.data = np.fromstring(self.data, dtype='c')

    def _generate_post(self):
        metadata = self._attrs
        chunks = [("notebook", self.data)]
        return (metadata, ("chunks", chunks))


class ImageCollection(object):

    def __init__(self, ds, name):
        self.ds = ds
        self.name = name
        self.images = []
        self.image_metadata = []

    def add_image(self, fn, descr):
        self.image_metadata.append(descr)
        self.images.append((os.path.basename(fn), np.fromfile(fn, dtype='c')))
