import  numpy as  np
from fractions import Fraction
import cv2
import numpy as np
import exifread
from exifread.utils import Ratio
import struct
import json
import torch
import time

Temp = np.ones([1536,2048]).astype(np.float32)
Timg = np.ones([768,1024,3]).astype(np.float32)

def apply_gamma(x):
    # return x ** (1.0 / 2.2)
    x = x.copy()
    idx = x <= 0.0031308
    x[idx] *= 12.92
    x[idx == False] = (x[idx == False] ** (1.0 / 2.4)) * 1.055 - 0.055
    return x

def binning(img,data):

    if data['cfa_pattern'] == [0,1,1,2]: 

      ch_R  = img[0::2, 0::2]
      ch_G  = (img[1::2, 0::2]+img[0::2,1::2])/2
      ch_B  = img[1::2, 1::2]
      out = np.dstack((ch_R, ch_G, ch_B))

    if data['cfa_pattern'] == [2,1,1,0]:      
    
      ch_R  = img[1::2, 1::2]
      ch_G  = (img[1::2, 0::2]+img[0::2,1::2])/2
      ch_B  = img[0::2, 0::2]
      out = np.dstack((ch_R, ch_G, ch_B))
    
    return out

def Four2One(img):
  Temp[0::2,0::2] = img[:,:,0]
  Temp[1::2,0::2] = img[:,:,1]
  Temp[0::2,1::2] = img[:,:,1]
  Temp[1::2,1::2] = img[:,:,2]
  
  return Temp
  
def One2Four(Temp):
  Timg[:,:,0] = Temp[0::2,0::2]
  Timg[:,:,1] = (Temp[1::2,0::2]+Temp[0::2,1::2])/2
  Timg[:,:,2] = Temp[1::2,1::2]
  
  return Timg
  
  
def white_balance(demosaic_img, as_shot_neutral):
    if type(as_shot_neutral[0]) is Ratio:
        as_shot_neutral = ratios2floats(as_shot_neutral)

    as_shot_neutral = np.asarray(as_shot_neutral)
    # transform vector into matrix
    if as_shot_neutral.shape == (3,):
        as_shot_neutral = np.diag(1. / as_shot_neutral)

    assert as_shot_neutral.shape == (3, 3)

    white_balanced_image = np.dot(demosaic_img, as_shot_neutral.T)
    white_balanced_image = np.clip(white_balanced_image, 0.0, 1.0)

    return white_balanced_image
    


    


def apply_color_space_transform(demosaiced_image, color_matrix):
    xyz2cam = np.reshape(np.asarray(color_matrix), (3, 3))
    # normalize rows (needed?)
    xyz2cam = xyz2cam / np.sum(xyz2cam, axis=1, keepdims=True)
    # inverse
    cam2xyz = np.linalg.inv(xyz2cam)
    # simplified matrix multiplication
    xyz_image = cam2xyz[np.newaxis, np.newaxis, :, :] * \
                demosaiced_image[:, :, np.newaxis, :]
    xyz_image = np.sum(xyz_image, axis=-1)
    xyz_image = np.clip(xyz_image, 0.0, 1.0)
    return xyz_image


    


def transform_xyz_to_srgb(xyz_image):
    xyz2srgb = np.array([[3.2404542, -1.5371385, -0.4985314],
                         [-0.9692660, 1.8760108, 0.0415560],
                         [0.0556434, -0.2040259, 1.0572252]])

    # normalize rows (needed?)
    xyz2srgb = xyz2srgb / np.sum(xyz2srgb, axis=-1, keepdims=True)

    srgb_image = xyz2srgb[np.newaxis, np.newaxis, :, :] * xyz_image[:, :, np.newaxis, :]
    srgb_image = np.sum(srgb_image, axis=-1)
    srgb_image = np.clip(srgb_image, 0.0, 1.0)
    return srgb_image




def fix_orientation(image, orientation):
    # 1 = Horizontal(normal)
    # 2 = Mirror horizontal
    # 3 = Rotate 180
    # 4 = Mirror vertical
    # 5 = Mirror horizontal and rotate 270 CW
    # 6 = Rotate 90 CW
    # 7 = Mirror horizontal and rotate 90 CW
    # 8 = Rotate 270 CW

    if type(orientation) is list:
        orientation = orientation[0]

    if orientation == "Horizontal(normal)":
        pass
    elif orientation == "Mirror horizonta":
        image = cv2.flip(image, 0)
    elif orientation == "Rotate 180":
        image = cv2.rotate(image, cv2.ROTATE_180)
    elif orientation == "Mirror vertical":
        image = cv2.flip(image, 1)
    elif orientation == "Mirror horizontal and rotate 270 CW":
        image = cv2.flip(image, 0)
        image = cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)
    elif orientation == "Rotate 90 CW":
        image = cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)
    elif orientation == "Mirror horizontal and rotate 90 CW":
        image = cv2.flip(image, 0)
        image = cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)
    elif orientation == "Rotate 270 CW":
        image = cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)

    return image