#!/usr/bin/env python3
"""Booru Lighting Tagger: lighting tags for anime-style illustrations.

Command line (run from anywhere; the model files are read from this script's folder):

    python inference.py image.png other.jpg            # print the tags
    python inference.py dataset/ --recursive --append-to-captions
    python inference.py dataset/ --jsonl results.jsonl  # every score, for your own thresholds
    python inference.py dataset/ -r --append-to-captions --preset recall   # more tags, a little less precise

Python:

    from inference import LightingTagger
    tagger = LightingTagger()                    # model files next to this script
    for result in tagger.tag_files(["image.png"]):
        print(result["tags"], result["scores"]["backlighting"])

The model is three fine-tuned copies ("seeds") of SmilingWolf/wd-eva02-large-tagger-v3. Each weight file
holds only the parameters that were trained (last 16 transformer blocks, final norm, new head); everything
else is loaded from the base model, pinned to the revision used for training. Scores are averaged over the
seeds. See README.md for the labels and their reliability.

Threshold presets: "default" emits as many tags as the teacher did on the validation images; "recall" emits
about 1.5 times as many (see README.md, "Threshold presets").
"""

from __future__ import annotations

import argparse
import contextlib
import copy
import json
import sys
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Iterable, Iterator

import numpy as np
import torch
from PIL import Image

CONFIG_FILE = "tagger_config.json"
PRESETS = ("default", "recall")
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif", ".tif", ".tiff"}
WHITE = (255, 255, 255)


# --- Preprocessing: identical to training, do not change ------------------------------------------------

def to_rgb(image: Image.Image) -> Image.Image:
    """RGB image; transparency is flattened onto white."""
    if image.mode not in ("RGB", "RGBA"):
        image = image.convert("RGBA" if "transparency" in image.info else "RGB")
    if image.mode == "RGBA":
        background = Image.new("RGBA", image.size, (255, 255, 255, 255))
        image = Image.alpha_composite(background, image).convert("RGB")
    return image


