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
from datasets import Dataset
from datasets.splits import NamedSplit
from PIL import Image, ImageDraw
from tqdm import tqdm


# https://drive.google.com/file/d/1xYyQ31CHFRnvTCTuuHdconlJCMk2SK7Z/view?usp=sharing
patient_data = {
    "TCGA-A7-A13E-01Z-00-DX1": "Breast",
    "TCGA-A7-A13F-01Z-00-DX1": "Breast",
    "TCGA-AR-A1AK-01Z-00-DX1": "Breast",
    "TCGA-AR-A1AS-01Z-00-DX1": "Breast",
    "TCGA-E2-A1B5-01Z-00-DX1": "Breast",
    "TCGA-E2-A14V-01Z-00-DX1": "Breast",
    "TCGA-B0-5711-01Z-00-DX1": "Kidney",
    "TCGA-HE-7128-01Z-00-DX1": "Kidney",
    "TCGA-HE-7129-01Z-00-DX1": "Kidney",
    "TCGA-HE-7130-01Z-00-DX1": "Kidney",
    "TCGA-B0-5710-01Z-00-DX1": "Kidney",
    "TCGA-B0-5698-01Z-00-DX1": "Kidney",
    "TCGA-18-5592-01Z-00-DX1": "Liver",
    "TCGA-38-6178-01Z-00-DX1": "Liver",
    "TCGA-49-4488-01Z-00-DX1": "Liver",
    "TCGA-50-5931-01Z-00-DX1": "Liver",
    "TCGA-21-5784-01Z-00-DX1": "Liver",
    "TCGA-21-5786-01Z-00-DX1": "Liver",
    "TCGA-G9-6336-01Z-00-DX1": "Prostate",
    "TCGA-G9-6348-01Z-00-DX1": "Prostate",
    "TCGA-G9-6356-01Z-00-DX1": "Prostate",
    "TCGA-G9-6363-01Z-00-DX1": "Prostate",
    "TCGA-CH-5767-01Z-00-DX1": "Prostate",
    "TCGA-G9-6362-01Z-00-DX1": "Prostate",
    "TCGA-DK-A2I6-01A-01-TS1": "Bladder",
    "TCGA-G2-A2EK-01A-02-TSB": "Bladder",
    "TCGA-AY-A8YK-01A-01-TS1": "Colon",
    "TCGA-NH-A8F7-01A-01-TS1": "Colon",
    "TCGA-KB-A93J-01A-01-TS1": "Stomach",
    "TCGA-RD-A8N9-01A-01-TS1": "Stomach",
}


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

    for region in ET.parse(path).getroot().findall("Annotation/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)

    return masks


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

    for file in tqdm(files):
        masks = get_masks(file, mask_size=(1000, 1000))
        tissue_path = Path(str(file).replace("Annotations", "Tissue Images"))
        image = np.asarray(Image.open(tissue_path.with_suffix(".tif")))

        yield {
            "patient": file.stem,
            "image": Image.fromarray(image.astype(np.uint8)),
            "instances": masks,
            "tissue": patient_data.get(file.stem, "Unknown"),
        }


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

    for file in tqdm(files):
        masks = get_masks(file, mask_size=(1000, 1000))
        image = np.asarray(Image.open(file.with_suffix(".tif")))

        yield {
            "patient": file.stem,
            "image": Image.fromarray(image.astype(np.uint8)),
            "instances": masks,
            "tissue": patient_data.get(file.stem, "Unknown"),
        }


features = datasets.Features(
    {
        "patient": datasets.Value("string"),
        "image": datasets.Image(mode="RGB"),
        "instances": datasets.Sequence(datasets.Image(mode="1")),
        "tissue": datasets.ClassLabel(
            names=[
                "Unknown",
                "Breast",
                "Kidney",
                "Liver",
                "Prostate",
                "Bladder",
                "Colon",
                "Stomach",
            ]
        ),
    }
)


if __name__ == "__main__":
    train = Dataset.from_generator(
        process_train,
        gen_kwargs={"src": "data/raw/MoNuSeg/MoNuSeg 2018 Training Data/Annotations"},
        features=features,
        split=NamedSplit("train"),
        keep_in_memory=True,
    )
    train.push_to_hub("RationAI/MoNuSeg")

    test = Dataset.from_generator(
        process_test,
        gen_kwargs={"src": "data/raw/MoNuSeg/MoNuSegTestData"},
        features=features,
        split=NamedSplit("test"),
        keep_in_memory=True,
    )
    test.push_to_hub("RationAI/MoNuSeg")