try:
    import cv2
    HAS_CV2 = True 
except:
    HAS_CV2 = False 

import numpy as np
import time
from typing import Optional, Tuple, Sequence
from .utils import find_files_open_binaries, init_ops

class VideoReader:
    """ Uses cv2 to read video files """
    def __init__(self, filenames: list):
        """ Uses cv2 to open video files and obtain their details for reading

        Parameters
        ------------
        filenames : int
            list of video files
        """
        cumframes = [0]
        containers = []
        Ly = []
        Lx = []
        for f in filenames:  # for each video in the list
            cap = cv2.VideoCapture(f)
            containers.append(cap)
            Lx.append(int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)))
            Ly.append(int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
            cumframes.append(cumframes[-1] + int(cap.get(cv2.CAP_PROP_FRAME_COUNT)))
        cumframes = np.array(cumframes).astype(int)
        Ly = np.array(Ly)
        Lx = np.array(Lx)
        if (Ly==Ly[0]).sum() < len(Ly) or (Lx==Lx[0]).sum() < len(Lx):
            raise ValueError("videos are not all the same size in y and x")
        else:
            Ly, Lx = Ly[0], Lx[0]

        self.filenames = filenames
        self.cumframes = cumframes 
        self.n_frames = cumframes[-1]
        self.Ly = Ly
        self.Lx = Lx
        self.containers = containers
        self.fs = containers[0].get(cv2.CAP_PROP_FPS)

    def close(self) -> None:
        """
        Closes the video files
        """
        for i in range(len(self.containers)):  # for each video in the list
            cap = self.containers[i]
            cap.release()

    def __enter__(self):
        return self

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

    @property
    def shape(self) -> Tuple[int, int, int]:
        """
        The dimensions of the data in the file

        Returns
        -------
        n_frames: int
            The number of frames
        Ly: int
            The height of each frame
        Lx: int
            The width of each frame
        """
        return self.n_frames, self.Ly, self.Lx

    def get_frames(self, cframes):
        """
        read frames "cframes" from videos

        Parameters
        ------------
        cframes : np.array
            start and stop of frames to read, or consecutive list of frames to read
        """
        cframes = np.maximum(0, np.minimum(self.n_frames - 1, cframes))
        cframes = np.arange(cframes[0], cframes[-1] + 1).astype(int)
        # find which video the frames exist in (ivids is length of cframes)
        ivids = (cframes[np.newaxis, :] >= self.cumframes[1:, np.newaxis]).sum(axis=0)
        nk = 0
        im = np.zeros((len(cframes), self.Ly, self.Lx), "uint8")
        for n in np.unique(ivids):  # for each video in cumframes
            cfr = cframes[ivids == n]
            start = cfr[0] - self.cumframes[n]
            end = cfr[-1] - self.cumframes[n] + 1
            nt0 = end - start
            capture = self.containers[n]
            if int(capture.get(cv2.CAP_PROP_POS_FRAMES)) != start:
                capture.set(cv2.CAP_PROP_POS_FRAMES, start)
            fc = 0
            ret = True
            while fc < nt0 and ret:
                ret, frame = capture.read()
                if ret:
                    im[nk + fc] = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                else:
                    print("img load failed, replacing with prev..")
                    im[nk + fc] = im[nk + fc - 1]
                fc += 1
            nk += nt0
        return im

def movie_to_binary(ops):
    """  finds movie files and writes them to binaries

    Parameters
    ----------
    ops : dictionary
        "nplanes", "data_path", "save_path", "save_folder", "fast_disk",
        "nchannels", "keep_movie_raw", "look_one_level_down" (optional: "subfolders")

    Returns
    -------
        ops : dictionary of first plane
            "Ly", "Lx", ops["reg_file"] or ops["raw_file"] is created binary

    """
    if not HAS_CV2:
        raise ImportError("cv2 is required for this file type, please 'pip install opencv-python-headless'")

    ops1 = init_ops(ops)

    nplanes = ops1[0]["nplanes"]
    nchannels = ops1[0]["nchannels"]
    
    # open all binary files for writing
    ops1, filenames, reg_file, reg_file_chan2 = find_files_open_binaries(ops1)
    
    ik = 0
    for j in range(ops["nplanes"]):
        ops1[j]["nframes_per_folder"] = np.zeros(len(filenames), np.int32)


    ncp = nplanes * nchannels
    nbatch = ncp * int(np.ceil(ops1[0]["batch_size"] / ncp))
    print(filenames)
    t0 = time.time()
    with VideoReader(filenames=filenames) as vr:
        if ops1[0]["fs"]<=0:
            for ops in ops1:
                ops["fs"] = vr.fs

        nframes_all = vr.cumframes[-1]
        nbatch = min(nbatch, nframes_all)
        nfunc = ops["functional_chan"] - 1 if nchannels > 1 else 0
        # loop over all video frames
        ik = 0
        while 1:
            irange = np.arange(ik, min(ik + nbatch, nframes_all), 1)
            if irange.size == 0:
                break
            im = vr.get_frames(irange).astype("int16")
            nframes = im.shape[0]
            for j in range(0, nplanes):
                if ik == 0:
                    ops1[j]["meanImg"] = np.zeros((im.shape[1], im.shape[2]),
                                                    np.float32)
                    if nchannels > 1:
                        ops1[j]["meanImg_chan2"] = np.zeros(
                            (im.shape[1], im.shape[2]), np.float32)
                    ops1[j]["nframes"] = 0
                i0 = nchannels * ((j) % nplanes)
                im2write = im[np.arange(int(i0) +
                                        nfunc, nframes, ncp), :, :].astype(
                                            np.int16)
                reg_file[j].write(bytearray(im2write))
                ops1[j]["meanImg"] += im2write.astype(np.float32).sum(axis=0)
                if nchannels > 1:
                    im2write = im[np.arange(int(i0) + 1 -
                                            nfunc, nframes, ncp), :, :].astype(
                                                np.int16)
                    reg_file_chan2[j].write(bytearray(im2write))
                    ops1[j]["meanImg_chan2"] += im2write.astype(
                        np.float32).sum(axis=0)
                ops1[j]["nframes"] += im2write.shape[0]
                #ops1[j]["nframes_per_folder"][ih5] += im2write.shape[0]
            ik += nframes
            if ik % (nbatch * 4) == 0:
                print("%d frames of binary, time %0.2f sec." %
                      (ik, time.time() - t0))

    # write ops files
    do_registration = ops1[0]["do_registration"]
    for ops in ops1:
        ops["Ly"] = im2write.shape[1]
        ops["Lx"] = im2write.shape[2]
        if not do_registration:
            ops["yrange"] = np.array([0, ops["Ly"]])
            ops["xrange"] = np.array([0, ops["Lx"]])
        ops["meanImg"] /= ops["nframes"]
        if nchannels > 1:
            ops["meanImg_chan2"] /= ops["nframes"]
        np.save(ops["ops_path"], ops)
    # close all binary files and write ops files
    for j in range(nplanes):
        reg_file[j].close()
        if nchannels > 1:
            reg_file_chan2[j].close()
    return ops1[0]
