#use sin cos to get better gradients (than nulayer)
#migth habe better gradients? (seems that way but not sure yet)

#should rename it, but who cares
#now also able to export the given matrix
from tensorflow.keras.layers import Layer
from tensorflow.keras import backend as K
from tensorflow import keras
import tensorflow as tf

import numpy as np



class ulayer(Layer):
    def __init__(self,siz,dex1,dex2, **kwargs):
        self.siz = siz
        self.dex1 = dex1
        self.dex2 = dex2
        super(ulayer, self).__init__(**kwargs)

    def build(self, input_shape):
        # Create a trainable weight variable for this layer.
        self.kernel = self.add_weight(name='kernel',
                                      shape=(1,),
                                      initializer=keras.initializers.RandomUniform(-0.5, 0.5),
                                      trainable=True)
        super(ulayer, self).build(input_shape)  # Be sure to call this at the end

    def numpify(self):
        mat=np.eye(self.siz)
        val=self.weights[0].numpy()[0]
        sin,cos=np.sin(val),np.cos(val)
        mat[self.dex1,self.dex2]=sin
        mat[self.dex2,self.dex1]=-sin
        mat[self.dex1,self.dex1]=cos
        mat[self.dex2,self.dex2]=cos
        return mat



    def call(self, x):
        kernel=self.kernel
        sin=K.sin(kernel)
        cos=K.cos(kernel)
        tan=sin/cos#that should diverge?
        rows=[tf.expand_dims(x[:,i],1) for i in range(self.siz)]
        #instead of ((1,a),(-a,1)), I want this to be
        #((1,a),(-a,1))/sqrt(1+a**2)
        #and with trigonometry, I can get the same result by
        #a=sin(kernel)?
        #multiply to make 1->cos(x) (aka *cos(x))
        #so a actually tan(kernel)
        z1=rows[self.dex2]*tan
        z2=rows[self.dex1]*tan
        rows[self.dex1]+=z1
        rows[self.dex2]-=z2
        rows[self.dex1]*=cos
        rows[self.dex2]*=cos
        rows=K.concatenate(rows,axis=1)
        return rows


        mat=tf.eye(self.siz)
        tf.assign(mat[self.dex1,self.dex2],self.kernel)
        #mat[self.dex2,self.dex1]=-self.kernel
        return K.dot(x, mat)

    def compute_output_shape(self, input_shape):
        return input_shape