import tensorflow as tf
# create Tensor from List with specific type
x = tf.constant([[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], dtype=tf.int32)

print("Type of every element:", x.dtype)
print("Number of axes (=Rank):", x.ndim)
print("Shape of tensor:", x.shape)
print("Total number of elements: ", tf.size(x).numpy())

print("To numpy array:\n", x.numpy()[0])
print("Back to Tensor:", tf.convert_to_tensor(x.numpy())[0])

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np
import matplotlib.pyplot as plt
import pprint
from timeit import default_timer as timer

"""
Common utils
"""
class DurationTime:
    pass

def gpu_config():
    pass

def load_dataset(dataname="cifar10", show_imgs=True):
    pass

def show_samples(images, labels, class_names, grid=(3,4)):
    pass

"""
Classifier
"""
def tf2_keras_classifier():
    gpu_config()
    train_data, test_data = load_dataset()
    model = create_model(train_data)
    train_model(model, train_data)
    test_model(model, test_data)

def create_model(dataset):
    pass

def train_model(model, train_data, split_ratio=0.8):
	pass

def test_model(model, test_data):
	pass

if __name__ == "__main__":
    tf2_keras_classifier()

def gpu_config():
    gpus = tf.config.experimental.list_physical_devices('GPU')
    if gpus:
        try:
            # Currently, memory growth needs to be the same across GPUs
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu, True)
            logical_gpus = tf.config.experimental.list_logical_devices('GPU')
            print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
        except RuntimeError as e:
            # Memory growth must be set before GPUs have been initialized
            print(e)

def load_dataset(dataname="cifar10", show_imgs=True):
    if dataname == "cifar10":
        dataset = tf.keras.datasets.cifar10
        class_names = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
    else:
        raise ValueError(f"Invalid dataset name: {dataname}")

    (x_train, y_train), (x_test, y_test) = dataset.load_data()
    x_train, x_test = x_train / 255.0, x_test / 255.0
    y_train, y_test = y_train[:, 0], y_test[:, 0]
    print(f"Load {dataname} dataset:", x_train.shape, y_train.shape, x_test.shape, y_test.shape)
    if show_imgs:
        show_samples(x_train, y_train, class_names)
    return (x_train, y_train), (x_test, y_test)

def show_samples(images, labels, class_names, grid=(3,4)):
    plt.figure(figsize=grid)
    num_samples = grid[0] * grid[1]
    for i in range(num_samples):
        plt.subplot(grid[0], grid[1], i+1)
        plt.xticks([])
        plt.yticks([])
        plt.grid(False)
        plt.imshow(images[i])
        plt.xlabel(class_names[labels[i]])
    plt.show()

def create_model(dataset, use_add=True):
    x, y = dataset
    input_shape = tf.shape(x)[1:].numpy()
    num_class = tf.reduce_max(y).numpy() + 1
    print(f"[create_model] input shape={input_shape}, num_class={num_class}")

    if use_add:
        model = keras.Sequential(name="tf-classifier")
        model.add(layers.Input(shape=input_shape))
        model.add(layers.Conv2D(filters=32, kernel_size=3, padding="same", activation="relu", name="conv1"))
        model.add(layers.MaxPool2D(pool_size=(2, 2), name="pooling1"))
        model.add(layers.Conv2D(filters=64, kernel_size=3, padding="same", activation="relu", name="conv2"))
        model.add(layers.MaxPool2D(pool_size=(2, 2), name="pooling2"))
        model.add(layers.Flatten(name="flatten"))
        model.add(layers.Dense(units=100, activation="relu", name="dense1"))
        model.add(keras.layers.Dropout(0.2))
        model.add(layers.Dense(units=num_class, activation="softmax", name="dense2"))
    else:
        model = keras.Sequential([
            layers.Conv2D(filters=32, kernel_size=3, padding="same", activation="relu", input_shape=input_shape, name="conv1"),
            layers.MaxPool2D(pool_size=(2, 2), name="pooling1"),
            layers.Conv2D(filters=64, kernel_size=3, padding="same", activation="relu", name="conv2"),
            layers.MaxPool2D(pool_size=(2, 2), name="pooling2"),
            layers.Flatten(name="flatten"),
            layers.Dense(units=100, activation="relu", name="dense1"),
            keras.layers.Dropout(0.2),
            layers.Dense(units=num_class, activation="softmax", name="dense2"),
            ],
            name="tf-classifier")
    model.summary()
    keras.utils.plot_model(model, "tf-clsf-model.png")
    return model

