"""Validate examples, group duplicates, and load target-specific labels."""

import hashlib
import json
import os
import random
import re
from collections import Counter, defaultdict
from functools import partial
from multiprocessing import Pool
from pathlib import Path

import torch
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms as T

from targets import ALIASES, BACKGROUND, target_name  # noqa: F401  (re-exported)


def read_rgb(path):
    with Image.open(path) as image:
        return image.convert("RGB")


def pixel_hash(path):
    image = read_rgb(path)
    return hashlib.sha256(str(image.size).encode() + image.tobytes()).hexdigest()


def average_hash(path):
    """Survives rescaling and pixel noise, so noisy copies of a tile group together."""
    with Image.open(path) as image:
        return _average_hash(image)


def _average_hash(image):
    pixels = image.convert("L").resize((8, 8), Image.Resampling.BOX).tobytes()
    mean = sum(pixels) / 64
    return "a" + format(int("".join("1" if v > mean else "0" for v in pixels), 2), "016x")


GRID_SIZES = {"3x3": 9, "4x4": 16}
IMAGE_NAMES = ("image.jpg", "image.jpeg", "image.png", "image.webp")


def grid_boxes(width, height, tiles):
    """Cell boxes, left to right then top to bottom, with edges rounded like the labeling tools."""
    n = int(tiles**0.5)
    xs = [round(width * i / n) for i in range(n + 1)]
    ys = [round(height * i / n) for i in range(n + 1)]
    return [(xs[c], ys[r], xs[c + 1], ys[r + 1]) for r in range(n) for c in range(n)]


def image_hash(image):
    return hashlib.sha256(str(image.size).encode() + image.tobytes()).hexdigest()


def record_image_path(root, record):
    """Locate a full image from either a published dataset or a legacy manifest."""
    if "image_path" in record:
        return Path(root) / record["image_path"]
    return Path(root) / record["id"] / record.get("image", "full_images.png")


def prepare_jsonl(root, output):
    """Import the published dataset's existing splits without reshuffling them."""
    root = Path(root).resolve()
    data = root / "data" if (root / "data" / "3x3").is_dir() else root
    records, counts = [], Counter()
    for grid, tiles in GRID_SIZES.items():
        for split in ("train", "validation", "test", "review", "excluded"):
            metadata = data / grid / split / "metadata.jsonl"
            if not metadata.is_file():
                raise ValueError(f"Missing dataset metadata: {metadata}")
            with metadata.open() as source:
                for line_number, line in enumerate(source, 1):
                    try:
                        row = json.loads(line)
                        if row["grid"] != grid or not isinstance(row["target"], str):
                            raise ValueError("invalid grid or target")
                        answers, uncertain = row["correct_answers"], row["uncertain_cells"]
                        if any(not isinstance(x, list) for x in (answers, uncertain)) or any(
                            type(i) is not int or not 0 <= i < tiles for i in answers + uncertain
                        ):
                            raise ValueError("invalid cell indices")
                        path = (metadata.parent / row["file_name"]).resolve()
                        if not path.is_relative_to(metadata.parent.resolve()) or not path.is_file():
                            raise ValueError("invalid or missing image")
                        counts[split] += 1
                        if split == "excluded":
                            continue
                        record = {
                            "id": row["id"], "tiles": tiles,
                            "target": target_name(row["target"]),
                            "answers": sorted(set(answers)), "uncertain": sorted(set(uncertain)),
                            "image_path": str(path.relative_to(data)), "tile_files": False,
                            "label_source": row.get("label_source"),
                            "split": "val" if split == "validation" else split,
                        }
                        if split == "review" and isinstance(row.get("machine_answers"), list):
                            record["labeler_answers"] = row["machine_answers"]
                        records.append(record)
                    except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
                        raise ValueError(f"Invalid record in {metadata}:{line_number}: {exc}") from exc
    if not records:
        raise ValueError(f"No labeled examples in {data}")
    manifest = {"root": str(data), "records": records}
    output = Path(output)
    output.mkdir(parents=True, exist_ok=True)
    (output / "split.json").write_text(json.dumps(manifest, indent=2) + "\n")
    summary = {"splits": dict(counts), "uncertain_cells": sum(len(r["uncertain"]) for r in records)}
    (output / "data_summary.json").write_text(json.dumps(summary, indent=2) + "\n")
    return manifest, summary


