import numpy as np
import PIL.Image
from PIL.PngImagePlugin import PngInfo


def read_depth_from_png(file_path):
    depth = PIL.Image.open(file_path) # Scale the saved image using the metadata
    max_depth = float(depth.info["max_value"])
    min_depth = float(depth.info["min_value"])
    depth = np.array(depth).astype(np.float32)

    # Scale from uint16 range
    depth = (depth / (2 ** 16 - 1)) * (max_depth - min_depth) + min_depth

    # replace the magic constant with positive infinity
    depth[depth == -1.0] = np.inf
    return depth


def write_depth_to_png(outpath_file, depth_pred):
    # MoGe predicts infinity for sky pixels
    depth_pred[np.isinf(depth_pred)] = -1.0

    min_depth = np.min(depth_pred)
    max_depth = np.max(depth_pred)

    # Normalize to [0, 1]
    depth_pred_normalized = (depth_pred - min_depth) / (max_depth - min_depth)

    # Scale to uint16 range
    depth_pred_scaled = np.round(depth_pred_normalized * (2 ** 16 - 1)).astype(np.uint16)

    # Save image with pred and error as uint16
    depth_pred_img = PIL.Image.fromarray(depth_pred_scaled)
    metadata = PngInfo()
    metadata.add_text("max_value", str(max_depth))
    metadata.add_text("min_value", str(min_depth))
    depth_pred_img.save(outpath_file,
                        format="PNG", mode="I", pnginfo=metadata)