"""Check physical luminance, PQ encoding, and the actual PNG byte stream."""

import struct
import tempfile
from pathlib import Path
import unittest
import xml.etree.ElementTree as ET
import zlib

import numpy as np

from avatar_color import BT2020_Y, SRGB_Y, SDR_PEAK_Y, TO_2020, equal_luminance, gradient_level, linear_srgb, sdr_color
from avatar_hdr import export_hdr, pq_decode, pq_encode, render_linear_hdr
from cmatrix_avatar import BLUE, GREEN, YELLOW, CMatrixAvatar


def read_png(path):
    data = Path(path).read_bytes()
    assert data[:8] == b'\x89PNG\r\n\x1a\n'
    position, chunks = 8, {}
    while position < len(data):
        size = struct.unpack_from('>I', data, position)[0]
        kind = data[position + 4:position + 8]
        payload = data[position + 8:position + 8 + size]
        crc = struct.unpack_from('>I', data, position + 8 + size)[0]
        assert crc == zlib.crc32(kind + payload) & 0xffffffff
        chunks[kind] = chunks.get(kind, b'') + payload
        position += 12 + size
    return chunks


class HDRTests(unittest.TestCase):
    def test_pq_reference_values_and_round_trip(self):
        self.assertAlmostEqual(float(pq_encode(100)), .50807842, places=7)
        self.assertAlmostEqual(float(pq_encode(1000)), .75182710, places=7)
        self.assertAlmostEqual(float(pq_encode(10000)), 1, places=7)
        nits = np.array([.02, 2, 100, 400, 1000, 4000])
        np.testing.assert_allclose(pq_decode(pq_encode(nits)), nits, rtol=1e-10)

    def test_equal_peaks_and_stronger_gradient(self):
        self.assertEqual(sdr_color(BLUE), BLUE)
        for base in (BLUE, GREEN, YELLOW):
            rgb = equal_luminance(base, 400, 10000, hdr=True)
            self.assertAlmostEqual(sum(a * b for a, b in zip(rgb, BT2020_Y)), 400)
            reference = np.asarray(TO_2020) @ np.asarray(linear_srgb(base))
            np.testing.assert_allclose(np.asarray(rgb) / sum(rgb), reference / reference.sum())
            self.assertLessEqual(max(equal_luminance(base, 2000, 10000, hdr=True)), 10000)
            sdr_y = sum(a * b for a, b in zip(linear_srgb(sdr_color(base)), SRGB_Y))
            self.assertAlmostEqual(sdr_y, SDR_PEAK_Y, delta=.004)
        levels = [gradient_level(i / 31) for i in range(32)]
        self.assertEqual(len(set(levels)), 32)
        self.assertEqual(levels, sorted(levels))
        self.assertEqual(levels[0], .005)
        self.assertEqual(levels[-1], 1)
        avatar = CMatrixAvatar(42)
        yellow = [c.level for row in avatar.cells for c in row if c.color == YELLOW]
        self.assertGreater(len(set(yellow)), 5)

    def test_hdr_file_encoding_and_decoded_luminance(self):
        svg = '<svg xmlns="http://www.w3.org/2000/svg" width="128" height="16" viewBox="0 0 128 16">'
        for index, (base, level) in enumerate(((GREEN, 1), (BLUE, 1), (YELLOW, 1), (GREEN, .005))):
            x = index * 32
            svg += f'<path data-base="{base}" data-level="{level}" d="M{x} 0h32v16h-32z"/>'
        svg += '</svg>'
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / 'hdr.png'
            export_hdr(svg, path, 400)
            chunks = read_png(path)
            self.assertEqual(struct.unpack('>IIBBBBB', chunks[b'IHDR']), (128, 16, 16, 2, 0, 0, 0))
            self.assertEqual(chunks[b'cICP'], bytes((9, 16, 0, 1)))
            self.assertNotIn(b'sRGB', chunks)
            self.assertNotIn(b'gAMA', chunks)
            raw = np.frombuffer(zlib.decompress(chunks[b'IDAT']), dtype=np.uint8).reshape(16, 128 * 6 + 1)
            self.assertTrue((raw[:, 0] == 0).all())
            signal = np.frombuffer(raw[:, 1:].copy().tobytes(), dtype='>u2').reshape(16, 128, 3)
            self.assertTrue((signal % 257 != 0).any())
            decoded = pq_decode(signal.astype(float) / 65535)
            y = decoded @ np.asarray(BT2020_Y)
            for column in (16, 48, 80):
                self.assertAlmostEqual(y[8, column], 400, delta=.1)
            self.assertAlmostEqual(y[8, 112], 2, delta=.01)
            max_cll, max_fall = struct.unpack('>II', chunks[b'cLLI'])
            self.assertAlmostEqual(max_cll / 10000, y.max(), delta=.0001)
            self.assertAlmostEqual(max_fall / 10000, y.mean(), delta=.0001)

    def test_invalid_hdr_peak_is_rejected(self):
        for value in (0, 3000, float('nan'), float('inf')):
            with self.assertRaises(ValueError):
                render_linear_hdr('<svg/>', value)


if __name__ == '__main__':
    unittest.main()