def _scan_labeled(directory, sam=False):
    """One <item>/{image.jpg, metadata.json} folder: (record, None) or (None, reason).
    With sam, 4x4 answers come from <item>/sam.json instead."""
    try:
        info = json.loads((directory / "metadata.json").read_text())
    except (OSError, json.JSONDecodeError):
        return None, "unreadable metadata.json"
    if not isinstance(info, dict):
        return None, "invalid metadata.json"
    if not isinstance(info.get("target"), str) or not info["target"].strip():
        return None, "invalid target"
    n = GRID_SIZES.get(str(info.get("grid", "")).lower())
    if n is None:
        return None, "unknown grid"
    if sam and n == 16:
        try:
            label = json.loads((directory / "sam.json").read_text())
        except (OSError, json.JSONDecodeError):
            return None, "no sam.json"
        if not isinstance(label, dict) or not {"correct_answers", "model"} <= label.keys():
            return None, "invalid sam.json"
        # SAM gives no uncertain cells; its cutoffs were tuned on reference grids.
        info = {
            **info,
            "correct_answers": label["correct_answers"],
            "uncertain_cells": [],
            "label_source": f"{label['model']}:sam",
        }
    if "correct_answers" not in info:
        return None, "not labeled"
    answers, uncertain = info["correct_answers"], info.get("uncertain_cells", [])
    if (
        not isinstance(answers, list)
        or not isinstance(uncertain, list)
        or any(type(i) is not int or not 0 <= i < n for i in answers + uncertain)
    ):
        return None, "invalid answers"
    path = next((directory / name for name in IMAGE_NAMES if (directory / name).exists()), None)
    if path is None:
        return None, "no image"
    try:
        image = read_rgb(path)
    except OSError:
        return None, "unreadable image"
    hashes = [image_hash(image)]
    if n == 9:
        # Separate photos get reused across grids; hashing each keeps reused photos in one split.
        hashes += [image_hash(image.crop(box)) for box in grid_boxes(*image.size, n)]
    return {
        "id": directory.name,
        "tiles": n,
        "target": target_name(info["target"]),
        "answers": sorted(set(answers)),
        "uncertain": sorted(set(uncertain)),
        "image": path.name,
        "tile_files": False,
        "label_source": info.get("label_source"),
        "hashes": hashes,
    }, None


def scan(root, workers=8, sam=False):
    """Answer-key folders (info.json + full_images.png + tile_N.png) and labeled
    <item>/{image.jpg, metadata.json} folders. Returns (records, {skip reason: count})."""
    if workers < 1:
        raise ValueError("Scanning workers must be positive")
    directories = sorted(
        p for p in Path(root).iterdir() if p.is_dir() and not p.name.startswith(".")
    )
    labeled = [d for d in directories if not (d / "info.json").exists()]
    records, skipped = [], Counter()
    if labeled:
        with Pool(workers) as pool:
            for record, reason in pool.imap(partial(_scan_labeled, sam=sam), labeled, chunksize=64):
                if record:
                    records.append(record)
                else:
                    skipped[reason] += 1
    for directory in directories:
        if not (directory / "info.json").exists():
            continue
        info = json.loads((directory / "info.json").read_text())
        indices = sorted(int(p.stem.split("_")[1]) for p in directory.glob("tile_*.png"))
        n = len(indices)
        if n not in (9, 16) or indices != list(range(n)):
            raise ValueError(f"Invalid tiles: {directory}")
        answers = info["correct_answers"]
        if (
            not isinstance(answers, list)
            or any(type(i) is not int or not 0 <= i < n for i in answers)
            or len(set(answers)) != len(answers)
        ):
            raise ValueError(f"Invalid answers: {directory}")
        paths = [directory / "full_images.png"] + [directory / f"tile_{i}.png" for i in range(n)]
        records.append(
            {
                "id": directory.name,
                "tiles": n,
                "target": target_name(info["keyword"]),
                "answers": sorted(answers),
                "hashes": [pixel_hash(p) for p in paths],
            }
        )
    if not records:
        raise ValueError(f"No labeled examples in {root} ({dict(skipped)})")
    records.sort(key=lambda r: r["id"])
    return records, dict(skipped)


