import numpy as np
import cv2

def bit_check(bit_num):
    return 2**bit_num-1, 2**(bit_num-4)

def rawread(name, pt = -1):
    bl = 64
    wh = 1023
    if name.endswith('dng'):
        import rawpy
        a = rawpy.imread(name).raw_image_visible.copy()
        h,w = a.shape[:2]
    elif name.endswith('npy'):
        a = np.load(name)
        if len(a.shape) > 2:
            a = rggb2bayer(a)
        h,w = a.shape[:2]
    else:
        a = np.fromfile(name, 'uint16')
    #depth
    wh, bl = bit_check(10)
    if a.max()>2*(2**14):
        wh, bl = bit_check(16)
    elif a.max()>2*(2**12):
        wh, bl = bit_check(14)
    elif a.max()>2*(2**10):
        wh, bl = bit_check(12)
    print('bl:', bl, 'wh:',wh, 'max=', a.max())
    #size
    if len(a) == 4512*6016:
        h = 4512
        w = 6016
    elif len(a) == 3000*4000:
        h = 3000
        w = 4000
    elif len(a) == 1824*2432:
        h = 1824
        w = 2432
    elif np.abs(len(a) - 2448*3264)<=1024:
        h = 2448
        w = 3264
    elif np.abs(len(a) - 2400*3200)<=1024:
        h = 2400
        w = 3200
    elif np.abs(len(a) - 1824*2432)<=1024:
        h = 1824
        w = 2432
    elif len(a) == 4640*3472:
        h = 3472
        w = 4640
    elif len(a) == 2632*3504:
        h = 2632
        w = 3504
    elif len(a) == 1940*2592:
        h = 1940
        w = 2592
    elif len(a) == 3472*4624:
        h = 3472
        w = 4624
    elif len(a) == 3072*4096:
        h = 3072
        w = 4096
    elif len(a) == 2720*3648:
        h = 2720
        w = 3648
    elif len(a) == 3072*4080:
        h = 3072
        w = 4080
    elif len(a) == 2304*4096:
        h = 2304
        w = 4096
    elif len(a) == 2304*1728:
        h = 1728
        w = 2304
    elif len(a) == 4160*3120:
        h = 3120
        w = 4160
    elif np.abs(len(a) - 3648*2736)<=1024:
        h = 2736
        w = 3648
    elif np.abs(len(a) - 3648*2736*4)<=1024:
        h = 2736*2
        w = 3648*2
    elif np.abs(len(a) - 4096*3072)<=1024:
        h = 3072
        w = 4096
    elif np.abs(len(a) - 4160*3120)<=1024:
        h = 3120
        w = 4160
    elif np.abs(len(a) - 3264*2432)<=1024:
        h = 2432
        w = 3264
    elif len(a)==4032*3024:
        h = 3024
        w = 4032
    elif len(a)==4208*3120:
        h = 3120
        w = 4208
    elif len(a)==2944*2208:
        h = 2208
        w = 2944
    elif len(a)==3840*2160:
        h = 2160
        w = 3840
    elif len(a)==2880*1616:
        h = 1616
        w = 2880
    elif len(a)==2880*1624:
        h = 1624
        w = 2880
    elif len(a)==2880*1620:
        h = 1620
        w = 2880
    elif len(a)==2688*1520:
        h = 1520
        w = 2688
    elif len(a)==1920*1080:
        h = 1080
        w = 1920
    print('h:',h,'w:',w)
    a = a[:h*w].reshape([h, w])
    m0 = a[::2, ::2].mean()
    m1 = a[::2, 1::2].mean()
    m2 = a[1::2, ::2].mean()
    m3 = a[1::2, 1::2].mean()
    m12 = max(m1, m2)/min(m1, m2)
    m03 = max(m0, m3)/min(m0, m3)
    if m12<m03:
        #XG
        #GX
        if m0>m3:
            #RGGB
            mode = 0
        else:#BGGR
            mode = 3
    else:
        #GX
        #XG
        if m1>m2:
            #GRBG
            mode = 1
        else:
            #GBRG
            mode = 2
    if pt>=0:
        mode = pt
    print('mode:', mode)
    #any to rggb
    if mode == 0:
        pass
    elif mode==1:
        a = a[:, ::-1]
    elif mode==2:
        a = a[::-1, :]
    elif mode==3:
        a = a[::-1, ::-1]
    return {'raw':a, 
            'h':h, 
            'w':w, 
            'bl':bl, 
            'wp':wh, 
            'mode':mode} 

def rgb2rggb(rgb):
    R = rgb[::2, ::2, 0:1]
    Gr = rgb[::2, 1::2, 1:2]
    Gb = rgb[1::2, ::2, 1:2]
    B = rgb[1::2, 1::2, 2:]
    rggb = np.concatenate((R, Gr, Gb, B), axis=2)
    return rggb

def rggb2bayer(rggb):
    h,w = rggb.shape[:2]
    bayer = rggb.reshape([h, w, 2, 2]).transpose([0, 2, 1, 3]).reshape([h*2, w*2])
    return bayer

def bayer2rggb(bayer):
    h,w = bayer.shape[:2]
    rggb = bayer.reshape([h//2, 2, w//2, 2]).transpose([0, 2, 1, 3]).reshape([h//2, w//2, 4])
    return rggb

def easydemoisac(bayer):
    if len(bayer.shape)==3:
        bayer = rggb2bayer(bayer)
    h,w = bayer.shape
    rgb = np.zeros([h//2, w//2, 3], bayer.dtype)
    rgb[:, :, 0] = bayer[::2, ::2]
    rgb[:, :, 1] = (bayer[::2, 1::2]+bayer[1::2, ::2])*0.5
    rgb[:, :, 2] = bayer[1::2, 1::2]
    rgb = cv2.GaussianBlur(rgb, (3, 3), 0)
    return rgb

def easyremosaic(rgb):
    bayer = np.zeros((rgb.shape[0], rgb.shape[1]))
    bayer[0::2, 0::2] = rgb[0::2, 0::2, 0]
    bayer[0::2, 1::2] = rgb[0::2, 1::2, 1]
    bayer[1::2, 0::2] = rgb[1::2, 0::2, 1]
    bayer[1::2, 1::2] = rgb[1::2, 1::2, 2]
    return bayer