from mynumpy import *

# def dist:
#     def p(z):
#         return exp(self.logp(z))

# class norm0(dist):
#     def logp(z):
#         return -z**2/2 - sqrt(2*pi)
#     def ppf(r):

# take a distribution that has a pdf cdf logpdf and logcdf WITH autograd
# and a ppf WITHOUT autograd
# and give it autograd

class mydist():

    def __init__(self,d,d0):
        self.d  = d
        self.d0 = d0

    def logpdf(self,z,*args):
        return self.d.logpdf(z,*args)
    
    def pdf(self,z,*args):
        return self.d.pdf(z,*args)
    
    def cdf(self,z,*args):
        return self.d.cdf(z,*args)

    def logcdf(self,z,*args):
        return self.d.logcdf(z,*args)
    
    def ppf_newt(self,r,*args):
        z = 0*r
        def f(z):
            return r-self.cdf(z,*args)
        #g = grad(f)
        def g(z):
            return -self.pdf(z,*args)

        for iter in range(10):
            print('inter',iter,'f',f(z),'z',z)
            if iter < 5:
                step = .5
            else:
                step = 1.0
            #z = z - step*f(z)/g(z)
            z = z - step*(r-self.cdf(z,*args))/(-self.pdf(z,*args))
        return z

    # here, we play a trick: we automatically start at the solution
    # but do a single iteration of Newton's method
    # this makes all the derivatives propagate correctly
    def ppf(self,r,*args):
        def f(z):
            return r-self.cdf(z,*args)
        def g(z):
            return -self.pdf(z,*args)
        args2 = list(map(getval,args))
        z = getval(self.d0.ppf(getval(r),*args2))
        #return z - f(z)/g(z)
        return z + (r-self.cdf(z,*args))/self.pdf(z,*args)

if __name__ == '__main__':
    d0 = autograd.scipy.stats.norm
    d1 = scipy.stats.norm

    d = mydist(d0,d1)

    r = .1
    loc = .7
    scale = .3
    d.ppf(r,loc,scale)

    from IPython import embed
    embed()