"""Some utility functions for working with PyTorch"""

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

# %% auto 0
__all__ = ['set_seed', 'pil_to_tensor', 'tensor_to_pil', 'iterate_modules', 'tensor_stats_df', 'get_torch_device',
           'denorm_img_tensor', 'move_data_to_device', 'ImageDataset', 'compute_mean_std']

# %% ../nbs/00_core.ipynb 3
# Import necessary modules from the standard library
from pathlib import Path  # For working with file paths
import logging  # For logging messages
import hashlib
import random

# Disable logging warnings
logging.disable(logging.WARNING)

import numpy as np  # For working with arrays

from PIL import Image  # For working with images

import torch  # PyTorch module for deep learning
from torchvision import transforms  # PyTorch module for image transformations

# %% ../nbs/00_core.ipynb 5
def set_seed(seed: int, # The seed value to be set for all random number generators.
             deterministic: bool = False # If True, uses deterministic algorithms in PyTorch where possible for reproducibility, at the cost of performance.
            ) -> None:
    """
    Sets the seed for generating random numbers in PyTorch, NumPy, and Python's random module.

    This function is used for reproducibility in stochastic operations, e.g. shuffling in data loaders,
    random initializations in neural networks, etc.

    Note: The deterministic flag does not guarantee complete reproducibility. Operations which rely on CUDA might
    still produce non-deterministic results.
    """
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.use_deterministic_algorithms(deterministic)

# %% ../nbs/00_core.ipynb 8
def pil_to_tensor(img:Image, # The input PIL image.
                  mean=[0.485, 0.456, 0.406], # The mean values for normalization.
                  std=[0.229, 0.224, 0.225] # The standard deviation values for normalization.
                 ):
    """
    Converts a PIL image to a normalized and batched PyTorch tensor.
    
    Returns:
        The normalized and batched tensor.
    """
    
    return transforms.Normalize(mean, std)(transforms.ToTensor()(img))[None]

# %% ../nbs/00_core.ipynb 16
def tensor_to_pil(tensor: torch.Tensor # The tensor to be converted
                 ):
    """
    Convert a tensor to a PIL image.
    
    Returns:
    img (PIL.Image): The PIL image
    """
    
    # Remove the first dimension if the tensor has 4 dimensions
    if len(tensor.shape) == 4: tensor.squeeze_(0)
        
    # Use the ToPILImage() function from the transforms module to convert the tensor to a PIL image
    return transforms.ToPILImage()(tensor)

# %% ../nbs/00_core.ipynb 20
def iterate_modules(module: torch.nn.Module): # A PyTorch module that contains child modules to be iterated over.
    """
    A generator function that yields the children and grandchildren of a PyTorch module.    
    """
    
    for child in module.children():
        yield child
        yield from iterate_modules(child)

# %% ../nbs/00_core.ipynb 23
import pandas as pd

# %% ../nbs/00_core.ipynb 24
def tensor_stats_df(tensor, # Input tensor for which statistics are to be calculated.
                    attrs = ["mean", "std", "min", "max"], # List of statistics to be calculated.
                    shape=True): # If True, include shape of the tensor in the output.
    """
    Calculate and return statistics of a given tensor as a pandas DataFrame.
    """
    
    attr_dict = {attr: getattr(tensor, attr)().item() for attr in attrs}
    if shape: attr_dict["shape"] = tensor.shape
    return pd.DataFrame.from_dict(attr_dict, orient='index')

# %% ../nbs/00_core.ipynb 27
def get_torch_device():
    """
    This function returns the device to be used for PyTorch computations.
    
    Returns:
    str: "mps" if Metal Performance Shaders (MPS) for MacOS is available, 
         "cuda" if CUDA is available, 
         "cpu" otherwise
    """
    device = (
        "mps"
        if torch.backends.mps.is_available()
        else "cuda"
        if torch.cuda.is_available()
        else "cpu"
    )
    return device

