import os
import cv2
from src.ss.ss import handle_ss
from src.sts.demo.sts import handle_sts
from src.ir.ir import handle_ir
from utils.utils import correct_rotation, create_folder


def gen_frames(video, dir_path, rotateCode):
    subpath = create_folder(dir_path)
    yield b'--frame\r\n'
    count = 0
    while True:
        ret, frame = video.read()
        # if frame is read correctly ret is True
        if rotateCode is not None:
            frame = correct_rotation(frame, rotateCode)
        if ret:
            print('Read a new frame: ', ret)
            img_path = os.path.join(subpath[0], "frames_%d.jpg" % count)
            cv2.imwrite(img_path, frame)
            segment_path, segment_array = handle_ss(img_path, subpath[1])
            output_path_box, output_path_text, output_path_visual, dict_box_sign_out, dict_rec_sign_out = handle_sts(
                img_path, segment_path, subpath[2])
            predicted = handle_ir(img_path, dict_rec_sign_out, subpath[3])
            print(predicted)

        else:
            print("Can't receive frame (stream end?). Exiting ...")
            break
        count = count + 1
        ret, jpeg = cv2.imencode('.jpg', frame)
        frame = jpeg.tobytes()
        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n\r\n')


def gen_videos(file_in, file_out):
    cap = cv2.VideoCapture(file_in)

    width = cap.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH)   # float `width`
    height = cap.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT)
    fps = cap.get(cv2.cv.CV_CAP_PROP_FPS)

    print(width, height, fps)

    # Define the codec and create VideoWriter object
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    out = cv2.VideoWriter(file_out, fourcc, 20.0, (640,  480))
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            print("Can't receive frame (stream end?). Exiting ...")
            break
        frame = cv2.flip(frame, 0)
        # write the flipped frame
        out.write(frame)
        cv2.imshow('frame', frame)
        if cv2.waitKey(1) == ord('q'):
            break
    # Release everything if job is finished
    cap.release()
    out.release()
    cv2.destroyAllWindows()