#AUTOGENERATED! DO NOT EDIT! File to edit: dev/10_train_wavedata.ipynb (unless otherwise specified).

__all__ = ['trainEpoch', 'validEpoch']

#Cell
# --- Must haves ---
import os, sys
sys.path.append('..')

import torch
from torch.utils.data import Dataset, DataLoader
from torch.utils.tensorboard import SummaryWriter
import torch.cuda as cuda
import torch.nn as nn
import torchvision
import torch.nn.functional as F

from .wave_equation import WaveDataset, wave_loss, batchLaplacian

from .utils import create_opt, create_one_cycle, find_lr, printNumModelParams, \
                                    rmse, writeMessage, plotSampleWprediction, plotSampleWpredictionByChannel, \
                                    plotSample, curl, jacobian, stream2uv, create_movie, convertSimToImage

from .models import Generator, Encoder, AE_no_P, AE_xhat_z, AE_xhat_zV2

import numpy as np
from tqdm import tqdm
from copy import deepcopy

#Cell
def trainEpoch(myDataLoader, tensorboard_writer, model, opt, p_loss, loss,
               wave_loss, metric, lr_scheduler, tensorboard_rate, device,
               tensorboard_recorder_step, total_steps):
    running_loss = 0.0
    running_rmse = 0.0
    running_ploss = 0.0
    running_waveloss = 0.0
    for i, (X,t0,s0) in enumerate(myDataLoader, start=1):

        # --- Main Training ---

        # gpu
        X = X.to(device)

        u0 = X[:,0,:,:]
        u1 = X[:,1,:,:]
        u2 = X[:,2,:,:]

        # zero the parameter gradients
        opt.zero_grad()

        X_hat, z = model(u1.unsqueeze(1))
        pl = p_loss(z,z) # fix later
        ll = loss(X_hat,u1.unsqueeze(1),device)
        wl = wave_loss(u0,X_hat.squeeze(),u2,device=device)
        combined_loss = pl + ll + wl
        combined_loss.backward()
        opt.step()

        # loss
        batch_loss = combined_loss.item()
        running_loss += batch_loss

        batch_ploss = pl.item()
        running_ploss += batch_ploss

        batch_waveloss = w1.item()
        running_waveloss += batch_waveloss

        # --- Metrics Recording ---

        # metrics
        r = metric(X_hat, X[:,1:2,:,:])
        running_rmse += r

        # record lr change
        total_steps += 1
        tensorboard_writer.add_scalar(tag="LR", scalar_value=opt.param_groups[0]['lr'], global_step=total_steps)
        lr_scheduler.step()

        # tensorboard writes
        if (i % tensorboard_rate == 0):
            tensorboard_recorder_step += 1
            avg_running_loss = running_loss/tensorboard_rate
            avg_running_rmse = running_rmse/tensorboard_rate
            avg_running_ploss = running_ploss/tensorboard_rate
            avg_running_waveloss = running_waveloss/tensorboard_rate
            tensorboard_writer.add_scalar(tag="Loss", scalar_value=avg_running_loss, global_step=tensorboard_recorder_step)
            tensorboard_writer.add_scalar(tag="p_loss", scalar_value=avg_running_ploss, global_step=tensorboard_recorder_step)
            tensorboard_writer.add_scalar(tag="wave_loss", scalar_value=avg_running_waveloss, global_step=tensorboard_recorder_step)
            tensorboard_writer.add_scalar(tag=metric.__name__, scalar_value=avg_running_rmse, global_step=tensorboard_recorder_step)
            # reset running_loss for the next set of batches. (tensorboard_rate number of batches)
            running_loss = 0.0
            running_rmse = 0.0
            running_ploss = 0.0
            running_waveloss = 0.0

    return batch_loss, tensorboard_recorder_step, total_steps


#Cell
def validEpoch(myDataLoader, tensorboard_writer, model, p_loss, loss, wave_loss, metric,
               device, tensorboard_recorder_step):
    running_loss = 0.0
    running_rmse = 0.0
    for i, (X,t0,s0) in enumerate(myDataLoader, start=1):

        # --- Metrics Recording ---

        # gpu
        X = X.to(device)

        u0 = X[:,0,:,:]
        u1 = X[:,1,:,:]
        u2 = X[:,2,:,:]

        perc = len(X)/len(myDataLoader.dataset)

        # forward, no gradient calculations
        with torch.no_grad():
            X_hat, z = model(u1.unsqueeze(1))

        # loss
        # fix later
        pl = p_loss(z,z) # fix later
        ll = loss(X_hat,u1.unsqueeze(1),device)
        wl = wave_loss(u0,X_hat.squeeze(),u2,device=device)
        combined_loss = pl + ll + wl

        running_loss += perc*(combined_loss.item())

        # metrics
        r = metric(X_hat, X[:,1:2,:,:])
        running_rmse += perc*r

    avg_running_loss = running_loss
    avg_running_rmse = running_rmse
    tensorboard_writer.add_scalar(tag="Loss", scalar_value=avg_running_loss, global_step=tensorboard_recorder_step)
    tensorboard_writer.add_scalar(tag=metric.__name__, scalar_value=avg_running_rmse, global_step=tensorboard_recorder_step)

    return running_loss