import numpy as np
import cv2
import sys
import os
import glob

def im2patch(im, pch_size, stride=1):
    '''
    Transform image to patches.
    Input:
        im: 3 x H x W or 1 X H x W image, numpy format
        pch_size: (int, int) tuple or integer
        stride: (int, int) tuple or integer
    '''
    if isinstance(pch_size, tuple):
        pch_H, pch_W = pch_size
    elif isinstance(pch_size, int):
        pch_H = pch_W = pch_size
    else:
        sys.exit('The input of pch_size must be a integer or a int tuple!')

    if isinstance(stride, tuple):
        stride_H, stride_W = stride
    elif isinstance(stride, int):
        stride_H = stride_W = stride
    else:
        sys.exit('The input of stride must be a integer or a int tuple!')


    C, H, W = im.shape
    num_H = len(range(0, H-pch_H+1, stride_H))
    num_W = len(range(0, W-pch_W+1, stride_W))
    num_pch = num_H * num_W
    pch = np.zeros((C, pch_H*pch_W, num_pch), dtype=im.dtype)
    kk = 0
    for ii in range(pch_H):
        for jj in range(pch_W):
            temp = im[:, ii:H-pch_H+ii+1:stride_H, jj:W-pch_W+jj+1:stride_W]
            pch[:, kk, :] = temp.reshape((C, num_pch))
            kk += 1

    return pch.reshape((C, pch_H, pch_W, num_pch))

def noise_estimate(im, pch_size=8):
    '''
    Implement of noise level estimation of the following paper:
    Chen G , Zhu F , Heng P A . An Efficient Statistical Method for Image Noise Level Estimation[C]// 2015 IEEE International Conference
    on Computer Vision (ICCV). IEEE Computer Society, 2015.
    Input:
        im: the noise image, H x W x 3 or H x W numpy tensor, range [0,1]
        pch_size: patch_size
    Output:
        noise_level: the estimated noise level
    '''

    if im.ndim == 3:
        im = im.transpose((2, 0, 1))
    else:
        im = np.expand_dims(im, axis=0)

    # image to patch
    pch = im2patch(im, pch_size, 3)  # C x pch_size x pch_size x num_pch tensor
    num_pch = pch.shape[3]
    pch = pch.reshape((-1, num_pch))  # d x num_pch matrix
    d = pch.shape[0]

    mu = pch.mean(axis=1, keepdims=True)  # d x 1
    X = pch - mu
    sigma_X = np.matmul(X, X.transpose()) / num_pch
    sig_value, _ = np.linalg.eigh(sigma_X)
    sig_value.sort()

    for ii in range(-1, -d-1, -1):
        tau = np.mean(sig_value[:ii])
        if np.sum(sig_value[:ii]>tau) == np.sum(sig_value[:ii] < tau):
            return np.sqrt(tau)
        
