# Copyright (C) 2021-2024, Mindee.

# This program is licensed under the Apache License 2.0.
# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.

import glob
import os
from typing import Any, Dict, List, Tuple, Union

import numpy as np
from PIL import Image
from scipy import io as sio
from tqdm import tqdm

from .datasets import VisionDataset
from .utils import convert_target_to_relative, crop_bboxes_from_image

__all__ = ["SynthText"]


class SynthText(VisionDataset):
    """SynthText dataset from `"Synthetic Data for Text Localisation in Natural Images"
    <https://arxiv.org/abs/1604.06646>`_ | `"repository" <https://github.com/ankush-me/SynthText>`_ |
    `"website" <https://www.robots.ox.ac.uk/~vgg/data/scenetext/>`_.

    .. image:: https://doctr-static.mindee.com/models?id=v0.5.0/svt-grid.png&src=0
        :align: center

    >>> from doctr.datasets import SynthText
    >>> train_set = SynthText(train=True, download=True)
    >>> img, target = train_set[0]

    Args:
    ----
        train: whether the subset should be the training one
        use_polygons: whether polygons should be considered as rotated bounding box (instead of straight ones)
        recognition_task: whether the dataset should be used for recognition task
        **kwargs: keyword arguments from `VisionDataset`.
    """

    URL = "https://thor.robots.ox.ac.uk/~vgg/data/scenetext/SynthText.zip"
    SHA256 = "28ab030485ec8df3ed612c568dd71fb2793b9afbfa3a9d9c6e792aef33265bf1"

    def __init__(
        self,
        train: bool = True,
        use_polygons: bool = False,
        recognition_task: bool = False,
        **kwargs: Any,
    ) -> None:
        super().__init__(
            self.URL,
            None,
            file_hash=None,
            extract_archive=True,
            pre_transforms=convert_target_to_relative if not recognition_task else None,
            **kwargs,
        )
        self.train = train
        self.data: List[Tuple[Union[str, np.ndarray], Union[str, Dict[str, Any]]]] = []
        np_dtype = np.float32

        # Load mat data
        tmp_root = os.path.join(self.root, "SynthText") if self.SHA256 else self.root
        # define folder to write SynthText recognition dataset
        reco_folder_name = "SynthText_recognition_train" if self.train else "SynthText_recognition_test"
        reco_folder_name = "Poly_" + reco_folder_name if use_polygons else reco_folder_name
        reco_folder_path = os.path.join(tmp_root, reco_folder_name)
        reco_images_counter = 0

        if recognition_task and os.path.isdir(reco_folder_path):
            self._read_from_folder(reco_folder_path)
            return
        elif recognition_task and not os.path.isdir(reco_folder_path):
            os.makedirs(reco_folder_path, exist_ok=False)

        mat_data = sio.loadmat(os.path.join(tmp_root, "gt.mat"))
        train_samples = int(len(mat_data["imnames"][0]) * 0.9)
        set_slice = slice(train_samples) if self.train else slice(train_samples, None)
        paths = mat_data["imnames"][0][set_slice]
        boxes = mat_data["wordBB"][0][set_slice]
        labels = mat_data["txt"][0][set_slice]
        del mat_data

        for img_path, word_boxes, txt in tqdm(
            iterable=zip(paths, boxes, labels), desc="Unpacking SynthText", total=len(paths)
        ):
            # File existence check
            if not os.path.exists(os.path.join(tmp_root, img_path[0])):
                raise FileNotFoundError(f"unable to locate {os.path.join(tmp_root, img_path[0])}")

            labels = [elt for word in txt.tolist() for elt in word.split()]
            # (x, y) coordinates of top left, top right, bottom right, bottom left corners
            word_boxes = (
                word_boxes.transpose(2, 1, 0)
                if word_boxes.ndim == 3
                else np.expand_dims(word_boxes.transpose(1, 0), axis=0)
            )

            if not use_polygons:
                # xmin, ymin, xmax, ymax
                word_boxes = np.concatenate((word_boxes.min(axis=1), word_boxes.max(axis=1)), axis=1)

            if recognition_task:
                crops = crop_bboxes_from_image(img_path=os.path.join(tmp_root, img_path[0]), geoms=word_boxes)
                for crop, label in zip(crops, labels):
                    if crop.shape[0] > 0 and crop.shape[1] > 0 and len(label) > 0:
                        # write data to disk
                        with open(os.path.join(reco_folder_path, f"{reco_images_counter}.txt"), "w") as f:
                            f.write(label)
                            tmp_img = Image.fromarray(crop)
                            tmp_img.save(os.path.join(reco_folder_path, f"{reco_images_counter}.png"))
                            reco_images_counter += 1
            else:
                self.data.append((img_path[0], dict(boxes=np.asarray(word_boxes, dtype=np_dtype), labels=labels)))

        if recognition_task:
            self._read_from_folder(reco_folder_path)

        self.root = tmp_root

    def extra_repr(self) -> str:
        return f"train={self.train}"

    def _read_from_folder(self, path: str) -> None:
        for img_path in glob.glob(os.path.join(path, "*.png")):
            with open(os.path.join(path, f"{os.path.basename(img_path)[:-4]}.txt"), "r") as f:
                self.data.append((img_path, f.read()))