# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_utils.ipynb.

# %% auto 0
__all__ = ['UnNormalize', 'is_sequential', 'params', 'is_frozen', 'is_unfrozen', 'freeze_params', 'unfreeze_params',
           'freeze_model', 'unfreeze_model', 'get_norm', 'is_norm', 'get_norm_id', 'del_norm', 'update_norm',
           'is_tensor', 'tensor_to_img', 'plt_show', 'show_img', 'plot_loss_and_acc']

# %% ../nbs/00_utils.ipynb 2
from .imports import *

# %% ../nbs/00_utils.ipynb 4
class UnNormalize(transforms.Normalize):
    def __init__(self,mean,std,*args,**kwargs):
        mean = [-m/s for m,s in zip(mean,std)]
        std = [1/s for s in std]
        super().__init__(mean, std, *args, **kwargs)

def is_sequential(x):
    return isinstance(x, nn.Sequential)

def params(m):
    return [p for p in m.parameters()]

def is_frozen(model):
    return np.array([not p.requires_grad for p in (params(model))]).all()

def is_unfrozen(model):
    return np.array([p.requires_grad for p in (params(model))]).all()

def freeze_params(params):
    for p in params:
        p.requires_grad = False

def unfreeze_params(params):
    for p in params:
        p.requires_grad = True

def freeze_model(model):
    freeze_params(model.parameters())

def unfreeze_model(model):
    unfreeze_params(model.parameters())

def get_norm(tfms):
    if tfms is None:
        return tfms
    try:
        tfms_list = list(tfms)
    except:
        tfms_list = list(tfms.transforms)
    for t in tfms_list:
        if is_norm(t):
            return t
    return None

def is_norm(x):
    return type(x).__name__ == 'Normalize'

def get_norm_id(tfms):
    try:
        tfms_list = list(tfms)
    except:
        tfms_list = tfms.transforms
    for i,t in enumerate(tfms_list):
        if is_norm(t):
            return t,i
    return None, None

def del_norm(tfms, idx=None):
    if idx is None:
        idx = get_norm_id(tfms)[1]
    if idx is not None:
        del tfms.transforms[idx]

def update_norm(tfms, mean, std, idx=None):
    if idx is None:
        idx = get_norm_id(tfms)[1]
    if idx is not None:
        tfms.transforms[idx].mean = mean
        tfms.transforms[idx].std = std
    else:
        tfms.transforms.append(transforms.Normalize(mean=mean, std=std))

def is_tensor(x):
    return isinstance(x, torch.Tensor)

def tensor_to_img(t):
    if t.dim() > 3:
        return [np.array(np.transpose(t_,(1,2,0))) for t_ in t]
    return np.array(np.transpose(t,(1,2,0)))

def plt_show(im, cmap=None, title='', figsize=(7,7)):
    if path_or_str(im):
        im = rgb_read(im)
    if is_tensor(im):
        im = tensor_to_img(im)
        if is_list(im): im = im[0]
    fig=plt.figure(figsize=figsize)
    plt.imshow(im, cmap=cmap)
    plt.title(title)
    plt.show()

def show_img(img, cmap=None, titles=[''], figsize=(7,7)):
    if not is_list(img):
        imgs = [img]
    else:
        imgs = img
    [plt_show(img, cmap=cmap, title=title, figsize=figsize) for img,title in zip(imgs,titles)]

def plot_loss_and_acc(log_dir, loss_ylim=(0.0, 0.9), acc_ylim=(0.7, 1.0), save_loss=None, save_acc=None):

    metrics = pd.read_csv(f"{log_dir}/metrics.csv")

    aggreg_metrics = []
    agg_col = "epoch"
    for i, dfg in metrics.groupby(agg_col):
        agg = dict(dfg.mean())
        agg[agg_col] = i
        aggreg_metrics.append(agg)

    df_metrics = pd.DataFrame(aggreg_metrics)
    df_metrics[["train_loss", "val_loss"]].plot(
        grid=True, legend=True, xlabel="Epoch", ylabel="Loss"
    )

    plt.ylim(loss_ylim)
    if save_loss is not None:
        plt.savefig(save_loss)

    df_metrics[["train_acc", "val_acc"]].plot(
        grid=True, legend=True, xlabel="Epoch", ylabel="ACC"
    )

    plt.ylim(acc_ylim)
    if save_acc is not None:
        plt.savefig(save_acc)
