import numpy as np
import torch
from util import pause

def computeMeanStd_RGB(dataloader, dataset_sizes, batch_sizeP, cuda):

    numBatches = np.round(dataset_sizes / batch_sizeP)

    pop_mean_R = []
    pop_mean_G = []
    pop_mean_B = []
    pop_std0_R = []
    pop_std0_G = []
    pop_std0_B = []

    for i, (data, y) in enumerate(dataloader):

        # display
        if i % 100 == 0:
            print("\tBatch n. {0} / {1}".format(i, int(numBatches)))

        if cuda:
            data = data.to('cuda')

        # shape (3,)
        batch_mean_R = torch.mean(data[:,0,:])
        batch_mean_G = torch.mean(data[:,1,:])
        batch_mean_B = torch.mean(data[:,2,:])

        batch_std0_R = torch.std(data[:,0,:])
        batch_std0_G = torch.std(data[:,1,:])
        batch_std0_B = torch.std(data[:,2,:])

        if cuda:
            batch_mean_R = batch_mean_R.detach().to('cpu')
            batch_mean_G = batch_mean_G.detach().to('cpu')
            batch_mean_B = batch_mean_B.detach().to('cpu')
            batch_std0_R = batch_std0_R.detach().to('cpu')
            batch_std0_G = batch_std0_G.detach().to('cpu')
            batch_std0_B = batch_std0_B.detach().to('cpu')

        pop_mean_R.append(batch_mean_R)
        pop_mean_G.append(batch_mean_G)
        pop_mean_B.append(batch_mean_B)
        pop_std0_R.append(batch_std0_R)
        pop_std0_G.append(batch_std0_G)
        pop_std0_B.append(batch_std0_B)

    # shape (num_iterations, 3) -> (mean across 0th axis) -> shape (3,)
    pop_mean_R = np.mean(pop_mean_R)
    pop_mean_G = np.mean(pop_mean_G)
    pop_mean_B = np.mean(pop_mean_B)
    pop_std0_R = np.mean(pop_std0_R)
    pop_std0_G = np.mean(pop_std0_G)
    pop_std0_B = np.mean(pop_std0_B)

    return pop_mean_R, pop_mean_G, pop_mean_B, pop_std0_R, pop_std0_G, pop_std0_B


def computeMeanStd(dataloader_all, dataset_sizes, batch_sizeP, cuda):

    numBatches = np.round(dataset_sizes / batch_sizeP)

    pop_mean = []
    pop_std0 = []

    for dataloader in dataloader_all:

        for i, (data, y) in enumerate(dataloader):

            ##################
            #if i > 0:
                #break
            ##################

            #print(data.size())
            #pause()

            # display
            #if i % 100 == 0:
                #print("\tBatch n. {0} / {1}".format(i, int(numBatches)))

            if cuda:
                data = data.to('cuda')

            # shape (3,)
            batch_mean = torch.mean(data)
            batch_std0 = torch.std(data)

            if cuda:
                batch_mean = batch_mean.detach().to('cpu')
                batch_std0 = batch_std0.detach().to('cpu')

            pop_mean.append(batch_mean)
            pop_std0.append(batch_std0)

            #print(pop_mean)
            #print(len(pop_mean))
            #pause()

            #if i > 100:
            #break

    # shape (num_iterations, 3) -> (mean across 0th axis) -> shape (3,)
    pop_mean = np.mean(pop_mean)
    pop_std0 = np.mean(pop_std0)

    return pop_mean, pop_std0