def split_records(records, seed, holdouts=("val", "test"), review=()):
    """Records whose id is in review go to a "review" split; unreviewed records sharing an image
    or tile with them get split "dropped" so nothing they show leaks into training."""
    # Connected components prevent leakage even when only one tile is repeated.
    review = set(review)
    parent = list(range(len(records)))

    def find(i):
        while parent[i] != i:
            parent[i] = parent[parent[i]]
            i = parent[i]
        return i

    seen = {}
    for i, record in enumerate(records):
        for digest in record["hashes"]:
            if digest in seen:
                parent[find(i)] = find(seen[digest])
            seen[digest] = i
    groups = defaultdict(list)
    for i, record in enumerate(records):
        groups[find(i)].append(record)
    strata = defaultdict(list)
    for group in groups.values():
        if review and any(r.get("id") in review for r in group):
            for record in group:
                record["split"] = "review" if record.get("id") in review else "dropped"
            continue
        strata[tuple(sorted({(r["tiles"], r["target"]) for r in group}))].append(group)
    rng = random.Random(seed)
    for key in sorted(strata):
        items = strata[key]
        rng.shuffle(items)
        # Fewer than three independent examples cannot support three splits.
        holdout = max(1, round(len(items) * 0.1)) if len(items) >= len(holdouts) + 1 else 0
        for i, group in enumerate(items):
            split = holdouts[i // holdout] if i < holdout * len(holdouts) else "train"
            for record in group:
                record["split"] = split
    return records


def prepare(root, output, seed=42, workers=8, sam=False, review_keys=()):
    """review_keys: reference answer files ({"items": {id: {"answers": [...]}}}). Those grids
    get the reviewed answers and their own "review" split, never used for training or selection."""
    output = Path(output)
    output.mkdir(parents=True, exist_ok=True)
    records, skipped = scan(root, workers, sam)
    review = {}
    for path in review_keys:
        review.update(json.loads(Path(path).read_text())["items"])
    for record in records:
        if record["id"] in review:
            # Keep the machine label so evaluation can score the labeler on the same grids.
            record.update(
                labeler_answers=record["answers"],
                answers=review[record["id"]]["answers"],
                uncertain=[],
                label_source="review",
            )
    # Quarantine contradictory answer keys without changing source files.
    duplicates = defaultdict(list)
    for record in records:
        if record["id"] in review:
            continue
        key = (record["hashes"][0], record["target"], record["tiles"])
        duplicates[key].append(record)
    excluded = [
        r
        for group in duplicates.values()
        if len({tuple(r["answers"]) for r in group}) > 1
        for r in group
    ]
    excluded_ids = {r["id"] for r in excluded}
    records = split_records(
        [r for r in records if r["id"] not in excluded_ids], seed, review=review
    )
    dropped = sum(r["split"] == "dropped" for r in records)
    records = [r for r in records if r["split"] != "dropped"]
    manifest = {
        "seed": seed,
        "root": str(Path(root).resolve()),
        "records": records,
        "excluded_conflicting_labels": excluded,
    }
    (output / "label_conflicts.json").write_text(json.dumps(excluded, indent=2) + "\n")
    (output / "split.json").write_text(json.dumps(manifest, indent=2) + "\n")
    counts = {}
    for n in (9, 16):
        counts[str(n)] = {
            target: dict(
                Counter(r["split"] for r in records if r["tiles"] == n and r["target"] == target)
            )
            for target in sorted({r["target"] for r in records if r["tiles"] == n})
        }
    counts["skipped"] = skipped
    counts["excluded_conflicting_labels"] = len(excluded)
    counts["dropped_sharing_images_with_review"] = dropped
    counts["uncertain_cells"] = sum(len(r.get("uncertain", [])) for r in records)
    (output / "data_summary.json").write_text(json.dumps(counts, indent=2) + "\n")
    return manifest, counts


def _file_hashes(path):
    with Image.open(path) as image:
        width, height = image.size
        thumbnail = image.convert("L").resize((32, 32), Image.Resampling.BOX).tobytes()
    return str(path), pixel_hash(path), average_hash(path), max(width, height), thumbnail


def _held_out_hashes(item):
    root, record = item
    directory = root / record["id"]
    image = read_rgb(record_image_path(root, record))
    pictures = [image]
    if record.get("tile_files", True):
        pictures.extend(read_rgb(directory / f"tile_{i}.png") for i in range(record["tiles"]))
    else:
        pictures.extend(image.crop(box) for box in grid_boxes(*image.size, record["tiles"]))
    return [
        (
            _average_hash(picture),
            picture.convert("L").resize((32, 32), Image.Resampling.BOX).tobytes(),
        )
        for picture in pictures
    ]


def same_picture(a, b, tolerance=12):
    """Average hashes collide for similar layouts; confirm with a 32x32 mean absolute difference."""
    return sum(abs(x - y) for x, y in zip(a, b)) / len(a) < tolerance


IMAGE_SUFFIXES = (".png", ".jpg", ".jpeg", ".webp")
# One YOLO box: class id, then centre x/y and width/height with six decimals. Some files
# lost their line breaks, gluing the next class id onto the previous number.
BOX = re.compile(r"(\d+)\s*(\d\.\d{6})\s+(\d\.\d{6})\s+(\d\.\d{6})\s+(\d\.\d{6})")


def label_file(path):
    """YOLO layout: <root>/images/<Class>/x.png pairs with <root>/labels/<Class>/x.txt."""
    parts = path.parts
    if "images" not in parts:
        return None
    i = len(parts) - 1 - parts[::-1].index("images")
    candidate = Path(*parts[:i], "labels", *parts[i + 1 :]).with_suffix(".txt")
    return candidate if candidate.exists() else None


def box_classes(path):
    return sorted({int(match[0]) for match in BOX.findall(path.read_text())})


def scan_folders(roots):
    """Images labeled by their parent folder; <split>/ and images/ levels are ignored.

    YOLO class ids are mapped to the folder where each id appears most often, so boxes
    for other classes add those targets to an image.
    """
    paths, labels, boxes = [], [], []
    for root in roots:
        for path in sorted(Path(root).rglob("*")):
            if path.suffix.lower() in IMAGE_SUFFIXES and not path.name.startswith("."):
                paths.append(path)
                labels.append(target_name(path.parent.name))
                label = label_file(path)
                boxes.append(box_classes(label) if label else [])
    votes = defaultdict(Counter)
    for label, ids in zip(labels, boxes):
        for i in ids:
            votes[i][label] += 1
    id_names = {i: counts.most_common(1)[0][0] for i, counts in votes.items()}
    targets = [sorted({label, *(id_names[i] for i in ids)}) for label, ids in zip(labels, boxes)]
    return paths, targets, {str(i): name for i, name in sorted(id_names.items())}


def prepare_folders(roots, old_manifest, output, seed=42, workers=8):
    """Merge exact copies, group near-copies, and re-split one or more class-folder datasets.

    The folders' own train/val divisions are ignored because many images appear in both.
    Images resembling held-out grids of the old dataset are dropped.
    """
    roots = [Path(r).resolve() for r in ([roots] if isinstance(roots, (str, Path)) else roots)]
    base = Path(os.path.commonpath(roots)) if len(roots) > 1 else roots[0]
    output = Path(output)
    output.mkdir(parents=True, exist_ok=True)
    paths, path_targets, id_names = scan_folders(roots)
    old = json.loads(Path(old_manifest).read_text())
    old_root = Path(old["root"])
    held_out = [(old_root, record) for record in old["records"] if record["split"] != "train"]
    with Pool(workers) as pool:
        hashed = pool.map(_file_hashes, paths, chunksize=64)
        held_out_hashes = defaultdict(list)
        for pictures in pool.map(_held_out_hashes, held_out, chunksize=64):
            for ahash, thumbnail in pictures:
                held_out_hashes[ahash].append(thumbnail)
    by_pixels, thumbnails = {}, {}
    sources = Counter()
    for path, labels, (_, digest, ahash, side, thumbnail) in zip(paths, path_targets, hashed):
        thumbnails[digest] = thumbnail
        relative = path.relative_to(base)
        if digest not in by_pixels:
            sources[relative.parts[0] if len(roots) > 1 else roots[0].name] += 1
        record = by_pixels.setdefault(
            digest,
            {
                "path": str(relative),
                "targets": [],
                "kind": "full" if side >= 400 else "tile",
                "tiles": "full" if side >= 400 else "tile",
                "hashes": [digest, ahash],
            },
        )
        record["targets"].extend(t for t in labels if t not in record["targets"])
    records, leaked, conflicts = [], [], []
    for record in by_pixels.values():
        record["targets"].sort()
        if any(
            same_picture(thumbnails[record["hashes"][0]], other)
            for other in held_out_hashes.get(record["hashes"][1], [])
        ):
            leaked.append(record["path"])
        else:
            if BACKGROUND in record["targets"] and len(record["targets"]) > 1:
                # Images moved out of "Other" into a target folder mostly do show that target.
                conflicts.append(record["path"])
                record["targets"].remove(BACKGROUND)
            record["target"] = record["targets"][0]
            records.append(record)
    if not records:
        raise ValueError("No folder images remain after excluding held-out images")
    split_records(records, seed, holdouts=("val",))
    for record in records:
        del record["target"], record["tiles"]
    summary = {
        "roots": [str(r) for r in roots],
        "files": len(paths),
        "unique_images": len(by_pixels),
        "unique_images_by_first_source": dict(sources),
        "yolo_class_ids": id_names,
        "exact_duplicates_merged": len(paths) - len(by_pixels),
        "multi_label_images": sum(len(r["targets"]) > 1 for r in records),
        "dropped_matching_old_heldout": len(leaked),
        "background_label_overridden_by_target": len(conflicts),
        "largest_near_duplicate_group": max(Counter(r["hashes"][1] for r in records).values()),
        "counts": {
            kind: {
                target: dict(
                    Counter(
                        r["split"] for r in records if r["kind"] == kind and target in r["targets"]
                    )
                )
                for target in sorted({t for r in records for t in r["targets"]})
            }
            for kind in ("tile", "full")
        },
    }
    manifest = {
        "seed": seed,
        "root": str(base),
        "sources": [str(r) for r in roots],
        "old_manifest": str(Path(old_manifest).resolve()),
        "records": records,
        "dropped_matching_old_heldout": leaked,
        "background_label_overridden_by_target": conflicts,
    }
    (output / "folders.json").write_text(json.dumps(manifest, indent=2) + "\n")
    (output / "folders_summary.json").write_text(json.dumps(summary, indent=2) + "\n")
    return manifest, summary


class AddNoise:
    """Some real tiles carry heavy pixel noise."""

    def __init__(self, probability=0.3, max_std=0.15):
        self.probability, self.max_std = probability, max_std

    def __call__(self, image):
        if random.random() >= self.probability:
            return image
        return (image + torch.randn_like(image) * random.uniform(0.02, self.max_std)).clamp(0, 1)


def transform(size, training=False, noise=False):
    operations = [T.Resize((size, size), interpolation=T.InterpolationMode.BILINEAR)]
    if training:
        operations.append(T.ColorJitter(brightness=0.15, contrast=0.15, saturation=0.1))
    operations.append(T.ToTensor())
    if training and noise:
        operations.append(AddNoise())
    return T.Compose(operations + [T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])


class GridDataset(Dataset):
    def __init__(
        self,
        root,
        records,
        classes,
        size,
        training=False,
        tiles=False,
        tile_grids=False,
        include_full_images=False,
    ):
        self.root = Path(root)
        self.records = records
        self.classes = classes
        self.transform = transform(size, training)
        self.tiles = tiles
        self.tile_grids = tile_grids
        self.items = [(r, i) for r in records for i in (range(r["tiles"]) if tiles else [None])]
        if include_full_images:
            if not tiles:
                raise ValueError("Whole-image classification examples require tile training")
            self.items += [(r, -1) for r in records if r["tiles"] == 16]

    def __len__(self):
        return len(self.items)

    def full_image(self, record):
        return read_rgb(record_image_path(self.root, record))

    def tile_images(self, record, indices):
        if record.get("tile_files", True):
            return [read_rgb(self.root / record["id"] / f"tile_{i}.png") for i in indices]
        image = self.full_image(record)
        boxes = grid_boxes(*image.size, record["tiles"])
        return [image.crop(boxes[i]) for i in indices]

    def __getitem__(self, index):
        record, tile = self.items[index]
        label = torch.tensor([float(i in record["answers"]) for i in range(record["tiles"])])
        if tile == -1:
            image = self.transform(self.full_image(record))
            label = label.any().float()
        elif tile is not None:
            image = self.transform(self.tile_images(record, [tile])[0])
            label = label[tile]
        elif record["tiles"] == 9 or self.tile_grids:
            image = torch.stack(
                [self.transform(t) for t in self.tile_images(record, range(record["tiles"]))]
            )
        else:
            image = self.transform(self.full_image(record))
        return image, self.classes.index(record["target"]), label


def supervision(classes, cells, target, cell_labels=None, any_label=None):
    """Labels and weights for (classes, cells) logits; zero weight means unknown."""
    labels, weights = torch.zeros(len(classes), cells), torch.zeros(len(classes), cells)
    any_labels, any_weights = torch.zeros(len(classes)), torch.zeros(len(classes))
    if cell_labels is not None:
        labels[target], weights[target] = cell_labels, 1
    if any_label is not None:
        any_labels[target], any_weights[target] = any_label, 1
    return labels, weights, any_labels, any_weights


class SupervisedGrids(Dataset):
    """Answer-key items in the per-class format: only the queried target is known.
    Cells listed as uncertain get zero weight unless keep_uncertain is set."""

    def __init__(self, grids, cells, keep_uncertain=False, flip=False):
        self.grids, self.cells, self.keep_uncertain, self.flip = grids, cells, keep_uncertain, flip

    def __len__(self):
        return len(self.grids)

    def __getitem__(self, index):
        record, tile = self.grids.items[index]
        image, target, label = self.grids[index]
        if tile == -1:
            return (image, *supervision(self.grids.classes, self.cells, target, any_label=label))
        uncertain = [] if self.keep_uncertain else record.get("uncertain", [])
        if self.flip and random.random() < 0.5:
            # Mirror left-right; a whole grid also mirrors its cell numbers.
            image = image.flip(-1)
            if tile is None:
                n = int(label.numel() ** 0.5)
                label = label.view(n, n).flip(-1).flatten()
                uncertain = [c - c % n + n - 1 - c % n for c in uncertain]
        labels, weights, any_labels, any_weights = supervision(
            self.grids.classes, self.cells, target, cell_labels=label
        )
        if tile is not None:
            if tile in uncertain:
                weights.zero_()
        else:
            for cell in uncertain:
                weights[target, cell] = 0
        return image, labels, weights, any_labels, any_weights


class FolderDataset(Dataset):
    """Single-label folder images: present for named targets, others assumed absent at low weight.

    A folder label says the image contains that target. Another target may also be visible,
    so absence is only weakly supervised, except for background images, which contain none.
    """

    def __init__(self, root, records, classes, size, cells=1, training=False, negative_weight=0.2):
        self.root = Path(root)
        self.records = records
        self.classes = classes
        self.cells = cells
        self.negative_weight = negative_weight
        self.transform = transform(size, training, noise=True)

    def __len__(self):
        return len(self.records)

    def __getitem__(self, index):
        record = self.records[index]
        image = self.transform(read_rgb(self.root / record["path"]))
        background = BACKGROUND in record["targets"]
        any_weights = torch.full((len(self.classes),), 1.0 if background else self.negative_weight)
        any_labels = torch.zeros(len(self.classes))
        for target in record["targets"]:
            if target in self.classes:
                any_labels[self.classes.index(target)] = any_weights[self.classes.index(target)] = 1
        empty = torch.zeros(len(self.classes), self.cells)
        return image, empty, empty.clone(), any_labels, any_weights