import time

from docarray import BaseDoc
from docarray.typing import TensorFlowTensor
import tensorflow as tf


class MyDoc(BaseDoc):
    x: TensorFlowTensor
    y: TensorFlowTensor


# 1. Eager execution with GradientTape
x = tf.Variable(3.0)  # trainable: Default: `True`, GradientTapes automatically watch uses of this variable.

with tf.GradientTape() as tape:
    y = x**2

doc = MyDoc(x=x, y=y)
doc.json()


# 2. GradientTape with tf.function
@tf.function
def square(a):
    return a**2


x = tf.Variable(tf.random.normal((3, 2)))

with tf.GradientTape() as tape:
    y = square(x)

doc = MyDoc(x=x, y=y)
doc.json()

# 3. TF constants
x = tf.Variable(tf.random.normal((3, 2)))
y = tf.constant(2.0)

doc = MyDoc(x=x, y=y)
doc.json()

# 4. Full example
w = tf.Variable(tf.random.normal((3, 2)), name="w")
b = tf.Variable(tf.zeros(2, dtype=tf.float32), name="b")
x = [[1.0, 2.0, 3.0]]

with tf.GradientTape(persistent=True) as tape:
    y = x @ w + b
    loss = tf.reduce_mean(y**2)

[dl_dw, dl_db] = tape.gradient(loss, [w, b])

doc = MyDoc(x=dl_dw, y=dl_db)
doc.json()