import pytorch_lightning as pl
from pytorch_lightning.loggers import CSVLogger
from week6.training.datasets import MedianFMNIST

from week9.evaluation.base_model import FashionMNISTDataModule, BATCH_SIZE, EPOCHS, ConvNeuralNetworkBase, \
    NeuralNetwork, RANDOM_SEED


def train_and_test(model, dataset, exp_name: str, seed: int):
    pl.seed_everything(seed)
    datamodule = FashionMNISTDataModule(dataset=dataset, batch_size=BATCH_SIZE)
    logger = CSVLogger('logs/median', name=exp_name)

    trainer = pl.Trainer(auto_lr_find=True, max_epochs=EPOCHS, weights_summary='full', logger=logger)

    trainer.tune(model, datamodule=datamodule)

    trainer.fit(model, datamodule=datamodule)

    trainer.test(model)


if __name__ == '__main__':
    models = [0, 1, 2, 3, 4, 5, 10, None]

    seed = RANDOM_SEED
    for model in models:
        for i in range(10):
            if model is None:
                network = NeuralNetwork()
                exp_name = f'median-fmnist-no-conv-layers-{seed}-seed'
            else:
                network = ConvNeuralNetworkBase(extra_layers=model)
                exp_name = f'median-fmnist-{model}-layers-{seed}-seed'

            dataset = MedianFMNIST
            train_and_test(model=network, dataset=dataset, exp_name=exp_name, seed=seed)

            seed += 1

import pytorch_lightning as pl
from pytorch_lightning.loggers import CSVLogger
from week6.training.datasets import MedianFMNIST

from week9.evaluation.base_model import FashionMNISTDataModule, BATCH_SIZE, EPOCHS, ConvNeuralNetworkBase, \
    NeuralNetwork, RANDOM_SEED


def train_and_test(model, dataset, exp_name: str, seed: int):
    pl.seed_everything(seed)
    datamodule = FashionMNISTDataModule(dataset=dataset, batch_size=BATCH_SIZE)
    logger = CSVLogger('logs/median', name=exp_name)

    trainer = pl.Trainer(auto_lr_find=True, max_epochs=EPOCHS, weights_summary='full', logger=logger)

    trainer.tune(model, datamodule=datamodule)

    trainer.fit(model, datamodule=datamodule)

    trainer.test(model)


if __name__ == '__main__':
    models = [0, 1, 2, 3, 4, 5, 10, None]

    seed = RANDOM_SEED
    for model in models:
        for i in range(10):
            if model is None:
                network = NeuralNetwork()
                exp_name = f'median-fmnist-no-conv-layers-{seed}-seed'
            else:
                network = ConvNeuralNetworkBase(extra_layers=model)
                exp_name = f'median-fmnist-{model}-layers-{seed}-seed'

            dataset = MedianFMNIST
            train_and_test(model=network, dataset=dataset, exp_name=exp_name, seed=seed)

            seed += 1
