# Copyright 2026 Toyota Research Institute.  All rights reserved.

import glob
import io
import json
import os
import pickle as pkl

import cv2
import numpy as np
import yaml
import zarr
from PIL import Image, ImageOps

from zarr.storage import ZipStore

from anydata.utils.data import to_namespace
from anydata.geometry.depth import decode_dense
from anydata.utils.decorators import iterate1
from anydata.utils.misc import frame_name, stack_sample, unjoin_dict
from anydata.utils.write import frame_name

try:
    import Imath  # @IgnoreException
    import OpenEXR as exr  # @IgnoreException
except:
    pass

try:
    from torchcodec.decoders import VideoDecoder
except:
    pass

try:
    from decord import VideoReader
    from decord.ndarray import NDArray
except:
    pass

# NOTE(bvh): decord's default num_threads=0 lets ffmpeg auto-size its decode pool per
# VideoReader; across many dataloader workers this explodes into thousands of threads
# on large-core nodes. Keep in sync with webdataset/autodecode.py mp4_loads.
DECORD_NUM_THREADS = int(os.environ.get('ANYDATA_DECORD_THREADS', 4))


def read_pickle(filename):
    """Read pickle file"""
    if not filename.endswith('.pkl'):
        filename += '.pkl'
    return pkl.load(open(filename, 'rb'))


def read_yaml(filename):
    """Opens and returns a yaml file"""
    return yaml.safe_load(open(filename))


def read_json(filename):
    """Opens and returns a json file"""
    with open(filename) as file:
        return json.load(file)


def read_txt(filename):
    """Opens and returns a json file"""
    with open(filename, 'r') as file:
        return [f.rstrip() for f in file]


@iterate1
def read_image(filename, mode='RGB', invert=False):
    """Read an image using PIL"""
    image = Image.open(filename)
    if mode != '':
        image = image.convert(mode)
    if invert:
        image = ImageOps.invert(image)
    return image


@iterate1
def read_numpy(filename, key=None):
    depth = np.load(filename, allow_pickle=True)
    if key is not None:
        depth = depth[key]
    return depth


def read_npz(filename):
    """Reads a npz file from disk and returns a dict or array"""
    loaded: np.lib.npyio.NpzFile = np.load(filename, allow_pickle=True)
    return {k: loaded[k][()] for k in loaded}


@iterate1
def read_depth(filename: str, div=256., key='depth'):
    """Load a depth map from file"""
    # If loading a .npz array
    if filename.endswith('npz'):
        data = read_npz(filename)
        encode_info = data.get('encode_info', None)
        return decode_dense(data[key], encode_info=encode_info)
    elif filename.endswith('npy'):
        return np.load(filename)
    # If loading a .png image
    elif filename.endswith('png'):
        depth_png = np.array(read_image(filename, mode=''), dtype=int)
        # assert (np.max(depth_png) > 255), 'Wrong .png depth file'
        return depth_png.astype(float) / div
    elif filename.endswith('exr'):
        exrfile = exr.InputFile(filename) # .as_posix())
        raw_bytes = exrfile.channel('B', Imath.PixelType(Imath.PixelType.FLOAT))
        depth = np.frombuffer(raw_bytes, dtype=np.float32)
        height = exrfile.header()['displayWindow'].max.y + 1 - exrfile.header()['displayWindow'].min.y
        width = exrfile.header()['displayWindow'].max.x + 1 - exrfile.header()['displayWindow'].min.x
        depth = np.reshape(depth, (height, width))
        return depth


def read_optical_flow(filename: str):
    """Read optical flow from .flo extension"""
    if filename.endswith('flo'):
        with open(filename, 'rb') as f:
            if np.fromfile(f, np.float32, count=1) == 202021.25:
                w = np.fromfile(f, np.int32, count=1)
                h = np.fromfile(f, np.int32, count=1)
                data = np.fromfile(f, np.float32, count=2 * int(w) * int(h))
                return np.resize(data, (int(h), int(w), 2))


