import yaml
from argparse import Namespace

import random
import numpy as np
import torch
import torch.nn as nn
import torchvision.transforms.functional as TF
from matplotlib.figure import Figure

def load_yaml(file: str) -> Namespace:
    with open(file) as f:
        config = yaml.safe_load(f)
    return Namespace(**config)

def set_seed(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)

def get_output_shape(model: nn.Module, input_shape: tuple):
    t = torch.rand(input_shape)
    return model(t).shape

def filename_to_width_height(filename: str) -> tuple:
    splitted = filename.split('_')
    return int(splitted[2]), int(splitted[3])

def show_valid_image_during_training(model: nn.Module, image, input_resolution: int, padding_mode: str, device: str) -> Figure:
    padding = ((input_resolution - image.sw) // 2, (input_resolution - image.sh) // 2)
    image_tensor = TF.pad(image.tensor, padding=padding, padding_mode=padding_mode).unsqueeze(0).unsqueeze(0).to(device)
    model = model.to(device)
    preds = model(image_tensor)
    segmentations = dict()
    for i, organ in enumerate(image.organs):
        pred = preds[:, i, :, :]
        segmentation = TF.center_crop(pred, output_size=(image.sh, image.sw))
        segmentations[organ] = segmentation.squeeze(0).detach().cpu().numpy()
    fig = image.show_segmented_images(segmentations=segmentations)
    return fig