from ....fit import HistContainer, IndexedContainer, UnbinnedContainer, XYContainer
from .. import _AVAILABLE_REPRESENTATIONS
from .._yaml_base import YamlReaderMixin, YamlWriterMixin
from ..error import common_error_tools
from ._base import DataContainerDReprBase

__all__ = ["DataContainerYamlReader", "DataContainerYamlWriter"]


class DataContainerYamlWriter(YamlWriterMixin, DataContainerDReprBase):
    DUMPER = common_error_tools.MatrixYamlDumper

    def __init__(self, data_container, output_io_handle):
        super(DataContainerYamlWriter, self).__init__(output_io_handle=output_io_handle, data_container=data_container)

    @classmethod
    def _make_representation(cls, container):
        """Create a dictionary representing a data container.

        :param container: The data container which will be converted.
        :type container: kafe2.fit._base.DataContainerBase
        """
        _yaml_doc = dict()
        _class = container.__class__

        # -- determine container type
        _type = cls._CLASS_TO_OBJECT_TYPE_NAME.get(_class, None)
        if _type is None:
            raise TypeError("Container type unknown or not supported: %s" % _class)
        _yaml_doc["type"] = _type

        # -- write representation for container types
        if _class is HistContainer:
            _yaml_doc["bin_edges"] = list(map(float, container.bin_edges))
            if container._manual_heights:
                _yaml_doc["bin_heights"] = list(map(float, container.data))  # float64 -> float
                _yaml_doc["underflow"] = float(container.underflow)
                _yaml_doc["overflow"] = float(container.underflow)
            else:
                _yaml_doc["raw_data"] = list(map(float, container.raw_data))  # float64 -> float
        elif _class is IndexedContainer or _class is UnbinnedContainer:
            _yaml_doc["data"] = container.data.tolist()
        elif _class is XYContainer:
            _yaml_doc["x_data"] = container.x.tolist()
            _yaml_doc["y_data"] = container.y.tolist()
        else:
            raise TypeError("Container type unknown or not supported: {}".format(_type))

        # -- write error representation for all container types
        if container.has_errors:
            common_error_tools.write_errors_to_yaml(container, _yaml_doc)

        # write labels for all container types
        if container.label is not None:
            _yaml_doc["label"] = container.label
        _x_axis_label, _y_axis_label = container.axis_labels
        if _x_axis_label is not None:
            _yaml_doc["x_label"] = _x_axis_label
        if _y_axis_label is not None:
            _yaml_doc["y_label"] = _y_axis_label

        return _yaml_doc


class DataContainerYamlReader(YamlReaderMixin, DataContainerDReprBase):
    LOADER = common_error_tools.MatrixYamlLoader

    def __init__(self, input_io_handle):
        super(DataContainerYamlReader, self).__init__(input_io_handle=input_io_handle, data_container=None)

    @classmethod
    def _get_required_keywords(cls, yaml_doc, container_class):
        if container_class is HistContainer:
            return []
        if container_class is IndexedContainer or container_class is UnbinnedContainer:
            return ["data"]
        if container_class is XYContainer:
            return ["x_data", "y_data"]
        raise TypeError("Unknown container type")

    @classmethod
    def _convert_yaml_doc_to_object(cls, yaml_doc):
        # -- determine container class from type
        _container_type = yaml_doc.pop("type")
        _class = cls._OBJECT_TYPE_NAME_TO_CLASS.get(_container_type)

        # -- read in representation for container types
        if _class is HistContainer:
            _bin_edges = yaml_doc.pop("bin_edges", None)
            _n_bins = yaml_doc.pop("n_bins", None)
            _bin_range = yaml_doc.pop("bin_range", None)
            if not _bin_edges and not (_n_bins and _bin_range):
                raise ValueError("When reading in a histogram dataset either " "bin_edges or n_bins and bin_range has to be specified!")
            if _bin_edges:
                _n_bins = len(_bin_edges) - 1
                _bin_range = (_bin_edges[0], _bin_edges[-1])
            _raw_data = yaml_doc.pop("raw_data", None)
            _bin_heights = yaml_doc.pop("bin_heights", None)
            if _raw_data and _bin_heights:
                raise ValueError("When reading in a histogram dataset only one out of " "raw_data and bin_heights can be specified!")
            _container_obj = HistContainer(n_bins=_n_bins, bin_range=_bin_range, bin_edges=_bin_edges, fill_data=_raw_data)
            if _bin_heights:
                _underflow = yaml_doc.pop("underflow", 0)
                _overflow = yaml_doc.pop("overflow", 0)
                _container_obj.set_bins(bin_heights=_bin_heights, underflow=_underflow, overflow=_overflow)
        elif _class is IndexedContainer:
            _data = yaml_doc.pop("data")
            _container_obj = IndexedContainer(_data)
        elif _class is UnbinnedContainer:
            _data = yaml_doc.pop("data")
            _container_obj = UnbinnedContainer(_data)
        elif _class is XYContainer:
            _x_data = yaml_doc.pop("x_data")
            _y_data = yaml_doc.pop("y_data")
            _container_obj = XYContainer(_x_data, _y_data)
        else:
            raise TypeError("Container type unknown or not supported: {}".format(_container_type))

        # get labels for all container types
        _container_obj.label = yaml_doc.pop("label", None)
        _container_obj.axis_labels = (yaml_doc.pop("x_label", None), yaml_doc.pop("y_label", None))

        # -- process error sources
        _container_obj, yaml_doc = common_error_tools.process_error_sources(container_obj=_container_obj, yaml_doc=yaml_doc)

        return _container_obj, yaml_doc


# register the above classes in the module-level dictionary
DataContainerYamlReader._register_class(_AVAILABLE_REPRESENTATIONS)
DataContainerYamlWriter._register_class(_AVAILABLE_REPRESENTATIONS)
