"""

DeepSlide

Computes the image statistics for normalization.

Authors: Naofumi Tomita

"""

import argparse

import json

from datetime import datetime

from pathlib import Path

from typing import (List, Tuple)

import torch

from PIL import Image

from torchvision.transforms import ToTensor

Image.MAX_IMAGE_PIXELS = None

def compute_stats(folderpath: Path,

                  image_ext: str) -> Tuple[List[float], List[float]]:

    """

    Compute the mean and standard deviation of the images found in folderpath.

    Args:

        folderpath: Path containing images.

        image_ext: Extension of the image files.

    Returns:

        A tuple containing the mean and standard deviation for the images over the channel, height, and width axes.

    This implementation is based on the discussion from: 

        https://discuss.pytorch.org/t/about-normalization-using-pre-trained-vgg16-networks/23560/9

    """

    class MyDataset(torch.utils.data.Dataset):

        """

        Creates a dataset by reading images.

        Attributes:

            data: List of the string image filenames.

        """

        def __init__(self, folder: Path) -> None:

            """

            Create the MyDataset object.

            Args:

                folder: Path to the images.

            """

            self.data = []

            for file in folder.rglob(f"*{image_ext}"):

                if not file.name.startswith("."):

                    self.data.append(file)

        def __getitem__(self, index: int) -> torch.Tensor:

            """

            Finds the specified image and outputs in correct format.

            Args:

                index: Index of the desired image.

            Returns:

                A PyTorch Tensor in the correct color space.

            """

            return ToTensor()(Image.open(self.data[index]).convert("RGB"))

        def __len__(self) -> int:

            return len(self.data)

    def online_mean_and_sd(

        loader: torch.utils.data.DataLoader, report_interval: int=1000

                           ) -> Tuple[List[float], List[float]]:

        """

        Computes the mean and standard deviation online.

            Var[x] = E[X^2] - (E[X])^2

        Args:

            loader: The PyTorch DataLoader containing the images to iterate over.

            report_interval: Report the intermediate results every N items. (N=0 to suppress reporting.)

        Returns:

            A tuple containing the mean and standard deviation for the images

            over the channel, height, and width axes.

        """

        cnt = 0

        fst_moment = torch.empty(3)

        snd_moment = torch.empty(3)

        for i, data in enumerate(loader, 1):

            b, __, h, w = data.shape

            nb_pixels = b * h * w

            fst_moment = (cnt * fst_moment +

                          torch.sum(data, dim=[0, 2, 3])) / (cnt + nb_pixels)

            snd_moment = (cnt * snd_moment + torch.sum(

                data**2, dim=[0, 2, 3])) / (cnt + nb_pixels)

            cnt += nb_pixels

            if report_interval != 0 and i % report_interval == 0:

                temp_mean = fst_moment.tolist()

                temp_std = torch.sqrt(snd_moment - fst_moment**2).tolist()

                print(f"Mean: {temp_mean}; STD: {temp_std} at iter: {i}")

        return fst_moment.tolist(), torch.sqrt(snd_moment -

                                               fst_moment**2).tolist()

    return online_mean_and_sd(

        loader=torch.utils.data.DataLoader(

            dataset=MyDataset(folder=folderpath),

            batch_size=1,

            num_workers=1,

            shuffle=False))

def save_stats(mean: List, std: List, datapath: str):

    data = {

        'mean': mean,

        'std': std,

        'datapath': datapath}

    data = json.dumps(data, indent=4)

    filename = f"stats_{datetime.now().strftime('%Y-%m-%d_%H:%M')}.json"

    with open(filename, 'w') as outfile:

        outfile.write(data)

    print(f"Results are saved in {filename}.")

def load_stats(jsonfile: str):

    """ Load a stats file in json and return mean and std in lists.

    """

    with open(jsonfile, 'r') as infile:        

        data = json.load(infile)

    print(f"Stats of \'{data['datapath']}\' are loaded from {jsonfile}.")

    return data['mean'], data['std']

if __name__ == '__main__':

    parser = argparse.ArgumentParser(

        description='Compute channel-wise patch color mean and std.')

    parser.add_argument('--datapath', '-i', type=str, required=True,

        help='Path containing images.')

    parser.add_argument('--image_ext', '-x', type=str, default='.png',

        help='Specify file extension of images. Default: .png')

    parser.add_argument('--report_interval', '-n', type=int, default=1000,

        help='Report the intermediate results every N items. Default: 1000')

    parser.add_argument('--save_results', '-d', action='store_true', default=False,

        help='Set this flag to save results.')

    args = parser.parse_args()

    mean, std = compute_stats(Path(args.datapath), args.image_ext,)

    print(f"Mean: {mean}; STD: {std}")

    if args.save_results:

        save_stats(mean=mean, std=std, datapath=args.datapath)