"""Tests for pylibjpeg.decode()."""

import os
from pathlib import Path
import pytest

import numpy as np

try:
    from pydicom.encaps import generate_frames
    from pydicom.pixels.utils import reshape_pixel_array

    HAS_PYDICOM = True
except ImportError:
    HAS_PYDICOM = False

from libjpeg import decode, decode_pixel_data
from libjpeg.data import get_indexed_datasets, JPEG_DIRECTORY


DIR_10918 = JPEG_DIRECTORY / "10918"
DIR_14495 = JPEG_DIRECTORY / "14495"


REF_DCM = {
    "1.2.840.10008.1.2.4.50": [
        # filename, (rows, columns, samples/px, bits/sample)
        ("JPEGBaseline_1s_1f_u_08_08.dcm", (100, 100, 1, 8)),
        ("SC_rgb_dcmtk_+eb+cy+np.dcm", (100, 100, 3, 8)),
        ("color3d_jpeg_baseline.dcm", (480, 640, 3, 8)),
        ("SC_rgb_dcmtk_+eb+cr.dcm", (100, 100, 3, 8)),
        ("SC_rgb_dcmtk_+eb+cy+n1.dcm", (100, 100, 3, 8)),
        ("SC_rgb_dcmtk_+eb+cy+s4.dcm", (100, 100, 3, 8)),
    ],
    "1.2.840.10008.1.2.4.51": [
        ("RG2_JPLY_fixed.dcm", (2140, 1760, 1, 12)),
        ("JPEGExtended_1s_1f_u_16_12.dcm", (1024, 256, 1, 12)),
        ("JPEGExtended_3s_1f_u_08_08.dcm", (576, 768, 3, 8)),
    ],
    "1.2.840.10008.1.2.4.57": [
        ("JPEGLossless_1s_1f_u_16_12.dcm", (1024, 1024, 1, 12)),
    ],
    "1.2.840.10008.1.2.4.70": [
        ("JPEG-LL.dcm", (1024, 256, 1, 16)),
        ("JPEGLosslessP14SV1_1s_1f_u_08_08.dcm", (768, 1024, 1, 8)),
        ("JPEGLosslessP14SV1_1s_1f_u_16_16.dcm", (535, 800, 1, 16)),
        ("MG1_JPLL.dcm", (4664, 3064, 1, 12)),
        ("RG1_JPLL.dcm", (1955, 1841, 1, 15)),
        ("RG2_JPLL.dcm", (2140, 1760, 1, 10)),
        ("SC_rgb_jpeg_gdcm.dcm", (100, 100, 3, 8)),
    ],
    "1.2.840.10008.1.2.4.80": [
        ("emri_small_jpeg_ls_lossless.dcm", (64, 64, 1, 12)),
        ("MR_small_jpeg_ls_lossless.dcm", (64, 64, 1, 16)),
        ("RG1_JLSL.dcm", (1955, 1841, 1, 16)),
        ("RG2_JLSL.dcm", (2140, 1760, 1, 10)),
    ],
    "1.2.840.10008.1.2.4.81": [
        ("CT1_JLSN.dcm", (512, 512, 1, 16)),
        ("MG1_JLSN.dcm", (4664, 3064, 1, 12)),
        ("RG1_JLSN.dcm", (1955, 1841, 1, 15)),
        ("RG2_JLSN.dcm", (2140, 1760, 1, 10)),
    ],
}


# TODO: convert using straight JPG data
@pytest.mark.skipif(not HAS_PYDICOM, reason="No pydicom")
def test_decode_bytes():
    """Test decode using bytes."""
    index = get_indexed_datasets("1.2.840.10008.1.2.4.50")
    ds = index["JPEGBaseline_1s_1f_u_08_08.dcm"]["ds"]
    nr_frames = ds.get("NumberOfFrames", 1)
    frame = next(generate_frames(ds.PixelData, number_of_frames=nr_frames))
    arr = decode(frame)
    assert arr.flags.writeable
    assert "uint8" == arr.dtype
    assert (ds.Rows, ds.Columns) == arr.shape

    # Reference values from GDCM handler
    assert 76 == arr[5, 50]
    assert 167 == arr[15, 50]
    assert 149 == arr[25, 50]
    assert 203 == arr[35, 50]
    assert 29 == arr[45, 50]
    assert 142 == arr[55, 50]
    assert 1 == arr[65, 50]
    assert 64 == arr[75, 50]
    assert 192 == arr[85, 50]
    assert 255 == arr[95, 50]


