import os
import logging
import warnings
import contextlib

import numpy as np
from ncempy.io.ser import fileSER

from libertem.common import Shape, Slice
from libertem.web.messages import MessageConverter
from .base import (
    DataSet, FileSet, BasePartition, DataSetException, DataSetMeta,
    _roi_to_indices, DataTile,
)

log = logging.getLogger(__name__)


class SERDatasetParams(MessageConverter):
    SCHEMA = {
        "$schema": "http://json-schema.org/draft-07/schema#",
        "$id": "http://libertem.org/SERDatasetParams.schema.json",
        "title": "SERDatasetParams",
        "type": "object",
        "properties": {
            "type": {"const": "SER"},
            "path": {"type": "string"},
            "nav_shape": {
                "type": "array",
                "items": {"type": "number", "minimum": 1},
                "minItems": 2,
                "maxItems": 2
            },
            "sig_shape": {
                "type": "array",
                "items": {"type": "number", "minimum": 1},
                "minItems": 2,
                "maxItems": 2
            },
            "sync_offset": {"type": "number"},
        },
        "required": ["type", "path"]
        }

    def convert_to_python(self, raw_data):
        data = {
            k: raw_data[k]
            for k in ["path"]
        }
        if "nav_shape" in raw_data:
            data["nav_shape"] = tuple(raw_data["nav_shape"])
        if "sig_shape" in raw_data:
            data["sig_shape"] = tuple(raw_data["sig_shape"])
        if "sync_offset" in raw_data:
            data["sync_offset"] = raw_data["sync_offset"]
        return data


class SERFile:
    def __init__(self, path, num_frames):
        self._path = path
        self._num_frames = num_frames

    def _get_handle(self):
        return fileSER(self._path)

    @contextlib.contextmanager
    def get_handle(self):
        with self._get_handle() as f:
            yield f

    @property
    def num_frames(self):
        return self._num_frames

    @property
    def start_idx(self):
        return 0

    @property
    def end_idx(self):
        return self.num_frames


class SERFileSet(FileSet):
    pass


