import numpy as np

from encord.common.bitmask_operations.bitmask_operations import (
    _mask_to_rle,
    _rle_to_mask,
    _rle_to_string,
    _string_to_rle,
    serialise_bitmask,
    transpose_bytearray,
)
from encord.objects.bitmask import (
    BitmaskCoordinates,
)


def test_rle_decode():
    resolution = 1914 * 2294

    encoded_str = "V[^i0g1Rj1101Df0D2N2G?H20OJ;J2O001I<J000J<J0O2J:J101J7M2O0N4M2N1L6M2OOK<I20ON4L3O000J9M001I8N0O1N201N1K7M2O0N4M2O0K7M2OOL6N100N4M2O0O1K8L100J60001I8N0O1N2O101J7M100K5O100N2O10001I8N0O1N200O1K5O2O0K7M2N001O2N0L501OOO3N10OK8M3N000J60000M211N1K402OOO2O2OOL5O2OOL6N10OO3N100K5O100O1N102O0J511OOK601N0O201N0L5O2OOO3N11NN1110OL6N011O0M211N0L3110OO2O2OOL5O2OOO2O2N00O1010K5O100O0O020O1K5O10OO10010M300O0L320O0O3O000J32100NO21000NO1200OK420O1N201N0L31010N01010NO2100KO5100NO2011NOM3110OO10010NM401OL041O00L41O0L05O10NL421OO0NN500OK06000M12O01NM5O00OM410OL0411NN11010NL51000OK150O0OL61N0O0110OL04010NL51O1ON2100KO510OOM4010O0O0110OL0410OO0110OOM41O0L05O01N02O10NL421OK212000ONM600ONO4O00L4001OOM410OO0110OL04010NL51000ON120O0LO61N0O0110OOM4010K31001ON2100NL5100K22100OO0M410OO0110OO1001OOM41O0L320O0OM5O10N0012ONO12000OK1500ON21O00OM50O0O0110OO01010KO5100NO21000NON51ONO120O1K31010NM401OO0110OOM411NN1110OL040100ON120O0O0110OOM410OL311O00O11O00L4010NO2100NO2100O0LO52OOOO22OOO0110OOM4001K32O01NO211ONL421NO001100ON1200ONO4O00O11O00O020O0OM50O0O0102ON001100ONN500OK4200ON120O0OM410OO01010N0102OMM5100NO21000ON031OONM60O1N02O01N01010NL5100NO211NK4110OO010100ONN50O0O11O00O1001OOM410OO0110OO01010N010100ONN500ON120O0O020O0O0110OO0110OOM402OON3000NO21000ON1200ONM61N0O020O0O0110OO01010N0102OMM5100NO21000NO121OO1M30O1NM5O00O1001OO0111N1N111NN1111NOM31100O1N200ON21O1O2N001OOM410O1N20O00O10100N2O100ON1200O1N20O0O1002N1K411N1N111N1N2O10M1102OON4O00O02M110O2M110OO0111N1N111N1N111N0O2O2OON30000M30000O2J410O2M111N1N111N1N111N1N102OON301OON30000O1N200O2L220N3L201N1N111N1N102O0N102O2K4O00ON4O00O02M201O2M11OO1N200O3L3O0O1N200O4K2O100M30003J30000000001N4K111N3L201N3L201N10O20N3L3N102L3N103J30001N3L201N2J6O1N6I201N3L201N3L3N106E5O106D60002M6E610M7F501N8B8N2N<ZO;01NgRhi1"

    r = _string_to_rle(encoded_str)
    assert sum(r) == resolution

    re_encoded_str = _rle_to_string(r)
    assert encoded_str == re_encoded_str


def test_basic_rle_to_mask():
    res = _rle_to_mask([5, 2, 5], size=12)
    assert res == b"\00\00\00\00\00\01\01\00\00\00\00\00"


def test_basic_rle_to_mask_zebra():
    res = _rle_to_mask([1, 1, 1, 1, 1, 1, 1, 1], size=8)
    assert res == b"\00\01\00\01\00\01\00\01"


def test_basic_mask_to_rle():
    mask = b"\00\00\00\00\00\01\01\00\00\00\00\00"
    rle = _mask_to_rle(mask)
    assert rle == [5, 2, 5]


def test_basic_mask_to_rle_zebra():
    mask = b"\00\01\00\01\00\01\00\01"
    rle = _mask_to_rle(mask)
    assert rle == [1, 1, 1, 1, 1, 1, 1, 1]


