from __future__ import absolute_import
from __future__ import division
from __future__ import print_function


import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from ..backend.common import get_session,addindent,get_time_suffix,get_class,format_time,get_terminal_size,snake2camel,camel2snake,get_function


_session = get_session()
_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
_epsilon = _session.epsilon

__all__ = ['l1_reg','l2_reg','orth_reg','get_reg','total_variation_norm_reg']


_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
_session=get_session()
_epsilon=_session.epsilon


def l1_reg(model:nn.Module,reg_weight=1e-6):
    #with torch.enable_grad():
    loss =torch.tensor(0,dtype=torch.float,requires_grad=True).to(_device)
    for name, param in model.named_parameters():
        if 'bias' not in name:
            loss = loss + (reg_weight * torch.sum(torch.abs(param)))
        return loss


def l2_reg(model:nn.Module,reg_weight=1e-6):
    loss =0
    for name, param in model.named_parameters():
        if 'bias' not in name:
            loss=loss+ reg_weight *torch.norm(param)
        return loss


def orth_reg(model:nn.Module,reg_weight=1e-6):
    loss =torch.tensor(0,dtype=torch.float,requires_grad=True).to(_device)
    for name, param in model.named_parameters():
        if 'bias' not in name:
            param_flat = param.view(param.shape[0], -1)
            sym = torch.mm(param_flat, torch.t(param_flat))
            sym -= torch.eye(param_flat.shape[0])
            loss = loss + (reg_weight * sym.abs().sum())
    return loss

def total_variation_norm_reg(output:torch.Tensor,reg_weight=1e-6):
    assert len(output.size())==4
    loss = reg_weight * (torch.sum(torch.abs(output[:, :, :, :-1] - output[:, :, :, 1:])) + torch.sum(torch.abs(output[:, :, :-1, :] - output[:, :, 1:, :])))
    return loss


def get_reg(reg_name):
    if reg_name is None:
        return None
    if '_reg' not in reg_name:
        reg_name=reg_name+'_reg'
    reg_modules = ['trident.optims.pytorch_regularizers']
    reg_fn = get_function(camel2snake(reg_name), reg_modules)
    return reg_fn