from ..media import copy_color_images
from ..models import geometry_adapter, moge2
from .background import remove_image_background
from .balloon import (
    write_camera_normal_texture,
    write_depth_analysis,
)
from .common import analysis_frames, input_frames
from .image_mask import (
    apply_selection_mask,
    create_selection_canvas,
    polygon_selection_mask,
)


def _write_field(frame, output_path):
    import numpy as np

    depth = np.asarray(frame.depth, dtype=np.float32)
    points = np.asarray(frame.points, dtype=np.float32)
    intrinsics = np.asarray(frame.intrinsics, dtype=np.float32)
    valid = np.asarray(frame.valid_mask, dtype=bool)
    if depth.ndim != 2 or points.shape != (*depth.shape, 3):
        raise ValueError("The depth model returned an invalid point field")
    valid &= np.isfinite(depth) & (depth > 0.0)
    valid &= np.isfinite(points).all(axis=-1)
    if not valid.any():
        raise ValueError("The depth model returned no valid geometry")
    reference_values = depth[valid]
    if frame.model_mask is not None:
        model_mask = np.asarray(frame.model_mask, dtype=bool)
        masked = depth[valid & model_mask]
        if len(masked):
            reference_values = masked
    np.savez_compressed(
        output_path,
        points=points,
        depth=depth,
        intrinsics=intrinsics,
        valid=valid,
        image_size=np.asarray((depth.shape[1], depth.shape[0]), dtype=np.int32),
        reference_depth=np.asarray(np.median(reference_values), dtype=np.float32),
    )


def _write_masked_color(input_path, polygon, output_path, invert=False):
    import numpy as np
    from PIL import Image, ImageChops, ImageFilter, ImageOps

    with Image.open(input_path) as opened:
        source = opened.convert("RGBA")
    selection = Image.fromarray(
        polygon_selection_mask(source.size, polygon, "NONZERO").astype(np.uint8) * 255,
        mode="L",
    ).filter(ImageFilter.GaussianBlur(radius=0.75))
    if invert:
        selection = ImageOps.invert(selection)
    alpha = ImageChops.multiply(source.getchannel("A"), selection)
    pixels = np.asarray(source, dtype=np.uint8).copy()
    pixels[:, :, 3] = np.asarray(alpha, dtype=np.uint8)
    pixels[pixels[:, :, 3] == 0, :3] = 0
    Image.fromarray(pixels, mode="RGBA").save(
        output_path,
        format="PNG",
        compress_level=1,
    )
    return output_path


def _write_refined_color(context, parameters, input_path, polygon, output_path):
    from PIL import Image, ImageChops

    canvas, selection_mask, bounds = create_selection_canvas(
        input_path, polygon, "NONZERO"
    )
    selection_path = context.directory / "selection-input.png"
    canvas.save(selection_path, format="PNG", compress_level=1)
    model_manager = context.resource("model_manager")
    foreground, _load_ms, _inference_ms = remove_image_background(
        selection_path,
        model_manager.directory(parameters["refine_model"]),
        parameters["device"],
        model_manager,
        context.check_cancelled,
    )
    refined = apply_selection_mask(foreground, selection_mask)
    with Image.open(input_path) as opened:
        source = opened.convert("RGBA")
    if parameters.get("invert"):
        mask = Image.new("L", source.size, 0)
        mask.paste(refined.getchannel("A"), bounds[:2])
        source.putalpha(
            ImageChops.multiply(source.getchannel("A"), ImageChops.invert(mask))
        )
        output = source
    else:
        output = Image.new("RGBA", source.size, (0, 0, 0, 0))
        output.paste(refined, bounds[:2])
    output.save(output_path, format="PNG", compress_level=1)
    return output_path


def _write_camera_normal_from_prediction(prediction, color_path, output_path):
    import numpy as np
    from PIL import Image

    normal = np.asarray(prediction.normal[0], dtype=np.float32)
    with Image.open(color_path) as source:
        alpha = source.convert("RGBA").getchannel("A")
        source_size = source.size
    valid = (
        np.asarray(
            alpha.resize((normal.shape[1], normal.shape[0]), Image.Resampling.BILINEAR)
        )
        > 127
    )
    return write_camera_normal_texture(
        normal,
        valid,
        (0, 0),
        source_size,
        source_size,
        output_path,
    )


def _write_camera_normal(context, parameters, color_path):
    normal_parameters = dict(parameters)
    normal_parameters.update(
        model=parameters["normal_model"],
        resolution_level=parameters["normal_resolution_level"],
    )
    prediction = moge2.infer(
        context.resource("model_manager"),
        normal_parameters,
        [color_path],
        include_points=False,
        cancel_check=context.check_cancelled,
    )
    return _write_camera_normal_from_prediction(
        prediction,
        color_path,
        context.directory / "normal.png",
    )