def test_decode_binaryio():
    """Test decode using binaryio."""
    with open(DIR_10918 / "p1" / "A1.JPG", "rb") as f:
        arr = decode(f)
        assert arr.flags.writeable
        assert "uint8" == arr.dtype
        assert arr.shape == (257, 255, 4)


def test_decode_raises():
    """Test decode raises if invalid type"""
    msg = (
        "Invalid type 'NoneType' - must be the path to a JPEG file, a "
        "buffer containing the JPEG data or an open JPEG file-like"
    )
    with pytest.raises(TypeError, match=msg):
        decode(None)


def test_v1_invalid_photometric_raises():
    """Test decode_pixel_data raises if invalid photometric interpretation"""
    msg = (
        r"The \(0028,0004\) Photometric Interpretation element is missing "
        "from the dataset"
    )
    with pytest.raises(ValueError, match=msg):
        decode_pixel_data(DIR_10918 / "p1" / "A1.JPG", version=1)


# TODO: convert to using straight JPG data
@pytest.mark.skipif(not HAS_PYDICOM, reason="No pydicom")
def test_invalid_colourspace_warns():
    """Test that using an unknown colourspace gives a warning."""
    index = get_indexed_datasets("1.2.840.10008.1.2.4.50")
    ds = index["JPEGBaseline_1s_1f_u_08_08.dcm"]["ds"]
    nr_frames = ds.get("NumberOfFrames", 1)
    frame = next(generate_frames(ds.PixelData, number_of_frames=nr_frames))
    msg = r"no colour transformation will be applied"
    ds.PhotometricInterpretation = "ANY"
    with pytest.warns(UserWarning, match=msg):
        arr = decode_pixel_data(frame, ds)

    arr = reshape_pixel_array(ds, arr)

    assert arr.flags.writeable
    assert "uint8" == arr.dtype
    assert (ds.Rows, ds.Columns) == arr.shape

    # Reference values from GDCM handler
    assert 76 == arr[5, 50]
    assert 167 == arr[15, 50]
    assert 149 == arr[25, 50]
    assert 203 == arr[35, 50]
    assert 29 == arr[45, 50]
    assert 142 == arr[55, 50]
    assert 1 == arr[65, 50]
    assert 64 == arr[75, 50]
    assert 192 == arr[85, 50]
    assert 255 == arr[95, 50]


@pytest.mark.skipif(not HAS_PYDICOM, reason="No pydicom")
def test_decode_adobe_v101():
    """Test decoding JPEG with Adobe v101 APP14 marker."""
    index = get_indexed_datasets("1.2.840.10008.1.2.4.50")
    ds = index["SC_jpeg_no_color_transform_2.dcm"]["ds"]
    frame = next(generate_frames(ds.PixelData, number_of_frames=1))
    arr = decode(frame, reshape=True)
    assert arr.shape == (256, 256, 3)


