import json
import unittest
from unittest.mock import patch

import numpy as np
import shapely
from affine import Affine

from samgis.prediction_api import predictors
from samgis.prediction_api.predictors import samexporter_predict
from tests import TEST_EVENTS_FOLDER


class TestPredictors(unittest.TestCase):
    @patch.object(predictors, "download_extent")
    def test_get_raster_inference(self, download_extent_mocked):
        name_fn = "samexporter_predict"

        with open(TEST_EVENTS_FOLDER / f"{name_fn}.json") as tst_json:
            inputs_outputs = json.load(tst_json)
            for k, input_output in inputs_outputs.items():
                input_payload = input_output["input"]
                prompt = input_payload["prompt"]
                model_name = input_payload["model_name"]
                bbox = input_payload["bbox"]
                zoom = input_payload["zoom"]
                print(f"k:{k}.")
                img = np.load(TEST_EVENTS_FOLDER / f"{name_fn}" / k / "img.npy")
                affine_transform = Affine.from_gdal(*input_payload["matrix"])
                download_extent_mocked.return_value = img, affine_transform
                expected_output = input_output["output"]

                output_dict = samexporter_predict(
                    bbox,
                    prompt,
                    zoom,
                    model_name
                )
                len_inference_out = output_dict["n_predictions"]
                geojson = output_dict["geojson"]
                n_shapes_geojson = output_dict["n_shapes_geojson"]

                assert isinstance(geojson, str)
                assert isinstance(n_shapes_geojson, int)
                assert len_inference_out == expected_output["n_predictions"]

                output_geojson = shapely.from_geojson(geojson)
                print("output_geojson::{}.".format(output_geojson))
                assert isinstance(output_geojson, shapely.GeometryCollection)
                assert len(output_geojson.geoms) > 0

    @patch.object(predictors, "get_raster_inference_with_embedding_from_dict")
    @patch.object(predictors, "SegmentAnythingONNX2")
    @patch.object(predictors, "download_extent")
    @patch.object(predictors, "get_vectorized_raster_as_geojson")
    def test_samexporter_predict_mocked(
            self,
            get_vectorized_raster_as_geojson_mocked,
            download_extent_mocked,
            segment_anything_onnx2_mocked,
            get_raster_inference_with_embedding_from_dict_mocked
    ):
        """
        model_instance = SegmentAnythingONNX()
        img, matrix = download_extent(DEFAULT_TMS, pt0[0], pt0[1], pt1[0], pt1[1], zoom)
        transform = get_affine_transform_from_gdal(matrix)
        mask, n_predictions = get_raster_inference(img, prompt, models_instance, model_name)
        get_vectorized_raster_as_geojson(mask, matrix)
        """
        aff = 1, 2, 3, 4, 5, 6
        segment_anything_onnx2_mocked.return_value = "SegmentAnythingONNX2_instance"
        input_downloaded = np.arange(0, 300, 1).reshape((10, 10, 3))
        download_extent_mocked.return_value = input_downloaded, aff
        mask_output = np.zeros((10, 10))
        mask_output[4:4, 6:6] = 255.0
        get_raster_inference_with_embedding_from_dict_mocked.return_value = mask_output, 1
        get_vectorized_raster_as_geojson_mocked.return_value = {"geojson": "{}", "n_shapes_geojson": 2}
        output = samexporter_predict(
            bbox=[[1, 2], [3, 4]], prompt=[{}], zoom=10, model_name="mobile_sam", source_name="localtest"
        )
        assert output == {"n_predictions": 1, "geojson": "{}", "n_shapes_geojson": 2}