import tensorflow as tf

tf.config.set_visible_devices([], 'GPU')
import math
import asyncio
import tensorflow_datasets as tfds
from sunstreaker.data import Dataloader
from sunstreaker.layers import Flatten
from sunstreaker.activations import Softmax
import jax.numpy as jnp
from sunstreaker.losses import categorical_crossentropy
from sunstreaker.metrics import categorical_accuracy
from sunstreaker.optimizers import RMSProp


def load(batch_size: int, func):
    async def tfds_load_data() -> Dataloader:
        ds, info = tfds.load(name="mnist", split=["train", "test"], as_supervised=True, with_info=True,
                             shuffle_files=True, batch_size=batch_size)
        train_ds, valid_ds = ds
        train_ds, valid_ds = func(train_ds), func(train_ds)
        train_ds, valid_ds = train_ds.cache().repeat(), valid_ds.cache().repeat()
        input_shape = tuple(list(info.features["image"].shape))
        num_train_batches = math.ceil(info.splits["train"].num_examples / batch_size)
        num_val_batches = math.ceil(info.splits["test"].num_examples / batch_size)
        return Dataloader(
            train_data=iter(tfds.as_numpy(train_ds)), val_data=iter(tfds.as_numpy(valid_ds)),
            input_shape=input_shape, batch_size=batch_size,
            num_train_batches=num_train_batches, num_val_batches=num_val_batches
        )

    return asyncio.run(tfds_load_data())


def load_dataset(batch_size: int):
    def func(ds):
        return ds.map(lambda x, y: (tf.divide(tf.cast(x, dtype=tf.float32), 255.0), tf.one_hot(y, depth=10)))

    return asyncio.run(load(batch_size, func))


def load_dataset_muti(batch_size: int):
    def func(ds):
        return ds.map(lambda x, y: ({"img": tf.divide(tf.cast(x, dtype=tf.float32), 255.0)}, {"out1": tf.one_hot(y, depth=10)}))

    return asyncio.run(load(batch_size, func))

from sunstreaker.engine.sequential import Model

data = load_dataset(batch_size=1024)
model = Model([Input(input_shape=(28, 28, 1)), Flatten(), Dense(100), Dense(10), Softmax()])

data = load_dataset_muti(batch_size=1024)
inputs = Input(input_shape=(28, 28, 1), name="img")
flatten = Flatten()(inputs)
dense1 = Dense(100, activation='relu')(flatten)
dense2 = Dense(10, use_bias=False)(dense1)
outputs = Softmax(name="out1")(dense2)

from sunstreaker.engine.functional import Model

model = Model(inputs=inputs, outputs=outputs)

from sunstreaker import Model

data = load_dataset(batch_size=1024)

class MyModel(Model):
    def build(self, rng=None):
        self.W = self.add_weight((784, 10))
        self.flatten = Flatten()
        self.softmax = Softmax()
        return (10,), [(self.W,)]

    def call(self, params, inputs, trainable=True, **kwargs):
        self.W, = params[0]
        x = self.flatten.forward(params=[], inputs=inputs)
        x = jnp.dot(x, self.W)
        y = self.softmax.forward(params=[], inputs=x)
        return y


model = MyModel()

model.compile(loss=categorical_crossentropy, optimizer=RMSProp(lr=0.001), metrics=[categorical_accuracy])
model.fit(data, epochs=10)
model.save("tfds_mnist_v2")

model.summary()
model.plot_model() 

model.plot_losses()
model.plot_accuracy()