import numpy as np
from astropy.modeling import Fittable1DModel, Fittable2DModel, Parameter


class Linear1D(Fittable1DModel):
    """
    One dimensional Line model.

    Parameters
    ----------
    slope : float
        Slope of the straight line
    intercept : float
        Intercept of the straight line
    """

    slope = Parameter(default=1, description="Slope of the straight line")
    intercept = Parameter(default=0, description="Intercept of the straight line")

    @staticmethod
    def evaluate(x, slope, intercept):
        return slope * x + intercept

    @staticmethod
    def fit_deriv(x, slope, intercept):
        d_slope = x
        d_intercept = np.ones_like(x)
        return [d_slope, d_intercept]


class FixedSlopeLine(Fittable1DModel):
    """
    Linear model, but fix the slope

    Parameters
    ----------
    slope : float
        Slope of the line
    intercept : float
        Intercept of the line
    """

    slope = Parameter()
    intercept = Parameter()

    @staticmethod
    def evaluate(x, slope, intercept):
        return slope * x + intercept

    @staticmethod
    def fit_deriv(x, slope, intercept):
        d_slope = np.zeros_like(x)
        d_intercept = np.ones_like(x)
        return [d_slope, d_intercept]


class pXLine(Fittable2DModel):
    """
    pX linear model.
    $Px = \\frac{k_2y - k_1x}{E}L+CL$.

    Parameters
    ----------
    L : float
        Slope of the straight line
    C : float
        Intercept / Slope of the straight line
    """

    linear = True

    L, C = Parameter(default=40), Parameter(default=0)
    k1, k2, b = Parameter(), Parameter(), Parameter()

    @staticmethod
    def evaluate(x, y, L, C, k1, k2, b):
        E = k1 * x + k2 * y + b
        return (k2 * y - k1 * x) / E * L + C * L

    @staticmethod
    def fit_deriv(x, y, L, C, k1, k2, b):
        E = k1 * x + k2 * y + b
        d_L = (k2 * y - k1 * x) / E + C
        d_C = np.full(x.shape, L)
        d_k1, d_k2, d_b = np.zeros_like(x), np.zeros_like(x), np.zeros_like(x)
        return [d_L, d_C, d_k1, d_k2, d_b]