from collections.abc import Iterable
from pathlib import Path
from typing import Any
from xml.etree import ElementTree as ET

import datasets
import numpy as np
import pyvips
from datasets import Dataset
from datasets.splits import NamedSplit
from PIL import Image, ImageDraw
from tqdm import tqdm


# https://drive.google.com/file/d/1kdOl3s6uQBRv0nToSIf1dPuceZunzL4N/view
patient_data = {
    "TCGA-55-1594": "Lung",
    "TCGA-69-7760": "Lung",
    "TCGA-69-A59K": "Lung",
    "TCGA-73-4668": "Lung",
    "TCGA-78-7220": "Lung",
    "TCGA-86-7713": "Lung",
    "TCGA-86-8672": "Lung",
    "TCGA-L4-A4E5": "Lung",
    "TCGA-MP-A4SY": "Lung",
    "TCGA-MP-A4T7": "Lung",
    "TCGA-5P-A9K0": "Kidney",
    "TCGA-B9-A44B": "Kidney",
    "TCGA-B9-A8YI": "Kidney",
    "TCGA-DW-7841": "Kidney",
    "TCGA-EV-5903": "Kidney",
    "TCGA-F9-A97G": "Kidney",
    "TCGA-G7-A8LD": "Kidney",
    "TCGA-MH-A560": "Kidney",
    "TCGA-P4-AAVK": "Kidney",
    "TCGA-SX-A7SR": "Kidney",
    "TCGA-UZ-A9PO": "Kidney",
    "TCGA-UZ-A9PU": "Kidney",
    "TCGA-A2-A0CV": "Breast",
    "TCGA-A2-A0ES": "Breast",
    "TCGA-B6-A0WZ": "Breast",
    "TCGA-BH-A18T": "Breast",
    "TCGA-D8-A1X5": "Breast",
    "TCGA-E2-A154": "Breast",
    "TCGA-E9-A22B": "Breast",
    "TCGA-E9-A22G": "Breast",
    "TCGA-EW-A6SD": "Breast",
    "TCGA-S3-AA11": "Breast",
    "TCGA-EJ-5495": "Prostate",
    "TCGA-EJ-5505": "Prostate",
    "TCGA-EJ-5517": "Prostate",
    "TCGA-G9-6342": "Prostate",
    "TCGA-G9-6499": "Prostate",
    "TCGA-J4-A67Q": "Prostate",
    "TCGA-J4-A67T": "Prostate",
    "TCGA-KK-A59X": "Prostate",
    "TCGA-KK-A6E0": "Prostate",
    "TCGA-KK-A7AW": "Prostate",
    "TCGA-V1-A8WL": "Prostate",
    "TCGA-V1-A9O9": "Prostate",
    "TCGA-X4-A8KQ": "Prostate",
    "TCGA-YL-A9WY": "Prostate",
    "TCGA-49-6743": "Lung",
    "TCGA-50-6591": "Lung",
    "TCGA-55-7570": "Lung",
    "TCGA-55-7573": "Lung",
    "TCGA-73-4662": "Lung",
    "TCGA-78-7152": "Lung",
    "TCGA-2Z-A9JG": "Kidney",
    "TCGA-2Z-A9JN": "Kidney",
    "TCGA-DW-7838": "Kidney",
    "TCGA-DW-7963": "Kidney",
    "TCGA-F9-A8NY": "Kidney",
    "TCGA-IZ-A6M9": "Kidney",
    "TCGA-MH-A55W": "Kidney",
    "TCGA-A2-A04X": "Breast",
    "TCGA-D8-A3Z6": "Breast",
    "TCGA-E2-A108": "Breast",
    "TCGA-EW-A6SB": "Breast",
    "TCGA-G9-6356": "Prostate",
    "TCGA-G9-6367": "Prostate",
    "TCGA-VP-A87E": "Prostate",
    "TCGA-VP-A87H": "Prostate",
    "TCGA-X4-A8KS": "Prostate",
    "TCGA-YL-A9WL": "Prostate",
}


features = datasets.Features(
    {
        "patient": datasets.Value("string"),
        "image": datasets.Image(mode="RGB"),
        "instances": datasets.Sequence(datasets.Image(mode="1")),
        "categories": datasets.Sequence(
            datasets.ClassLabel(
                names=[
                    "Ambiguous",
                    "Epithelial",
                    "Lymphocyte",
                    "Macrophage",
                    "Neutrophil",
                ],
            )
        ),
        "tissue": datasets.ClassLabel(
            names=[
                "Breast",
                "Kidney",
                "Lung",
                "Prostate",
            ]
        ),
    }
)


def get_masks(
    path: Path, mask_size: tuple[int, int]
) -> tuple[list[Image.Image], list[str]]:
    masks = []
    categories = []

    root = ET.parse(path).getroot()

    for annotation in root.findall("Annotation"):
        for region in annotation.findall("Regions/Region"):
            polygon = [
                (float(vertex.attrib["X"]), float(vertex.attrib["Y"]))
                for vertex in region.findall("Vertices/Vertex")
            ]

            if len(polygon) < 2:
                continue

            mask = Image.new("1", size=mask_size)
            canvas = ImageDraw.Draw(mask)
            canvas.polygon(xy=polygon, outline=True, fill=True)
            masks.append(mask)
            categories.append(annotation.find("Attributes/Attribute").attrib["Name"])

    return masks, categories


def process(src: str) -> Iterable[dict[str, Any]]:
    files = list(Path(src).rglob("*.xml"))

    for file in tqdm(files):
        try:
            image = np.asarray(Image.open(file.with_suffix(".tif")))
        except FileNotFoundError:
            image = pyvips.Image.new_from_file(file.with_suffix(".svs"))
            image = image.numpy()

        masks, categories = get_masks(file, mask_size=(image.shape[1], image.shape[0]))
        patient_id = file.parent.stem[:12]

        yield {
            "patient": patient_id,
            "image": Image.fromarray(image.astype(np.uint8)),
            "instances": masks,
            "categories": categories,
            "tissue": patient_data[patient_id],
        }


if __name__ == "__main__":
    train = Dataset.from_generator(
        process,
        gen_kwargs={"src": "data/raw/MoNuSAC/MoNuSAC_images_and_annotations"},
        features=features,
        split=NamedSplit("train"),
        keep_in_memory=True,
    )
    train.push_to_hub("RationAI/MoNuSAC")

    test = Dataset.from_generator(
        process,
        gen_kwargs={"src": "data/raw/MoNuSAC/MoNuSAC Testing Data and Annotations"},
        features=features,
        split=NamedSplit("test"),
        keep_in_memory=True,
    )
    test.push_to_hub("RationAI/MoNuSAC")