@iterate1
def read_float3(filename):
    f = open(filename, 'rb')
    if (f.readline().decode('utf-8')) != 'float\n':
        print('asdfasdfasdf')
    dim = int(f.readline())
    dims = []
    count = 1
    for i in range(0, dim):
        d = int(f.readline())
        dims.append(d)
        count *= d
    dims = list(reversed(dims))
    return np.fromfile(f, np.float32, count).reshape(dims)


def read_png32(filename):
    """Read 4 uint8 channel image as float32"""
    return cv2.imread(filename, cv2.IMREAD_UNCHANGED).view("<f4").squeeze()


def read_png8(filename):
    """Read 1 uint16 image as float32 [0,65535]"""
    return cv2.imread(filename, cv2.IMREAD_UNCHANGED).astype(np.float32) / 1000


def read_config(path):
    """Read configuration from file"""
    with open(path) as cfg:
        config = yaml.load(cfg, Loader=yaml.FullLoader)
    cfg = to_namespace(config)
    return cfg


def read_image_seq(folder: str, index: int = None, ext: str = "jpg", pil: bool = False, frame_name_fn=frame_name):
    """Reads a sequence of images from disk, one per timestep, and stacks into an array"""
    if index is not None:
        rgb = read_image(f"{folder}/{frame_name_fn(index)}.{ext}")
        return rgb if pil else np.asarray(rgb)
    else:
        seq_len = len(glob.glob(f"{folder}/*.{ext}"))
        indices = np.arange(seq_len)[index] if index is not None else np.arange(seq_len)
        rgbs = [read_image(f"{folder}/{frame_name_fn(i)}.{ext}") for i in indices]
        return rgbs if pil else np.stack([np.asarray(img) for img in rgbs])


def _read_video_seq_torchcodec(source, index):
    decoder = VideoDecoder(source)
    torch_frames = decoder.get_frames_at(index).data if isinstance(index, list) else decoder[index]
    rgb_array = torch_frames.movedim(-3, -1).numpy().astype('uint8')
    return rgb_array


def _read_video_seq_decord(source, index):
    decoder = VideoReader(source, num_threads=DECORD_NUM_THREADS)
    decord_frames: NDArray = decoder.get_batch(index) if isinstance(index, list) else decoder[index]
    rgb_array = decord_frames.asnumpy().astype('uint8')
    return rgb_array


def read_video_seq(src: str | bytes, index=slice(None), ext: str = "mp4", pil: bool = False, decoder: str = 'torchcodec'):
    """Reads a video file from disk and returns as a sequence of images"""
    source = f"{src.removesuffix(f'.{ext}')}.{ext}" if isinstance(src, str) else io.BytesIO(src)
    if decoder == 'torchcodec':
        try:
            rgb_array = _read_video_seq_torchcodec(source, index)
        except Exception:
            rgb_array = _read_video_seq_decord(source, index)
    else:
        rgb_array = _read_video_seq_decord(source, index)
    ndims = rgb_array.ndim
    if pil:
        return Image.fromarray(rgb_array) if ndims == 3 else [Image.fromarray(rgb) for rgb in rgb_array]
    return rgb_array


def read_npzs_seq(folder: str, key: str = None, index=slice(None), frame_name_fn=frame_name):
    """Reads a sequence of npz files from disk, one per timestep, and stacks into an array"""
    if isinstance(index, int):
        # NOTE(bvh): use raw np.load (not read_npz + unjoin_dict) to preserve 0-d numpy arrays.
        # Callers like load_lowdim depend on .item() working on object arrays.
        loaded = np.load(f"{folder}/{frame_name_fn(index)}.npz", allow_pickle=True)
        frame_dict = {k: loaded[k] for k in loaded}
        return frame_dict[key] if key is not None else frame_dict
    
    else:
        files = sorted(glob.glob(f"{folder}/*.npz"))
        indices = np.arange(len(files))[index]
        dict_list = []
        for i in indices:
            frame_dict = unjoin_dict(read_npz(f"{folder}/{frame_name_fn(i)}.npz"))
            if key is not None:
                frame_dict = read_npz_node(frame_dict, key=key.split('/'))
            dict_list.append(frame_dict)
        combiner = stack_sample if dict_list and isinstance(dict_list[0], dict) else np.stack
        return combiner(dict_list) if dict_list else np.array([])


