import os
from pathlib import Path

import torch


def set_bn_eval(module):
    for submodule in module.modules():
        if "batchnorm" in submodule.__class__.__name__.lower():
            submodule.train(False)


def set_train(model):
    """Disable batch normalization when training."""
    model.train()
    set_bn_eval(model)


def set_eval(model):
    model.eval()


def save_model(model, savename):
    Path(os.path.dirname(savename)).mkdir(parents=True, exist_ok=True)
    torch.save(model.state_dict(), savename + ".tmp")
    os.rename(savename + ".tmp", savename)
    print(f"saved {savename}")


def r1_reg(d_out, x_in):
    # zero-centered gradient penalty for real images
    # from https://github.com/clovaai/stargan-v2/blob/875b70a150609e8a678ed8482562e7074cdce7e5/core/solver.py#L282
    batch_size = x_in.size(0)
    # TODO ideally we should scale here : https://pytorch.org/docs/stable/notes/amp_examples.html#gradient-penalty
    # maybe scale d_out.sum() and unscale . there's a default scaler
    grad_dout = torch.autograd.grad(
        outputs=d_out.sum(),
        inputs=x_in,
        create_graph=True,
        retain_graph=True,
        only_inputs=True,
        # allow_unused=True,
    )[0]
    grad_dout2 = grad_dout.pow(2)
    assert grad_dout2.size() == x_in.size()
    reg = 0.5 * grad_dout2.reshape(batch_size, -1).sum(1).mean(0)
    return reg
