import numpy as np
#TODO documentation

__all__ = ['linear_model', 'linear_model_derivative', 'quadratic_model', 'quadratic_model_derivative',
           'cubic_model', 'cubic_model_derivative', 'exponential_model', 'normal_distribution']


def linear_model(x, a, b):
    return a * x + b


def linear_model_derivative(x, a, b):
    return np.ones_like(x) * a


linear_model.expression_format_string = "{a} * {x} + {b}"
linear_model.latex_expression_format_string = r"{a} \, {x} + {b}"


def quadratic_model(x, a, b, c):
    return a * x ** 2 + b * x + c


def quadratic_model_derivative(x, a, b, c):
    return 2 * a * x + b


quadratic_model.expression_format_string = "{a} * {x} ** 2 + {b} * {x} + {c}"
quadratic_model.latex_expression_format_string = r"{a} \, {x}^2 + {b} \, {x} + {c}"


def cubic_model(x, a, b, c, d):
    return a * x ** 3 + b * x ** 2 + c * x + d


def cubic_model_derivative(x, a, b, c, d):
    return 3 * a * x ** 2 + 2 * b * x + c


cubic_model.expression_format_string = "{a} * {x} ** 3 + {b} * {x} ** 2 + {c} * {x} + {d}"
cubic_model.latex_expression_format_string = r"{a} \, {x}^3 + {b} \, {x}^2 + {c} \, {x} + {d}"


def exponential_model(x, A_0, x_0):
    return A_0 * np.exp(x / x_0)


def exponential_model_derivative(x, A_0, x_0):
    return A_0 * np.exp(x / x_0)


exponential_model.expression_format_string = "{A_0} * exp({x} / {x_0})"
exponential_model.latex_expression_format_string = r"{A_0} \, e^{{{x} \ / {x_0}}}"


def normal_distribution(x, mu, sigma):
    return np.exp(-0.5 * ((x - mu) / sigma) ** 2) / np.sqrt(2.0 * np.pi * sigma ** 2)


normal_distribution.expression_format_string = \
    "1 / sqrt(2 * pi * {sigma}) * exp(-0.5 * (({x} - {mu}) / {sigma}) ** 2)"
normal_distribution.latex_expression_format_string = \
    (r"\frac{{1}}{{2 \pi {sigma}}} e^{{- \frac{{1}}{{2}}"
    r"\left( \frac{{{x} - {mu}}}{{{sigma}}} \right)^2 }}")

normal_distribution_pdf = normal_distribution  # Backwards compatibility


STRING_TO_FUNCTION = {
    'line': linear_model,
    'linear': linear_model,
    'linear_model': linear_model,
    'quadratic': quadratic_model,
    'quadratic_model': quadratic_model,
    'cubic': cubic_model,
    'cubic_model': cubic_model,
    'exp': exponential_model,
    'exponential': exponential_model,
    'exponential_model': exponential_model,
    'normal': normal_distribution,
    'normal_distribution': normal_distribution,
    'normal_distribution_pdf': normal_distribution,
}