def read_npz_seq(path: str, key: str = None, index=slice(None)):
    """Reads a full npz file from disk and returns a nested indexed dict"""
    loaded = read_npz(path + ".npz")
    return read_npz_node(loaded, key=key.split('/') if key is not None else None, index=index)


def read_npz_node(node: dict | np.ndarray, key: list | str = None, index=slice(None)):
    """Recursively read a npz file or group."""
    if isinstance(node, np.ndarray):
        return node[index]
    if isinstance(key, str):
        key = key.split('/')
    if key is not None and len(key):
        return read_npz_node(node[key[0]], key=key[1:], index=index)
    if isinstance(node, list):
        if len(node) == 0:  # Empty list
            return node
        elif isinstance(index, int):  # Single index
            return node[index]
        else:  # Slice index
            return [node[i] for i in range(index.start, index.stop, index.step)]
    return node if index == slice(None) else {k: read_npz_node(node[k], index=index) for k in node}


def read_zarr_seq(filename: str, key: str = None, index=slice(None)):
    """
    Read zarr data from disk. Note that the key can be None to read the full hierarchy recursively, or a 
    string path (eg "foo/bar/baz") to read a specific array/group. The index is applied to array leaves.
    """
    store = ZipStore(f"{filename.removesuffix('.zarr')}.zarr", mode='r')
    try:
        root = zarr.open(store, mode='r')
        node = root if key is None else root[key]
        val = _read_zarr_node(node, index=index)
        if 'encode_info' in root:  # decode zarr data
            encode_info = {key: np.array(root['encode_info'][key]) for key in root['encode_info'].keys()}
            val = decode_dense(val, encode_info=encode_info)
            if len(val.shape) == 3:  # expand if there is no channel dimension
                val = np.expand_dims(val, 1)
    finally:
        store.close()
    return val


def _read_zarr_node(node: zarr.Group | zarr.Array, index=slice(None)):
    """Recursively read a zarr array or group."""
    if isinstance(node, zarr.Array):
        return node[index]
    return {k: _read_zarr_node(node[k], index=index) for k in node.keys()}


def read_seq(path: str, ext: str = None, *, key: str = None, index=slice(None), sto=None, mp4_decoder='torchcodec', **kwargs):
    """General function to read a sequence of files from disk, one per timestep, and stacks into an array"""
    if ext is None and "." in os.path.basename(path):
        ext = os.path.splitext(path)[-1][1:]
    if ext in ['jpg', 'png']:
        return [read_image_seq(path, index=ind, ext=ext, **kwargs) for ind in index]
    elif ext == 'mp4':
        return read_video_seq(path, index=index, ext=ext, decoder=mp4_decoder, **kwargs)
    elif ext == 'zarr':
        return read_zarr_seq(path, key=key, index=index, **kwargs)
    elif ext == 'npzs' or os.path.isdir(path):
        # NOTE(bvh): npz with per-frame directory layout is equivalent to npzs.
        # npzs may be deprecated in favor of npz handling both cases.
        return read_npzs_seq(path, key=key, index=index, **kwargs)
    elif ext == 'npz':
        read_npz_seq_fn = read_npz_seq if sto in ['sequence','videos'] else read_npzs_seq
        return read_npz_seq_fn(path, key=key, index=index, **kwargs)
    else:
        print(f"Unsupported file extension {ext} for read_seq")
        return None