import numpy as np
from scipy.optimize import curve_fit


class MODELS:

    LINEAR = {
        'id': 'linear',
        'f': lambda x,a,b: a * x + b,
        'str': '{}*x + {}',
        'min_values': 2
    }

    LOG = {
        'id': 'logarithmic',
        'f': lambda x,a,b: a * np.log(x) + b,
        'str': '{}*log(x) + {}',
        'min_values': 2
    }

    POLY2DEG = {
        'id': '2 degree polynomial',
        'f': lambda x,a,b,c: a * x + b * (x ** 2) + c,
        'str': '{}*x + {}*x^2 + {}',
        'min_values': 3
    }

    POLY3DEG = {
        'id': '3 degree polynomial',
        'f': lambda x,a,b,c,d: a * x + b * (x ** 2) + c * (x ** 3) + d,
        'str': '{}*x + {}*x^2 + {}*x^3 + {}',
        'min_values': 4
    }

    POLY4DEG = {
        'id': '4 degree polynomial',
        'f': lambda x,a,b,c,d,e: a*x+b*(x**2)+c*(x**3)+d*(x**4)+e,
        'str': '{}*x + {}*x^2 + {}*x^3 + {}*x^4 + {}',
        'min_values': 5
    }

    MAPPER = {
        'lin': LINEAR,
        'log': LOG,
        'poly2deg': POLY2DEG,
        'poly3deg': POLY3DEG,
        'poly4deg': POLY4DEG
    }


def build_regression_model(func, raw_data):
    x = np.arange(1, len(raw_data) + 1)
    y = np.array(raw_data)
    popt, pcov = curve_fit(func, x, y)

    x1 = np.arange(len(raw_data) + 1, len(raw_data) + 4)

    data = func(np.append(x, x1), *popt)
    residuals = y - func(x, *popt)
    ss_res = np.sum(residuals ** 2)
    ss_tot = np.sum((y - np.mean(y)) ** 2)
    r_squared = 1 - (ss_res / ss_tot)

    return data, r_squared, popt


def _safe_value(value):
    if np.isnan(value) or value in [np.inf, np.NINF]:
        return str(value)
    else:
        return value


def predictions(raw_data, models=[]):
    result = {}

    for model in models:
        if len(raw_data) < model['min_values']:
            result[model['id']] = {
                'values': ('Not enough input values. '
                           'Expected={}, current={}'.format(model['min_values'],
                                                            len(raw_data))),
                'r2': '',
                'function': '',
                'params': ''
            }
        else:
            values, r2, params = build_regression_model(model['f'], raw_data)
            result[model['id']] = {
                'values': values.tolist(),
                'r2': _safe_value(r2),
                'function': model['str'].format(*params),
                'params': params.tolist()
            }

    return result
