import os
from argparse import ArgumentParser

import tensorflow as tf
import tensorflow_addons as tfa
import tensorflow_datasets as tfds
from tensorflow.keras.callbacks import TensorBoard

from model import VisionTransformer
import wandb
from wandb.keras import WandbCallback

AUTOTUNE = tf.data.experimental.AUTOTUNE

if __name__ == "__main__":  
    parser = ArgumentParser()
    parser.add_argument("--logdir", default="logs")
    parser.add_argument("--image-size", default=32, type=int)
    parser.add_argument("--patch-size", default=4, type=int)
    parser.add_argument("--num-layers", default=4, type=int)
    parser.add_argument("--d-model", default=64, type=int)
    parser.add_argument("--num-heads", default=4, type=int)
    parser.add_argument("--mlp-dim", default=128, type=int)
    parser.add_argument("--lr", default=3e-4, type=float)
    parser.add_argument("--weight-decay", default=1e-4, type=float)
    parser.add_argument("--batch-size", default=32, type=int)
    parser.add_argument("--epochs", default=100, type=int)
    args = parser.parse_args()
    
    wandb.init(project='vit', entity='dadajonjurakuziev')

    # 2. Save model inputs and hyperparameters for wandb
    config = wandb.config
    config.logdir = args.logdir
    config.image_size=args.image_size
    config.patch_size=args.patch_size
    config.num_layers=args.num_layers
    config.d_model=args.d_model
    config.num_heads=args.num_heads
    config.mlp_dim=args.mlp_dim
    config.learning_rate = args.lr
    config.weight_decay = args.weight_decay
    config.batch_size = args.batch_size
    config.epochs = args.epochs
    config.num_classes=10
    config.channels=3
    config.dropout=0.1

    ds = tfds.load("cifar10", as_supervised=True)
    ds_train = (
        ds["train"]
        .cache()
        .shuffle(5 * args.batch_size)
        .batch(args.batch_size)
        .prefetch(AUTOTUNE)
    )
    ds_test = (
        ds["test"]
        .cache()
        .batch(args.batch_size)
        .prefetch(AUTOTUNE)
    )

    strategy = tf.distribute.MirroredStrategy()

    with strategy.scope():
        model = VisionTransformer(
            image_size=args.image_size,
            patch_size=args.patch_size,
            num_layers=args.num_layers,
            num_classes=10,
            d_model=args.d_model,
            num_heads=args.num_heads,
            mlp_dim=args.mlp_dim,
            channels=3,
            dropout=0.1,
        )
        model.compile(
            loss=tf.keras.losses.SparseCategoricalCrossentropy(
                from_logits=True
            ),
            optimizer=tfa.optimizers.AdamW(
                learning_rate=args.lr, weight_decay=args.weight_decay
            ),
            metrics=["accuracy"],
        )

    model.fit(
        ds_train,
        validation_data=ds_test,
        epochs=args.epochs,
        callbacks=[WandbCallback()],
    )
    model.save_weights(os.path.join(args.logdir, "vit"))

import os
from argparse import ArgumentParser

import tensorflow as tf
import tensorflow_addons as tfa
import tensorflow_datasets as tfds
from tensorflow.keras.callbacks import TensorBoard

from model import VisionTransformer
import wandb
from wandb.keras import WandbCallback

AUTOTUNE = tf.data.experimental.AUTOTUNE

if __name__ == "__main__":  
    parser = ArgumentParser()
    parser.add_argument("--logdir", default="logs")
    parser.add_argument("--image-size", default=32, type=int)
    parser.add_argument("--patch-size", default=4, type=int)
    parser.add_argument("--num-layers", default=4, type=int)
    parser.add_argument("--d-model", default=64, type=int)
    parser.add_argument("--num-heads", default=4, type=int)
    parser.add_argument("--mlp-dim", default=128, type=int)
    parser.add_argument("--lr", default=3e-4, type=float)
    parser.add_argument("--weight-decay", default=1e-4, type=float)
    parser.add_argument("--batch-size", default=32, type=int)
    parser.add_argument("--epochs", default=100, type=int)
    args = parser.parse_args()
    
    wandb.init(project='vit', entity='dadajonjurakuziev')

    # 2. Save model inputs and hyperparameters for wandb
    config = wandb.config
    config.logdir = args.logdir
    config.image_size=args.image_size
    config.patch_size=args.patch_size
    config.num_layers=args.num_layers
    config.d_model=args.d_model
    config.num_heads=args.num_heads
    config.mlp_dim=args.mlp_dim
    config.learning_rate = args.lr
    config.weight_decay = args.weight_decay
    config.batch_size = args.batch_size
    config.epochs = args.epochs
    config.num_classes=10
    config.channels=3
    config.dropout=0.1

    ds = tfds.load("cifar10", as_supervised=True)
    ds_train = (
        ds["train"]
        .cache()
        .shuffle(5 * args.batch_size)
        .batch(args.batch_size)
        .prefetch(AUTOTUNE)
    )
    ds_test = (
        ds["test"]
        .cache()
        .batch(args.batch_size)
        .prefetch(AUTOTUNE)
    )

    strategy = tf.distribute.MirroredStrategy()

    with strategy.scope():
        model = VisionTransformer(
            image_size=args.image_size,
            patch_size=args.patch_size,
            num_layers=args.num_layers,
            num_classes=10,
            d_model=args.d_model,
            num_heads=args.num_heads,
            mlp_dim=args.mlp_dim,
            channels=3,
            dropout=0.1,
        )
        model.compile(
            loss=tf.keras.losses.SparseCategoricalCrossentropy(
                from_logits=True
            ),
            optimizer=tfa.optimizers.AdamW(
                learning_rate=args.lr, weight_decay=args.weight_decay
            ),
            metrics=["accuracy"],
        )

    model.fit(
        ds_train,
        validation_data=ds_test,
        epochs=args.epochs,
        callbacks=[WandbCallback()],
    )
    model.save_weights(os.path.join(args.logdir, "vit"))

# 2. Save model inputs and hyperparameters for wandb