"""Round 84: Kaleidostrobe, the 2002 iTunes visualizer, as the fourth style.

``kaleidostrobe_golden.npz`` was recorded from the *original* PowerPC plugin,
run under emulation through its own VisualPluginHandler: every frame's two
uploaded textures (as digests) and every triangle it handed OpenGL, for a
fixed run of spectra; plus the output of each warp (calcDelta) and each
wave-table entry. The port has to reproduce them exactly.

1. The engine against the original: textures byte for byte, geometry.
2. The pieces: warps, wave shapes, transform's quirks, Darwin's rand().
3. The renderer: 25 fps, the GL picture, dim mode, pause.
4. The style in the app.
"""
import hashlib
import math
from pathlib import Path

import numpy as np
import pytest

from PyQt6.QtCore import QObject, QRectF, pyqtSignal
from PyQt6.QtGui import QColor, QImage, QPainter

from lintunes.kaleidostrobe import engine as ks
from lintunes.gui import visualizer_styles as vs
from lintunes.gui.visualizer_kaleido import KaleidostrobeRenderer, itunes_spectrum
from lintunes.preferences import Preferences


GOLDEN = np.load(Path(__file__).with_name("kaleidostrobe_golden.npz"))
RATE = 44100


def _digest(a):
    return hashlib.sha256(np.ascontiguousarray(a).tobytes()).hexdigest()[:16]


def _pattern(seed):
    y, x = np.mgrid[0:128, 0:128]
    return ((x * (7 + seed) + y * (13 + 2 * seed)) ^ (x * y + seed)).astype(np.uint8)