def train_model(model, train_data, split_ratio=0.8):
    x, y = train_data
    trainlen = x.shape[0] * split_ratio
    x_train, y_train = x[:trainlen], y[:trainlen]
    x_val, y_val = x[trainlen:], y[trainlen:]

    model.compile(
        loss=keras.losses.SparseCategoricalCrossentropy(),
        optimizer=keras.optimizers.RMSprop(),
        metrics=[keras.metrics.SparseCategoricalAccuracy()],
    )

    with DurationTime("** training time") as duration:
        history = model.fit(x_train, y_train, batch_size=32, epochs=5, validation_data=(x_val, y_val))
    history = {key: np.array(values) for key, values in history.history.items()}
    np.set_printoptions(precision=4, suppress=True)
    pp = pprint.PrettyPrinter(indent=2, width=100, compact=True)
    print("[train_model] training history:")
    pp.pprint(history)

class DurationTime:
    def __init__(self, context):
        self.start = 0
        self.context = context

    def __enter__(self):        # entering 'with' context
        self.start = timer()
        return self             # pass object by 'as'

    def __exit__(self, type, value, trace_back):    # exiting 'with' context
        print(f"{self.context}: {timer() - self.start:1.2f}")

def train_model(model, train_data, split_ratio=0.8):
    ...
    with DurationTime("** training time") as duration:
        history = model.fit(x_train, y_train, batch_size=32, epochs=5, validation_data=(x_val, y_val))

def train_model(model, train_data, split_ratio=0.8):
    ...
    start = timer()
    history = model.fit(x_train, y_train, batch_size=32, epochs=5, validation_data=(x_val, y_val))
    print(f"** trainig time: {timer() - start:1.2f}")

def some_func():
    start_total = timer()
    start = timer()
    sub_func1()
    print(f"elapsed time1:", timer() - start)
    start = timer()
    sub_func1()
    print(f"elapsed time2:", timer() - start)
    print(f"total elapsed time:", timer() - start_total)

def some_func():
    with DurationTime("total elapsed time"):
        with DurationTime("elased time1"):
            sub_func1()
        with DurationTime("elased time1"):
            sub_func1()

def test_model(model, test_data):
    x_test, y_test = test_data
    loss, accuracy = model.evaluate(x_test, y_test)
    print("[test_model] evaluate by model.evaluate()")
    print(f"  test loss: {loss:1.4f}")
    print(f"  test accuracy: {accuracy:1.4f}\n")

    predicts = model.predict(x_test)
    print("[test_model] predict by model.predict()")
    print("  prediction shape:", predicts.shape, y_test.shape)
    print("  first 5 predicts:\n", predicts[:5])
    print("  check probability:", np.sum(predicts[:5], axis=1))
    print("  manual accuracy:", np.mean(np.argmax(predicts, axis=1) == y_test))

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np
from timeit import default_timer as timer

"""
Common utils
"""
class DurationTime:
    pass

def load_dataset(dataname="cifar10", show_imgs=True):
    pass

def show_samples(images, labels, class_names, grid=(3,4)):
    pass

"""
Classifier
"""
class AdvancedClassifier:
    def __init__(self, batch_size=32, val_ratio=0.2):
        pass
    
    def build_model(self, x, y):
        pass
    
    def train(self, x, y, epochs, eager_mode=True):
        pass
    
    @tf.function
    def train_batch_graph(self, x_batch, y_batch):
        pass

    def evaluate(self, x, y_true, verbose=True):
        pass

def tf2_advanced_classifier():
    gpu_config()
    (x_train, y_train), (x_test, y_test) = load_dataset("cifar10")
    clsf = AdvancedClassifier()
    clsf.build_model(x_train, y_train)
    clsf.train(x_train, y_train, 5)
    clsf.evaluate(x_test, y_test)

