import re

import numpy as np
import yaml

from ....core.error import MatrixGaussianError, SimpleGaussianError
from ...xy import XYContainer

__all__ = [
    "add_error_to_container",
    "write_errors_to_yaml",
    "process_error_sources",
    "MatrixYamlDumper",
    "MatrixYamlLoader",
]


_yaml_error_section_for_axis = {0: "x_errors", 1: "y_errors", None: "errors"}


def add_error_to_container(err_type, container_obj, **kwargs):
    # TODO: check kwargs explicitly
    if err_type == "simple":
        container_obj.add_error(**kwargs)
    elif err_type == "matrix":
        container_obj.add_matrix_error(**kwargs)
    else:
        raise TypeError("Unknown error type '{}'. " "Valid: {}".format(err_type, ("simple", "matrix")))
    return container_obj


def write_errors_to_yaml(container, yaml_doc):
    # TODO: create public error retrieval interface
    for _err_name, _err_dict in container._error_dicts.items():
        _err_obj = _err_dict["err"]
        _err_axis = _err_dict.get("axis", None)

        # get the relevant error section for the current axis, creating it if necessary
        _yaml_section = yaml_doc.setdefault(_yaml_error_section_for_axis[_err_axis], [])

        # -- check for relative errors
        _is_relative = _err_obj.relative
        if _is_relative:
            _err_val = _err_obj.error_rel
        else:
            _err_val = _err_obj.error

        # -- collapse identical errors to one float
        if np.allclose(_err_val[0], _err_val):
            _err_val = float(_err_val[0])
        else:
            _err_val = _err_val.tolist()

        # -- handle different error types
        # TODO shouldn't each error be wrapped inside an 'error' namespace?
        if _err_obj.__class__ is SimpleGaussianError:
            _yaml_section.append(
                dict(
                    name=_err_name,
                    type="simple",
                    error_value=_err_val,
                    relative=_is_relative,
                    correlation_coefficient=_err_obj._corr_coeff,
                    # TODO: public interface for _corr_coeff!
                )
            )
        elif _err_obj.__class__ is MatrixGaussianError:
            _mtype = _err_obj._matrix_type_at_construction  # TODO: public interface!
            _yaml_section.append(
                dict(
                    name=_err_name,
                    type="matrix",
                    matrix_type=_mtype,
                    relative=_is_relative,
                )
            )
            if _mtype == "covariance":
                if _is_relative:
                    _yaml_section[-1]["matrix"] = _err_obj.cov_mat_rel  # .tolist()
                else:
                    _yaml_section[-1]["matrix"] = _err_obj.cov_mat  # .tolist()
            elif _mtype == "correlation":
                _yaml_section[-1]["matrix"] = _err_obj.cor_mat  # .tolist()
                _yaml_section[-1]["error_value"] = _err_val
            else:
                raise TypeError("Unknown error matrix type '{}'. " "Valid: 'correlation' or 'covariance'.")
        else:
            raise TypeError("No representation for error type {} " "implemented!".format(type(_err_obj)))

    return yaml_doc