# %% ../nbs/00_core.ipynb 30
def denorm_img_tensor(img_tensor:torch.Tensor, # The tensor representing the normalized image.
                      mean:list, # The mean values used for normalization.
                      std:list): # The standard deviation values used for normalization.
    """
    Denormalize an image tensor.
    
    Returns:
        torch.Tensor: The tensor representing the denormalized image.
    """
    # Convert the mean and standard deviation values to tensors
    mean_tensor = torch.Tensor(mean).view(1,1,-1).permute(2, 0, 1)
    std_tensor = torch.Tensor(std).view(1,1,-1).permute(2, 0, 1)
    # Denormalize the image tensor
    return img_tensor*std_tensor+mean_tensor

# %% ../nbs/00_core.ipynb 34
def move_data_to_device(data, # Data to move to the device.
                        device:torch.device # The PyTorch device to move the data to.
                       ): # Moved data with the same structure as the input but residing on the specified device.
    """
    Recursively move data to the specified device.

    This function takes a data structure (could be a tensor, list, tuple, or dictionary)
    and moves all tensors within the structure to the given PyTorch device.
    """
    
    # If the data is a tuple, iterate through its elements and move each to the device.
    if isinstance(data, tuple):
        return tuple(move_data_to_device(d, device) for d in data)
    
    # If the data is a list, iterate through its elements and move each to the device.
    if isinstance(data, list):
        return list(move_data_to_device(d, device) for d in data)
    
    # If the data is a dictionary, iterate through its key-value pairs and move each value to the device.
    elif isinstance(data, dict):
        return {k: move_data_to_device(v, device) for k, v in data.items()}
    
    # If the data is a tensor, directly move it to the device.
    elif isinstance(data, torch.Tensor):
        return data.to(device)
    
    # If the data type is not a tensor, list, tuple, or dictionary, it remains unchanged.
    else:
        return data

# %% ../nbs/00_core.ipynb 38
import os
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
from pathlib import Path
from tqdm.auto import tqdm

from typing import List

# %% ../nbs/00_core.ipynb 39
class ImageDataset(Dataset):
    """
    A PyTorch Dataset for RGB images.
    """
    def __init__(self, image_paths, transform=None):
        self.image_paths = image_paths
        self.transform = transform or transforms.ToTensor()

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        # Open image and ensure it's in RGB format
        img = Image.open(self.image_paths[idx]).convert('RGB')
        # Apply transformation
        img = self.transform(img)
        return img

# %% ../nbs/00_core.ipynb 40
def compute_mean_std(image_paths:List[Path], # List of image file paths.
                     batch_size:int=64, # Number of images to process in a batch.
                     num_workers:int=os.cpu_count(), # Number of subprocesses to use for data loading.
                     device:str='cpu', # Device to use for computation (e.g., 'cpu', 'cuda')
                     image_size:int=224, # Size to resize images to.
                     transform:transforms.Compose=None # Torchvision transforms to apply to the images.
                    )->dict: # Dictionary containing 'mean' and 'std' values.
    """
    Computes the mean and standard deviation of images provided in image_paths.
    """

    if not transform:
        # Define transformation (without normalization)
        transform = transforms.Compose([
            transforms.Resize((image_size, image_size)),  # Resize images
            transforms.ToTensor(),
        ])

    # Create custom dataset
    dataset = ImageDataset(image_paths, transform=transform)

    # Create DataLoader
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        num_workers=num_workers,
        shuffle=False
    )

    # Initialize accumulators
    # Move accumulators to the specified device
    mean = torch.zeros(3, device=device)
    std = torch.zeros(3, device=device)
    total_pixels = 0

    # Iterate over DataLoader
    for data in tqdm(loader):
        # data shape: (batch_size, channels, height, width)
        batch_samples = data.size(0)
        # Flatten height and width
        data = data.view(batch_samples, data.size(1), -1).to(device)
        total_pixels += batch_samples * data.size(2)

        mean += data.sum(dim=[0, 2])
        std += (data ** 2).sum(dim=[0, 2])

    # Finalize mean and std computation
    mean /= total_pixels
    std = torch.sqrt(std / total_pixels - mean ** 2)

    return {'mean': mean.tolist(), 'std': std.tolist()}
