import numpy as np
from contextlib import contextmanager
from typing import Generator, Dict, Union
import io

#only scalar gradient
#op must be tree. 그래프 구현할려면, 위상정렬해서 순회해야하기 때문에 그렇게 하지 않음.

def broadcasting_be(a,b):
    i = len(a)-1
    j = len(b)-1
    abroad = []
    bbroad = []
    while i >= 0 and j >= 0:
        if a[i] == b[j]:
            abroad.insert(0,1)
            bbroad.insert(0,1)
        elif a[i] == 1 or b[j] == 1:
            abroad.insert(0,b[j])
            bbroad.insert(0,a[i])
        else:
            raise ValueError
        i -= 1
        j -= 1
    while i >= 0:
        bbroad.insert(0,a[i])
        i -= 1
    while j >= 0:
        abroad.insert(0,b[j])
        j -= 1
    return abroad, bbroad


class NonExistVarableError(ValueError):
    pass

def make_mermaid_graph(result):
    with io.StringIO("") as graph:
        graph.write("graph TD\n")
        result.mermaid_graph(graph)
        graph.write(f"{id(result)}-->Result\n")
        return graph.getvalue()

class OpTree:
    def __init__(self):
        super().__init__()
    def __matmul__(self,a):
        return MatMulOp(self,a)
    def __add__(self,a):
        return AddOp(self,a)
    
    @property
    def T(self):
        return TransposeOp(self)

class MatMulOp(OpTree):
    def __init__(self,a,b):
        super().__init__()
        self.a = a
        self.b = b
        va = self.a.numpy() if isinstance(self.a,OpTree) else self.a
        vb = self.b.numpy() if isinstance(self.b,OpTree) else self.b
        self.v = va @ vb
    
    def __str__(self):
        return f"MatmulOp"

    def mermaid_graph(self,writer):
        if isinstance(self.a,OpTree):
            self.a.mermaid_graph(writer)
            writer.write(f'{id(self.a)}-->{id(self)}[MatmulOp]\n')
        if isinstance(self.b,OpTree):
            self.b.mermaid_graph(writer)
            writer.write(f'{id(self.b)}-->{id(self)}[MatmulOp]\n')

    def numpy(self):
        return self.v

    def backprop(self,seed):
        #a @ b
        a = self.a.numpy() if isinstance(self.a,OpTree) else self.a
        b = self.b.numpy() if isinstance(self.b,OpTree) else self.b
        if isinstance(self.a,OpTree):
            s = seed * np.transpose(b) if seed.shape == () else  (seed) @ np.transpose(b)
            self.a.backprop((s))
        if isinstance(self.b,OpTree):
            s = np.transpose(a) * seed if seed.shape == () else np.transpose(a) @ seed
            self.b.backprop(s)

def matmul(a,b):
    return MatMulOp(a,b)

class AddOp(OpTree):
    def __init__(self,a,b):
        super().__init__()
        self.a = a
        self.b = b
        va = self.a.numpy() if isinstance(self.a,OpTree) else self.a
        vb = self.b.numpy() if isinstance(self.b,OpTree) else self.b
        self.v = va + vb
    
    def __str__(self):
        return f"AddOp"

    def mermaid_graph(self,writer):
        if isinstance(self.a,OpTree):
            self.a.mermaid_graph(writer)
            writer.write(f'{id(self.a)}-->{id(self)}[AddOp]\n')
        if isinstance(self.b,OpTree):
            self.b.mermaid_graph(writer)
            writer.write(f'{id(self.b)}-->{id(self)}[AddOp]\n')

    def numpy(self):
        return self.v

    def backprop(self,seed):
        #a + b
        ashape, bshape = broadcasting_be(self.a.numpy().shape,self.b.numpy().shape)
        aai = np.where(np.array(ashape) != 1)
        bbi = np.where(np.array(bshape) != 1)
        if isinstance(self.a,OpTree):
            self.a.backprop(np.sum(seed,axis=tuple(aai[0])))
        if isinstance(self.b,OpTree):
            self.b.backprop(np.sum(seed,axis=tuple(bbi[0])))


def addmul(a,b):
    return AddOp(a,b)

class FunctionOp(OpTree):
    def __init__(self,f, f_grad, f_name, i):
        super().__init__()
        self.f = np.vectorize(f)
        self.f_grad = np.vectorize(f_grad)
        self.f_name = f_name
        self.i = i
        self.v = self.f(i.numpy())
    
    def __str__(self):
        return f"Function{self.f_name}Op"

    def mermaid_graph(self,writer):
        self.i.mermaid_graph(writer)
        writer.write(f'{id(self.i)}-->{id(self)}[Function{self.f_name}Op]\n')

    def numpy(self):
        return self.v

    def backprop(self,seed):
        self.i.backprop(seed * (self.f_grad(self.i.numpy())))

class TransposeOp(OpTree):
    def __init__(self, i):
        super().__init__()
        self.i = i
        self.v = np.transpose(i.numpy())
    
    def __str__(self):
        return f"TransposeOp"

    def mermaid_graph(self,writer):
        self.i.mermaid_graph(writer)
        writer.write(f'{id(self.i)}-->{id(self)}[TransposeOp]\n')

    def numpy(self):
        return self.v

    def backprop(self,seed):
        self.i.backprop(np.transpose(seed))

def transposemul(a):
    return TransposeOp(a)

def relu(v):
    relu_f = lambda x: np.max([x,0])
    relu_diff = lambda x: 1 if x > 0 else 0
    return FunctionOp(relu_f,relu_diff,"Relu",v)
#row vector
def softmaxHelp(i):
    m = np.max(i,axis=i.ndim-1)
    e = np.exp(i-m.reshape(*m.shape,1))
    sumofe = np.sum(e,axis=i.ndim - 1)
    sumofe = sumofe.reshape(*sumofe.shape,1)
    return e / sumofe
class SoftmaxWithNegativeLogLikelihood(OpTree):
    #row vector
    def __init__(self, i, y):
        super().__init__()
        epsilon = 1e-15
        self.i = i
        self.s = softmaxHelp(i.numpy())
        self.y = y
        self.v = -y*np.log(self.s+epsilon)
        self.v = np.sum(self.v,axis=self.v.ndim-1)
    
    def __str__(self):
        return f"SoftmaxWithNegativeLogLikelihoodOp"

    def mermaid_graph(self,writer):
        self.i.mermaid_graph(writer)
        writer.write(f'{id(self.i)}-->{id(self)}[SoftmaxWithNegativeLogLikelihoodOp]\n')

    def numpy(self):
        return self.v
    
    def softmax_numpy(self):
        return self.s

    def backprop(self,seed):
        self.i.backprop(seed * (self.s-self.y))

class Variable(OpTree):
    def __init__(self,x):
        super().__init__()
        self.x = x
        self.grad = None

    def numpy(self):
        return self.x
    def mermaid_graph(self,writer):
        writer.write(f'{id(self)}["Variable{self.x.shape}"]\n')
    def backprop(self,seed):
        self.grad = seed