import numpy as np
from copy import deepcopy as copy

class AIlib:
	def sigmoid(x):
		return 1/(1 + np.exp(-x))

	def correctFunc(inp:np.array): # generates the correct answer for the AI 
		return np.asarray( [1.0 - inp[0], 1.0 - inp[1], 1.0 - inp[2]] ) # basically invert the rgb values

	def calcCost( predicted:np.array, correct:np.array ): # cost function, lower -> good, higher -> bad, bad bot, bad
		costSum = 0
		maxLen = len(correct)

		for i in range(maxLen):
			costSum += (predicted[i] - correct[i])**2

		return costSum / maxLen

	def getThinkCost( inp:np.array, predicted:np.array ):
		corr = AIlib.correctFunc(inp)
		return AIlib.calcCost( predicted, corr )

	def genRandomMatrix( x:int, y:int, min: float=0.0, max: float=1.0 ): # generate a matrix with x, y dimensions with random values from min-max in it
		# apply ranger with * and -
		mat = np.random.rand(x, y) - 0.25
		return mat

	def think( inp:np.array, obj, layerIndex: int=0 ): # recursive thinking, hehe
		maxLayer = len(obj.weights) - 1
		weightedLayer = np.dot( inp, obj.weights[layerIndex] ) # dot multiply the input and the weights
		layer = AIlib.sigmoid( np.add(weightedLayer, obj.bias[layerIndex]) ) # add the biases

		if( layerIndex < maxLayer ):
			return AIlib.think( layer, obj, layerIndex + 1 )
		else:
			out = np.squeeze(np.asarray(layer))
			return out

	def propDer( dCost, dProp ):
		# Calculate the partial derivative for that prop
		return dCost / dProp

	def compareAIobjects( obj1, obj2 ):
		# Compare the two instances
		res1 = AIlib.think( inp, obj1 )
		cost1 = AIlib.getThinkCost( inp, res1 ) # get the cost

		res2 = AIlib.think( inp, obj2 )
		cost2 = AIlib.getThinkCost( inp, res2 ) # get the second cost

		# Actually calculate stuff 
		dCost = cost2 - cost1
		return dCost

	def compareInstance( obj, theta, layerIndex, neuronIndex_X=0, neuronIndex_Y=0 ):
		# Create new instances of the object
		obj2_w = copy(obj) # annoying way to create a new instance of the object
		obj2_b = copy(obj)

		obj2_w.weights[layerIndex][neuronIndex_X][neuronIndex_Y] += theta # mutate the second objects neuron
		dCost_weight = AIlib.compareAIobjects( obj, obj2_w ) # compare the two and get the dCost with respect to the weights

		obj2_b.bias[layerIndex][neuronIndex_X][neuronIndex_Y] += theta # do the same thing for the bias
		dCost_bias = AIlib.compareAIobjects( obj, obj2_b )

		return dCost_weight, dCost_bias

	def getChangeInCost( obj, theta, layerIndex ):
		mirrorObj = copy(obj)

		# Fill the buffer with None so that the dCost can replace it later
		mirrorObj.weights[layerIndex].fill(None)
		mirrorObj.bias[layerIndex].fill(None)

	def gradient( inp:np.array, obj, theta:float, maxLayer:int, layerIndex: int=0, grads=None, obj1=None, obj2=None ): # Calculate the gradient for that prop
		# Check if grads exists, if not create the buffer
		if( not grads ):
			grads = [None] * (maxLayer+1)

		# Calculate the gradient for the layer
		weightDer = AIlib.propDer( dCost_W, theta )
		biasDer = AIlib.propDer( dCost_B, theta )

		# Append the gradients to the list
		grads[layerIndex] = {
			"weight": weightDer,
			"bias": biasDer
		}

		newLayer = layerIndex + 1
		if( newLayer <= maxLayer ):
			return AIlib.gradient( inp, obj, theta, maxLayer, newLayer, grads, obj1, obj2 )
		else:
			return grads, res1, cost1

	def mutateProps( inpObj, maxLen:int, gradient:list ):
		obj = copy(inpObj)
		for i in range(maxLen):
			obj.weights[i] -= obj.learningrate * gradient[i]["weight"] # mutate the weights
			obj.bias[i] -= obj.learningrate * gradient[i]["bias"]

		return obj

	def learn( inputNum:int, targetCost:float, obj, theta:float, curCost: float=None ):
		# Calculate the derivative for:
		# Cost in respect to weights
		# Cost in respect to biases

		# i.e. : W' = W - lr * gradient (respect to W in layer i) = W - lr*[ dC / dW[i] ... ]
		# So if we change all the weights with i.e. 0.01 = theta, then we can derive the gradient with math and stuff

		inp = np.asarray(np.random.rand( 1, inputNum ))[0] # create a random learning sample

		while( not curCost or curCost > targetCost ): # targetCost is the target for the cost function
			maxLen = len(obj.bias)
			grads, res, curCost = AIlib.gradient( inp, obj, theta, maxLen - 1 )

			obj = AIlib.mutateProps( obj, maxLen, grads ) # mutate the props for next round
			print("Cost:", curCost, "|", inp, res)


		print("DONE\n")
		print(obj.weights)
		print(obj.bias)