def run(imgPath, patchSize, internalNumPatches, dirOut, saveResults=True):
    """
    Estimates the standard deviation of (additive white gaussian) noise of image patches.
    The noise is estimated patch by patch.
    Based on: "An Efficient Statistical Method for Image Noise Level Estimation" (2015)
    :param imgPath: Path to the input image.
    :param patchSize: Image patch size. 
    :param internalNumPatches: Internal number of sub-image-patches.
    :param dirOut: Directory where to save the noise estimation results.
    :param saveResults: Whether to save the estimation results or not.
    :return: None
    """
    # Load image
    img = np.array(cv2.imread(imgPath))
    try:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        img = img / 255.0
        
        h, w = img.shape
        psize = min(min(patchSize, h), w)
        psize -= psize % 2
        patch_step = psize
        shift_factor = 2
    
        # Result array
        estimatedNoiseMap = np.zeros([h, w], dtype=np.int8)
        rangex = range(0, w, patch_step)
        rangey = range(0, h, patch_step)
        for start_x in rangex:
            for start_y in rangey:
    
                end_x = start_x + psize
                end_y = start_y + psize
                if end_x > w:
                    end_x = w
                    end_x = shift_factor * ((end_x) // shift_factor)
                    start_x = end_x - psize
                if end_y > h:
                    end_y = h
                    end_y = shift_factor * ((end_y) // shift_factor)
                    start_y = end_y - psize
    
                tileM = img[start_y:end_y, start_x:end_x]
                h_, w_ = tileM.shape
                sigma = noise_estimate(tileM, internalNumPatches) * 255.0
                estimatedNoiseMap[start_y :start_y + h_, start_x : start_x + w_] = sigma   
                
        if saveResults:
            if dirOut is not None:
                imgName = imgPath.split(os.sep)[-1].split(".")[0]
                dirOut = os.path.join(dirOut)
                if not os.path.exists(dirOut):
                    os.makedirs(dirOut)
                    
            noiseMapPath = os.path.join(dirOut, imgName + ".npz")
            if not os.path.exists(noiseMapPath):
                np.savez_compressed(noiseMapPath, estimatedNoiseMap)
            
        return estimatedNoiseMap
    except:
        return None


# Example
# if __name__ == '__main__':
#     dirIn = r"../../../data/udacity/img/GT"
#     dirOut = r"../../../data/udacity/labels_noise_patchwise/PCA"
#     imgFileEnding = ".jpg"
#     for imgPath in glob.glob(os.path.join(dirIn, "*" + imgFileEnding)):
#         run(imgPath, 128, 8, dirOut)
[((34, 10, 34, 61), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((70, 19, 70, 42), 'numpy.linalg.eigh', 'np.linalg.eigh', ({(70, 34, 70, 41): 'sigma_X'}, {}), '(sigma_X)', True, 'import numpy as np\n'), ((59, 13, 59, 39), 'numpy.expand_dims', 'np.expand_dims', (), '', True, 'import numpy as np\n'), ((74, 14, 74, 37), 'numpy.mean', 'np.mean', ({(74, 22, 74, 36): 'sig_value[:ii]'}, {}), '(sig_value[:ii])', True, 'import numpy as np\n'), ((91, 19, 91, 38), 'cv2.imread', 'cv2.imread', ({(91, 30, 91, 37): 'imgPath'}, {}), '(imgPath)', False, 'import cv2\n'), ((93, 14, 93, 51), 'cv2.cvtColor', 'cv2.cvtColor', ({(93, 27, 93, 30): 'img', (93, 32, 93, 50): 'cv2.COLOR_RGB2GRAY'}, {}), '(img, cv2.COLOR_RGB2GRAY)', False, 'import cv2\n'), ((103, 28, 103, 59), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((20, 8, 20, 75), 'sys.exit', 'sys.exit', ({(20, 17, 20, 74): '"""The input of pch_size must be a integer or a int tuple!"""'}, {}), "('The input of pch_size must be a integer or a int tuple!')", False, 'import sys\n'), ((27, 8, 27, 73), 'sys.exit', 'sys.exit', ({(27, 17, 27, 72): '"""The input of stride must be a integer or a int tuple!"""'}, {}), "('The input of stride must be a integer or a int tuple!')", False, 'import sys\n'), ((75, 11, 75, 37), 'numpy.sum', 'np.sum', ({(75, 18, 75, 36): '(sig_value[:ii] > tau)'}, {}), '(sig_value[:ii] > tau)', True, 'import numpy as np\n'), ((75, 41, 75, 69), 'numpy.sum', 'np.sum', ({(75, 48, 75, 68): '(sig_value[:ii] < tau)'}, {}), '(sig_value[:ii] < tau)', True, 'import numpy as np\n'), ((76, 19, 76, 31), 'numpy.sqrt', 'np.sqrt', ({(76, 27, 76, 30): 'tau'}, {}), '(tau)', True, 'import numpy as np\n'), ((132, 27, 132, 65), 'os.path.join', 'os.path.join', ({(132, 40, 132, 46): 'dirOut', (132, 48, 132, 64): "imgName + '.npz'"}, {}), "(dirOut, imgName + '.npz')", False, 'import os\n'), ((128, 25, 128, 45), 'os.path.join', 'os.path.join', ({(128, 38, 128, 44): 'dirOut'}, {}), '(dirOut)', False, 'import os\n'), ((133, 19, 133, 47), 'os.path.exists', 'os.path.exists', ({(133, 34, 133, 46): 'noiseMapPath'}, {}), '(noiseMapPath)', False, 'import os\n'), ((134, 16, 134, 68), 'numpy.savez_compressed', 'np.savez_compressed', ({(134, 36, 134, 48): 'noiseMapPath', (134, 50, 134, 67): 'estimatedNoiseMap'}, {}), '(noiseMapPath, estimatedNoiseMap)', True, 'import numpy as np\n'), ((129, 23, 129, 45), 'os.path.exists', 'os.path.exists', ({(129, 38, 129, 44): 'dirOut'}, {}), '(dirOut)', False, 'import os\n'), ((130, 20, 130, 39), 'os.makedirs', 'os.makedirs', ({(130, 32, 130, 38): 'dirOut'}, {}), '(dirOut)', False, 'import os\n')]