def pad_square(image: Image.Image) -> Image.Image:
    """Pads to a square on white, centered, without cropping: the whole lighting stays visible."""
    side = max(image.size)
    square = Image.new("RGB", (side, side), WHITE)
    square.paste(image, ((side - image.width) // 2, (side - image.height) // 2))
    return square


def prepare(image: Image.Image, transform) -> torch.Tensor:
    """WD preprocessing: white square padding, the base model's timm transform, RGB -> BGR."""
    return transform(pad_square(to_rgb(image)))[[2, 1, 0]]


# --- Model ----------------------------------------------------------------------------------------------

class _Student(torch.nn.Module):
    """WD backbone with a new linear head (parameter names match the weight files)."""

    def __init__(self, backbone, n_outputs: int):
        super().__init__()
        self.backbone = backbone
        self.head = torch.nn.Linear(backbone.num_features, n_outputs)

    def forward(self, x):
        return self.head(self.backbone.forward_head(self.backbone.forward_features(x), pre_logits=True))


def _load_base(repo: str, revision: str):
    """Base WD tagger at the pinned revision, on the CPU, and its timm eval transform."""
    import timm

    ref = f"{repo}@{revision}"
    model = timm.create_model(f"hf-hub:{ref}").eval()
    model.load_state_dict(timm.models.load_state_dict_from_hf(ref))
    transform = timm.data.create_transform(**timm.data.resolve_data_config(model.pretrained_cfg, model=model))
    return model, transform


def _trained_parameters(model: _Student, blocks: int) -> set[str]:
    depth = len(model.backbone.blocks)
    prefixes = tuple(f"backbone.blocks.{i}." for i in range(depth - blocks, depth)) + ("backbone.fc_norm.", "head.")
    return {name for name, _ in model.named_parameters() if name.startswith(prefixes)}


def _load_seed(base, path: Path, n_outputs: int, blocks: int) -> _Student:
    """A copy of the base model with the trained parameters of one seed loaded on top."""
    from safetensors.torch import load_file

    model = _Student(copy.deepcopy(base), n_outputs)
    state = load_file(str(path), device="cpu")
    expected = _trained_parameters(model, blocks)
    if set(state) != expected:
        missing, extra = sorted(expected - set(state)), sorted(set(state) - expected)
        raise ValueError(f"{path.name}: unexpected parameters (missing {missing[:5]}, extra {extra[:5]})")
    model.load_state_dict(state, strict=False)
    return model.eval()


def _decode(logits: np.ndarray, n_labels: int, steps: int) -> tuple[np.ndarray, np.ndarray]:
    """Raw outputs -> (scores 0-10, tonal key probabilities).

    Ordinal target: for each label, `steps` outputs give P(score >= t) for t = 1..steps; the predicted score
    is their sum. The last outputs are the tonal key logits (softmax)."""
    width = n_labels * steps
    sigmoid = 1 / (1 + np.exp(-logits[:, :width].astype(np.float64)))
    scores = sigmoid.reshape(-1, n_labels, steps).sum(-1)
    keys = logits[:, width:].astype(np.float64)
    keys = np.exp(keys - keys.max(1, keepdims=True))
    return scores.astype(np.float32), (keys / keys.sum(1, keepdims=True)).astype(np.float32)


def _resolve_device(device: str) -> torch.device:
    if device == "auto":
        if torch.cuda.is_available():
            device = "cuda"
        elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
            device = "mps"
        else:
            device = "cpu"
    return torch.device(device)


def _autocast_dtype(device: torch.device, precision: str):
    """Mixed precision on CUDA only (bf16, as in training; fp16 on GPUs without bf16)."""
    if device.type != "cuda" or precision == "fp32":
        return None
    if precision == "bf16":
        return torch.bfloat16
    if precision == "fp16":
        return torch.float16
    return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16


class LightingTagger:
    """Loads the seeds once; `tag_images` and `tag_files` return one result dict per image."""

    def __init__(self, model_dir: str | Path | None = None, device: str = "auto", seeds: int | None = None,
                 precision: str = "auto", preset: str = "default"):
        self.model_dir = Path(model_dir) if model_dir else Path(__file__).resolve().parent
        config_path = self.model_dir / CONFIG_FILE
        if not config_path.exists():
            raise FileNotFoundError(f"{config_path} not found: pass the folder that holds the model files")
        self.config = json.loads(config_path.read_text(encoding="utf-8"))
        self.version = self.config["version"]
        self.labels = [label["name"] for label in self.config["labels"]]
        # Configs older than v1.0.1 have no recall thresholds: the default ones are used instead.
        self.presets = {
            "default": np.array([label["threshold"] for label in self.config["labels"]], dtype=np.float32),
            "recall": np.array([label.get("threshold_recall", label["threshold"]) for label in self.config["labels"]],
                               dtype=np.float32),
        }
        if preset not in self.presets:
            raise ValueError(f"preset must be one of {PRESETS}")
        self.preset = preset
        self.thresholds = self.presets[preset]
        self.status = {label["name"]: label["reliability"]["status"] for label in self.config["labels"]}
        self.keys = [key["name"] for key in self.config["tonal_keys"]]
        self.steps = int(self.config["ordinal_steps"])
        self.device = _resolve_device(device)
        self.dtype = _autocast_dtype(self.device, precision)

        all_seeds = self.config["seeds"]
        if seeds is not None and not 1 <= seeds <= len(all_seeds):
            raise ValueError(f"seeds must be between 1 and {len(all_seeds)}")
        base_model = self.config["base_model"]
        base, self.transform = _load_base(base_model["repo"], base_model["revision"])
        n_outputs = len(self.labels) * self.steps + len(self.keys)
        self.models = [_load_seed(base, self.model_dir / seed["file"], n_outputs, self.config["trained_blocks"])
                       .to(self.device) for seed in all_seeds[:seeds]]
        del base

    @classmethod
    def from_pretrained(cls, repo_id: str, revision: str | None = None, **kwargs) -> LightingTagger:
        """Downloads the model from the Hugging Face Hub (cached), then loads it."""
        from huggingface_hub import snapshot_download

        return cls(snapshot_download(repo_id, revision=revision), **kwargs)

    # --- Inference --------------------------------------------------------------------------------------

    @torch.inference_mode()
    def predict(self, batch: torch.Tensor) -> tuple[np.ndarray, np.ndarray]:
        """Preprocessed batch -> (scores: images x labels, tonal key probabilities: images x 4), seed average."""
        batch = batch.to(self.device)
        scores, keys = [], []
        for model in self.models:
            autocast = (torch.autocast(device_type="cuda", dtype=self.dtype) if self.dtype is not None
                        else contextlib.nullcontext())
            with autocast:
                logits = model(batch).float().cpu().numpy()
            s, k = _decode(logits, len(self.labels), self.steps)
            scores.append(s)
            keys.append(k)
        return np.mean(scores, axis=0), np.mean(keys, axis=0)

    def format(self, scores: np.ndarray, key_probs: np.ndarray, validated_only: bool = False,
               tonal_key_tag: bool = True, preset: str | None = None) -> dict:
        """One image: tags (by descending score, then the tonal key), all scores and the tonal key.
        `preset` overrides the tagger's threshold preset for this call ("default" or "recall")."""
        thresholds = self.thresholds if preset is None else self.presets[preset]
        tagged = [k for k in np.flatnonzero(scores >= thresholds)
                  if not validated_only or self.status[self.labels[k]] == "validated"]
        tags = [self.labels[k] for k in sorted(tagged, key=lambda k: -scores[k])]
        key = self.keys[int(np.argmax(key_probs))]
        if tonal_key_tag and key != "normal":
            tags.append(key)
        return {
            "tags": tags,
            "scores": {name: round(float(s), 3) for name, s in zip(self.labels, scores)},
            "tonal_key": key,
            "tonal_key_probabilities": {name: round(float(p), 4) for name, p in zip(self.keys, key_probs)},
        }

    def tag_images(self, images: list[Image.Image], **options) -> list[dict]:
        """PIL images -> one result per image (see `format`)."""
        if not images:
            return []
        scores, keys = self.predict(torch.stack([prepare(image, self.transform) for image in images]))
        return [self.format(s, k, **options) for s, k in zip(scores, keys)]

    def tag_files(self, paths: Iterable[str | Path], batch_size: int = 8, workers: int = 4,
                  **options) -> Iterator[dict]:
        """Image files -> one result per file, in order, as they are computed. An unreadable file gets an
        `error` entry and does not stop the run."""
        def load(path: Path):
            try:
                with Image.open(path) as image:
                    return prepare(image, self.transform)
            except Exception as exc:  # noqa: BLE001: report and go on
                return f"{type(exc).__name__}: {exc}"

        paths = [Path(p) for p in paths]
        with ThreadPoolExecutor(max_workers=max(1, workers)) as pool:
            for start in range(0, len(paths), batch_size):
                chunk = paths[start:start + batch_size]
                loaded = list(pool.map(load, chunk))
                ok = [k for k, item in enumerate(loaded) if not isinstance(item, str)]
                results = {}
                if ok:
                    scores, keys = self.predict(torch.stack([loaded[k] for k in ok]))
                    results = {k: self.format(s, p, **options) for k, s, p in zip(ok, scores, keys)}
                for k, path in enumerate(chunk):
                    if k in results:
                        yield {"file": str(path), **results[k], "error": None}
                    else:
                        yield {"file": str(path), "tags": [], "scores": None, "tonal_key": None,
                               "tonal_key_probabilities": None, "error": loaded[k]}


# --- Captions -------------------------------------------------------------------------------------------

def _normalize(tag: str) -> str:
    return tag.strip().lower().replace("_", " ")


def caption_path(image: Path, extension: str) -> Path:
    return image.parent / (image.stem + extension)


def append_to_caption(path: Path, tags: list[str]) -> int:
    """Adds the tags missing from a comma-separated caption file (created if needed). Returns the count."""
    text = path.read_text(encoding="utf-8") if path.exists() else ""
    present = {_normalize(tag) for tag in text.split(",") if tag.strip()}
    new = [tag for tag in tags if _normalize(tag) not in present]
    if new:
        text = text.rstrip()
        path.write_text((text + ", " if text else "") + ", ".join(new), encoding="utf-8")
    return len(new)


# --- Command line ---------------------------------------------------------------------------------------

def find_images(inputs: list[str], recursive: bool) -> list[Path]:
    paths = []
    for item in map(Path, inputs):
        if item.is_dir():
            pattern = "**/*" if recursive else "*"
            paths += sorted(p for p in item.glob(pattern) if p.is_file() and p.suffix.lower() in IMAGE_EXTENSIONS)
        elif item.is_file():
            paths.append(item)
        else:
            print(f"not found: {item}", file=sys.stderr)
    return paths


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Booru Lighting Tagger: adds lighting tags (backlighting, rim_light, sunset...) to images.")
    parser.add_argument("inputs", nargs="+", help="image files and/or folders")
    parser.add_argument("-r", "--recursive", action="store_true", help="look into subfolders")
    parser.add_argument("--jsonl", type=Path, help="write every result (all scores) to this JSON Lines file")
    parser.add_argument("--append-to-captions", nargs="?", const=".txt", metavar="EXT",
                        help="add the missing lighting tags to each image's caption file "
                             "(<image name>EXT, default .txt; created if missing)")
    parser.add_argument("--write-captions", nargs="?", const=".lighting.txt", metavar="EXT",
                        help="write the lighting tags alone to <image name>EXT (default .lighting.txt), "
                             "replacing that file")
    parser.add_argument("--spaces", action="store_true",
                        help="write tags with spaces instead of underscores (rim light, not rim_light)")
    parser.add_argument("--validated-only", action="store_true",
                        help="only emit the labels validated against human answers (see README)")
    parser.add_argument("--preset", choices=PRESETS, default="default",
                        help="threshold preset: default (as many tags as the teacher) or recall (about 1.5 times "
                             "as many tags, a little less precise; see README)")
    parser.add_argument("--no-tonal-key", action="store_true",
                        help="do not emit the tonal key tag (dark, dim_lighting, high_key)")
    parser.add_argument("--seeds", type=int, help="number of seeds to use, 1 to 3 (default 3; fewer is faster "
                                                   "and slightly less accurate)")
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--workers", type=int, default=4, help="image decoding threads")
    parser.add_argument("--device", default="auto", help="auto, cuda, cuda:1, mps or cpu")
    parser.add_argument("--precision", default="auto", choices=("auto", "bf16", "fp16", "fp32"),
                        help="CUDA only (auto: bf16 when supported, as in training)")
    parser.add_argument("--model-dir", type=Path, help="folder with the model files (default: this script's)")
    parser.add_argument("--repo", help="download the model from this Hugging Face repo instead")
    parser.add_argument("--revision", help="with --repo: a tag such as v1.0.0, a branch or a commit")
    parser.add_argument("-q", "--quiet", action="store_true", help="do not print a line per image")
    args = parser.parse_args(argv)
    for stream in (sys.stdout, sys.stderr):
        with contextlib.suppress(AttributeError, ValueError):
            stream.reconfigure(encoding="utf-8", errors="replace")

    paths = find_images(args.inputs, args.recursive)
    if not paths:
        print("no image found", file=sys.stderr)
        return 1
    options = {"device": args.device, "seeds": args.seeds, "precision": args.precision, "preset": args.preset}
    tagger = (LightingTagger.from_pretrained(args.repo, revision=args.revision, **options) if args.repo
              else LightingTagger(args.model_dir, **options))
    print(f"booru-lighting-tagger {tagger.version}: {len(tagger.models)} seed(s) on {tagger.device}, "
          f"{tagger.preset} thresholds, {len(paths)} image(s)", file=sys.stderr)

    jsonl = args.jsonl.open("w", encoding="utf-8") if args.jsonl else None
    done = errors = tagged_images = added = 0
    try:
        for result in tagger.tag_files(paths, batch_size=args.batch_size, workers=args.workers,
                                       validated_only=args.validated_only, tonal_key_tag=not args.no_tonal_key):
            done += 1
            if jsonl:
                jsonl.write(json.dumps(result, ensure_ascii=False) + "\n")
            if result["error"]:
                errors += 1
                print(f"{result['file']}: error: {result['error']}", file=sys.stderr)
                continue
            tags = [tag.replace("_", " ") for tag in result["tags"]] if args.spaces else result["tags"]
            tagged_images += bool(tags)
            image = Path(result["file"])
            if args.append_to_captions:
                added += append_to_caption(caption_path(image, args.append_to_captions), tags)
            if args.write_captions:
                caption_path(image, args.write_captions).write_text(", ".join(tags), encoding="utf-8")
            if not args.quiet:
                print(f"{result['file']}: {', '.join(tags) if tags else '(no lighting tag)'}")
    finally:
        if jsonl:
            jsonl.close()
    summary = f"{done} image(s), {tagged_images} with at least one tag, {errors} error(s)"
    if args.append_to_captions:
        summary += f", {added} tag(s) added to captions"
    print(summary, file=sys.stderr)
    return 0 if errors == 0 else 2


if __name__ == "__main__":
    sys.exit(main())