import random

import matplotlib.pyplot as plt
import wandb
import hydra
import torch
import torch.utils.data as data_utils
from model import ChessPiecePredictor
from torch import nn, optim
from google.cloud import storage
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import ImageFolder


@hydra.main(config_path="../conf", config_name="config")
def train(cfg):
    print(f"Training started with parameters: {cfg}")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    wandb.init()

    torch.manual_seed(cfg.seed)

    model = ChessPiecePredictor(
        image_size=cfg.image_size,
        patch_size=cfg.patch_size,
        in_channels=cfg.in_channels,
        embed_dim=cfg.embed_dim,
        num_heads=cfg.num_heads,
    )
    wandb.watch(model)

    t = transforms.Compose(
        [
            transforms.Resize((cfg.image_size, cfg.image_size)),
            transforms.Grayscale(num_output_channels=cfg.in_channels),
            transforms.ToTensor(),
        ]
    )

    train_data = ImageFolder(f"{cfg.data_path}/train", transform=t)
    validation_data = ImageFolder(f"{cfg.data_path}/test", transform=t)


    indices_train = random.sample(range(1, 60000), 5000)
    indices_valid = random.sample(range(1, 30000), 1000)

    train_data = data_utils.Subset(train_data, indices_train)
    validation_data = data_utils.Subset(validation_data, indices_valid)

    train_loader = DataLoader(train_data, batch_size=cfg.batch_size, shuffle=True)
    validation_loader = DataLoader(validation_data, batch_size=cfg.batch_size, shuffle=True)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=cfg.lr)

    print("Training started...")
    train_losses = []
    validation_losses = []

    batch_count = len(train_loader)
    epochs = 2
    for e in range(epochs):
        train_loss = 0
        train_correct = 0

        validation_loss = 0
        validation_correct = 0

        i = 0
        for images, labels in train_loader:
            # in case we use cuda to train on gpu
            images = images.to(device)
            labels = labels.to(device)

            optimizer.zero_grad()

            preds = model(images)
            loss = criterion(preds, labels)
            loss.backward()
            optimizer.step()
            train_loss += loss.item()

            # accuracy
            _, preds_indices = torch.max(preds, dim=1)
            train_correct += (preds_indices == labels).sum()

            i += 1
            if i % 100 == 0:
                print(
                    f"Epoch: {e+1} / {epochs}"
                    f" - progress: {i} / {batch_count}"
                    f" - loss: {loss.data.mean()}"
                )

        for images, labels in validation_loader:
            images = images.to(device)
            labels = labels.to(device)

            preds = model(images)
            loss = criterion(preds, labels)

            validation_loss += loss.item()

            # accuracy
            _, preds_indices = torch.max(preds, dim=1)
            validation_correct += (preds_indices == labels).sum()

        train_accuracy = float(train_correct / (len(train_loader) * cfg.batch_size))
        validation_accuracy = float(validation_correct / (len(validation_loader) * cfg.batch_size))

        wandb.log({
            "train_loss": train_loss,
            "validation_loss": validation_loss,
            "train_accuracy": train_accuracy,
            "validation_accuracy": validation_accuracy,
        })

        train_losses.append(train_loss / len(train_loader))
        validation_losses.append(validation_loss / len(validation_loader))

    # plotting
    plt.plot(list(range(1, len(train_losses) + 1)), train_losses, label="Training loss")
    print("Train losses:", train_losses)

    plt.plot(list(range(1, len(validation_losses) + 1)), validation_losses, label="Validation loss")
    print("Validation losses:", validation_losses)

    plt.xlabel("epoch")
    plt.ylabel("loss")
    plt.legend()

    fig_path = "training_run.png"
    plt.savefig(fig_path)
    print(f"Saved training loss figure to {fig_path}")

    model_path = "trained_model.pth"
    torch.save(model.state_dict(), model_path)
    print(f"Saved trained model to {model_path}")

    storage_client = storage.Client()
    bucket = storage_client.bucket("chess_predictor")
    blob = bucket.blob("model_blob")

    blob.upload_from_filename("outputs/model_0.pt")


if __name__ == "__main__":
    train()

