from __future__ import annotations
from abc import ABC, abstractmethod
from typing import Dict, Any

import numpy as np
from numpy.typing import NDArray

from ...core.hashing import list_to_hash
from ...methods.functionals.base_functional import Functional
from ..regularizers.base_regularizer import Regularizer


class LossFunction(ABC):

    """This is the base class from which specific loss functions are derived.

    Parameters
    ----------
    functional
        A class derived from :class:`Functional <mumott.methods.functionals.base_functional.Functional>`
    use_weights
        Whether to multiply residuals with weights before calculating the residual norm. The calculation
        is also applied to the gradient.
    preconditioner
        A preconditioner to be applied to the gradient. Must have the same shape as
        :attr:`functional.coefficients <mumott.methods.functionals.base_functional.Functional.coefficients>`
        or it must be possible to broadcast by multiplication.
    residual_norm_multiplier
        A multiplier that is applied to the residual norm and gradient. Useful in cases where
        a very small or large loss function value changes the optimizer behaviour.
    """

    def __init__(self,
                 functional: Functional,
                 use_weights: bool = False,
                 preconditioner: NDArray[float] = None,
                 residual_norm_multiplier: float = 1):
        self._functional = functional
        self.use_weights = use_weights
        self.preconditioner = preconditioner
        self.residual_norm_multiplier = residual_norm_multiplier
        self._regularizers = {}
        self._regularization_weights = {}

    def get_loss(self,
                 coefficients: NDArray[float] = None,
                 get_gradient: bool = False):
        """Returns loss function value and possibly gradient based on the given :attr:`coefficients`.

        Notes
        -----
        This method simply calls the methods :meth:`get_residual_norm` and :meth:`get_regularization_norm`
        and sums up their respective contributions.

        Parameters
        ----------
        coefficients
            An ``NDArray`` of values of the same shape as :attr:`functional.coefficients
            <mumott.methods.functionals.base_functional.Functional.coefficients>`. Default value
            is ``None``, which leaves it to :meth:`get_residual_norm`
            to handle the choice of coefficients, which in general defaults to using the coefficients
            of the attached :attr:`functional`.
        get_gradient
            If ``True``, returns a ``'gradient'`` of the same shape as :attr:`functional.coefficients
            <mumott.methods.functionals.base_functional.Functional.coefficients>`.
            Otherwise, the entry ``'gradient'`` will be ``None``.

        Returns
        -------
            A dictionary with at least two entries, ``loss`` and ``gradient``.
        """
        residual_norm = self.get_residual_norm(coefficients, get_gradient)
        regularization = self.get_regularization_norm(self._functional.coefficients, get_gradient)

        result = dict(loss=0., gradient=None)
        result['loss'] += residual_norm['residual_norm'] * self.residual_norm_multiplier
        for name in regularization:
            result['loss'] += regularization[name]['regularization_norm'] * self.regularization_weights[name]

        if get_gradient:
            result['gradient'] = residual_norm['gradient'] * self.residual_norm_multiplier
            if self.preconditioner is not None:
                if self._functional.coefficients.shape[:-1] != self.preconditioner.shape[:-1]:
                    raise ValueError('The first three dimensions of the preconditioner must'
                                     ' have the same size as the coefficients of the functional,'
                                     ' and the last index must be the same or 1, but the'
                                     ' functional coefficients have shape'
                                     f' {self._functional.coefficients.shape},'
                                     ' while the preconditioner has shape'
                                     f' {self.preconditioner.shape}!')
                result['gradient'] *= self.preconditioner
            for name in regularization:
                result['gradient'] += regularization[name]['gradient'] * self.regularization_weights[name]
        return result

    def get_residual_norm(self,
                          coefficients: NDArray[float] = None,
                          get_gradient: bool = False) -> Dict:
        """Returns residual norm and possibly gradient based on the attached
        :attr:`functional <mumott.methods.functionals.base_functional.Functional>`.
        If :attr:`coefficients` is given, :attr:`functional.coefficients
        <mumott.methods.functionals.base_functional.Functional.coefficients>` will be updated
        with these new values, otherwise, the residual norm and possibly the gradient
        will just be calculated using the current coefficients.

        Parameters
        ----------
        coefficients
            An ``NDArray`` of values of the same shape as :attr:`functional.coefficients
            <mumott.methods.functionals.base_functional.Functional.coefficients>`.
        get_gradient
            If ``True``, returns a ``'gradient'`` of the same shape as :attr:`functional.coefficients
            <mumott.methods.functionals.base_functional.Functional.coefficients>`.
            Otherwise, the entry ``'gradient'`` will be ``None``.

        Returns
        -------
            A dictionary with at least two entries, ``residual_norm`` and ``gradient``.
        """
        if coefficients is not None:
            self._functional.coefficients = coefficients
        result = self._get_residual_norm_internal(get_gradient)
        return result

    @abstractmethod
    def _get_residual_norm_internal(self, get_gradient: bool) -> Dict:
        """ Method that implements the actual calculation of the residual norm. """
        pass

    @property
    def use_weights(self) -> bool:
        """ Whether to use weights or not in calculating the residual
        and gradient. """
        return self._use_weights

    @use_weights.setter
    def use_weights(self, val: bool) -> None:
        self._use_weights = val

    @property
    def residual_norm_multiplier(self) -> float:
        """ Multiplicative factor by which the residual norm will be scaled. Can be used,
        together with any :attr:`regularization_weights`, to scale the loss function,
        in order to address unexpected behaviour that arises when some optimizers are given
        very small or very large loss functions. """
        return self._residual_norm_multiplier

    @residual_norm_multiplier.setter
    def residual_norm_multiplier(self, val: float) -> None:
        self._residual_norm_multiplier = val

    @property
    def initial_values(self) -> NDArray[float]:
        """ Initial coefficient values for optimizer; defaults to zeros. """
        return np.zeros_like(self._functional.coefficients)

    @property
    def preconditioner(self) -> NDArray[float]:
        """ Preconditioner that is applied to the gradient by multiplication. """
        return self._preconditioner

    @property
    def preconditioner_hash(self) -> str:
        """ Hash of the preconditioner. """
        if self._preconditioner is None:
            return None
        return list_to_hash([self._preconditioner])[:6]

    @preconditioner.setter
    def preconditioner(self, value: NDArray[float]) -> None:
        self._preconditioner = value

    @property
    @abstractmethod
    def _function_as_str(self) -> str:
        """ Should return a string representation of the associated loss function
        of the residual in Python idiom, e.g. 'L(r) = 0.5 * r ** 2' for squared loss. """
        pass

    @property
    @abstractmethod
    def _function_as_tex(self) -> str:
        """ Should return a string representation of the associated loss function
        of the residual in MathJax-renderable TeX, e.g. $L(r) = \frac{r^2}{2}$ for squared loss"""
        pass

    def get_regularization_norm(self,
                                coefficients: NDArray[float],
                                get_gradient: bool = False) -> Dict[str, Dict[str, Any]]:
        """ Returns the regularization norm, and if requested, the gradient, from all
        regularizers attached to this instance, based on the provided :attr`coefficients`.
        If no coefficients are provided, the ones from the attached :attr:`functional` are used.

        Parameters
        ----------
        coefficients
            An ``np.ndarray`` of values of the same shape as :attr:`functional.coefficients
            <mumott.methods.functionals.base_functional.Functional.coefficients>`.
        get_gradient
            Whether to compute and return the gradient. Default is ``False``

        Returns
        -------
            A dictionary with one entry for each regularizer in :attr:`regularizers`, containing
            ``'regularization_norm'`` and ``'gradient'`` as entries.
        """
        regularization = dict()

        if coefficients is None:
            coefficients = self._functional.coefficients

        for name in self._regularizers:
            reg = self._regularizers[name].get_regularization_norm(coefficients=coefficients,
                                                                   get_gradient=get_gradient)
            sub_dict = dict(gradient=None)
            sub_dict['regularization_norm'] = reg['regularization_norm']

            if get_gradient:
                sub_dict['gradient'] = reg['gradient']
            regularization[name] = sub_dict
        return regularization

    def add_regularizer(self,
                        name: str,
                        regularizer: Regularizer,
                        regularization_weight: float) -> None:
        r""" Add a regularizer to the loss function.

        Parameters
        ----------
        name
            Name of the regularizer, to be used as its key.
        regularizer
            The :class:`Regularizer` instance to be attached.
        regularization_weight
            The regularization weight (often denoted :math:`\lambda`),
            by which the residual norm and gradient will be scaled.
        """
        self._regularizers[name] = regularizer
        self._regularization_weights[name] = regularization_weight

    @property
    def regularizers(self) -> Dict[str, Regularizer]:
        """ The dictionary of regularizers appended to this loss function."""
        return self._regularizers

    @property
    def regularization_weights(self) -> Dict[str, float]:
        """ The dictionary of regularization weights appended to this
        loss function. """
        return self._regularization_weights

    def __str__(self) -> str:
        s = []
        wdt = 74
        s += ['=' * wdt]
        s += [self.__class__.__name__.center(wdt)]
        s += ['-' * wdt]
        with np.printoptions(threshold=4, precision=5, linewidth=60, edgeitems=1):
            s += ['{:18} : {}'.format('Functional', self._functional.__class__.__name__)]
            s += ['{:18} : {}'.format('Uses weights', self.use_weights)]
            s += ['{:18} : {}'.format('preconditioner_hash', self.preconditioner_hash)]
            s += ['{:18} : {}'.format('residual_norm_multiplier',
                                      self._residual_norm_multiplier)]
            s += ['{:18} : {}'.format('Function of residual', self._function_as_str)]
            s += ['{:18} : {}'.format('hash', hex(hash(self))[2:8])]
        s += ['-' * wdt]
        return '\n'.join(s)

    def _repr_html_(self) -> str:
        s = []
        s += [f'<h3>{__class__.__name__}</h3>']
        s += ['<table border="1" class="dataframe">']
        s += ['<thead><tr><th style="text-align: left;">Field</th><th>Size</th><th>Data</th></tr></thead>']
        s += ['<tbody>']
        with np.printoptions(threshold=4, edgeitems=2, precision=2, linewidth=40):
            s += ['<tr><td style="text-align: left;">Functional</td>']
            s += [f'<td>{1}</td><td>{self._functional.__class__.__name__}</td></tr>']
            s += ['<tr><td style="text-align: left;">use_weights</td>']
            s += [f'<td>{1}</td><td>{self.use_weights}</td></tr>']
            s += ['<tr><td style="text-align: left;">preconditioner_hash</td>']
            s += [f'<td>{1}</td><td>{self.preconditioner_hash}</td></tr>']
            s += ['<tr><td style="text-align: left;">residual_norm_multiplier</td>']
            s += [f'<td>{1}</td><td>{self._residual_norm_multiplier}</td></tr>']
            s += ['<tr><td style="text-align: left;">Function of residual r</td>']
            s += [f'<td>1</td><td>{self._function_as_tex}</td></tr>']
            s += ['<tr><td style="text-align: left;">Hash</td>']
            h = hex(hash(self))
            s += [f'<td>{len(h)}</td><td>{h[2:8]}</td></tr>']
        s += ['</tbody>']
        s += ['</table>']
        return '\n'.join(s)

    def __hash__(self) -> int:
        to_hash = [hash(self._functional),
                   self.use_weights,
                   self._residual_norm_multiplier,
                   self._preconditioner]
        for r in self.regularizers:
            to_hash.append(hash(self.regularizers[r]))
            to_hash.append(hash(self.regularization_weights[r]))
        return int(list_to_hash(to_hash), 16)
