# -*- coding: utf-8 -*-

#

# Licensed under the terms of the BSD 3-Clause

# (see plotpy/LICENSE for details)

"""Tests around image transforms: rotation, translation, ..."""

# guitest: show

from __future__ import annotations

import os

import numpy as np

import pytest

from guidata.env import execenv

from guidata.qthelpers import qt_app_context

from qtpy import QtCore as QC

from qtpy import QtGui as QG

from plotpy import io

from plotpy.builder import make

from plotpy.constants import LUTAlpha

from plotpy.items import TrImageItem, assemble_imageitems

from plotpy.tests import vistools as ptv

from plotpy.tests.data import gen_image4

DEFAULT_CHARS = "".join([chr(c) for c in range(32, 256)])

def get_font_array(sz: int, chars: str = DEFAULT_CHARS) -> np.ndarray | None:

"""Return array of font characters

Args:

sz: Font size

chars: Characters to include (default: all printable characters)

Returns:

Array of font characters

"""

font = QG.QFont()

font.setFixedPitch(True)

font.setPixelSize(sz)

font.setStyleStrategy(QG.QFont.NoAntialias)

dummy = QG.QImage(10, 10, QG.QImage.Format_ARGB32)

pnt = QG.QPainter(dummy)

pnt.setFont(font)

metric = pnt.fontMetrics()

rct = metric.boundingRect(chars)

pnt.end()

h, w = rct.height(), rct.width()

img = QG.QImage(w, h, QG.QImage.Format_ARGB32)

paint = QG.QPainter()

paint.begin(img)

paint.setFont(font)

paint.setBrush(QG.QColor(255, 255, 255))

paint.setPen(QG.QColor(255, 255, 255))

paint.drawRect(0, 0, w + 1, h + 1)

paint.setPen(QG.QColor(0, 0, 0))

paint.setBrush(QG.QColor(0, 0, 0))

paint.drawText(0, paint.fontMetrics().ascent(), chars)

paint.end()

try:

data = img.bits().asstring(h * w * 4)

except AttributeError:

data = img.bits()

npy: np.ndarray = np.frombuffer(data, np.uint8)

return npy.reshape(h, w, 4)[:, :, 0]

def write_text_on_array(

data: np.ndarray,

x: int,

y: int,

sz: int,

txt: str,

range: tuple[int, int] | None = None,

) -> None:

"""Write text in image (in-place)

Args:

data: Image data

x: X-coordinate of top-left corner

y: Y-coordinate of top-left corner

sz: Font size

txt: Text to write

range: Range of values to map to 0-255 (default: None)

"""

arr = get_font_array(sz, txt)

if arr is None:

return

if range is None:

m, M = data.min(), data.max()

else:

m, M = range

z = (float(M) - float(m)) * np.array(arr, float) / 255.0 + m

arr = np.array(z, data.dtype)

dy, dx = arr.shape

data[y : y + dy, x : x + dx] = arr

def make_items(N: int) -> list[TrImageItem]:

"""Make test TrImageItem items

Args:

N: Image size (N x N)

Returns:

List of image items

"""

data = gen_image4(N, N)

m = data.min()

M = data.max()

items = [make.trimage(data, alpha_function=LUTAlpha.LINEAR, colormap="jet")]

for dtype in (np.uint8, np.uint16, np.int8, np.int16):

info = np.iinfo(dtype().dtype) # pylint: disable=no-value-for-parameter

s = float((info.max - info.min))

a1 = s * (data - m) / (M - m)

img = np.array(a1 + info.min, dtype)

write_text_on_array(img, 0, 0, int(N / 15.0), dtype.__name__)

items.append(make.trimage(img, colormap="jet"))

nc = int(np.sqrt(len(items)) + 1.0)

maxy, x, y = 0, 0, 0

w = None

for index, item in enumerate(items):

h = item.boundingRect().height()

if index % nc == 0:

x = 0

y += maxy

maxy = h

else:

x += w

maxy = max(maxy, h)

w = item.boundingRect().width()

item.set_transform(x, y, 0.0)

# item.set_selectable(False)

return items

def save_image(name: str, data: np.ndarray) -> None:

"""Save image to file

Args:

name: Base name of file

data: Image data

"""

for fname in (name + ".u16.tif", name + ".u8.png"):

if os.path.exists(fname):

os.remove(fname)

size = int(data.nbytes / 1024.0)

print(f"Saving image: {data.shape[0]} x {data.shape[1]} ({size} KB):")

print(" --> uint16")

io.imwrite(name + ".u16.tif", data, dtype=np.uint16, max_range=True)

print(" --> uint8")

io.imwrite(name + ".u8.png", data, dtype=np.uint8, max_range=True)

def get_bbox(items: list[TrImageItem]) -> QC.QRectF:

"""Get bounding box of items

Args:

items: List of image items

Returns:

Bounding box of items

"""

rectf = QC.QRectF()

for item in items:

rectf = rectf.united(item.boundingRect())

return rectf

def build_image(items: list[TrImageItem]) -> None:

"""Build image from items

Args:

items: List of image items

"""

r = get_bbox(items)

_x, _y, w, h = r.getRect()

print("-" * 80)

print(f"Assemble test1: {int(w)} x {int(h)}")

dest = assemble_imageitems(items, r, w, h)

if not execenv.unattended:

save_image("test1", dest)

print("-" * 80)

print(f"Assemble test1: {int(w / 4)} x {int(h / 4)}")

dest = assemble_imageitems(items, r, w / 4, h / 4)

if not execenv.unattended:

save_image("test2", dest)

print("-" * 80)

@pytest.mark.parametrize("N", [500])

@pytest.mark.parametrize("assemble_images", [False, True])

def test_transform(N: int, assemble_images: bool) -> None:

"""Test image transforms

Args:

N: Image size (N x N)

assemble_images: If True, assemble images (default: False)

"""

with qt_app_context(exec_loop=True):

items = make_items(N)

_win = ptv.show_items(

items,

wintitle="Transform test ({}x{} images)".format(N, N),

plot_type="image",

show_itemlist=True,

winsize=(1000, 600),

)

if assemble_images:

build_image(items)

if __name__ == "__main__":

test_transform(N=500, assemble_images=True)