from pathlib import Path

from .surface import write_normal_texture


def write_da3_depth(prediction, output_path):
    import numpy as np

    values = {"depth": np.round(prediction.depth, 8)}
    for name in ("conf", "extrinsics", "intrinsics"):
        value = getattr(prediction, name, None)
        if value is not None:
            values[name] = np.round(value, 2) if name == "conf" else value
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output_path, **values)
    return output_path


def write_moge_depth(output_directory, prediction):
    import numpy as np

    output_directory = Path(output_directory)
    np.savez(
        output_directory / "results.npz",
        **prediction.archive_values(),
    )
    for index, (depth, normal, mask) in enumerate(
        zip(prediction.depth, prediction.normal, prediction.mask)
    ):
        _write_moge_maps(
            output_directory,
            None if len(prediction.depth) == 1 else index,
            depth,
            normal,
            mask,
        )


def _write_moge_maps(output_directory, index, depth, normal, mask):
    import numpy as np
    from PIL import Image

    suffix = "" if index is None else f".{index + 1:04d}"
    valid = mask & np.isfinite(depth) & (depth > 0)
    preview = np.zeros(depth.shape, dtype=np.uint8)
    if np.any(valid):
        low, high = np.percentile(depth[valid], (2.0, 98.0))
        span = max(float(high - low), 1e-6)
        preview[valid] = np.clip(
            (depth[valid] - low) / span * 255.0,
            0.0,
            255.0,
        ).astype(np.uint8)
    Image.fromarray(preview, mode="L").save(
        output_directory / f"depth{suffix}.png"
    )
    write_normal_texture(
        normal,
        output_directory / f"normal{suffix}.png",
    )
    Image.fromarray(mask.astype(np.uint8) * 255, mode="L").save(
        output_directory / f"mask{suffix}.png"
    )