@pytest.mark.skipif(not HAS_PYDICOM, reason="No pydicom")
class TestDecodeDCM:
    """Tests for get_parameters() using DICOM datasets."""

    def generate_frames(self, ds):
        """Return a generator object with the dataset's pixel data frames."""
        nr_frames = ds.get("NumberOfFrames", 1)
        return generate_frames(ds.PixelData, number_of_frames=nr_frames)

    @pytest.mark.parametrize("fname, info", REF_DCM["1.2.840.10008.1.2.4.50"])
    def test_baseline(self, fname, info):
        """Test get_parameters() for the baseline datasets."""
        # info: (rows, columns, spp, bps)
        index = get_indexed_datasets("1.2.840.10008.1.2.4.50")
        ds = index[fname]["ds"]

        frame = next(self.generate_frames(ds))
        arr = decode(frame, reshape=True)

        if info[2] == 1:
            assert (info[0], info[1]) == arr.shape
        else:
            assert (info[0], info[1], info[2]) == arr.shape

        if 1 <= info[3] <= 8:
            assert arr.dtype == "uint8"
        if 9 <= info[3] <= 16:
            assert arr.dtype == "uint16"

    @pytest.mark.parametrize("fname, info", REF_DCM["1.2.840.10008.1.2.4.51"])
    def test_extended(self, fname, info):
        """Test get_parameters() for the baseline datasets."""
        # info: (rows, columns, spp, bps)
        index = get_indexed_datasets("1.2.840.10008.1.2.4.51")
        ds = index[fname]["ds"]

        frame = next(self.generate_frames(ds))
        arr = decode(frame, reshape=True)

        if info[2] == 1:
            assert (info[0], info[1]) == arr.shape
        else:
            assert (info[0], info[1], info[2]) == arr.shape

        if 1 <= info[3] <= 8:
            assert arr.dtype == "uint8"
        if 9 <= info[3] <= 16:
            assert arr.dtype == "uint16"

    @pytest.mark.parametrize("fname, info", REF_DCM["1.2.840.10008.1.2.4.57"])
    def test_lossless(self, fname, info):
        """Test get_parameters() for the lossless datasets."""
        # info: (rows, columns, spp, bps)
        index = get_indexed_datasets("1.2.840.10008.1.2.4.57")
        ds = index[fname]["ds"]

        frame = next(self.generate_frames(ds))
        arr = decode(frame, reshape=True)

        if info[2] == 1:
            assert (info[0], info[1]) == arr.shape
        else:
            assert (info[0], info[1], info[2]) == arr.shape

        if 1 <= info[3] <= 8:
            assert arr.dtype == "uint8"
        if 9 <= info[3] <= 16:
            assert arr.dtype == "uint16"

    @pytest.mark.parametrize("fname, info", REF_DCM["1.2.840.10008.1.2.4.70"])
    def test_lossless_sv1(self, fname, info):
        """Test get_parameters() for the lossless SV1 datasets."""
        # info: (rows, columns, spp, bps)
        index = get_indexed_datasets("1.2.840.10008.1.2.4.70")
        ds = index[fname]["ds"]

        frame = next(self.generate_frames(ds))
        arr = decode(frame, reshape=True)

        if info[2] == 1:
            assert (info[0], info[1]) == arr.shape
        else:
            assert (info[0], info[1], info[2]) == arr.shape

        if 1 <= info[3] <= 8:
            assert arr.dtype == "uint8"
        if 9 <= info[3] <= 16:
            assert arr.dtype == "uint16"

    @pytest.mark.parametrize("fname, info", REF_DCM["1.2.840.10008.1.2.4.80"])
    def test_jls_lossless(self, fname, info):
        """Test get_parameters() for the LS lossless datasets."""
        # info: (rows, columns, spp, bps)
        index = get_indexed_datasets("1.2.840.10008.1.2.4.80")
        ds = index[fname]["ds"]

        frame = next(self.generate_frames(ds))
        arr = decode(frame, reshape=True)

        if info[2] == 1:
            assert (info[0], info[1]) == arr.shape
        else:
            assert (info[0], info[1], info[2]) == arr.shape

        if 1 <= info[3] <= 8:
            assert arr.dtype == "uint8"
        if 9 <= info[3] <= 16:
            assert arr.dtype == "uint16"

    @pytest.mark.parametrize("fname, info", REF_DCM["1.2.840.10008.1.2.4.81"])
    def test_jls(self, fname, info):
        """Test get_parameters() for the LS lossy datasets."""
        # info: (rows, columns, spp, bps)
        index = get_indexed_datasets("1.2.840.10008.1.2.4.81")
        ds = index[fname]["ds"]

        frame = next(self.generate_frames(ds))
        arr = decode(frame, reshape=True)

        if info[2] == 1:
            assert (info[0], info[1]) == arr.shape
        else:
            assert (info[0], info[1], info[2]) == arr.shape

        if 1 <= info[3] <= 8:
            assert arr.dtype == "uint8"
        if 9 <= info[3] <= 16:
            assert arr.dtype == "uint16"

    def test_v2_non_conformant_raises(self):
        """Test that a non-conformant JPEG image raises an exception."""
        ds_list = get_indexed_datasets("1.2.840.10008.1.2.4.51")
        # Image has invalid Se value in the SOS marker segment
        item = ds_list["JPEG-lossy.dcm"]
        assert 0xC000 == item["Status"][1]
        ds = item["ds"]
        nr_frames = ds.get("NumberOfFrames", 1)
        frame = next(generate_frames(ds.PixelData, number_of_frames=nr_frames))
        msg = (
            r"libjpeg error code '-1038' returned from Decode\(\): A "
            r"misplaced marker segment was found - scan start must be zero "
            r"and scan stop must be 63 for the sequential operating modes"
        )
        with pytest.raises(RuntimeError, match=msg):
            decode_pixel_data(frame, version=2)


