"""DeepMosaics restorers — real generative mosaic removal, in-process.

The DeepMosaics network code (GPL-3.0) is vendored under ``_deepmosaics/`` (see its
NOTICE/LICENSE). We load the models **once** and run in-process — far faster than
spawning a subprocess per frame (which reloaded the models every time). Only the model
*weights* are user-supplied. Both engines locate the mosaic themselves (BiSeNet
``mosaic_position.pth``); detections are not passed to them.

Two engines:
  - :class:`DeepMosaicsRestorer` (per-frame): reproduces ``cleanmosaic_img_server`` —
    locate mosaic → run the image generator on the crop → feather it back. Image weights
    ``clean_youknow_resnet_9blocks.pth``.
  - :class:`DeepMosaicsVideoRestorer` (temporal/BVDNet): reproduces
    ``cleanmosaic_video_fusion`` — a window of neighbouring frames + recurrence, for
    temporal coherence. Video weights ``clean_youknow_video.pth``; needs a contiguous
    sequence (see :meth:`Restorer.restore_sequence`).

Setup (see README → Восстановление): drop the chosen ``clean_*.pth`` + ``mosaic_position.pth``
into one folder (``models/deepmosaics``) and pick it in the restore dialog.
"""

from __future__ import annotations

import sys
from pathlib import Path
from types import SimpleNamespace

import numpy as np

from ..detection.types import Detection
from .base import (
    CancelCheck,
    Cancelled,
    DetGetter,
    FrameGetter,
    Restorer,
    ResultSink,
)

_VENDOR = Path(__file__).parent / "_deepmosaics"
# Default place to drop DeepMosaics clean weights (gitignored — see models/).
DEFAULT_WEIGHTS_DIR = Path(__file__).resolve().parents[3] / "models" / "deepmosaics"


def discover_models(
    extra_dir: str | None = None, include_video: bool = False
) -> list[tuple[str, str]]:
    """Find usable clean models: (display_name, full_path).

    Scans the bundled ``models/deepmosaics`` folder (plus ``extra_dir`` if given)
    for ``clean_*.pth``. By default the video model is skipped — it can't run
    per-frame; pass ``include_video=True`` for the temporal (BVDNet) engine, which
    *needs* ``clean_*_video.pth``.
    """
    dirs = [DEFAULT_WEIGHTS_DIR]
    if extra_dir:
        dirs.insert(0, Path(extra_dir))
    out: list[tuple[str, str]] = []
    seen: set[str] = set()
    for d in dirs:
        if not d.is_dir():
            continue
        for p in sorted(d.glob("clean_*.pth")):
            if p.name in seen:
                continue
            if "video" in p.name.lower() and not include_video:
                continue
            seen.add(p.name)
            out.append((p.stem, str(p)))
    return out


def _find_mosaic_position(model: Path, dm_dir: str | None) -> Path | None:
    """Locate ``mosaic_position.pth`` (the BiSeNet mosaic locator) for ``model``."""
    candidates = [model.parent / "mosaic_position.pth"]
    if dm_dir:
        candidates.append(Path(dm_dir) / "pretrained_models" / "mosaic" / "mosaic_position.pth")
    return next((p for p in candidates if p.is_file()), None)


def _netg_kind(model_name: str) -> str:
    """Pick DeepMosaics' netG type from the weights filename (see their options.py)."""
    n = model_name.lower()
    if "video" in n:
        raise ValueError(
            "Видеомодель (clean_*_video.pth) не работает покадрово — ей нужен соседний "
            "кадр.\nУкажите картиночную модель clean_youknow_resnet_9blocks.pth."
        )
    if "unet_128" in n:
        return "unet_128"
    if "hd" in n:
        return "HD"
    return "resnet_9blocks"


