import cv2
import numpy as np

def gen_2to1(batch_size=100+1,video_path):
    "split video into 2 lowQ frams with timestamps followed by an HD frame"
    vidcap = cv2.VideoCapture(video_path)
    success ,image = vidcap.read()
    count_x=0
    count_y=0
    xtempC=0
    xtemp=np.empty((2,144,176,3),dtype=np.uint8)
    xB=np.empty((batch_size,2,144,176,3),dtype=np.uint8)
    yB=np.empty((batch_size,1080,1920,3),dtype=np.uint8)
    while success:
      success,image = vidcap.read()
      xtempC+=1
      if xtempC==1:
          image_lr=cv2.resize(image,(176,144))
          image_lr=np.expand_dims(image_lr,axis=0)
          xtemp[0]=image_lr
      if xtempC==2:
          image_lr=cv2.resize(image,(176,144))
          image_lr=np.expand_dims(image_lr,axis=0)
          xtemp[1]=image_lr
          xB[count_x]=xtemp
          count_x+=1
      if xtempC==3:
          yB[count_y]=image
          count_y+=1
          xtempC=0
      if count_x==batch_size:
          count_x=0
          count_y=0
          yield xB,yB

def gen_1to1(batch_size=100+1,video_path):
    "split video into lowQ followed by HD Frame "
    vidcap = cv2.VideoCapture(video_path)
    success ,image = vidcap.read()
    count_x=0
    count_y=0
    xtempC=0
    xB=np.empty((batch_size,144,176,3),dtype=np.uint8)
    yB=np.empty((batch_size,1080,1920,3),dtype=np.uint8)
    while success:
      success,image = vidcap.read()
      image_lr=cv2.resize(image,(176,144))
      xB[count_x]=image_lr
      yB[count_y]=image 
      count_x+=1
      count_y+=1
      if batch_size==count_x:
        yield xB,yB


def gen_2to2(batch_size=100+1,video_path):
  "split the video into 2 HD frames followed by 2 lowQ frames with timestamps"
    vidcap = cv2.VideoCapture(video_path)
    s ,image = vidcap.read()
    count_x=0
    count_y=0
    xtempC=0
    xtemp=np.empty((2,144,176,3),dtype=np.uint8)
    ytemp=np.empty((2,1080,1920,3),dtype=np.uint8)
    xB=np.empty((batch_size,2,144,176,3),dtype=np.uint8)
    yB=np.empty((batch_size,2,1080,1920,3),dtype=np.uint8)
    while s:
      s,image = vidcap.read()
      xtempC+=1
      if xtempC==1:
          image_lr=cv2.resize(image,(176,144))
          image_lr=np.expand_dims(image_lr,axis=0)
          xtemp[0]=image_lr
      if xtempC==2:
          image_lr=cv2.resize(image,(176,144))
          image_lr=np.expand_dims(image_lr,axis=0)
          xtemp[1]=image_lr
          xB[count_x]=xtemp
          count_x+=1
      if xtempC==3:
          image_hr=np.expand_dims(image,axis=0)
          ytemp[0]=image_hr
      if xtempC==4:
          xtempC=0
          image_hr=np.expand_dims(image,axis=0)
          ytemp[1]=image_hr
          yB[count_y]=ytemp
          count_y+=1
      if count_x==batch_size:
          count_x=0
          count_y=0
          yield xB,yB