r"""Change detection with a Mahalanobis-type metric"""
import numpy as np
from numpy.linalg import inv

from ruptures.base import BaseCost
from ruptures.exceptions import NotEnoughPoints


class CostMl(BaseCost):

    r"""Mahalanobis-type cost function."""

    model = "mahalanobis"

    def __init__(self, metric=None):
        """Create a new instance.

        Args:
            metric (ndarray, optional): PSD matrix that defines a Mahalanobis-type pseudo distance. If None, defaults to the Mahalanobis matrix. Shape (n_features, n_features).
        """
        self.metric = metric
        self.gram = None
        self.min_size = 2

    def fit(self, signal) -> "CostMl":
        """Set parameters of the instance.

        Args:
            signal (array): signal. Shape (n_samples,) or (n_samples, n_features)

        Returns:
            self
        """

        s_ = signal.reshape(-1, 1) if signal.ndim == 1 else signal

        # Mahalanobis metric if self.metric is None
        if self.metric is None:
            covar = np.cov(s_.T)
            self.metric = inv(covar.reshape(1, 1) if covar.size == 1 else covar)

        self.gram = s_.dot(self.metric).dot(s_.T)
        self.signal = s_

        return self

    def error(self, start, end):
        """Return the approximation cost on the segment [start:end].

        Args:
            start (int): start of the segment
            end (int): end of the segment

        Returns:
            float: segment cost

        Raises:
            NotEnoughPoints: when the segment is too short (less than ``'min_size'`` samples).
        """
        if end - start < self.min_size:
            raise NotEnoughPoints
        sub_gram = self.gram[start:end, start:end]
        val = np.diagonal(sub_gram).sum()
        val -= sub_gram.sum() / (end - start)
        return val