REF_JPG = {
    "10918": {
        # Reference values using libjpeg's `jpeg` app + imageJ
        # Very circular but having issues with other programs interpreting
        #   it as CYMK and doing weird stuff
        "p1": [
            # ([top left pixel], [bottom right pixel])
            ("A1.JPG", (257, 255, 4, 8), ([138, 76, 239, 216], [155, 217, 191, 115])),
            # ('B1.JPG', (1, 1, 1, 1)),
            # ('B2.JPG', (1, 1, 1, 1)), missing DHT marker (its in B1.JPG)
        ],
        "p2": [
            ("A1.JPG", (257, 255, 4, 8), ([138, 76, 239, 216], [155, 217, 191, 115])),
            # ('B1.JPG', ()),
            # ('B2.JPG', ()), missing DHT (in B1)
            ("C1.JPG", (257, 255, 4, 8), ([138, 76, 239, 216], [155, 217, 191, 115])),
            ("C2.JPG", (257, 255, 4, 8), ([138, 76, 239, 216], [155, 217, 191, 115])),
        ],
        "p4": [
            ("A1.JPG", (257, 255, 4, 8), ([138, 76, 239, 216], [155, 217, 191, 115])),
            # ('B1.JPG', ()),
            # ('B2.JPG', ()),
            ("C1.JPG", (257, 255, 4, 8), ([138, 76, 239, 216], [155, 217, 191, 115])),
            ("C2.JPG", (257, 255, 4, 8), ([138, 76, 239, 216], [155, 217, 191, 115])),
            (
                "E1.JPG",
                (257, 255, 4, 12),
                ([2119, 1183, 3907, 3487], [2502, 3402, 3041, 1872]),
            ),
            (
                "E2.JPG",
                (257, 255, 4, 12),
                ([2119, 1183, 3907, 3487], [2502, 3402, 3041, 1872]),
            ),
        ],
        "p14": [
            ("O1.JPG", (257, 255, 4, 8), ([132, 76, 245, 218], [156, 211, 191, 116])),
            (
                "O2.JPG",
                (257, 255, 4, 16),
                ([33792, 19456, 62720, 55808], [39936, 53888, 48768, 29696]),
            ),
        ],
    },
    "14495": {
        "JLS": [  # Lossless
            (
                # Reference values using GDCM
                "T8C0E0.JLS",
                (256, 256, 3, 8),
                ([161, 122, 108], [101, 99, 95]),
            ),
            (
                # GDCM and CharPyLS fail on this one - both use CharLS
                # Circular reference
                "T8C1E0.JLS",
                (256, 256, 3, 8),
                ([161, 122, 108], [101, 99, 95]),
            ),
            (
                # CharLS fails
                # Circular reference
                "T8C2E0.JLS",
                (256, 256, 3, 8),
                ([161, 122, 108], [101, 99, 95]),
            ),
            (
                # Reference values using GDCM
                "T8NDE0.JLS",
                (128, 128, 1, 8),
                (108, 231),
            ),
            (
                # CharLS fails
                # Circular reference
                "T8SSE0.JLS",
                (256, 256, 3, 8),
                ([161, 122, 108], [101, 171, 231]),
            ),
            (
                # Reference values using GDCM
                "T16E0.JLS",
                (256, 256, 1, 12),
                (1963, 1596),
            ),
        ],
        "JNL": [  # Near lossless
            (
                # Reference values using GDCM
                "T8C0E3.JLS",
                (256, 256, 3, 8),
                ([161, 119, 105], [98, 96, 93]),
            ),
            (
                # CharLS fails
                # Circular reference
                "T8C1E3.JLS",
                (256, 256, 3, 8),
                ([161, 119, 105], [101, 100, 97]),
            ),
            (
                # CharLS fails
                # Circular reference
                "T8C2E3.JLS",
                (256, 256, 3, 8),
                ([161, 119, 105], [98, 96, 94]),
            ),
            (
                # Reference values using GDCM
                "T8NDE3.JLS",
                (128, 128, 1, 8),
                (105, 229),
            ),
            (
                # CharLS fails
                # Circular reference
                "T8SSE3.JLS",
                (256, 256, 3, 8),
                ([161, 119, 105], [102, 169, 234]),
            ),
            (
                # Reference values using GDCM
                "T16E3.JLS",
                (256, 256, 1, 12),
                (1960, 1593),
            ),
        ],
    },
    "15444": {},
}


