import os
import subprocess
import numpy as np
import torch
from apex import amp
from torch.optim.lr_scheduler import LambdaLR

def set_seed_everywhere(seed, cuda):
    np.random.seed(seed)
    torch.manual_seed(seed)
    if cuda:
        torch.cuda.manual_seed_all(seed)

def handle_dirs(dirpath):
    if not os.path.exists(dirpath):
        os.makedirs(dirpath)

def make_train_state(args):
    return {'stop_early': False,
            'early_stopping_step': 0,
            'early_stopping_best_val': 1e8,
            'learning_rate': args.learning_rate,

            'epoch_index': 0,
            'cursor_train': 0,
            'cursor_val': 0,

            'train_running_loss': [],
            'val_running_loss': [],
            'test_running_loss': [],
            
            'model_filename': args.model_state_file,
    }

def update_train_state(args, model, optimizer, train_state):
    # Save one model at least
    if train_state['epoch_index'] == 0:
        checkpoint = {
            'model': model.state_dict(),
            'optimizer': optimizer.state_dict(),
        }
        if args.fp16: checkpoint['amp'] = amp.state_dict()
        torch.save(checkpoint, train_state['model_filename'])
        
        train_state['stop_early'] = False

    # Save model if performance improved
    elif train_state['epoch_index'] >= 1:
        loss_tm1, loss_t = train_state['val_running_loss'][-2:]
         
        # If loss worsened
        if loss_t >= loss_tm1:
            # Update step
            train_state['early_stopping_step'] += 1
        # Loss decreased
        else:
            # Save the best model
            if loss_t < train_state['early_stopping_best_val']:
                
                checkpoint = {
                    'model': model.state_dict(),
                    'optimizer': optimizer.state_dict(),
                    'cursor_train': train_state['cursor_train'],
                    'cursor_val': train_state['cursor_val'],
                }
                if args.fp16: checkpoint['amp'] = amp.state_dict()
                torch.save(checkpoint, train_state['model_filename'])

                train_state['early_stopping_best_val'] = loss_t

            # Reset early stopping step
            train_state['early_stopping_step'] = 0

        # Stop early ?
        train_state['stop_early'] = \
            train_state['early_stopping_step'] >= args.early_stopping_criteria

    return train_state

def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):
    """ Create a schedule with a learning rate that decreases linearly after
    linearly increasing during a warmup period.
    """

    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        return max(
            0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))
            )

    return LambdaLR(optimizer, lr_lambda, last_epoch)

def rm_rf(obj_dir):
    if not os.path.exists(obj_dir): return False
    for root, dirs, files in os.walk(obj_dir, topdown=False):
        for name in files:
            os.remove(os.path.join(root, name))
        for name in dirs:
            os.rmdir(os.path.join(root, name))

if __name__ == "__main__":
    pass