from ..abstract_model import AbstractModel
from scipy.optimize import curve_fit
import numpy as np


class PolynomialFit(AbstractModel):
    """A model that defines the best fitting polynomial for a given dataset.
    This polynomial is defined as p[0] * x ** n + p[1] * x ** (n-1) + ... + p[n-2] * x + p[n-1].

    """
    def __init__(self, order=1):
        self._order = order
        self._weights = np.random.rand(self._order + 1)

    def accuracy(self, x, y):
        """

        Args:
            x (array_like, shape (`M`,)):
            y (array_like, shape (`M`,) or (`M`, `K`)):

        Returns:
            Mean-Squared Error.
        """
        return np.mean(np.diff(self.predict(x), y)**2)

    def predict(self, x):
        """Predicts an output based on a curve generated by the input, x, and the weights.

        Args:
            x (array_like, shape (`M`,)): Input used for predicting the output.

        Returns:
            An array_like, shape (`M`,) or (`M`, `K`).
        """
        return np.polyval(self._weights, x)

    def train(self, x, y):
        """Overrides the previously set weights with new weights based on the
        fit of the curve to the x and y data parameters.

        Args:
            x (array_like, shape (`M`,)): Input used to fit the polynomial with target, y.
            y (array_like, shape (`M`,) or (`M`, `K`)): Targeted output for fitted polynomial.
        """
        self._weights = np.polyfit(x, y, self._order)


class CurveFit(AbstractModel):
    """A model that defines the best fitting exponential curve for a given dataset.
    This exponential curve is defined as a * exp(b * x) + c.

    """
    def __init__(self):
        self._weights = np.random.rand(3)

    @staticmethod
    def __curve__(x, a, b, c):
        """Defines the general exponential curve.

        Args:
            x: Input
            a: First parameter
            b: Second parameter
            c: Third parameter

        Returns:
            a * exp(b * x) + c
        """
        return a * np.exp(b * x) + c

    def accuracy(self, x, y):
        """

        Args:
            x (M-length sequence or a (k,M)-shaped array):
            y (M-length sequence):

        Returns:
            Mean-Squared Error.
        """
        return np.mean(np.diff(self.predict(x), y) ** 2)

    def predict(self, x):
        """Predicts an output based on a curve generated by the input, x, and the weights.

        Args:
            x (M-length sequence or a (k,M)-shaped array): Input used for predicting the output.

        Returns:
            M-length sequence.
        """
        return self.__curve__(x, self._weights[0], self._weights[1], self._weights[2])

    def train(self, x, y, initial=(1., 1e-6, 1.)):
        """Overrides the previously set weights with new weights based on the
        fit of the curve to the x and y data parameters.

        Args:
            x (M-length sequence or a (k,M)-shaped array): Input used to fit the curve with target, y.
            y (M-length sequence): Targeted output for fitted curve.
            initial (tuple): Sets the initial weights Default: (1, 1e-6, 1)
        """
        self._weights = curve_fit(self.__curve__, x, y, p0=initial)[0]
