import os
import sys
import warnings
warnings.filterwarnings("ignore")

import numpy as np
import matplotlib.pyplot as plt

# the following line is not required if BatchFlow is installed as a python package.
sys.path.append("../..")
from batchflow import Pipeline, B, C, D, F, V
from batchflow.opensets import MNIST, CIFAR10, CIFAR100
from batchflow.models.torch import VGG7

"""
Explanation: Train a ready to use PyTorch model with a simple pipeline
End of explanation
"""


BATCH_SIZE = 64

"""
Explanation: BATCH_SIZE might be increased for modern GPUs with lots of memory (4GB and higher).
End of explanation
"""


dataset = MNIST(bar=True)

"""
Explanation: Create a dataset
MNIST is a dataset of handwritten digits frequently used as a baseline for machine learning tasks.
Downloading MNIST database might take a few minutes to complete.
End of explanation
"""


config = dict(model=VGG7)

"""
Explanation: There are also predefined CIFAR10 and CIFAR100 datasets.
Define a pipeline config
Config allows to create flexible pipelines which take parameters.
For instance, if you put a model type into config, you can run a pipeline against different models.
See the list of available models to choose the one which fits you best.
End of explanation
"""


train_template = (Pipeline()
                .init_variable('loss_history', [])
                .init_model('dynamic', C('model'), 'conv_nn',
                            config={'inputs/images/shape': B.image_shape,
                                    'inputs/labels/classes': D.num_classes,
                                    'initial_block/inputs': 'images',
                                    #'device': 'gpu:0',
                                   })
                .to_array(channels='first', dtype='float32')
                .train_model('conv_nn', B.images, B.labels,
                             fetches='loss', save_to=V('loss_history', mode='a'), use_lock=True)
)

"""
Explanation: Create a template pipeline
A template pipeline is not linked to any dataset. It's just an abstract sequence of actions, so it cannot be executed, but it serves as a convenient building block.
Note: Uncomment device option to train on GPU
End of explanation
"""


train_pipeline = (train_template << dataset.train) << config

"""
Explanation: Train the model
Apply a dataset and a config to a template pipeline to create a runnable pipeline:
End of explanation
"""


train_pipeline.run(BATCH_SIZE, shuffle=True, n_epochs=1, drop_last=True, bar=True, prefetch=1)

"""
Explanation: Run the pipeline (it might take from a few minutes to a few hours depending on your hardware)
End of explanation
"""


plt.figure(figsize=(15, 5))
plt.plot(train_pipeline.v('loss_history'))
plt.xlabel("Iterations"), plt.ylabel("Loss")
plt.show()

"""
Explanation: Note that the progress bar often increments by 2 at a time - that's prefetch in action.
It does not give much here, though, since almost all time is spent in model training which is performed under a thread-lock one batch after another without any parallelism (otherwise the model would not learn anything as different batches would rewrite one another's model weights updates).
End of explanation
"""


test_pipeline = (dataset.test.p
                .import_model('conv_nn', train_pipeline)
                .init_variable('predictions') 
                .init_variable('metrics') 
                .to_array(channels='first', dtype='float32')
                .predict_model('conv_nn', B.images,
                               fetches='predictions', save_to=V('predictions'))
                .gather_metrics('class', targets=B.labels, predictions=V.predictions,
                                fmt='logits', axis=-1, save_to=V('metrics', mode='w'))
                .run(BATCH_SIZE, shuffle=True, n_epochs=1, drop_last=False, bar=True)
)

"""
Explanation: Test the model
It is much faster than training, but if you don't have GPU it would take some patience.
End of explanation
"""


metrics = test_pipeline.get_variable('metrics')

"""
Explanation: Let's get the accumulated metrics information
End of explanation
"""


metrics.evaluate('accuracy')

metrics.evaluate(['false_positive_rate', 'false_negative_rate'], multiclass=None)

"""
Explanation: Or a shorter version: metrics = test_pipeline.v('metrics')
Now we can easiliy calculate any metrics we need
End of explanation
"""


train_pipeline.save_model_now('conv_nn', path='path/to/model.torch')

"""
Explanation: Save the model
After learning the model, you may need to save it. It's easy to do this.
End of explanation
"""