class DeepMosaicsRestorer(Restorer):
    def __init__(
        self,
        deepmosaics_dir: str | None,   # optional extra dir to find mosaic_position.pth
        model_path: str | None,
        gpu_id: str = "0",
    ) -> None:
        if not model_path or not Path(model_path).is_file():
            discovered = discover_models()  # fall back to a bundled model
            if discovered:
                model_path = discovered[0][1]
            else:
                raise ValueError(
                    "Не найдены веса DeepMosaics (clean_*.pth).\n"
                    "Положите clean_youknow_resnet_9blocks.pth + mosaic_position.pth в "
                    "models/deepmosaics (или выберите в «Восстановление…»). См. README."
                )
        model = Path(model_path)
        self._netg = _netg_kind(model.name)  # raises on a video model
        pos = _find_mosaic_position(model, deepmosaics_dir)
        if pos is None:
            raise ValueError(
                "Рядом с clean-моделью не найден mosaic_position.pth.\n"
                "Положите mosaic_position.pth в ту же папку, что и clean_*.pth. См. README."
            )
        self._model = str(model)
        self._pos = str(pos)
        self._gpu = gpu_id
        self._loaded = False  # models loaded lazily on first restore

    @property
    def name(self) -> str:
        return f"DeepMosaics(gpu={self._gpu})"

    # ------------------------------------------------------------------ engine
    def _ensure_loaded(self) -> None:
        if self._loaded:
            return
        if str(_VENDOR) not in sys.path:
            sys.path.insert(0, str(_VENDOR))  # so vendored `from models/util import …` resolve

        # Fall back to CPU when CUDA isn't available: DeepMosaics calls `.cuda()`
        # whenever gpu_id != "-1", which raises "Torch not compiled with CUDA enabled"
        # on a CPU-only torch build.
        import torch
        if self._gpu != "-1" and not torch.cuda.is_available():
            self._gpu = "-1"

        import util.image_processing as impro  # type: ignore

        from models import loadmodel, runmodel  # type: ignore

        self._runmodel = runmodel
        self._impro = impro
        self._opt = SimpleNamespace(
            gpu_id=self._gpu,
            netG=self._netg,
            model_path=self._model,
            mosaic_position_model_path=self._pos,
            mask_threshold=64,
            all_mosaic_area=False,
            ex_mult=1.5,
            no_feather=False,
            traditional=False,
        )
        self._netM = loadmodel.bisenet(self._opt, "mosaic")
        self._netG = loadmodel.pix2pix(self._opt)
        self._loaded = True

    def restore(
        self,
        image: np.ndarray,
        detections: list[Detection],
        should_cancel: CancelCheck | None = None,
    ) -> np.ndarray:
        if should_cancel is not None and should_cancel():
            raise Cancelled("Восстановление отменено")
        self._ensure_loaded()
        rm, impro, opt = self._runmodel, self._impro, self._opt

        # DeepMosaics' cleanmosaic_img_server, faithfully reproduced.
        x, y, size, mask = rm.get_mosaic_position(image, self._netM, opt)
        if size <= 100:
            return image.copy()  # no mosaic located — leave the frame untouched
        if should_cancel is not None and should_cancel():
            raise Cancelled("Восстановление отменено")
        work = image.copy()
        img_mosaic = work[y - size:y + size, x - size:x + size]
        img_fake = rm.run_pix2pix(img_mosaic, self._netG, opt)
        return impro.replace_mosaic(work, img_fake, mask, x, y, size, opt.no_feather)