[((346, 401), 'hydra.main', 'hydra.main', ([], {'config_path': '"""../conf"""', 'config_name': '"""config"""'}), "(config_path='../conf', config_name='config')\n", (356, 401), False, 'import hydra\n'), ((552, 564), 'wandb.init', 'wandb.init', ([], {}), '()\n', (562, 564), False, 'import wandb\n'), ((570, 597), 'torch.manual_seed', 'torch.manual_seed', (['cfg.seed'], {}), '(cfg.seed)\n', (587, 597), False, 'import torch\n'), ((611, 772), 'model.ChessPiecePredictor', 'ChessPiecePredictor', ([], {'image_size': 'cfg.image_size', 'patch_size': 'cfg.patch_size', 'in_channels': 'cfg.in_channels', 'embed_dim': 'cfg.embed_dim', 'num_heads': 'cfg.num_heads'}), '(image_size=cfg.image_size, patch_size=cfg.patch_size,\n    in_channels=cfg.in_channels, embed_dim=cfg.embed_dim, num_heads=cfg.\n    num_heads)\n', (630, 772), False, 'from model import ChessPiecePredictor\n'), ((815, 833), 'wandb.watch', 'wandb.watch', (['model'], {}), '(model)\n', (826, 833), False, 'import wandb\n'), ((1078, 1128), 'torchvision.datasets.ImageFolder', 'ImageFolder', (['f"""{cfg.data_path}/train"""'], {'transform': 't'}), "(f'{cfg.data_path}/train', transform=t)\n", (1089, 1128), False, 'from torchvision.datasets import ImageFolder\n'), ((1151, 1200), 'torchvision.datasets.ImageFolder', 'ImageFolder', (['f"""{cfg.data_path}/test"""'], {'transform': 't'}), "(f'{cfg.data_path}/test', transform=t)\n", (1162, 1200), False, 'from torchvision.datasets import ImageFolder\n'), ((1335, 1379), 'torch.utils.data.Subset', 'data_utils.Subset', (['train_data', 'indices_train'], {}), '(train_data, indices_train)\n', (1352, 1379), True, 'import torch.utils.data as data_utils\n'), ((1402, 1451), 'torch.utils.data.Subset', 'data_utils.Subset', (['validation_data', 'indices_valid'], {}), '(validation_data, indices_valid)\n', (1419, 1451), True, 'import torch.utils.data as data_utils\n'), ((1472, 1535), 'torch.utils.data.DataLoader', 'DataLoader', (['train_data'], {'batch_size': 'cfg.batch_size', 'shuffle': '(True)'}), '(train_data, batch_size=cfg.batch_size, shuffle=True)\n', (1482, 1535), False, 'from torch.utils.data import DataLoader\n'), ((1560, 1628), 'torch.utils.data.DataLoader', 'DataLoader', (['validation_data'], {'batch_size': 'cfg.batch_size', 'shuffle': '(True)'}), '(validation_data, batch_size=cfg.batch_size, shuffle=True)\n', (1570, 1628), False, 'from torch.utils.data import DataLoader\n'), ((1646, 1667), 'torch.nn.CrossEntropyLoss', 'nn.CrossEntropyLoss', ([], {}), '()\n', (1665, 1667), False, 'from torch import nn, optim\n'), ((4020, 4039), 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""epoch"""'], {}), "('epoch')\n", (4030, 4039), True, 'import matplotlib.pyplot as plt\n'), ((4044, 4062), 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""loss"""'], {}), "('loss')\n", (4054, 4062), True, 'import matplotlib.pyplot as plt\n'), ((4067, 4079), 'matplotlib.pyplot.legend', 'plt.legend', ([], {}), '()\n', (4077, 4079), True, 'import matplotlib.pyplot as plt\n'), ((4119, 4140), 'matplotlib.pyplot.savefig', 'plt.savefig', (['fig_path'], {}), '(fig_path)\n', (4130, 4140), True, 'import matplotlib.pyplot as plt\n'), ((4353, 4369), 'google.cloud.storage.Client', 'storage.Client', ([], {}), '()\n', (4367, 4369), False, 'from google.cloud import storage\n'), ((3369, 3528), 'wandb.log', 'wandb.log', (["{'train_loss': train_loss, 'validation_loss': validation_loss,\n    'train_accuracy': train_accuracy, 'validation_accuracy':\n    validation_accuracy}"], {}), "({'train_loss': train_loss, 'validation_loss': validation_loss,\n    'train_accuracy': train_accuracy, 'validation_accuracy':\n    validation_accuracy})\n", (3378, 3528), False, 'import wandb\n'), ((509, 534), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (532, 534), False, 'import torch\n'), ((885, 936), 'torchvision.transforms.Resize', 'transforms.Resize', (['(cfg.image_size, cfg.image_size)'], {}), '((cfg.image_size, cfg.image_size))\n', (902, 936), False, 'from torchvision import transforms\n'), ((950, 1007), 'torchvision.transforms.Grayscale', 'transforms.Grayscale', ([], {'num_output_channels': 'cfg.in_channels'}), '(num_output_channels=cfg.in_channels)\n', (970, 1007), False, 'from torchvision import transforms\n'), ((1021, 1042), 'torchvision.transforms.ToTensor', 'transforms.ToTensor', ([], {}), '()\n', (1040, 1042), False, 'from torchvision import transforms\n'), ((2449, 2472), 'torch.max', 'torch.max', (['preds'], {'dim': '(1)'}), '(preds, dim=1)\n', (2458, 2472), False, 'import torch\n'), ((3084, 3107), 'torch.max', 'torch.max', (['preds'], {'dim': '(1)'}), '(preds, dim=1)\n', (3093, 3107), False, 'import torch\n')]