import json

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, 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)
    pixels = np.asarray(source, dtype=np.uint8).copy()
    pixels[:, :, 3] = np.asarray(selection, 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_unbounded_selection(
    input_path,
    polygon,
    output_path,
    invert=False,
):
    import numpy as np
    from PIL import Image, ImageFilter

    if invert:
        _write_masked_color(input_path, polygon, output_path, invert=True)
        with Image.open(input_path) as source:
            return output_path, (0, 0, source.width, source.height)
    canvas, selection_mask, bounds = create_selection_canvas(
        input_path, polygon, 'NONZERO'
    )
    selection = Image.fromarray(
        selection_mask.astype(np.uint8) * 255,
        mode='L',
    ).filter(ImageFilter.GaussianBlur(radius=0.75))
    pixels = np.asarray(canvas, dtype=np.uint8).copy()
    pixels[:, :, 3] = np.asarray(selection, 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, bounds


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_cutout_images(
    input_path,
    selection_path,
    output_directory,
    selection_bounds=None,
):
    import numpy as np
    from PIL import Image

    with Image.open(input_path) as opened:
        source = opened.convert("RGBA")
    with Image.open(selection_path) as opened:
        selection = opened.convert("RGBA")
    if selection_bounds is None and selection.size != source.size:
        raise ValueError("The Cutout selection does not match the source image")
    if selection_bounds is None:
        selection_bounds = (0, 0, source.width, source.height)
    selection_pixels = np.asarray(selection, dtype=np.uint8).copy()
    alpha = selection_pixels[:, :, 3]
    alpha[alpha <= 1] = 0
    selection_pixels[alpha == 0, :3] = 0
    selection = Image.fromarray(selection_pixels, mode="RGBA")
    visible_y, visible_x = np.nonzero(alpha >= 128)
    if not len(visible_x):
        raise ValueError("The Cutout selection contains no visible pixels")
    bounds = (
        int(visible_x.min()),
        int(visible_y.min()),
        int(visible_x.max()) + 1,
        int(visible_y.max()) + 1,
    )
    content_bounds = bounds
    selection_left, selection_top, _selection_right, _selection_bottom = (
        selection_bounds
    )
    bounds = (
        selection_left + content_bounds[0],
        selection_top + content_bounds[1],
        selection_left + content_bounds[2],
        selection_top + content_bounds[3],
    )
    color_canvas = Image.new('RGBA', selection.size, (0, 0, 0, 0))
    color_canvas.paste(source, (-selection_left, -selection_top))
    color = color_canvas.crop(content_bounds)
    color_path = output_directory / "image.0001.png"
    color.save(color_path, format="PNG", compress_level=1)

    selection_crop = selection.crop(content_bounds)
    selection_path = output_directory / "selection.png"
    selection_crop.save(selection_path, format="PNG", compress_level=1)

    analysis_pixels = np.asarray(color, dtype=np.uint8).copy()
    mask_alpha = np.asarray(selection_crop, dtype=np.uint8)[:, :, 3]
    analysis_pixels[:, :, 3] = np.minimum(
        analysis_pixels[:, :, 3],
        mask_alpha,
    )
    analysis_pixels[mask_alpha < 128] = 0
    analysis_path = output_directory / "analysis.png"
    Image.fromarray(analysis_pixels, mode="RGBA").save(
        analysis_path,
        format="PNG",
        compress_level=1,
    )
    bounds_path = output_directory / "bounds.json"
    bounds_path.write_text(
        json.dumps(bounds, separators=(",", ":")),
        encoding="utf-8",
    )
    return color_path, selection_path, analysis_path, bounds_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,
    include_normal=False,
):
    import numpy as np
    from PIL import Image

    with Image.open(input_path) as source:
        rgba = np.asarray(source.convert("RGBA"), dtype=np.uint8)
        source_size = source.size
    selection_mask = rgba[:, :, 3] >= 128
    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")
        analysis_path = None
        result = {}
        if polygon:
            selection_bounds = None
            full_selection_path = context.directory / "selection-full.png"
            if parameters.get("refine_selection"):
                _write_refined_color(
                    context,
                    parameters,
                    images[0],
                    polygon,
                    full_selection_path,
                )
            else:
                _selection_path, selection_bounds = _write_unbounded_selection(
                    images[0],
                    polygon,
                    full_selection_path,
                    parameters.get("invert", False),
                )
            color_path, selection_path, analysis_path, bounds_path = (
                _write_cutout_images(
                    images[0],
                    full_selection_path,
                    context.directory,
                    selection_bounds,
                )
            )
            result.update(
                selection=selection_path.name,
                bounds=bounds_path.name,
            )
        else:
            copy_color_images(images, context.directory)
            color_path = min(context.directory.glob("image.0001.*"))

        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")
            model_path = analysis_path or 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, model_path
                ).name
            elif shape == "DEPTH_BALLOON":
                context.progress(0.4, "Estimating Balloon depth")
                depth_path, normal_path = _write_balloon_depth(
                    context,
                    parameters,
                    model_path,
                    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,
                    model_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, model_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