###########################################################
# Functions to read/write/extend movies from jpg-files.
# References:
#   https://learnopencv.com/reading-and-writing-videos-using-opencv/#read-video-from-file
# Needed: (use pip3 if you don't use it under conda)
#   conda install -c fastai opencv-python-headless
#------------------------------------------------------------------------
# Hilversum   , 2022/11/07       License GPL3          (c) Rob Alblas
############################################################################################

import os
import cv2
import numpy as np

def codec2fourcc(codec):
  scodec=''
  for i in codec:
    scodec=scodec+i
  return scodec

# codec: 'mp4v', 'DIVX'; 'VP09' voor webm werkt niet???
def make_movie(jpg_files,ofile,fps=2,codec='DIVX'):
  scodec=codec2fourcc(codec)

  img_array=[]
  for fn in jpg_files:
    print('file='+fn)
    img=cv2.imread(fn)
    h,w,l=img.shape
    size=(w,h)
    img_array.append(img)

  fourcc=cv2.VideoWriter_fourcc(*scodec)
  out=cv2.VideoWriter(ofile,fourcc,fps,size)

  print('Write ' + str(ofile)+' in '+codec+', '+str(len(img_array))+' framse, fps='+str(fps))
  for i in range(len(img_array)):
    out.write(img_array[i])
  out.release()

def read_movie(ifile):
  oframe=[]
  vid_capture = cv2.VideoCapture(ifile)
  if (vid_capture.isOpened() == False):
    print('Error opening '+ifile)
    return oframe

  fps = int(vid_capture.get(5))
  while(vid_capture.isOpened()):
    ret, iframe = vid_capture.read()
    if ret==False:
      break
    h,w,l=iframe.shape
    size=(w,h)
    if ret == True:
      oframe.append(iframe)
    else:
      break
  vid_capture.release()

  return oframe,size,fps

def write_movie(ofile,oframe,size,fps=2,codec='DIVX'):
  scodec=codec2fourcc(codec)
  fourcc=cv2.VideoWriter_fourcc(*scodec)
  out=cv2.VideoWriter(ofile,fourcc,fps,size)

  print('Write ' + str(ofile)+' in '+codec+', '+str(len(oframe))+' framse, fps='+str(fps))

  for i in range(len(oframe)):
    out.write(oframe[i])
  out.release()

def extend_movie(ifile,efile,ofile,maxnrfrm=0,fps=0,codec='DIVX'):
  scodec=codec2fourcc(codec)
  if os.path.isfile(ifile):
    (oframe,size,ifps)=read_movie(ifile)
    if fps==0:
      fps=ifps
    nr_iframes=len(oframe)

    if ((maxnrfrm>0) and (nr_iframes>=maxnrfrm)):
      nr_skipframes=nr_iframes-maxnrfrm+1
      ooframe=[]
      for i in range(len(oframe)):
        if i >= nr_skipframes:
          ooframe.append(oframe[i])
      oframe=ooframe

    eframe=cv2.imread(efile)
    h,w,l=eframe.shape

    #size will be rounded to multiple of 4 by movie generator
    w=int((w/4))*4
    esize=(w,h)
    if esize != size:
      print('ERR: size of new file='+str(esize)+' != size of movie='+str(size))
    else:
      oframe.append(eframe)

    write_movie(ofile,oframe,size,fps)

  else:
    jpg_files=[]
    jpg_files.append(efile)
    if fps==0:
      fps=1
    make_movie(jpg_files,ofile,fps=fps,codec=codec)

  return True