def process_error_sources(container_obj, yaml_doc):
    # -- process error sources
    # errors can be specified as a single float, a list of floats, or a kafe2 error object
    # lists of the above are also valid, if the error object is not a list
    if isinstance(container_obj, XYContainer):  # also applies for XYParamModel
        _xerrs = yaml_doc.pop("x_errors", [])
        if isinstance(_xerrs, (int, float, str)):
            _xerrs = [_xerrs] * container_obj.size
        if len(_xerrs) > 0 and isinstance(_xerrs[0], (int, float, str)):
            _xerrs = [_xerrs]
        _yerrs = yaml_doc.pop("y_errors", [])
        if not isinstance(_yerrs, list):
            _yerrs = [_yerrs] * container_obj.size
        if len(_yerrs) > 0 and isinstance(_yerrs[0], (int, float, str)):
            _yerrs = [_yerrs]
        _errs = _xerrs + _yerrs
        _axes = [0] * len(_xerrs) + [1] * len(_yerrs)  # 0 for 'x', 1 for 'y'
    else:
        _errs = yaml_doc.pop("errors", [])
        if not isinstance(_errs, list):
            _errs = [_errs] * container_obj.size
        if len(_errs) > 0 and isinstance(_errs[0], float):
            _errs = [_errs]
        _axes = [None] * len(_errs)

    # add error sources, if any
    for _err, _axis in zip(_errs, _axes):
        # If error is a list, check for mixture of relative/absolute errors.
        # Add it as a simple error and don't try to interpret it as a kafe2 error object
        if isinstance(_err, list):
            _rel = np.zeros(len(_err))
            _abs = np.zeros(len(_err))
            for i, _val in enumerate(_err):
                if isinstance(_val, str) and _val.endswith("%"):
                    try:
                        _rel_err_percent = float(_val[:-1])
                    except ValueError:
                        raise ValueError("Cannot convert string to relative error: %s" % _err)
                    _rel[i] = _rel_err_percent
                else:
                    try:
                        _abs_err = float(_val)
                    except ValueError:
                        raise ValueError("Cannot convert {} to error value.".format(_val))
                    _abs[i] = _val
            if _axis is not None:
                container_obj = add_error_to_container("simple", container_obj, err_val=0.01 * _rel, relative=True, axis=_axis)
                container_obj = add_error_to_container("simple", container_obj, err_val=_abs, relative=False, axis=_axis)
            else:
                container_obj = add_error_to_container("simple", container_obj, err_val=0.01 * _rel, relative=True)
                container_obj = add_error_to_container("simple", container_obj, err_val=_abs, relative=False)
            continue
        elif isinstance(_err, (float, int, str)):
            raise ValueError("Failed to read in errors: {}".format(_err))

        _add_kwargs = dict()
        # translate and check that all required keys are present
        try:
            _err_type = _err.get("type", "simple")

            _add_kwargs["name"] = _err.get("name", None)

            if _err_type == "simple":
                _add_kwargs["err_val"] = _err["error_value"]
                _add_kwargs["correlation"] = _err.get("correlation_coefficient", 0)
            elif _err_type == "matrix":
                _add_kwargs["err_matrix"] = _err["matrix"]
                _add_kwargs["matrix_type"] = _err["matrix_type"]
                # default None only mandatory for cor mats; check done later
                _add_kwargs["err_val"] = _err.get("error_value", None)
            else:
                raise TypeError("Unknown error type '{}'. " "Valid: {}".format(_err_type, ("simple", "matrix")))

            _add_kwargs["relative"] = _err.get("relative", False)

            # if needed, specify the axis (only for 'xy' containers)
            if _axis is not None:
                _add_kwargs["axis"] = _axis
        except KeyError as e:
            # KeyErrors mean the YAML is incomplete -> raise
            raise ValueError("Missing required key '%s' for error specification" % e.args[0])

        # add error to data container
        container_obj = add_error_to_container(_err_type, container_obj, **_add_kwargs)

    return container_obj, yaml_doc


class MatrixYamlDumper(yaml.Dumper):
    """Custom directives for writing out matrices"""

    _regex_space_after_newline = re.compile(r"\n\s+")  # match all spaces immediately after newline
    # match exactly 'n' sequences of non-whitespace followed by optional whitespace
    _regex_format_n_numbers = r"((\S+\s*){{}})"

    def matrix(self, numpy_matrix):
        """Represent a matrix as a space/endline separated string.
        Note: symmetric matrices are represented as a lower triangular matrix by default
        """
        _is_symmetric = np.allclose(numpy_matrix, numpy_matrix.T)

        # remove brackets
        _string_repr = str(numpy_matrix).replace("[[", " [").replace("]]", "] ").replace("[", "").replace("]", "").strip()
        # remove all spaces immediately after newline
        _string_repr = re.sub(self.__class__._regex_space_after_newline, "\n", _string_repr)

        # if symmetric, remove everything above upper diagonal
        if _is_symmetric and False:
            _rows = _string_repr.split("\n")
            for _irow, _row in enumerate(_rows):
                _rows[_irow] = re.sub(r"^((\S+\s*){{{}}}).*$".format(_irow + 1), r"\1", _row).strip()
            _string_repr = "\n".join(_rows)
            # write lower triangular matrix using the '|'-style
            return self.represent_scalar("!symmetric_matrix", _string_repr, style="|")

        # write full matrix using the '|'-style
        return self.represent_scalar("!matrix", _string_repr, style="|")


# representers for covariance matrices errors
MatrixYamlDumper.add_representer(np.array, MatrixYamlDumper.matrix)


class MatrixYamlLoader(yaml.Loader):
    """custom directives for reading in matrices"""

    def matrix(self, node):
        """construct a matrix from a tab/endline separated string"""
        _rows = node.value.split("\n")
        _mat = [_row.split() for _row in _rows if _row]  # last row may be empty -> leave out
        return np.array(_mat, dtype=float)

    def symmetric_matrix(self, node):
        """construct a lower triangular matrix from a tab/endline separated string"""
        _rows = node.value.split("\n")[:-1]  # last endline ends the matrix -> no new row
        _mat = [_row.split() for _row in _rows]
        _np_mat = np.array(np.zeros((len(_rows), len(_rows))), dtype=float)

        for _i, _row in enumerate(_mat):
            # check shape -> row index must match row length
            if len(_row) != _i + 1:
                raise ValueError(
                    "Cannot parse lower triangular matrix: " "row #{} should have length {}, got {} " "instead!".format(_i + 1, _i + 1, len(_row))
                )
            # fill matrix
            _np_mat[_i, 0 : len(_row)] = _row  # fill below diagonal

        _np_mat = _np_mat + np.tril(_np_mat, -1).T  # symmetrize

        return _np_mat


MatrixYamlLoader.add_constructor("!matrix", MatrixYamlLoader.matrix)
MatrixYamlLoader.add_constructor("!symmetric_matrix", MatrixYamlLoader.symmetric_matrix)
