import os
import pickle
import torch
import matplotlib.pyplot as plt
import numpy as np
import yaml
from einops import rearrange
from roboimi.utils.constants import CONFIG_DIR

class IOUtils:
    @staticmethod
    def save_checkpoint(model_state_dict, epoch, ckpt_dir, seed,tag=None,min_val_loss=None):
        if tag == 'last':
            ckpt_path = f"{ckpt_dir}/policy_last.ckpt"
            torch.save(model_state_dict, ckpt_path)
            print(f"Last checkpoint saved at {ckpt_path}")
        elif tag == 'best':
            ckpt_path = f"{ckpt_dir}/policy_best.ckpt"
            torch.save(model_state_dict, ckpt_path)
            print(f"Best ckpt with val loss {min_val_loss:.6f} at epoch{epoch} saved!")
        else:
            ckpt_path = f"{ckpt_dir}/checkpoint_epoch_{epoch}_seed_{seed}.ckpt"
            torch.save(model_state_dict, ckpt_path)
            print(f"Checkpoint saved at {ckpt_path} for epoch {epoch}")


    @staticmethod
    def load_stats(ckpt_dir):
        stats_path = os.path.join(ckpt_dir, 'dataset_stats.pkl')
        with open(stats_path, 'rb') as f:
            return pickle.load(f)

    @staticmethod
    def save_results(ckpt_dir, ckpt_name, summary_str, episode_returns, highest_rewards):
        result_file_name = 'result_' + ckpt_name.split('.')[0] + '.txt'
        with open(os.path.join(ckpt_dir, result_file_name), 'w') as f:
            f.write(summary_str)
            f.write(repr(episode_returns))
            f.write('\n\n')
            f.write(repr(highest_rewards))

    @staticmethod
    def plot_history(train_history, validation_history, num_epochs, ckpt_dir, seed):
        for key in train_history[0]:
            plot_path = os.path.join(ckpt_dir, f'train_val_{key}_seed_{seed}.png')
            plt.figure()
            train_values = [summary[key].item() for summary in train_history]
            val_values = [summary[key].item() for summary in validation_history]
            plt.plot(np.linspace(0, num_epochs-1, len(train_history)), train_values, label='train')
            plt.plot(np.linspace(0, num_epochs-1, len(validation_history)), val_values, label='validation')
            plt.tight_layout()
            plt.legend()
            plt.title(key)
            plt.savefig(plot_path)
        print(f'Saved plots to {ckpt_dir}')

    @staticmethod
    def load_config():
        with open(CONFIG_DIR, 'r') as file:
            config = yaml.safe_load(file)
        return config


    @staticmethod
    def save_stats(ckpt_dir, stats):
        stats_path = os.path.join(ckpt_dir, f'dataset_stats.pkl')
        with open(stats_path, 'wb') as f:
            pickle.dump(stats, f)

    @staticmethod
    def load_stats(ckpt_dir):
        stats_path = os.path.join(ckpt_dir, 'dataset_stats.pkl')
        with open(stats_path, 'rb') as f:
            return pickle.load(f)
        

    # @staticmethod
    # def get_image(ts, camera_names):
    #     curr_images = []
    #     for cam_name in camera_names:
    #         curr_image = rearrange(ts.observation['images'][cam_name], 'h w c -> c h w')
    #         curr_images.append(curr_image)
    #     curr_image = np.stack(curr_images, axis=0)
    #     curr_image = torch.from_numpy(curr_image / 255.0).float().cuda().unsqueeze(0)
    #     return curr_image
    
    # @staticmethod
    # def load_data(dataset_dir, num_episodes, camera_names, batch_size, val_batch_size):
    #     from utils import load_data
    #     return load_data(dataset_dir, num_episodes, camera_names, batch_size, val_batch_size)

    @staticmethod
    def load_policy(config, ckpt_name):
        from roboimi.utils.model_interface import ModelInterface
        model_interface = ModelInterface(config)
        policy = model_interface.make_policy()
        ckpt_path = os.path.join(config['ckpt_dir'], ckpt_name)
        policy.load_state_dict(torch.load(ckpt_path))
        policy.cuda()
        policy.eval()
        return policy

    # @staticmethod
    # def make_env(config):
    #     if config['real_robot']:
    #         from aloha_scripts.real_env import make_real_env
    #         return make_real_env(init_node=True)
    #     else:
    #         from sim_env import make_sim_env
    #         return make_sim_env(config['task_name'])