if __name__ == "__main__":
    tf2_advanced_classifier()

    def __init__(self, batch_size=32, val_ratio=0.2):
        self.model = keras.Model()
        self.loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
        self.optimizer = tf.keras.optimizers.RMSprop()
        self.batch_size = batch_size
        self.val_ratio = val_ratio

    def build_model(self, x, y):
        input_shape = x.shape[1:]
        num_class = tf.reduce_max(y).numpy() + 1
        input_tensor = layers.Input(shape=input_shape)
        x = layers.Conv2D(filters=32, kernel_size=3, padding="same", activation="relu", name="conv1")(input_tensor)
        x = layers.MaxPool2D(pool_size=(2, 2), name="pooling1")(x)
        x = layers.Conv2D(filters=64, kernel_size=3, padding="same", activation="relu", name="conv2")(x)
        x = layers.MaxPool2D(pool_size=(2, 2), name="pooling2")(x)
        x = layers.Flatten(name="flatten")(x)
        x = layers.Dense(units=100, activation="relu", name="dense1")(x)
        x = layers.Dropout(0.2)(x)
        output_tensor = layers.Dense(units=num_class, activation="softmax", name="dense2")(x)
        self.model = keras.Model(inputs=input_tensor, outputs=output_tensor, name="tf-classifier")
        self.model.summary()
        keras.utils.plot_model(self.model, "tf-clsf-model-adv.png")

    def train(self, x, y, epochs):
        trainlen = int(x.shape[0] * (1 - self.val_ratio))
        x_train, y_train = x[:trainlen], y[:trainlen]
        x_val, y_val = x[trainlen:], y[trainlen:]

        dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
        dataset = dataset.shuffle(200).batch(self.batch_size)
        with DurationTime("** training time") as duration:
            for epoch in range(epochs):
                for x_batch, y_batch in dataset:
                    self.train_batch_graph(x_batch, y_batch)
                loss, accuracy = self.evaluate(x_val, y_val, verbose=False)
                print(f"[Training] epoch={epoch}, val_loss={loss}, val_accuracy={accuracy}")

    @tf.function
    def train_batch_graph(self, x_batch, y_batch):
        with tf.GradientTape() as tape:
            # training=True is only needed if there are layers with different
            # behavior during training versus inference (e.g. Dropout).
            predictions = self.model(x_batch, training=True)
            loss = self.loss_object(y_batch, predictions)
        gradients = tape.gradient(loss, self.model.trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, self.model.trainable_variables))

import tensorflow as tf

def main():
    # create Tensor from List with specific type
    x = tf.constant([[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], dtype=tf.int32)
    shape_out = print_tensor_shape(x, "eager")
    print("eager 6) shape output:", shape_out)
    shape_out = print_tensor_shape_graph(x, "graph")
    print("graph 6) shape output:", shape_out)
    shape_out = print_tensor_shape_graph(x, "graph")

def print_tensor_shape(tensor, title):
    print(f"{title} 1) Tensor.shape:", tensor.shape, type(tensor.shape))
    print(f"{title} 2) Tensor.get_shape():", tensor.get_shape())
    print(f"{title} 3) tf.shape():", tf.shape(tensor))
    h, w = tensor[0, 0, 1], tensor[0, 0, 2]
    zeros = tf.zeros((h, w))
    print(f"{title} 4) Tensor.shape:", zeros.shape)
    print(f"{title} 5) tf.shape():", tf.shape(zeros))
    return tf.shape(zeros)

@tf.function
def print_tensor_shape_graph(tensor, title):
    return print_tensor_shape(tensor, title)

if __name__ == "__main__":
    main()

    def evaluate(self, x, y_true, verbose=True):
        if verbose:
            print("[evaluate] predict by model.__call__()")
        y_pred = self.model(x)
        accuracy = np.mean(np.argmax(y_pred, axis=1) == y_true)
        loss = self.loss_object(y_true, y_pred)
        if verbose:
            np.set_printoptions(precision=4, suppress=True)
            print("  prediction shape:", y_pred.shape, y_true.shape)
            print("  first 5 predicts:\n", y_pred[:5].numpy())
            print("  check probability:", np.sum(y_pred[:5], axis=1))
            print(f"  loss={loss:1.4f}, accuracy={accuracy:1.4f}")
        return loss, accuracy

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
from timeit import default_timer as timer
"""
Common utils
"""
class DurationTime:
    pass

def load_dataset(dataname="cifar10", show_imgs=False):
    pass

def show_samples(images, labels, class_names, grid=(3,4)):
    pass
"""
Classifier
"""
class TorchClsfModel(nn.Module):
    def __init__(self):
        pass

    def forward(self, x):
        pass

class TorchClassifier:
    def __init__(self, model, batch_size=32, val_ratio=0.2):
        pass
    
    def train(self, x, y, epochs):
        pass

    def train_batch(self, x_batch, y_batch):
        pass

    def evaluate(self, x, y_true, verbose=True):
        pass

def torch_classifier():
    (x_train, y_train), (x_test, y_test) = load_dataset("cifar10")
    model = TorchClsfModel()
    clsf = TorchClassifier(model)
    clsf.train(x_train, y_train, 5)
    clsf.evaluate(x_test, y_test)

if __name__ == "__main__":
    torch_classifier()

def load_dataset(dataname="cifar10", show_imgs=False):
    if dataname == "cifar10":
        trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True)
        testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True)
        class_names = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
    else:
        raise ValueError(f"Invalid dataset name: {dataname}")

    # pytorch uses image shape like (batch, channel, height, width)
    x_train = np.transpose(np.array(trainset.data), (0, 3, 1, 2))
    y_train = np.array(trainset.targets)
    x_test = np.transpose(np.array(testset.data), (0, 3, 1, 2))
    y_test = np.array(testset.targets)
    x_train, x_test = x_train / 255.0 * 2. - 1., x_test / 255.0 * 2. - 1.
    print(f"Load {dataname} dataset:", x_train.shape, y_train.shape, x_test.shape, y_test.shape)
    if show_imgs:
        show_samples(trainset.data, trainset.targets, class_names)
    return (x_train, y_train), (x_test, y_test)