def test_rle_to_mask():
    resolution = 1914 * 2294
    encoded_str = "V[^i0g1Rj1101Df0D2N2G?H20OJ;J2O001I<J000J<J0O2J:J101J7M2O0N4M2N1L6M2OOK<I20ON4L3O000J9M001I8N0O1N201N1K7M2O0N4M2O0K7M2OOL6N100N4M2O0O1K8L100J60001I8N0O1N2O101J7M100K5O100N2O10001I8N0O1N200O1K5O2O0K7M2N001O2N0L501OOO3N10OK8M3N000J60000M211N1K402OOO2O2OOL5O2OOL6N10OO3N100K5O100O1N102O0J511OOK601N0O201N0L5O2OOO3N11NN1110OL6N011O0M211N0L3110OO2O2OOL5O2OOO2O2N00O1010K5O100O0O020O1K5O10OO10010M300O0L320O0O3O000J32100NO21000NO1200OK420O1N201N0L31010N01010NO2100KO5100NO2011NOM3110OO10010NM401OL041O00L41O0L05O10NL421OO0NN500OK06000M12O01NM5O00OM410OL0411NN11010NL51000OK150O0OL61N0O0110OL04010NL51O1ON2100KO510OOM4010O0O0110OL0410OO0110OOM41O0L05O01N02O10NL421OK212000ONM600ONO4O00L4001OOM410OO0110OL04010NL51000ON120O0LO61N0O0110OOM4010K31001ON2100NL5100K22100OO0M410OO0110OO1001OOM41O0L320O0OM5O10N0012ONO12000OK1500ON21O00OM50O0O0110OO01010KO5100NO21000NON51ONO120O1K31010NM401OO0110OOM411NN1110OL040100ON120O0O0110OOM410OL311O00O11O00L4010NO2100NO2100O0LO52OOOO22OOO0110OOM4001K32O01NO211ONL421NO001100ON1200ONO4O00O11O00O020O0OM50O0O0102ON001100ONN500OK4200ON120O0OM410OO01010N0102OMM5100NO21000ON031OONM60O1N02O01N01010NL5100NO211NK4110OO010100ONN50O0O11O00O1001OOM410OO0110OO01010N010100ONN500ON120O0O020O0O0110OO0110OOM402OON3000NO21000ON1200ONM61N0O020O0O0110OO01010N0102OMM5100NO21000NO121OO1M30O1NM5O00O1001OO0111N1N111NN1111NOM31100O1N200ON21O1O2N001OOM410O1N20O00O10100N2O100ON1200O1N20O0O1002N1K411N1N111N1N2O10M1102OON4O00O02M110O2M110OO0111N1N111N1N111N0O2O2OON30000M30000O2J410O2M111N1N111N1N111N1N102OON301OON30000O1N200O2L220N3L201N1N111N1N102O0N102O2K4O00ON4O00O02M201O2M11OO1N200O3L3O0O1N200O4K2O100M30003J30000000001N4K111N3L201N3L201N10O20N3L3N102L3N103J30001N3L201N2J6O1N6I201N3L201N3L3N106E5O106D60002M6E610M7F501N8B8N2N<ZO;01NgRhi1"
    rle = _string_to_rle(encoded_str)
    mask = _rle_to_mask(rle, size=resolution)
    assert len(mask) == resolution


def test_mask_roundtrip():
    orig_mask = np.zeros((2000, 3000), dtype=bool)
    orig_mask[:1000, 1000:] = 1
    orig_mask[1000:, :1000] = 1

    mask_coordinates = BitmaskCoordinates(orig_mask)

    new_mask = mask_coordinates.to_numpy_array()
    assert np.allclose(orig_mask, new_mask)

    new_mask_using_constructor = np.array(mask_coordinates)
    assert np.allclose(orig_mask, new_mask_using_constructor)


def test_mask_rle_mask_roundtrip():
    orig_mask = np.zeros((2000, 3000), dtype=bool)
    orig_mask[:1000, 1000:] = 1
    orig_mask[1000:, :1000] = 1

    mask_coordinates = BitmaskCoordinates(orig_mask)
    mask_dict = mask_coordinates.to_dict()

    new_mask = BitmaskCoordinates.from_dict(mask_dict).to_numpy_array()
    assert np.allclose(orig_mask, new_mask)


def test_transpose_bytearray():
    mask = b"\00\00\00\00\00\01\01\00\00\00\00\00"
    shape = (3, 4)
    expected = b"\x00\x00\x00\x00\x01\x00\x00\x01\x00\x00\x00\x00"
    ans = transpose_bytearray(mask, shape)
    assert ans == expected


def test_serialise_bitmask_starting_with_true():
    mask = np.array(
        [
            [True, True, True],
            [False, False, False],
        ]
    )
    shape = mask.shape
    rle_string = serialise_bitmask(mask.tobytes())
    mask_encoded = BitmaskCoordinates.EncodedBitmask(
        top=0, left=0, height=shape[0], width=shape[1], rle_string=rle_string
    )
    mask_decoded = BitmaskCoordinates(mask_encoded).to_numpy_array()

    assert np.array_equal(mask, mask_decoded)


def test_serialise_bitmask_empty():
    mask = np.array([[]], dtype=bool)
    shape = mask.shape

    rle_string = serialise_bitmask(mask.tobytes())
    mask_encoded = BitmaskCoordinates.EncodedBitmask(
        top=0, left=0, height=shape[0], width=shape[1], rle_string=rle_string
    )
    mask_decoded = BitmaskCoordinates(mask_encoded).to_numpy_array()

    assert np.array_equal(mask, mask_decoded)  # Mask is inverted and the test fails.
