import os
import re
import cv2
import numpy as np

from ..utils.common import Notify

def read_list(list_path):
    """Read list."""
    if list_path is None or not os.path.exists(list_path):
        print(Notify.FAIL, 'Not exist', list_path, Notify.ENDC)
        exit(-1)
    content = open(list_path).read().splitlines()
    return content


def load_pfm(pfm_path):
    with open(pfm_path, 'rb') as fin:
        color = None
        width = None
        height = None
        scale = None
        data_type = None
        header = str(fin.readline().decode('UTF-8')).rstrip()

        if header == 'PF':
            color = True
        elif header == 'Pf':
            color = False
        else:
            raise Exception('Not a PFM file.')

        dim_match = re.match(r'^(\d+)\s(\d+)\s$',
                             fin.readline().decode('UTF-8'))
        if dim_match:
            width, height = map(int, dim_match.groups())
        else:
            raise Exception('Malformed PFM header.')
        scale = float((fin.readline().decode('UTF-8')).rstrip())
        if scale < 0:  # little-endian
            data_type = '<f'
        else:
            data_type = '>f'  # big-endian
        data_string = fin.read()
        data = np.fromstring(data_string, data_type)
        shape = (height, width, 3) if color else (height, width)
        data = np.reshape(data, shape)
        data = np.flip(data, 0)
    return data


def _parse_img(img_paths, idx, config):
    img_path = img_paths[idx]
    img = cv2.imread(img_path)[:, :, ::-1]
    if config['resize'] > 0:
        img = cv2.resize(
            img, (config['resize'], config['resize']))
    return img


def _parse_depth(depth_paths, idx, config):
    depth = load_pfm(depth_paths[idx])

    if config['resize'] > 0:
        target_size = config['resize']
    if config['input_type'] == 'raw':
        depth = cv2.resize(depth, (int(target_size/2), int(target_size/2)))
    else:
        depth = cv2.resize(depth, (target_size, target_size))
    return depth


def _parse_kpts(kpts_paths, idx, config):
    kpts = np.load(kpts_paths[idx])['pts']
    # output: [N, 2] (W first H last)
    return kpts