from pathlib import Path
import json
import numpy as np
import imageio
from .utils import blend_rgba


def load_blender(split, scene="lego", half_res=False):
    assert split in ("train", "val", "test")

    env_fname = Path(__file__).resolve().parents[1] / "env.json"
    with env_fname.open("r") as f:
        root = json.load(f)['data_root']
    root = Path(root) / scene

    with open(root / f'transforms_{split}.json', "r") as f:
        meta = json.load(f)

    imgs, poses = [], []

    for frame in meta['frames']:
        file_name = root / f"{frame['file_path']}.png"
        im = imageio.imread(file_name)
        c2w = frame['transform_matrix']

        imgs.append(im)
        poses.append(c2w)

    imgs = (np.array(imgs) / 255.).astype(np.float32)  # (RGBA) imgs
    imgs = blend_rgba(imgs)
    poses = np.array(poses).astype(float)
    # print(imgs.shape)
    H, W = imgs[0].shape[:2]
    W = 64
    H = 64
    camera_angle_x = float(meta['camera_angle_x'])
    f = 1 / np.tan(camera_angle_x / 2) * (W / 2)

    if half_res:
        raise NotImplementedError()
    
    K = np.array([
        [f, 0, -(W/2 - 0.5)],
        [0, -f, -(H/2 - 0.5)],
        [0, 0, -1]
    ])  # note OpenGL -ve z convention;

    return imgs, K, poses