import numpy as np
import unittest
import weakref
import contextlib

@contextlib.contextmanager
def using_config(name,value):
    old_value = getattr(config,name)
    setattr(Config,name,value)
    try:
        yield
    finally:
        setattr(Config,name,old_value)

class Variable:
    def __init__(self,data,name=None):
        
        if data is not None:
            if not isinstance(data, np.ndarray):
                raise TypeError('{} is not supported'.format(type(data)))

        self.data = data
        self.grad  = None
        self.name = name
        self.creator = None
        self.generation = 0

    def set_creator(self,func):
        self.creator = func
        self.generation = func.generation + 1

    def cleargrad(self):
        self.grad = None
    
    def __len__(self):
        return len(self.data)
    
    def __repr__(self):
        if self.data is None:
            return 'Variable(None)'
        p = str(self.data).replace('\n','\n'+' '*9)
        return 'variable(' + p +')'

    @property
    def shape(self):
        return self.data.shape

    @property
    def ndim(self):
        return self.data.ndim
    @property
    def size(self):
        return self.data.size
    
    @property
    def dtype(self):
        return self.data.dtype

    def backward(self,retain_grad = False):
        if self.grad is None:
            self.grad = np.ones_like(self.data)

        funcs = []
        seen_set = set()

        def add_func(f):
            if f not in seen_set:
                funcs.append(f)
                seen_set.add(f)
                funcs.sort(key=lambda x: x.generation)

        add_func(self.creator)


        while funcs:
            f = funcs.pop() #関数を取得
            #gys = [output.grad for output in f.outputs]
            gys = [output().grad for output in f.outputs]
            
            gxs = f.backward(*gys)
            if not isinstance(gxs,tuple):
                gxs = (gxs,)

            for x, gx in zip(f.inputs, gxs):
                if x.grad is None:
                    x.grad = gx
                else:
                    x.grad = x.grad + gx

                if x.creator is not None:
                    add_func(x.creator)
                
            if not retain_grad:
                for y in f.outputs:
                    y().grad = None
            

class Function:
    def __call__(self,*inputs):
        xs = [x.data for x in inputs]
        ys = self.forward(*xs)
        if not isinstance(ys,tuple):
            ys = (ys,)
        outputs = [Variable(as_array(y)) for y in ys]

        if Config.enable_backprop:
            self.generation = max([x.generation for x in inputs])

            for output in outputs:
                output.set_creator(self)
            self.inputs = inputs
            self.outputs = [weakref.ref(output) for output in outputs]
        return outputs if len(outputs) > 1 else outputs[0]

    def forward(self, xs):
        raise NotImplementedError

    def backward(self,gys):
        raise NotImplementedError

class Config:
    enable_backprop = True

"""
class Mul(Function):
    def forward(self,x0,x1):
        y = x0*x1
        return y
    
    def backward(self,gy):
        x0,x1 = self.inputs[0].data, self.inputs[1].data
        return gy * x1, gy * x0
"""

def __mul__(self,other):
    return mul(self,other)




class Square(Function):
    def forward(self,x):
        y = x **2
        return y

    def backward(self,gy):
        x = self.inputs[0].data
        gx = 2*x*gy
        return gx

class Exp(Function):
    def forward(self,x):
        y = np.exp(x)
        return y
    
    def backward(self,gy):
        x = self.input.data
        gx = np.exp(x)*gy
        return gx

class Add(Function):
    def forward(self,x0,x1):
        y = x0 + x1

        return y

    def backward(self,gy):
        return gy, gy

def add(x0,x1):
    return Add()(x0,x1)

def square(x):
    return Square()(x)

def exp(x):
    return Exp()(x)

def mul(x0,x1):
    return Mul()(x0,x1)


def as_array(x):
    if np.isscalar(x):
        return np.array(x)
    return x

def numercial_diff(f,x,eps=1e-4):
    x0 = Variable(x.data - eps)
    x1 = Variable(x.data + eps)
    y0 = f(x0)
    y1 = f(x1)

    return (y1.data - y0.data) / (2*eps)

class SquareTest(unittest.TestCase):
    def test_forward(self):
        x = Variable(np.array(2.0))
        y = square(x)
        expected = np.array(4.0)
        self.assertEqual(y.data, expected)

    def test_backward(self):
        x = Variable(np.array(3.0))
        y = square(x)
        y.backward()
        expected = np.array(6.0)
        self.assertEqual(x.grad, expected)
    
    def test_gradient_check(self):
        x = Variable(np.random.rand(1))
        y = square(x)
        y.backward()
        num_grad = numercial_diff(square,x)
        flg = np.allclose(x.grad, num_grad)
        self.assertTrue(flg)

def no_grad():
    return using_config('enable_backprop',False)


a = Variable(np.array(3.0))
b = Variable(np.array(2.0))
c = Variable(np.array(1.0))

y = add(mul(a,b),c)
y.backward()

print(y)
print(a.grad)
print(b.grad)