import os
import getpass
from pathlib import Path
import torch
import torch.nn as nn
import torch.nn.functional as F

import numpy as np

  
class GeneralizedCELoss(nn.Module):

    def __init__(self, q=0.7):
        super(GeneralizedCELoss, self).__init__()
        self.q = q
             
    def forward(self, logits, targets):
        p = F.softmax(logits, dim=1)
        if np.isnan(p.mean().item()):
            raise NameError('GCE_p')
        Yg = torch.gather(p, 1, torch.unsqueeze(targets, 1))
        # modify gradient of cross entropy
        loss_weight = (Yg.squeeze().detach()**self.q)*self.q
        if np.isnan(Yg.mean().item()):
            raise NameError('GCE_Yg')

        loss = F.cross_entropy(logits, targets, reduction='none') * loss_weight

        return loss


def has_checkpoint(suffix = ''):
    slurm_job_id = os.environ.get('SLURM_JOB_ID')
    if slurm_job_id is not None and Path(f'/checkpoint/{getpass.getuser()}/{slurm_job_id}/chkpt{suffix}').exists():
        return True
    return False      

def load_checkpoint(suffix = ''):   
    slurm_job_id = os.environ.get('SLURM_JOB_ID')
    fname = Path(f'/checkpoint/{getpass.getuser()}/{slurm_job_id}/chkpt{suffix}')
    if slurm_job_id is not None and fname.exists():
        return torch.load(fname)

def save_checkpoint(obj, suffix = ''):   
    slurm_job_id = os.environ.get('SLURM_JOB_ID')        
    
    if slurm_job_id is not None and Path(f'/checkpoint/{getpass.getuser()}/{slurm_job_id}/').exists():        
        torch.save(obj, 
                   Path(f'/checkpoint/{getpass.getuser()}/{slurm_job_id}/chkpt{suffix}').open('wb')                  
                  )           