class TorchClsfModel(nn.Module):
    def __init__(self):
        super(TorchClsfModel, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, 
                               padding=1, padding_mode='zeros')
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1, padding_mode='zeros')
        self.pool2 = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 8 * 8, 100)
        self.fc2 = nn.Linear(100, 10)
        self.softmax = nn.Softmax()

    def forward(self, x):
        x = self.pool1(F.relu(self.conv1(x)))
        x = self.pool2(F.relu(self.conv2(x)))
        x = x.view(-1, 64 * 8 * 8)
        x = F.relu(self.fc1(x))
        x = F.dropout(x)
        x = F.softmax(self.fc2(x), dim=1)
        return x

seq_modules = nn.Sequential(
    flatten,
    layer1,
    nn.ReLU(),
    nn.Linear(20, 10)
)

class TorchClsfModelUsingSequential(nn.Module):
    def __init__(self):
        super(TorchClsfModelUsingSequential, self).__init__()
        self.conv_relu_pool1 = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, 
                      padding=1, padding_mode='zeros'),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.ReLU(),
        )
        self.conv_relu_pool2 = nn.Sequential(
            nn.Conv2d(32, 64, 3, padding=1, padding_mode='zeros'),
            nn.MaxPool2d(2, 2),
            nn.ReLU(),
        )
        self.fc1 = nn.Linear(64 * 8 * 8, 100)
        self.fc2 = nn.Linear(100, 10)
        self.softmax = nn.Softmax()

    def forward(self, x):
        x = self.conv_relu_pool1(x)
        x = self.conv_relu_pool2(x)
        x = x.view(-1, 64 * 8 * 8)
        x = F.relu(self.fc1(x))
        x = F.dropout(x)
        x = F.softmax(self.fc2(x), dim=1)
        return x

class TorchClassifier:
    def __init__(self, model, batch_size=32, val_ratio=0.2):
        self.model = model
        self.loss_object = nn.CrossEntropyLoss()
        self.optimizer = optim.RMSprop(self.model.parameters(), lr=0.001)
        self.batch_size = batch_size
        self.val_ratio = val_ratio

    def train(self, x, y, epochs):
        x, y = torch.from_numpy(x).float(), torch.from_numpy(y)
        trainlen = int(x.shape[0] * (1 - self.val_ratio))
        x_train, y_train = x[:trainlen], y[:trainlen]
        trainset = torch.utils.data.TensorDataset(x_train, y_train)
        trainloader = torch.utils.data.DataLoader(trainset, batch_size=self.batch_size, shuffle=True, num_workers=2)
        x_val, y_val = x[trainlen:], y[trainlen:]

        with DurationTime("** training time"):
            for epoch in range(epochs):
                for x_batch, y_batch in trainloader:
                    self.train_batch(x_batch, y_batch)

                loss, accuracy = self.evaluate(x_val, y_val, verbose=False)
                print(f"[Training] epoch={epoch}, val_loss={loss:1.4f}, val_accuracy={accuracy:1.4f}")

    def train_batch(self, x_batch, y_batch):
        # zero the parameter gradients
        self.optimizer.zero_grad()
        # forward + backward + optimize
        y_pred = self.model(x_batch)
        loss = self.loss_object(y_pred, y_batch)
        loss.backward()
        self.optimizer.step()

    def evaluate(self, x, y_true, verbose=True):
        if isinstance(x, np.ndarray):
            x, y_true = torch.from_numpy(x).float(), torch.from_numpy(y_true)

        y_pred = self.model(x)
        accuracy = torch.mean((torch.argmax(y_pred, dim=1) == y_true).float())
        loss = self.loss_object(y_pred, y_true)
        if verbose:
            np.set_printoptions(precision=4, suppress=True)
            print("  prediction shape:", y_pred.shape, y_true.shape)
            print("  first 5 predicts:\n", y_pred[:5].detach().numpy())
            print("  check probability:", torch.sum(y_pred[:5], dim=1))
            print(f"  loss={loss.detach().numpy():1.4f}, accuracy={accuracy.detach().numpy():1.4f}")
        return loss, accuracy