import torch
import numpy as np 
import os
import logging


class CheckpointSaver:
    def __init__(self, dirpath, top_n=5):
        """
        dirpath: Directory path where to store all model weights 
        decreasing: If decreasing is `True`, then lower metric is better
        top_n: Total number of models to track based on validation metric value
        """
        if not os.path.exists(dirpath): os.makedirs(dirpath)
        self.dirpath = dirpath
        self.top_n = top_n
        self.top_model_paths = []
        
    def __call__(self, model, iteration, optimizer, scaler):
        model_path = os.path.join(self.dirpath, f'iteration_{iteration}.pt')
        torch.save({ 'model': model.state_dict(), 'optimizer': optimizer, 'scaler': scaler }, model_path)
        self.top_model_paths.append({'path': model_path})
        if len(self.top_model_paths)>self.top_n: 
            self.cleanup()
    
    def cleanup(self):
        to_remove = self.top_model_paths[:len(self.top_model_paths)-self.top_n]
        logging.info(f"Removing extra models.. {to_remove}")
        for o in to_remove:
            os.remove(o['path'])
        self.top_model_paths = self.top_model_paths[len(self.top_model_paths)-self.top_n:]


def load_model(model: torch.nn.Module, checkpoint, optimizer: torch.nn.Module=None, scaler: torch.nn.Module=None, iteration=''):
    # list of checkpoint files in the folder checkpoints
    if iteration == '':
        checkpoints = [os.path.join(checkpoint, f) for f in os.listdir(checkpoint) if f.endswith('.pt')]
    else:
        checkpoints = [os.path.join(checkpoint, f) for f in os.listdir(checkpoint) if f.endswith(f'iteration_{iteration}.pt')]
        assert len(checkpoints) == 1
    iterations = np.asarray([int(f.split('_')[-1].split('.')[0]) for f in checkpoints])
    last_iteration = np.argmax(iterations)
    checkpoint = torch.load(checkpoints[last_iteration])
    iteration = iterations[last_iteration]
    model.load_state_dict(checkpoint['model'], strict=False)
    if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer'].state_dict())
    if scaler is not None and checkpoint['scaler'] is not None: scaler.load_state_dict(checkpoint['scaler'])
    print(f'Loaded model from {checkpoints[last_iteration]} at iteration {iteration}')
    return int(iteration) if iteration != '' else 0