# flake8: noqa F401 (imported but unused)
from collections import OrderedDict
from copy import deepcopy

import numpy as np

from ...tools import print_dict_as_table
from .._base import DataContainerBase, FitBase
from .._base.cost import CostFunction_Chi2
from ..util import collect
from .container import IndexedContainer
from .model import IndexedModelFunction, IndexedParametricModel
from .plot import IndexedPlotAdapter

__all__ = ["IndexedFit"]


class IndexedFit(FitBase):
    CONTAINER_TYPE = IndexedContainer
    MODEL_TYPE = IndexedParametricModel
    MODEL_FUNCTION_TYPE = IndexedModelFunction
    PLOT_ADAPTER_TYPE = IndexedPlotAdapter
    RESERVED_NODE_NAMES = {
        "data",
        "model",
        "cost",
        "data_error",
        "model_error",
        "total_error",
        "data_cov_mat",
        "model_cov_mat",
        "total_cov_mat",
        "data_cor_mat",
        "model_cor_mat",
        "total_cor_mat",
    }
    _BASIC_ERROR_NAMES = {"data_error", "model_error", "data_cov_mat", "model_cov_mat"}

    def __init__(
        self,
        data,
        model_function,
        cost_function="chi2",
        minimizer=None,
        minimizer_kwargs=None,
        dynamic_error_algorithm="nonlinear",
    ):
        """
        Construct a fit of a model to a series of indexed measurements.

        :param data: the measurement values
        :type data: iterable of float
        :param model_function: the model function
        :type model_function: :py:class:`~kafe2.fit.indexed.IndexedModelFunction` or unwrapped native Python function
        :param cost_function: the cost function
        :type cost_function: :py:class:`~kafe2.fit._base.CostFunctionBase`-derived or unwrapped native Python function
        :param minimizer: the minimizer to use for fitting.
        :type minimizer: None, "iminuit", "tminuit", or "scipy".
        :param minimizer_kwargs: dictionary with kwargs for the minimizer.
        :type minimizer_kwargs: dict
        """
        super(IndexedFit, self).__init__(
            data=data,
            model_function=model_function,
            cost_function=cost_function,
            minimizer=minimizer,
            minimizer_kwargs=minimizer_kwargs,
            dynamic_error_algorithm=dynamic_error_algorithm,
        )

    # -- private methods

    def _init_nexus(self):
        super(IndexedFit, self)._init_nexus()

        self._nexus.add_alias("y_data", alias_for="data")
        self._nexus.add_alias("y_model", alias_for="model")
        self._nexus.add_alias("y_total_cov_mat", alias_for="total_cov_mat")

        self._nexus.add_dependency("model", depends_on=("parameter_values"))

    def _report_data(self, output_stream, indent, indentation_level):
        output_stream.write(indent * indentation_level + "########\n")
        output_stream.write(indent * indentation_level + "# Data #\n")
        output_stream.write(indent * indentation_level + "########\n\n")

        _data_table_dict = OrderedDict()
        _data_table_dict["Index"] = range(self.data_size)
        _data_table_dict["Data"] = self.data
        if self.has_data_errors:
            _data_table_dict["Data Error"] = self.data_error
            _data_table_dict["Data Correlation Matrix"] = self.data_cor_mat

        print_dict_as_table(_data_table_dict, output_stream=output_stream, indent_level=indentation_level + 1)
        output_stream.write("\n")

    def _report_model(self, output_stream, indent, indentation_level):
        # call base method to show header and model function
        super(IndexedFit, self)._report_model(output_stream, indent, indentation_level)
        _model_table_dict = OrderedDict()
        _model_table_dict["Index"] = range(self.data_size)
        _model_table_dict["Model"] = self.model
        if self.has_model_errors:
            _model_table_dict["Model Error"] = self.model_error
            _model_table_dict["Model Correlation Matrix"] = self.model_cor_mat

        print_dict_as_table(_model_table_dict, output_stream=output_stream, indent_level=indentation_level + 1)
        output_stream.write("\n")

        if self._param_model.get_matching_errors({"relative": True}):
            output_stream.write(indent * (indentation_level + 1))
            output_stream.write("Model covariance matrix was calculated dynamically relative to y model values.\n")
            output_stream.write("\n")

    # -- public properties

    def _set_new_data(self, new_data):
        if isinstance(new_data, self.CONTAINER_TYPE):
            self._data_container = deepcopy(new_data)
        elif isinstance(new_data, DataContainerBase):
            raise TypeError("Incompatible container type '%s' (expected '%s')" % (type(new_data), self.CONTAINER_TYPE))
        else:
            self._data_container = IndexedContainer(new_data, dtype=float)
        self._data_container._on_error_change_callback = self._on_error_change

        self._nexus.get("data").mark_for_update()

    def _set_new_parametric_model(self):
        self._param_model = IndexedParametricModel(self._model_function, self.parameter_values, shape_like=self.data)

    def _get_default_x_ticks(self):
        if self.data_size <= 10:
            return np.arange(self.data_size)
        else:
            return []

    # -- public properties

    @property
    def model(self):
        """array of model predictions for the data points"""
        self._param_model.parameters = self.parameter_values  # this is lazy, so just do it
        return self._param_model.data
