"""
Define a machine learning model executed by ONNX Runtime (https://ai/)
for Segment Anything (https://segment-anything.com).
Modified from https://github.com/vietanhdev/samexporter/

Copyright (c) 2023 Viet Anh Nguyen

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""
from copy import deepcopy

from cv2 import INTER_LINEAR, warpAffine
from numpy import array as np_array, uint8, zeros, ndarray
from numpy import concatenate, float32, linalg, matmul, ones
from onnxruntime import get_available_providers, InferenceSession

from samgis_core import app_logger, MODEL_FOLDER
from samgis_core.utilities.constants import DEFAULT_INPUT_SHAPE, MODEL_ENCODER_NAME, MODEL_DECODER_NAME
from samgis_core.utilities.type_hints import PIL_Image, list_dict, tuple_ndarr_int


class SegmentAnythingONNX:
    """Segmentation model using SegmentAnything"""

    def __init__(self, encoder_model_path, decoder_model_path) -> None:
        self.target_size = DEFAULT_INPUT_SHAPE[1]
        self.input_size = DEFAULT_INPUT_SHAPE

        # Load models
        providers = get_available_providers()

        # Pop TensorRT Runtime due to crashing issues
        # TODO: Add back when TensorRT backend is stable
        providers = [p for p in providers if p != "TensorrtExecutionProvider"]

        if providers:
            app_logger.info(
                "Available providers for ONNXRuntime: %s", ", ".join(providers)
            )
        else:
            app_logger.warning("No available providers for ONNXRuntime")
        self.encoder_session = InferenceSession(
            encoder_model_path, providers=providers
        )
        self.encoder_input_name = self.encoder_session.get_inputs()[0].name
        self.decoder_session = InferenceSession(
            decoder_model_path, providers=providers
        )

    @staticmethod
    def get_input_points(prompt):
        """Get input points"""
        points = []
        labels = []
        for mark in prompt:
            if mark["type"] == "point":
                points.append(mark["data"])
                labels.append(mark["label"])
            elif mark["type"] == "rectangle":
                points.append([mark["data"][0], mark["data"][1]])  # top left
                points.append(
                    [mark["data"][2], mark["data"][3]]
                )  # bottom right
                labels.append(2)
                labels.append(3)
        points, labels = np_array(points), np_array(labels)
        return points, labels

    def run_encoder(self, encoder_inputs):
        """Run encoder"""
        output = self.encoder_session.run(None, encoder_inputs)
        image_embedding = output[0]
        return image_embedding

    @staticmethod
    def get_preprocess_shape(old_h: int, old_w: int, long_side_length: int):
        """
        Compute the output size given input size and target long side length.
        """
        scale = long_side_length * 1.0 / max(old_h, old_w)
        new_h, new_w = old_h * scale, old_w * scale
        new_w = int(new_w + 0.5)
        new_h = int(new_h + 0.5)
        return new_h, new_w

    def apply_coords(self, coords: ndarray, original_size, target_length):
        """
        Expects a numpy np_array of length 2 in the final dimension. Requires the
        original image size in (H, W) format.
        """
        old_h, old_w = original_size
        new_h, new_w = self.get_preprocess_shape(
            original_size[0], original_size[1], target_length
        )
        coords = deepcopy(coords).astype(float)
        coords[..., 0] = coords[..., 0] * (new_w / old_w)
        coords[..., 1] = coords[..., 1] * (new_h / old_h)
        return coords

    def run_decoder(
        self, image_embedding, original_size, transform_matrix, prompt
    ):
        """Run decoder"""
        input_points, input_labels = self.get_input_points(prompt)

        # Add a batch index, concatenate a padding point, and transform.
        onnx_coord = concatenate(
            [input_points, np_array([[0.0, 0.0]])], axis=0
        )[None, :, :]
        onnx_label = concatenate([input_labels, np_array([-1])], axis=0)[
            None, :
        ].astype(float32)
        onnx_coord = self.apply_coords(
            onnx_coord, self.input_size, self.target_size
        ).astype(float32)

        # Apply the transformation matrix to the coordinates.
        onnx_coord = concatenate(
            [
                onnx_coord,
                ones((1, onnx_coord.shape[1], 1), dtype=float32),
            ],
            axis=2,
        )
        onnx_coord = matmul(onnx_coord, transform_matrix.T)
        onnx_coord = onnx_coord[:, :, :2].astype(float32)

        # Create an empty mask input and an indicator for no mask.
        onnx_mask_input = zeros((1, 1, 256, 256), dtype=float32)
        onnx_has_mask_input = zeros(1, dtype=float32)

        decoder_inputs = {
            "image_embeddings": image_embedding,
            "point_coords": onnx_coord,
            "point_labels": onnx_label,
            "mask_input": onnx_mask_input,
            "has_mask_input": onnx_has_mask_input,
            "orig_im_size": np_array(self.input_size, dtype=float32),
        }
        masks, _, _ = self.decoder_session.run(None, decoder_inputs)

        # Transform the masks back to the original image size.
        inv_transform_matrix = linalg.inv(transform_matrix)
        transformed_masks = self.transform_masks(
            masks, original_size, inv_transform_matrix
        )

        return transformed_masks

    @staticmethod
    def transform_masks(masks, original_size, transform_matrix):
        """Transform masks
        Transform the masks back to the original image size.
        """
        output_masks = []
        for batch in range(masks.shape[0]):
            batch_masks = []
            for mask_id in range(masks.shape[1]):
                mask = masks[batch, mask_id]
                try:
                    try:
                        app_logger.debug(f"mask_shape transform_masks:{mask.shape}, dtype:{mask.dtype}.")
                    except Exception as e_mask_shape_transform_masks:
                        app_logger.error(f"e_mask_shape_transform_masks:{e_mask_shape_transform_masks}.")
                    mask = warpAffine(
                        mask,
                        transform_matrix[:2],
                        (original_size[1], original_size[0]),
                        flags=INTER_LINEAR,
                    )
                except Exception as e_warp_affine1:
                    app_logger.error(f"e_warp_affine1 mask shape:{mask.shape}, dtype:{mask.dtype}.")
                    app_logger.error(f"e_warp_affine1 transform_matrix:{transform_matrix}, [:2] {transform_matrix[:2]}.")
                    app_logger.error(f"e_warp_affine1 original_size:{original_size}.")
                    raise e_warp_affine1
                batch_masks.append(mask)
            output_masks.append(batch_masks)
        return np_array(output_masks)

    def encode(self, cv_image):
        """
        Calculate embedding and metadata for a single image.
        """
        original_size = cv_image.shape[:2]

        # Calculate a transformation matrix to convert to self.input_size
        scale_x = self.input_size[1] / cv_image.shape[1]
        scale_y = self.input_size[0] / cv_image.shape[0]
        scale = min(scale_x, scale_y)
        transform_matrix = np_array(
            [
                [scale, 0, 0],
                [0, scale, 0],
                [0, 0, 1],
            ]
        )
        try:
            cv_image = warpAffine(
                cv_image,
                transform_matrix[:2],
                (self.input_size[1], self.input_size[0]),
                flags=INTER_LINEAR,
            )
        except Exception as e_warp_affine2:
            app_logger.error(f"e_warp_affine2:{e_warp_affine2}.")
            np_cv_image = np_array(cv_image)
            app_logger.error(f"e_warp_affine2 cv_image shape:{np_cv_image.shape}, dtype:{np_cv_image.dtype}.")
            app_logger.error(f"e_warp_affine2 transform_matrix:{transform_matrix}, [:2] {transform_matrix[:2]}")
            app_logger.error(f"e_warp_affine2 self.input_size:{self.input_size}.")
            raise e_warp_affine2

        encoder_inputs = {
            self.encoder_input_name: cv_image.astype(float32),
        }
        image_embedding = self.run_encoder(encoder_inputs)
        return {
            "image_embedding": image_embedding,
            "original_size": original_size,
            "transform_matrix": transform_matrix,
        }

    def predict_masks(self, embedding, prompt):
        """
        Predict masks for a single image.
        """
        masks = self.run_decoder(
            embedding["image_embedding"],
            embedding["original_size"],
            embedding["transform_matrix"],
            prompt,
        )

        return masks


def get_raster_inference(
        img: PIL_Image or ndarray, prompt: list_dict, models_instance: SegmentAnythingONNX, model_name: str
     ) -> tuple_ndarr_int:
    """
    Get inference output for a given image using a SegmentAnythingONNX model

    Args:
        img: input PIL Image
        prompt: list of prompt dict
        models_instance: SegmentAnythingONNX instance model
        model_name: model name string

    Returns:
        raster prediction mask, prediction number
    """
    np_img = np_array(img)
    app_logger.info(f"img type {type(np_img)}, prompt:{prompt}.")
    app_logger.debug(f"onnxruntime input shape/size (shape if PIL) {np_img.size}.")
    try:
        app_logger.debug(f"onnxruntime input shape (NUMPY) {np_img.shape}.")
    except Exception as e_shape:
        app_logger.error(f"e_shape:{e_shape}.")
    app_logger.info(f"instantiated model {model_name}, ENCODER {MODEL_ENCODER_NAME}, "
                    f"DECODER {MODEL_DECODER_NAME} from {MODEL_FOLDER}: Creating embedding...")
    embedding = models_instance.encode(np_img)
    app_logger.debug(f"embedding created, running predict_masks with prompt {prompt}...")
    inference_out = models_instance.predict_masks(embedding, prompt)
    len_inference_out = len(inference_out[0, :, :, :])
    app_logger.info(f"Created {len_inference_out} prediction_masks,"
                    f"shape:{inference_out.shape}, dtype:{inference_out.dtype}.")
    mask = zeros((inference_out.shape[2], inference_out.shape[3]), dtype=uint8)
    for n, m in enumerate(inference_out[0, :, :, :]):
        app_logger.debug(f"{n}th of prediction_masks shape {inference_out.shape}"
                         f" => mask shape:{mask.shape}, {mask.dtype}.")
        mask[m > 0.0] = 255
    return mask, len_inference_out