class TestDecodeJPG:
    """Tests for get_parameters() using JPEG compliance data."""

    @pytest.mark.parametrize("fname, info, ref", REF_JPG["10918"]["p1"])
    def test_baseline(self, fname, info, ref):
        """Test decoding the baseline compliance images."""
        # info: (rows, columns, spp, bps)
        with open(os.path.join(DIR_10918, "p1", fname), "rb") as fp:
            data = fp.read()

        arr = decode(data, reshape=True)
        assert arr.flags.writeable

        if info[2] == 1:
            assert (info[0], info[1]) == arr.shape
        else:
            assert (info[0], info[1], info[2]) == arr.shape

        # Process 1 is always 8-bit
        assert arr.dtype == "uint8"

        assert ref[0] == arr[0, 0, :].tolist()
        assert ref[1] == arr[-1, -1, :].tolist()

        buffer = decode_pixel_data(data, version=2)
        assert isinstance(buffer, bytearray)
        arr = np.frombuffer(buffer, dtype="u1")
        assert arr.flags.writeable
        if info[2] == 1:
            arr = arr.reshape((info[0], info[1]))
        else:
            arr = arr.reshape((info[0], info[1], info[2]))

        assert ref[0] == arr[0, 0, :].tolist()
        assert ref[1] == arr[-1, -1, :].tolist()

    @pytest.mark.parametrize("fname, info, ref", REF_JPG["10918"]["p2"])
    def test_extended_p2(self, fname, info, ref):
        """Test decoding the extended p2 compliance images."""
        # info: (rows, columns, spp, bps)
        with open(os.path.join(DIR_10918, "p2", fname), "rb") as fp:
            data = fp.read()

        arr = decode(data, reshape=True)
        assert arr.flags.writeable

        if info[2] == 1:
            assert (info[0], info[1]) == arr.shape
        else:
            assert (info[0], info[1], info[2]) == arr.shape

        # Process 2 is always 8-bit
        assert arr.dtype == "uint8"

        assert ref[0] == arr[0, 0, :].tolist()
        assert ref[1] == arr[-1, -1, :].tolist()

    @pytest.mark.parametrize("fname, info, ref", REF_JPG["10918"]["p4"])
    def test_extended_p4(self, fname, info, ref):
        """Test decoding the extended p4 compliance images."""
        # info: (rows, columns, spp, bps)
        with open(os.path.join(DIR_10918, "p4", fname), "rb") as fp:
            data = fp.read()

        arr = decode(data, reshape=True)
        assert arr.flags.writeable

        if info[2] == 1:
            assert (info[0], info[1]) == arr.shape
        else:
            assert (info[0], info[1], info[2]) == arr.shape

        if 1 <= info[3] <= 8:
            assert arr.dtype == "uint8"
        if 9 <= info[3] <= 16:
            assert arr.dtype == "uint16"

        assert ref[0] == arr[0, 0, :].tolist()
        assert ref[1] == arr[-1, -1, :].tolist()

    @pytest.mark.parametrize("fname, info, ref", REF_JPG["10918"]["p14"])
    def test_lossless_p14(self, fname, info, ref):
        """Test decoding the extended p14 compliance images."""
        # info: (rows, columns, spp, bps)
        with open(os.path.join(DIR_10918, "p14", fname), "rb") as fp:
            data = fp.read()

        arr = decode(data, reshape=True)
        assert arr.flags.writeable

        if info[2] == 1:
            assert (info[0], info[1]) == arr.shape
        else:
            assert (info[0], info[1], info[2]) == arr.shape

        if 1 <= info[3] <= 8:
            assert arr.dtype == "uint8"
        if 9 <= info[3] <= 16:
            assert arr.dtype == "uint16"

        assert ref[0] == arr[0, 0, :].tolist()
        assert ref[1] == arr[-1, -1, :].tolist()

    @pytest.mark.parametrize("fname, info, ref", REF_JPG["14495"]["JLS"])
    def test_jls(self, fname, info, ref):
        """Test decoding the JPEG-LS lossless compliance images."""
        # info: (rows, columns, spp, bps)
        with open(os.path.join(DIR_14495, "JLS", fname), "rb") as fp:
            data = fp.read()

        arr = decode(data, reshape=True)
        assert arr.flags.writeable

        if info[2] == 1:
            assert (info[0], info[1]) == arr.shape
        else:
            assert (info[0], info[1], info[2]) == arr.shape

        if 1 <= info[3] <= 8:
            assert arr.dtype == "uint8"
        if 9 <= info[3] <= 16:
            assert arr.dtype == "uint16"

        if info[2] == 1:
            assert ref[0] == arr[0, 0].tolist()
            assert ref[1] == arr[-1, -1].tolist()
        else:
            assert ref[0] == arr[0, 0, :].tolist()
            assert ref[1] == arr[-1, -1, :].tolist()

    @pytest.mark.parametrize("fname, info, ref", REF_JPG["14495"]["JNL"])
    def test_jnl(self, fname, info, ref):
        """Test decoding the JPEG-LS near lossless compliance images."""
        # info: (rows, columns, spp, bps)
        with open(os.path.join(DIR_14495, "JNL", fname), "rb") as fp:
            data = fp.read()

        arr = decode(data, reshape=True)
        assert arr.flags.writeable

        if info[2] == 1:
            assert (info[0], info[1]) == arr.shape
        else:
            assert (info[0], info[1], info[2]) == arr.shape

        if 1 <= info[3] <= 8:
            assert arr.dtype == "uint8"
        if 9 <= info[3] <= 16:
            assert arr.dtype == "uint16"

        if info[2] == 1:
            assert ref[0] == arr[0, 0].tolist()
            assert ref[1] == arr[-1, -1].tolist()
        else:
            assert ref[0] == arr[0, 0, :].tolist()
            assert ref[1] == arr[-1, -1, :].tolist()

    def test_str(self):
        """Test decoding using a str path."""
        p = DIR_10918 / "p1" / "A1.JPG"
        p = os.fspath(p)
        assert isinstance(p, str)
        arr = decode(p)
        assert arr.shape == (257, 255, 4)

    def test_path(self):
        """Test decoding using a Path path."""
        p = DIR_10918 / "p1" / "A1.JPG"
        assert isinstance(p, Path)
        arr = decode(p)
        assert arr.shape == (257, 255, 4)