class SERDataSet(DataSet):
    """
    Read TIA SER files.

    Examples
    --------

    >>> ds = ctx.load("ser", path="/path/to/file.ser")  # doctest: +SKIP

    Parameters
    ----------
    path: str
        Path to the .ser file

    nav_shape: tuple of int, optional
        A n-tuple that specifies the size of the navigation region ((y, x), but
        can also be of length 1 for example for a line scan, or length 3 for
        a data cube, for example)

    sig_shape: tuple of int, optional
        Signal/detector size (height, width)

    sync_offset: int, optional
        If positive, number of frames to skip from start
        If negative, number of blank frames to insert at start
    """
    def __init__(self, path, emipath=None, nav_shape=None,
                 sig_shape=None, sync_offset=0, io_backend=None):
        super().__init__(io_backend=io_backend)
        if io_backend is not None:
            raise ValueError("SERDataSet currently doesn't support alternative I/O backends")
        self._path = path
        self._meta = None
        self._filesize = None
        self._num_frames = None
        if emipath is not None:
            warnings.warn(
                "emipath is not used anymore, as it was removed from ncempy", DeprecationWarning
            )
        self._nav_shape = tuple(nav_shape) if nav_shape else nav_shape
        self._sig_shape = tuple(sig_shape) if sig_shape else sig_shape
        self._sync_offset = sync_offset

    def _do_initialize(self):
        self._filesize = os.stat(self._path).st_size
        reader = SERFile(path=self._path, num_frames=None)

        with reader.get_handle() as f1:
            self._num_frames = f1.head['ValidNumberElements']
            if f1.head['ValidNumberElements'] == 0:
                raise DataSetException("no data found in file")

            data, meta_data = f1.getDataset(0)
            dtype = f1._dictDataType[meta_data['DataType']]
            nav_dims = tuple(
                reversed([
                    int(dim['DimensionSize'])
                    for dim in f1.head['Dimensions']
                ])
            )
            self._image_count = int(self._num_frames)
            if self._nav_shape is None:
                self._nav_shape = nav_dims
            if self._sig_shape is None:
                self._sig_shape = tuple(data.shape)
            elif int(np.prod(self._sig_shape)) != int(np.prod(data.shape)):
                raise DataSetException(
                    "sig_shape must be of size: %s" % int(np.prod(data.shape))
                )
            self._nav_shape_product = int(np.prod(self._nav_shape))
            self._sync_offset_info = self.get_sync_offset_info()
            self._shape = Shape(self._nav_shape + self._sig_shape, sig_dims=len(self._sig_shape))
            self._meta = DataSetMeta(
                shape=self._shape,
                raw_dtype=dtype,
                sync_offset=self._sync_offset,
                image_count=self._image_count,
            )
        return self

    def initialize(self, executor):
        return executor.run_function(self._do_initialize)

    @classmethod
    def get_msg_converter(cls):
        return SERDatasetParams

    @classmethod
    def get_supported_extensions(cls):
        return set(["ser"])

    @classmethod
    def detect_params(cls, path, executor):
        if path.lower().endswith(".ser"):
            ds = cls(path)
            ds = ds.initialize(executor)
            return {
                "parameters": {
                    "path": path,
                    "nav_shape": tuple(ds.shape.nav),
                    "sig_shape": tuple(ds.shape.sig),
                },
                "info": {
                    "image_count": int(np.prod(ds.shape.nav)),
                    "native_sig_shape": tuple(ds.shape.sig),
                }
            }
        return False

    @property
    def dtype(self):
        return self._meta.raw_dtype

    @property
    def shape(self):
        return self._meta.shape

    def check_valid(self):
        try:
            with fileSER(self._path) as f1:
                if f1.head['ValidNumberElements'] == 0:
                    raise DataSetException("no data found in file")
                if f1.head['DataTypeID'] not in (0x4120, 0x4122):
                    raise DataSetException("unknown datatype id: %s" % f1.head['DataTypeID'])
            return True
        except (IOError, OSError) as e:
            raise DataSetException("invalid dataset: %s" % e) from e

    def get_cache_key(self):
        return {
            "path": self._path,
            "shape": tuple(self.shape),
            "sync_offset": self._sync_offset,
        }

    def get_num_partitions(self):
        """
        returns the number of partitions the dataset should be split into
        """
        # let's try to aim for 512MB per partition
        res = max(self._cores, self._filesize // (64*1024*1024))
        return res

    def _get_fileset(self):
        assert self._num_frames is not None
        return SERFileSet([
            SERFile(
                path=self._path,
                num_frames=self._num_frames,
            )
        ])

    def get_partitions(self):
        fileset = self._get_fileset()
        for part_slice, start, stop in self.get_slices():
            yield SERPartition(
                path=self._path,
                meta=self._meta,
                partition_slice=part_slice,
                fileset=fileset,
                start_frame=start,
                num_frames=stop - start,
                io_backend=self.get_io_backend(),
            )

    def __repr__(self):
        return "<SERDataSet for %s>" % (self._path,)


class SERPartition(BasePartition):
    def __init__(self, path, *args, **kwargs):
        self._path = path
        super().__init__(*args, **kwargs)

    def validate_tiling_scheme(self, tiling_scheme):
        supported = (1,) + tuple(self.shape.sig)
        if tuple(tiling_scheme.shape) != supported:
            raise ValueError(
                "invalid tiling scheme: only supports %r, not %r" % (supported, tiling_scheme.shape)
            )

    def adjust_tileshape(self, tileshape, roi):
        # force single-frame tiles
        return (1,) + tileshape[1:]

    def get_base_shape(self, roi):
        return (1,) + tuple(self.shape.sig)

    def _preprocess(self, tile_data, tile_slice):
        if self._corrections is None:
            return
        self._corrections.apply(tile_data, tile_slice)

    def get_tiles(self, tiling_scheme, dest_dtype="float32", roi=None):
        sync_offset = self.meta.sync_offset
        shape = Shape((1,) + tuple(self.shape.sig), sig_dims=self.shape.sig.dims)

        self.validate_tiling_scheme(tiling_scheme)

        start = self._start_frame
        if start < self.meta.image_count:
            stop = min(start + self._num_frames, self.meta.image_count)
            if roi is None:
                indices = np.arange(max(0, start), stop)
                # in case of a negative sync_offset, 'start' can be negative
                if start < 0:
                    offset = abs(sync_offset)
                else:
                    offset = start - sync_offset
            else:
                indices = _roi_to_indices(roi, max(0, start), stop, sync_offset)
                # in case of a negative sync_offset, 'start' can be negative
                if start < 0:
                    offset = np.count_nonzero(roi.reshape((-1,))[:abs(sync_offset)])
                else:
                    offset = np.count_nonzero(roi.reshape((-1,))[:start - sync_offset])

            with fileSER(self._path) as f:
                for num, idx in enumerate(indices):
                    origin = (num + offset,) + tuple([0] * self.shape.sig.dims)
                    tile_slice = Slice(origin=origin, shape=shape)
                    data, metadata = f.getDataset(int(idx))
                    if data.dtype != np.dtype(dest_dtype):
                        data = data.astype(dest_dtype)
                    data = data.reshape(shape)
                    self._preprocess(data, tile_slice)
                    yield DataTile(
                        data,
                        tile_slice=tile_slice,
                        scheme_idx=0,
                    )
