# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_io.ipynb.

# %% auto 0
__all__ = ['load_avi', 'moviefps', 'save_avi']

# %% ../nbs/04_io.ipynb 6
import warnings
from string import Template
from enum import StrEnum, auto
from inspect import signature as getsig, getfullargspec as getspec

from dataclasses import dataclass, field

# %% ../nbs/04_io.ipynb 8
from types import TracebackType
from typing import Any, Dict, Tuple, Optional, Mapping, Literal, TypeAlias

# %% ../nbs/04_io.ipynb 11
#| export

# %% ../nbs/04_io.ipynb 13
#| export


# %% ../nbs/04_io.ipynb 15
try: import torchvision.io as tvio, torchvision.transforms.functional as F
except ImportError: ...

# %% ../nbs/04_io.ipynb 17
#| export


# %% ../nbs/04_io.ipynb 19
from atyp import StrQ, DictQ, FloatQ, PathLike, ImgSizeQ, Tensor, TensorQ
from chck import isnone, notnone
from nchr import NIL, DOT
from nlit import EXT_AVI
from psrc import insert_exts
from fpos import val1st
from putl import takekeys

# %% ../nbs/04_io.ipynb 21
from .enum import Channels

# %% ../nbs/04_io.ipynb 24
def load_avi(
    path: PathLike, resize: ImgSizeQ = None, channel_format: Channels = Channels.THWC, 
    ignore_exported: bool = False, **kwargs
):
    warnings.filterwarnings('ignore', category=UserWarning, module='torchvision')
    if notnone(resize): channel_format = Channels.TCHW
    
    kwargs.setdefault('pts_unit', 'sec')
    readkeys = ('start_pts', 'end_pts', 'pts_unit')
    
    infile = insert_exts(path, 'torch', idx=-2)
    if infile.exists() and not ignore_exported: path = infile

    video, audio, mdata = tvio.read_video(str(path), output_format=channel_format, **takekeys(kwargs, readkeys))

    vsize = video.count_nonzero().item()
    asize = audio.count_nonzero().item()
    
    if notnone(resize) and vsize > 0: video = F.resize(video, resize)
    if notnone(resize) and asize > 0: audio = F.resize(audio, resize)
    
    return video, audio, mdata

# %% ../nbs/04_io.ipynb 26
def moviefps(__default: float = 30.0, **kwargs) -> FloatQ:
    # check if metadata is provided
    mkeys, fkeys = ('meta', 'video_meta', 'metadata',), ('fps', 'video_fps', )
    metadata = {**val1st(kwargs, mkeys , default={}), **kwargs}    
    videofps = int(val1st(metadata, fkeys, default=__default))
    return videofps

# %% ../nbs/04_io.ipynb 27
def save_avi(
    path: PathLike, video: Tensor, exists_ok: bool = True, **kwargs
):
    outfile = insert_exts(path, 'torch', idx=-2)    
    fps = moviefps(**kwargs)
    
    if outfile.is_file() and not exists_ok: 
        raise FileExistsError(outfile)
    
    writekeys = ('video_codec', 'options', 'audio_array', 'audio_fps', 'audio_codec', 'audio_options', )
    tvio.write_video(str(outfile), Channels.to(video, Channels.THWC), fps = fps, **takekeys(kwargs, writekeys))