def _write_balloon_depth(
    context,
    parameters,
    input_path,
    polygon,
    selection_size=None,
    include_normal=False,
):
    from PIL import Image

    with Image.open(input_path) as source:
        source_size = source.size
    selection_size = selection_size or source_size
    scale_x = source_size[0] / selection_size[0]
    scale_y = source_size[1] / selection_size[1]
    analysis_polygon = tuple(
        (point[0] * scale_x, point[1] * scale_y) for point in polygon
    )
    selection_mask = polygon_selection_mask(
        source_size,
        analysis_polygon,
        "NONZERO",
    )
    prediction = moge2.infer(
        context.resource("model_manager"),
        parameters,
        [input_path],
        include_points=False,
        cancel_check=context.check_cancelled,
    )
    depth_path = write_depth_analysis(
        prediction,
        selection_mask,
        (0, 0),
        source_size,
        context.directory / "depth.npz",
    )
    normal_path = None
    if include_normal:
        normal_path = _write_camera_normal_from_prediction(
            prediction,
            input_path,
            context.directory / "normal.png",
        )
    return depth_path, normal_path


def _write_surface_depth(
    context, parameters, color_path, include_normal=False
):
    model_manager = context.resource("model_manager")
    adapter = geometry_adapter(parameters["model"])
    if adapter is moge2:
        prediction = adapter.infer(
            model_manager,
            parameters,
            [color_path],
            include_points=True,
            cancel_check=context.check_cancelled,
        )
    else:
        prediction = adapter.infer(
            model_manager,
            parameters,
            [color_path],
            cancel_check=context.check_cancelled,
        )
    context.check_cancelled()
    field_path = context.directory / "field.npz"
    _write_field(adapter.surface_frame(prediction), field_path)
    normal_path = None
    if include_normal and adapter is moge2:
        normal_path = _write_camera_normal_from_prediction(
            prediction, color_path, context.directory / "normal.png"
        )
    return field_path, normal_path


def run(context, parameters):
    context.progress(0.02, "Preparing mesh input")
    context.directory.mkdir(parents=True, exist_ok=True)
    with input_frames(context, parameters) as images:
        if len(images) != 1:
            raise ValueError("Empty Image to Mesh currently supports one image")
        shape = parameters["shape"]
        polygon = parameters.get("polygon")
        if polygon:
            if parameters.get("refine_selection"):
                color_path = _write_refined_color(
                    context,
                    parameters,
                    images[0],
                    polygon,
                    context.directory / "image.0001.png",
                )
            else:
                color_path = _write_masked_color(
                    images[0],
                    polygon,
                    context.directory / "image.0001.png",
                    parameters.get("invert", False),
                )
        else:
            copy_color_images(images, context.directory)
            color_path = min(context.directory.glob("image.0001.*"))
        result = {}
        normal_analysis = bool(parameters.get("normal_analysis"))
        with analysis_frames(context, parameters, images) as model_images:
            if len(model_images) != 1:
                raise ValueError("Empty Image to Mesh currently supports one image")
            analysis_path = model_images[0]

            if shape in {"CUTOUT", "BALLOON"} and normal_analysis:
                context.progress(0.4, "Estimating surface normal")
                result["normal"] = _write_camera_normal(
                    context, parameters, analysis_path
                ).name
            elif shape == "DEPTH_BALLOON":
                from PIL import Image

                with Image.open(images[0]) as selection_source:
                    selection_size = selection_source.size
                context.progress(0.4, "Estimating Balloon depth")
                depth_path, normal_path = _write_balloon_depth(
                    context,
                    parameters,
                    analysis_path,
                    polygon,
                    selection_size,
                    include_normal=normal_analysis,
                )
                result["depth"] = depth_path.name
                if normal_path is not None:
                    result["normal"] = normal_path.name
            elif shape in {"DEPTH_PLANE", "DEPTH_SURFACE"}:
                context.progress(0.4, "Estimating image depth")
                field_path, normal_path = _write_surface_depth(
                    context,
                    parameters,
                    analysis_path,
                    include_normal=normal_analysis,
                )
                result["field"] = field_path.name
                if (
                    normal_analysis
                    and normal_path is None
                    and shape == "DEPTH_SURFACE"
                ):
                    normal_path = _write_camera_normal(
                        context, parameters, analysis_path
                    )
                if normal_path is not None:
                    result["normal"] = normal_path.name

    context.check_cancelled()
    context.progress(1.0, "Mesh data is ready")
    return result