#! /usr/bin/env python
# -*- coding: utf-8 -*-
'''
Created on 2022/4/24 21:29:16

@author: LiuKuan
@copyright: Apache License, Version 2.0
'''
import os
from laok.third_lib import cv2, np
from ..cvt.cv2_ import keep_bgr_cv, keep_gray_cv, image_size_cv, is_color_cv, swap_rb_cv
#===============================================================================
'''     
'''
#===============================================================================
__all__ = ['fix_cv_read', 'read_img_cv', 'read_color_cv', 'read_gray_cv',
           'write_img_cv',
           'read_video_cv', 'read_camera_cv', 'VideoWriterCv',
           ]

####################    读取图像

def fix_cv_read():
    _org_read = cv2.imread
    cv2.imread = read_img_cv


def read_img_cv(filename, flags = 1):
    '''
    :param filename:
    :param flags: cv2.IMREAD_COLOR
    :return:
    '''
    return cv2.imdecode(np.fromfile(filename, dtype=np.uint8), flags)

def read_color_cv(filename, bgr=True):
    img = read_img_cv(filename, cv2.IMREAD_COLOR)
    if not bgr:
        return swap_rb_cv(img)
    return img

def read_gray_cv(filename):
    return read_img_cv(filename, cv2.IMREAD_GRAYSCALE)

def write_img_cv(data, filename):
    name, ext = os.path.splitext(filename)
    retval, buf = cv2.imencode(ext, data)
    return buf.tofile(filename)

####################    读取视频和相机
def _read_capture_cv(fileOrId):
    try:
        cap = cv2.VideoCapture(fileOrId)
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            yield frame
    finally:
        cap.release()

def read_video_cv(video_file):
    yield from _read_capture_cv(video_file)

def read_camera_cv(camId = 0):
    yield from _read_capture_cv(camId)


class VideoWriterCv:
    def __init__(self, filename, fourcc="mp4v", fps=24, frameSize=None, isColor=None):
        self.filename_ = filename
        self.frameSize_ = frameSize
        self.isColor_ = isColor
        self.fps_ = fps
        self.fourcc_ =  cv2.VideoWriter_fourcc(*fourcc)
        self.writer_ = None

    def _make_video(self, img):
        if self.frameSize_ is None:
            self.frameSize_ = image_size_cv(img)

        if self.isColor_ is None:
            self.isColor_ = is_color_cv(img)
        self.writer_ = cv2.VideoWriter(filename=self.filename_, fourcc=self.fourcc_,
                                      fps=self.fps_, frameSize=self.frameSize_, isColor=self.isColor_)

    def write(self, img):
        if self.writer_ is None:
            self._make_video(img)

        if self.isColor_ :
            img = keep_bgr_cv(img)
        else:
            img = keep_gray_cv(img)

        if image_size_cv(img) != self.frameSize_:
            img = cv2.resize(img, dsize=(self.frameSize_[0], self.frameSize_[1]),
                             interpolation=cv2.INTER_LINEAR)

        self.writer_.write(img)

    def close(self):
        if self.writer_ is not None:
            self.writer_.release()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
