import numpy as np

from .__prelude__ import dynamic
from .CurveParam import CurveParam
from .Curve import Curve

class HeavisideCurve(Curve):
    step = dynamic.variable()

    def __init__(self, curve_after):
        super().__init__()
        self.step = 0.0
        self.curve_after = curve_after

    __mime_type__ = "application/x-insarviz/HeavisideCurve"

    def to_dict(self):
        return {
            "step": self.step,
            "curve_after": self.curve_after.to_dict()
        }
    def init_from_dict(self, dct):
        self.step = dct["step"]
        self.curve_after.init_from_dict(dct["curve_after"])

    @property
    def nparams(self):
        return self.curve_after.nparams
    @property
    def params(self):
        return self.curve_after.params

    def func(self, xs, *params):
        return np.where(xs <= self.step, 0.0, self.curve_after.func(xs, *params))
    def jacobian(self, xs, *params):
        return np.where(xs <= self.step, 0.0, self.curve_after.jacobian(xs, *params))
    def hessian_epsilon(self, xs, err, grad, *params):
        xs_filter = xs >= self.step
        return self.curve_after.hessian_epsilon(xs[xs_filter], err[xs_filter], grad, *params)
