import itertools
import pandas as pd
import torch

from mnist_ten.data import train_labeled_loader, train_rotoflip_loader
from mnist_ten.models import weights_path, classifier_head, classifier, rotoflip as rotoflip_classifier

from tqdm.auto import tqdm, trange
import plotly.graph_objects as go

criterion = torch.nn.CrossEntropyLoss()

def train_step(model, optimizer, images, labels):
    predictions = model(images)
    loss = criterion(predictions, labels)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()

def train(model, optimizer, data_loader, epochs):
    losses = []
    for epoch in epochs:
        for images, labels in data_loader:
            losses.append(train_step(
                model=model,
                optimizer=optimizer,
                images=images,
                labels=labels
            ))
    return losses

def plot_losses(losses, title):
    losses_smooth = pd.Series(losses).rolling(window=16).mean()
    return go.Figure(
        layout=dict(
            title=title,
            xaxis_title='Batch number',
            yaxis_title='Loss',
            showlegend=False
        ),
        data=[
            go.Scatter(
                y=losses,
                line=dict(color='rgba(128,206,255,0.25)')
            ),
            go.Scatter(
                y=losses_smooth,
                line=dict(color='rgba(128,206,255,1)')
            )
        ]
    )

torch.save(classifier.state_dict(), weights_path)