def _rotz(deg):
    r = math.radians(deg)
    c, s = math.cos(r), math.sin(r)
    return np.array([[c, -s, 0, 0], [s, c, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])


def _lift(z):
    m = np.eye(4)
    m[2, 3] = z
    return m


def _eye_space(frame):
    """Each triangle as (s, t, x, y, z) × 3 in eye space, the way the
    golden file recorded the original's GL calls."""
    out = []
    for layer in frame.layers:
        base = _lift(-ks.EYE_Z) @ _lift(layer.lift) @ _rotz(layer.turn)
        for tri in layer.triangles:
            m = base @ _rotz(tri.turn)
            out.append([(*tri.st[i], *(m @ np.array((*tri.xyz[i], 1.0)))[:3])
                        for i in range(3)])
    return np.array(out)


# --------------------------------------------------------------------------
# 1. The engine against the original
# --------------------------------------------------------------------------

class TestAgainstTheOriginal:
    def test_every_frame_matches(self):
        engine = ks.Kaleidostrobe(seed=1000, ticks=1000)   # srand(TickCount())
        engine.set_size(640, 480)
        counts = GOLDEN["tri_counts"]
        offsets = np.concatenate([[0], np.cumsum(counts)])
        for f, (spectrum, ticks) in enumerate(zip(GOLDEN["spectra"], GOLDEN["ticks"])):
            frame = engine.step(spectrum, int(ticks))
            for tex in (0, 1):
                assert _digest(frame.textures[tex]) == GOLDEN["tex_digest"][f][tex], \
                    f"frame {f}, texture {tex}"
            ours = _eye_space(frame)
            theirs = GOLDEN["tris"][offsets[f]:offsets[f + 1]]
            assert ours.shape == theirs.shape, f"frame {f}"
            np.testing.assert_allclose(ours[..., :2], theirs[..., :2], atol=1e-4)
            np.testing.assert_allclose(ours[..., 2:], theirs[..., 2:], atol=1e-3)

    def test_the_run_saw_beats_and_switches(self):
        """The recorded run must exercise more than the opening state."""
        engine = ks.Kaleidostrobe(seed=1000, ticks=1000)
        deltas = set()
        for spectrum, ticks in zip(GOLDEN["spectra"], GOLDEN["ticks"]):
            engine.step(spectrum, int(ticks))
            deltas.add(tuple(engine.cur_delta))
        assert engine.beat_count >= 8
        assert len(deltas) > 1


# --------------------------------------------------------------------------
# 2. The pieces
# --------------------------------------------------------------------------

class TestPieces:
    @pytest.mark.parametrize("i, n", list(enumerate(ks.DELTA_TABLE)))
    def test_every_warp_matches(self, i, n):
        state = {}
        for it in range(3):
            out = ks.delta(n, _pattern(n * 3 + it), state)
            assert _digest(out) == GOLDEN["delta_digest"][i][it], f"iteration {it}"

    @pytest.mark.parametrize("idx", range(len(ks.WAVE_TABLE)))
    def test_every_wave_matches(self, idx):
        engine = ks.Kaleidostrobe(seed=1, ticks=1000)
        engine.rng = ks.ParkMiller(4321)
        engine.cur_wave[0] = idx
        engine._init_wave(0)
        assert _digest(engine.buffers[0, 0]) == GOLDEN["wave_digest"][idx]

    def test_rand_is_darwins(self):
        rng = ks.ParkMiller(1)
        assert [rng.rand() for _ in range(3)] == [16807, 282475249, 1622650073]

    def test_flip_diag_keeps_its_off_by_one(self):
        img = np.arange(128 * 128, dtype=np.int64).reshape(128, 128) % 251
        img = img.astype(np.uint8)
        before = img.copy()
        ks.flip_diag(img)
        assert img[10, 3] == before[3, 10]         # transposed…
        assert img[10, 9] == before[10, 9]         # …except beside the diagonal
        assert img[9, 10] == before[9, 10]

    def test_inside_out_flips_each_half(self):
        img = np.repeat(np.arange(128, dtype=np.uint8)[:, None], 128, 1)
        ks.inside_out(img)
        assert img[0, 0] == 63 and img[63, 0] == 0
        assert img[64, 0] == 127 and img[127, 0] == 64

    def test_a_reselected_warp_resumes_its_own_counter(self):
        engine = ks.Kaleidostrobe(seed=3, ticks=1000)
        engine.cur_delta = [0, 0]          # calcDelta000 on both layers
        for _ in range(5):
            engine._update_textures()
        engine.cur_delta[0] = 1            # another warp: its own counter, 0
        engine._update_textures()
        assert engine.delta_state[0][1] == {'h': 1, 'v': 1}
        assert engine.delta_state[0][0] == {'h': 5}


# --------------------------------------------------------------------------
# 3. The renderer
# --------------------------------------------------------------------------

def _frame(dt=0.016, seconds=4096 / RATE, hz=60.0):
    t = np.arange(int(seconds * RATE)) / RATE
    samples = (0.4 * np.sin(2 * np.pi * hz * t)).astype(np.float32)
    return vs.Frame(samples, RATE, np.zeros(0, np.float32), dt)


def _paint(renderer, w, h, background="white", color="blue"):
    img = QImage(w, h, QImage.Format.Format_RGB32)
    img.fill(QColor(background))
    painter = QPainter(img)
    renderer.paint(painter, QRectF(0, 0, w, h), QColor(color), QColor("gray"))
    painter.end()
    a = np.frombuffer(img.constBits().asstring(img.sizeInBytes()), np.uint8)
    return a.reshape(h, w, 4)[..., [2, 1, 0]].astype(float)


def _bilinear(tex, s, t):
    u, v = s * 128 - 0.5, t * 128 - 0.5
    u0, v0 = np.floor(u).astype(int), np.floor(v).astype(int)
    fu, fv = (u - u0)[..., None], (v - v0)[..., None]
    T = tex.astype(float)

    def at(y, x):
        return T[y % 128, x % 128]
    return (at(v0, u0) * (1 - fu) * (1 - fv) + at(v0, u0 + 1) * fu * (1 - fv)
            + at(v0 + 1, u0) * (1 - fu) * fv + at(v0 + 1, u0 + 1) * fu * fv)


def _gl_reference(frame, w, h):
    """What OpenGL drew: cast a ray through each pixel centre, take the
    nearest triangle, interpolate (s, t), GL_LINEAR + GL_REPEAT, blend."""
    f = 1 / math.tan(math.radians(ks.FOVY) / 2)
    px, py = np.meshgrid(np.arange(w) + 0.5, np.arange(h) + 0.5)
    d = np.stack([(2 * px / w - 1) * (w / h) / f, (1 - 2 * py / h) / f,
                  -np.ones_like(px)], -1)
    out = np.zeros((h, w, 3))
    for layer in frame.layers:
        depth = np.full((h, w), -np.inf)
        col = np.zeros((h, w, 3))
        hit = np.zeros((h, w), bool)
        base = _lift(-ks.EYE_Z) @ _lift(layer.lift) @ _rotz(layer.turn)
        for tri in layer.triangles:
            m = base @ _rotz(tri.turn)
            p = np.array([(m @ np.array((*v, 1.0)))[:3] for v in tri.xyz])
            e1, e2 = p[1] - p[0], p[2] - p[0]
            n = np.cross(e1, e2)
            with np.errstate(divide='ignore', invalid='ignore'):
                x = d * ((p[0] @ n) / (d @ n))[..., None]
                d00, d01, d11 = e1 @ e1, e1 @ e2, e2 @ e2
                v2 = x - p[0]
                den = d00 * d11 - d01 * d01
                b1 = (d11 * (v2 @ e1) - d01 * (v2 @ e2)) / den
                b2 = (d00 * (v2 @ e2) - d01 * (v2 @ e1)) / den
            b0 = 1 - b1 - b2
            take = (b0 >= 0) & (b1 >= 0) & (b2 >= 0) & (x[..., 2] > depth)
            st = np.array(tri.st)
            s = b0 * st[0, 0] + b1 * st[1, 0] + b2 * st[2, 0]
            t = b0 * st[0, 1] + b1 * st[1, 1] + b2 * st[2, 1]
            c = _bilinear(frame.textures[layer.texture][..., :3], s, t)
            col[take] = c[take]
            depth[take] = x[..., 2][take]
            hit |= take
        out[hit] = out[hit] * (1 - layer.alpha) + col[hit] * layer.alpha
    return out


class TestRenderer:
    def test_runs_at_the_plugins_25_fps(self, qapp):
        r = KaleidostrobeRenderer()
        for _ in range(125):                 # 2 s of 16 ms widget ticks
            r.advance(_frame())
        assert r._engine.frame_count == 50

    def test_a_stall_is_not_replayed(self, qapp):
        r = KaleidostrobeRenderer()
        r.advance(_frame())
        r.advance(_frame(dt=5.0))           # the GUI thread froze for 5 s
        assert r._engine.frame_count <= 1 + 3

    @pytest.mark.parametrize("frames", [3, 30, 90])
    def test_paints_what_opengl_drew(self, qapp, frames):
        r = KaleidostrobeRenderer()
        r.advance(_frame())
        r._engine.set_size(96, 60)
        for _ in range(frames):
            r.advance(_frame(dt=0.04))
        ours = _paint(r, 96, 60)
        gl = _gl_reference(r._frame, 96, 60)
        assert np.abs(ours - gl).max() <= 3

    def test_fills_any_shape(self, qapp):
        r = KaleidostrobeRenderer()
        for _ in range(5):
            r.advance(_frame(dt=0.04))
        for w, h in ((140, 36), (300, 900), (4, 4)):
            img = _paint(r, w, h, background="magenta")
            magenta = (img[..., 0] == 255) & (img[..., 1] == 0) & (img[..., 2] == 255)
            assert not magenta.any()        # the cone reaches every corner

    def test_dim_mode_is_gray(self, qapp):
        r = KaleidostrobeRenderer()
        for _ in range(5):
            r.advance(_frame(dt=0.04))
        r.dimmed = True                     # the widget then passes its gray
        img = _paint(r, 120, 80, color="gray")
        assert np.abs(img[..., 0] - img[..., 2]).max() <= 2   # gray tint

    def test_pause_blanks_but_keeps_the_plugin(self, qapp):
        r = KaleidostrobeRenderer()
        for _ in range(5):
            r.advance(_frame(dt=0.04))
        engine, count = r._engine, r._engine.frame_count
        r.clear()
        img = _paint(r, 40, 30, background="white")
        assert (img == 255).all()           # nothing drawn
        r.advance(_frame(dt=0.04))
        assert r._engine is engine and engine.frame_count == count + 1

    @pytest.mark.parametrize("rate", [44100, 48000, 96000])
    def test_spectrum_bins_are_itunes_43_hz(self, rate):
        t = np.arange(4096) / rate
        tone = (0.5 * np.sin(2 * np.pi * 1000 * t)).astype(np.float32)
        spec = itunes_spectrum(tone, rate)
        assert spec.shape == (2, 512) and spec.dtype == np.uint8
        assert abs(int(spec[0].argmax()) - round(1000 / (44100 / 1024))) <= 1
        assert (spec[0] == spec[1]).all()

    def test_silence_is_zero(self):
        assert not itunes_spectrum(np.zeros(4096, np.float32), RATE).any()


# --------------------------------------------------------------------------
# 4. In the app
# --------------------------------------------------------------------------

class _FakePlayer(QObject):
    audio_buffer = pyqtSignal(object)
    playing_changed = pyqtSignal(bool)
    track_changed = pyqtSignal(object)

    def is_playing(self):
        return False


class _LastFm(QObject):
    login_finished = pyqtSignal(bool, str)
    status_message = pyqtSignal(str)

    def is_logged_in(self):
        return False


class TestInTheApp:
    def test_is_a_style(self, qapp, tmp_path):
        from lintunes.gui.visualizer import VisualizerWidget
        prefs = Preferences(tmp_path)
        prefs.set("visualizer_style", "kaleidostrobe")
        vis = VisualizerWidget(_FakePlayer(), prefs)
        assert isinstance(vis._renderer, KaleidostrobeRenderer)

    def test_preferences_offers_it(self, qapp, tmp_path):
        from lintunes.gui.preferences_dialog import PreferencesDialog
        prefs = Preferences(tmp_path)
        dialog = PreferencesDialog(prefs, _LastFm())
        dialog._viz_radios["kaleidostrobe"].setChecked(True)
        assert prefs.get("visualizer_style") == "kaleidostrobe"
        assert dialog._viz_pages.currentIndex() == 3