class DeepMosaicsVideoRestorer(Restorer):
    """Temporal DeepMosaics (BVDNet) — un-censors using *neighbouring* frames.

    Reproduces DeepMosaics' ``cleanmosaic_video_fusion`` per target frame: for frame
    ``i`` it feeds the network a temporal window of ``T`` frames (sampled at step ``S``
    around ``i``) plus its own previous output (recurrent), so the reconstruction is
    temporally coherent. Because of that recurrence the frames MUST be processed in
    order over a contiguous range — see :meth:`restore_sequence` (the batch run).

    Needs the **video** weights ``clean_youknow_video.pth`` + ``mosaic_position.pth``
    (beside it). Single-frame :meth:`restore` degrades to a window of the same frame.
    """

    temporal = True

    # DeepMosaics fusion window: N before/after at step S → T = 2N+1 frames, INPUT_SIZE px.
    _N, _T, _S = 2, 5, 3
    _INPUT_SIZE = 256

    def __init__(
        self,
        deepmosaics_dir: str | None,
        model_path: str | None,
        gpu_id: str = "0",
        feed_restored: bool = True,
    ) -> None:
        # When True, already-restored PAST frames are fed into the temporal window
        # (instead of the original mosaic frames). Future neighbours and the centre
        # frame stay original — they aren't restored yet / are what we're cleaning.
        self._feed_restored = feed_restored
        chosen: Path | None = None
        if model_path and Path(model_path).is_file() and "video" in Path(model_path).name.lower():
            chosen = Path(model_path)
        else:  # configured model missing or not a video model → auto-pick a video model
            vids = [p for _n, p in discover_models(include_video=True) if "video" in Path(p).name.lower()]
            if vids:
                chosen = Path(vids[0])
        if chosen is None:
            raise ValueError(
                "Не найдены веса видеомодели DeepMosaics (clean_*_video.pth).\n"
                "Положите clean_youknow_video.pth + mosaic_position.pth в models/deepmosaics "
                "(или выберите в «Восстановление…»). См. README."
            )
        pos = _find_mosaic_position(chosen, deepmosaics_dir)
        if pos is None:
            raise ValueError(
                "Рядом с видеомоделью не найден mosaic_position.pth.\n"
                "Положите mosaic_position.pth в ту же папку, что и clean_*_video.pth. См. README."
            )
        self._model = str(chosen)
        self._pos = str(pos)
        self._gpu = gpu_id
        self._loaded = False

    @property
    def name(self) -> str:
        return f"DeepMosaicsVideo(gpu={self._gpu})"

    # ------------------------------------------------------------------ engine
    def _ensure_loaded(self) -> None:
        if self._loaded:
            return
        if str(_VENDOR) not in sys.path:
            sys.path.insert(0, str(_VENDOR))

        import torch
        if self._gpu != "-1" and not torch.cuda.is_available():
            self._gpu = "-1"  # CPU fallback (see DeepMosaicsRestorer for why)

        import util.data as data  # type: ignore
        import util.image_processing as impro  # type: ignore

        from models import loadmodel, runmodel  # type: ignore

        self._torch = torch
        self._runmodel = runmodel
        self._data = data
        self._impro = impro
        self._opt = SimpleNamespace(
            gpu_id=self._gpu,
            model_path=self._model,
            mosaic_position_model_path=self._pos,
            mask_threshold=64,
            all_mosaic_area=False,
            ex_mult=1.5,
            no_feather=False,
        )
        self._netM = loadmodel.bisenet(self._opt, "mosaic")
        self._netG = loadmodel.video(self._opt)  # BVDNet
        self._loaded = True

    def restore(
        self,
        image: np.ndarray,
        detections: list[Detection],
        should_cancel: CancelCheck | None = None,
    ) -> np.ndarray:
        """Single-frame restore — no neighbours, so the window is the same frame."""
        out: dict[int, np.ndarray] = {}
        self.restore_sequence(
            1,
            lambda _i: image,
            lambda _i: detections,
            lambda i, r: out.__setitem__(i, r),
            should_cancel,
        )
        return out.get(0, image.copy())

    def restore_sequence(
        self,
        count: int,
        get_frame: FrameGetter,
        get_dets: DetGetter,
        emit: ResultSink,
        should_cancel: CancelCheck | None = None,
    ) -> None:
        self._ensure_loaded()
        torch, data, impro, opt = self._torch, self._data, self._impro, self._opt
        N, T, S, SZ = self._N, self._T, self._S, self._INPUT_SIZE
        reach = N * S  # how far back/forward the window samples (frames)

        # Rolling cache of already-restored frames, used as window neighbours when
        # ``feed_restored`` is on. Only the last ``reach`` frames are ever needed.
        restored: dict[int, np.ndarray] = {}

        def remember(idx: int, frame: np.ndarray) -> None:
            if not self._feed_restored:
                return
            restored[idx] = frame
            for old in [k for k in restored if k < idx - reach]:
                restored.pop(old, None)

        previous = None  # recurrent state: the network's previous output (a tensor)
        for i in range(count):
            if should_cancel is not None and should_cancel():
                raise Cancelled("Восстановление отменено")
            img_origin = get_frame(i)
            x, y, size, mask = self._runmodel.get_mosaic_position(img_origin, self._netM, opt)
            if size <= 50:
                clean = img_origin.copy()  # no mosaic here; recurrence carries over
                emit(i, clean)
                remember(i, clean)
                continue

            stream = []
            for k in range(T):
                j = min(max(i + (k - N) * S, 0), count - 1)  # clamp window to range edges
                if j == i:
                    frame = img_origin
                elif self._feed_restored and j < i and j in restored:
                    frame = restored[j]   # already-restored past neighbour
                else:
                    frame = get_frame(j)  # original (future neighbour / not yet cached)
                crop = frame[y - size:y + size, x - size:x + size]
                stream.append(impro.resize(crop, SZ)[:, :, ::-1])  # BGR→RGB, SZ×SZ

            if previous is None:  # seed recurrence with the (centre) input crop
                previous = data.im2tensor(stream[N], bgr2rgb=False, gpu_id=opt.gpu_id)

            arr = np.array(stream).reshape(1, T, SZ, SZ, 3).transpose((0, 4, 1, 2, 3))
            tensor = data.to_tensor(data.normalize(arr), gpu_id=opt.gpu_id)
            with torch.no_grad():
                pred = self._netG(tensor, previous)
            previous = pred
            img_fake = data.tensor2im(pred, rgb2bgr=True)
            result = impro.replace_mosaic(img_origin.copy(), img_fake, mask, x, y, size, opt.no_feather)
            emit(i, result)
            remember(i, result)