import os
import sys

import torch

import torch.optim as optim
import torch.nn as nn

from torch.utils.data import DataLoader
from data.dataset import Dataset
from utils.utils import Accumulator
from models.depth_estim import DepthEstimation


def train_step(batch):
    ### Training
    model.train() # train mode is on by default after construction

    # 1. Forward pass
    y_pred = model(batch['mesh_posed'])

    # 2. Calculate loss
    loss = loss_fn(y_pred, batch['mesh_posed'])

    # 3. Zero grad optimizer
    optimizer.zero_grad()

    # 4. Loss backward
    loss.backward()

    # 5. Step the optimizer
    optimizer.step()

    return { 'train/loss': loss.detach().cpu().numpy() }


def valid_step(batch):
    ### Testing
    model.eval() # put the model in evaluation mode for testing (inference)
    # 1. Forward pass
    with torch.inference_mode():
        pred = model(batch['mesh_posed'])
    
        # 2. Calculate the loss
        loss = loss_fn(pred, batch['mesh_posed'])

    return { 'valid/loss': loss.detach().cpu().numpy() }


def batch_to_device(batch, device):
    filter = {'mesh_posed'}
    for key in batch:
        if key in filter:
            batch[key] = batch[key].to(device)


if __name__ == '__main__':
    device = "cuda" if torch.cuda.is_available() else "cpu"


    train_dataset = Dataset('/home/workspace/local-data/train')
    valid_dataset = Dataset('/home/workspace/local-data/valid')
    train_dataloader = DataLoader(train_dataset, batch_size=4, drop_last=True, num_workers=2)
    valid_dataloader = DataLoader(valid_dataset, batch_size=4, drop_last=True, num_workers=2)
    train_accumulator = Accumulator()
    valid_accumulator = Accumulator()

    model = DepthEstimation()

    loss_fn = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    iteration = 0
    while True:

        for batch in train_dataloader:
            train_loss = train_step(batch)
            train_accumulator.update(train_loss)

            if iteration % 1000 == 0:
                print(f"Iteration: {iteration} | Train loss: {train_accumulator.get()['train/loss']}")
                train_accumulator.reset()

            iteration += 1

        for batch in valid_dataloader:
            valid_loss = valid_step(batch)
            valid_accumulator.update(valid_loss)

            print(f"Iteration: {iteration} | Valid loss: {valid_accumulator.get()['valid/loss']}")
            valid_accumulator.reset()



    pass