import torch
import os
from tqdm import tqdm
import numpy as np
from copy import deepcopy
from itertools import repeat
import matplotlib.pyplot as plt
import time
from roboimi.utils.utils import set_seed, compute_dict_mean, detach_dict, load_data
from roboimi.utils.io_utils import IOUtils
from roboimi.utils.model_interface import ModelInterface
import matplotlib.pyplot as plt

def train_bc(config):
    num_epochs = config['num_epochs']
    ckpt_dir = config['ckpt_dir']
    seed = config['seed']

    os.makedirs(ckpt_dir, exist_ok=True)

    set_seed(seed)

    model_interface = ModelInterface(config)
    model_interface.setup()

    policy = model_interface.make_policy()
    policy.cuda()
    optimizer = model_interface.make_optimizer(policy)
    # print("cam names=====",config['camera_names'])
    train_dataloader, val_dataloader, stats, _ = load_data(
        config['dataset_dir'], 
        config['num_episodes'], 
        config['camera_names'], 
        config['batch_size'], 
        config['batch_size'])

    IOUtils.save_stats(ckpt_dir, stats)

    train_history = []
    validation_history = []
    min_val_loss = np.inf
    min_train_loss = np.inf
    best_ckpt_info = None

    plt.ion()
    fig, ax = plt.subplots()
    train_losses, val_losses = [], []
    train_line, = ax.plot([], [], label='Train Loss')
    val_line, = ax.plot([], [], label='Validation Loss')
    ax.autoscale_view()
    ax.set_xlabel('Epoch')
    ax.set_ylabel('Loss')
    ax.legend()
    ax.grid(True)
    

    train_annotation = ax.annotate('', xy=(0, 0), textcoords='offset points')
    val_annotation = ax.annotate('', xy=(0, 0), textcoords='offset points')
    

    min_train_text = ax.text(0.85, 0.5, '', transform=ax.transAxes, fontsize=10, verticalalignment='center', horizontalalignment='left', bbox=dict(facecolor='white', alpha=0.5))
    min_val_text = ax.text(0.85, 0.45, '', transform=ax.transAxes, fontsize=10, verticalalignment='center', horizontalalignment='left', bbox=dict(facecolor='white', alpha=0.5))

    for epoch in tqdm(range(num_epochs)):
        print(f'\nEpoch {epoch}')

        # Validation
        epoch_val_loss, epoch_summary = validate(policy, val_dataloader)
        validation_history.append(epoch_summary)
        val_losses.append(epoch_val_loss.cpu().item()) 

        if epoch_val_loss < min_val_loss:
            min_val_loss = epoch_val_loss
            min_val_epoch = epoch
            best_ckpt_info = (epoch, min_val_loss,
                              deepcopy(policy.state_dict()))

        print(f'Val loss:   {epoch_val_loss:.5f}')
        print_summary(epoch_summary)

        # Training
        epoch_train_loss, epoch_summary = train_epoch(
            policy, optimizer, train_dataloader)
        train_history.append(epoch_summary)
        train_losses.append(epoch_train_loss.cpu().item()) 

        if epoch_train_loss < min_train_loss:
            min_train_loss = epoch_train_loss
            min_train_epoch = epoch

        print(f'Train loss: {epoch_train_loss:.5f}')
        print_summary(epoch_summary)

        # Update the plot with the new data
        train_line.set_xdata(range(len(train_losses)))
        train_line.set_ydata(train_losses)
        val_line.set_xdata(range(len(val_losses)))
        val_line.set_ydata(val_losses)
        
        # Update annotations with the latest loss values at their respective positions
        train_annotation.set_position((len(train_losses)-1, train_losses[-1]))
        train_annotation.xy = (len(train_losses)-1, train_losses[-1])
        train_annotation.set_text(f'{train_losses[-1]:.5f}')
        
        val_annotation.set_position((len(val_losses)-1, val_losses[-1]))
        val_annotation.xy = (len(val_losses)-1, val_losses[-1])
        val_annotation.set_text(f'{val_losses[-1]:.5f}')
        
        # Update text objects with the minimum loss values, fixed on the right side
        min_train_text.set_text(f'Min Train Loss: {min_train_loss:.5f} (Epoch {min_train_epoch})')
        min_val_text.set_text(f'Min Val Loss: {min_val_loss:.5f} (Epoch {min_val_epoch})')
        
        ax.relim()
        ax.autoscale_view()
        plt.draw()
        plt.pause(0.1)


    plt.ioff() 
    IOUtils.save_checkpoint(policy, 'last', ckpt_dir, seed, 'last')

    best_epoch, min_val_loss, best_state_dict = best_ckpt_info
    IOUtils.save_checkpoint(best_state_dict, best_epoch,
                            ckpt_dir, seed, 'best', min_val_loss)
    print(
        f'Training finished:\nSeed {seed}, val loss {min_val_loss:.6f} at epoch {best_epoch}')

    IOUtils.plot_history(train_history, validation_history,
                         num_epochs, ckpt_dir, seed)

    return best_ckpt_info






def validate(policy, dataloader):
    policy.eval()
    epoch_dicts = []
    with torch.inference_mode():
        for data in dataloader:
            forward_dict = forward_pass(data, policy)
            epoch_dicts.append(forward_dict)
    epoch_summary = compute_dict_mean(epoch_dicts)
    return epoch_summary['loss'], epoch_summary


def train_epoch(policy, optimizer, dataloader):
    policy.train()
    epoch_dicts = []
    for data in dataloader:
        optimizer.zero_grad()
        forward_dict = forward_pass(data, policy)
        loss = forward_dict['loss']
        loss.backward()
        optimizer.step()
        epoch_dicts.append(detach_dict(forward_dict))
    epoch_summary = compute_dict_mean(epoch_dicts)
    return epoch_summary['loss'], epoch_summary


def forward_pass(data, policy):
    image_data, qpos_data, action_data, is_pad = data
    image_data, qpos_data, action_data, is_pad = image_data.cuda(
    ), qpos_data.cuda(), action_data.cuda(), is_pad.cuda()
    return policy(qpos_data, image_data, action_data, is_pad)


def print_summary(summary):
    summary_string = ' '.join(
        [f'{k}: {v.item():.3f}' for k, v in summary.items()])
    print(summary_string)


if __name__ == '__main__':
    io_utils = IOUtils()
    config = io_utils.load_config()
    train_bc(config)
