# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.
from __future__ import annotations

from functools import partial

import pytest

from sisl.utils.ranges import *

pytestmark = pytest.mark.utils


@pytest.mark.ranges
class TestRanges:
    def test_strseq_colon(self):
        ranges = strseq(int, "1:2:5")
        assert ranges == (1, 2, 5)
        ranges = strseq(int, "1:2")
        assert ranges == (1, 2)
        ranges = strseq(int, "1:")
        assert ranges == (1, None)
        ranges = strseq(int, ":4")
        assert ranges == (None, 4)
        ranges = strseq(int, ":4:")
        assert ranges == (None, 4, None)

    def test_strseq_minus(self):
        ranges = strseq(int, "1-2-5")
        assert ranges == (1, 2, 5)
        ranges = strseq(int, "1-2")
        assert ranges == (1, 2)
        ranges = strseq(int, "1-")
        assert ranges == (1, None)
        ranges = strseq(int, "-4")
        assert ranges == -4
        ranges = strseq(int, "-4", end=5)
        assert ranges == 1
        ranges = strseq(int, "-4-")
        assert ranges == (-4, None)

    def test_strmap1(self):
        assert strmap(int, "1") == [1]
        assert strmap(int, "") == [None]
        assert strmap(int, "1-") == [(1, None)]
        assert strmap(int, "-") == [(None, None)]
        assert strmap(int, "-1") == [-1]
        assert strmap(int, "-1", start=1) == [-1]
        assert strmap(int, "-", start=1, end=2) == [(1, 2)]
        assert strmap(int, "-1:2") == [(-1, 2)]
        assert strmap(int, "1,2") == [1, 2]
        assert strmap(int, "1,2[0,2-]") == [1, (2, [0, (2, None)])]
        assert strmap(int, "1,2-[0,-2]") == [1, ((2, None), [0, -2])]
        assert strmap(int, "1,2[0,2]") == [1, (2, [0, 2])]
        assert strmap(int, "1,2-3[0,2]") == [1, ((2, 3), [0, 2])]
        assert strmap(int, "1,[2,3][0,2]") == [1, (2, [0, 2]), (3, [0, 2])]
        assert strmap(int, "[82][10]") == [(82, [10])]
        assert strmap(int, "[82,83][10]") == [(82, [10]), (83, [10])]
        assert strmap(int, "[82,83][10-13]") == [(82, [(10, 13)]), (83, [(10, 13)])]

    def test_strmap2(self):
        with pytest.raises(ValueError):
            strmap(int, "1", sep="*")

    def test_strmap3(self):
        sm = partial(strmap, sep="c")
        assert sm(int, "1") == [1]
        assert sm(int, "1,2") == [1, 2]
        assert sm(int, "1,2{0,2}") == [1, (2, [0, 2])]
        assert sm(int, "1,2-3{0,2}") == [1, ((2, 3), [0, 2])]
        assert sm(int, "1,{2,3}{0,2}") == [1, (2, [0, 2]), (3, [0, 2])]
        assert sm(int, "{82}{10}") == [(82, [10])]
        assert sm(int, "{82,83}{10}") == [(82, [10]), (83, [10])]
        assert sm(int, "{82,83}{10-13}") == [(82, [(10, 13)]), (83, [(10, 13)])]

    def test_strmap4(self):
        with pytest.raises(ValueError):
            strmap(int, "1[oestuh]]")

    def test_strmap5(self):
        r = strmap(int, "1-", end=5)
        assert r == [(1, 5)]
        r = strmap(int, "1-", start=0, end=5)
        assert r == [(1, 5)]
        r = strmap(int, "-4", start=0, end=5)
        assert r == [1]
        r = strmap(int, "-", start=0, end=5)
        assert r == [(0, 5)]

    def test_lstranges1(self):
        ranges = strmap(int, "1,2-3[0,2]")
        assert lstranges(ranges) == [1, [2, [0, 2]], [3, [0, 2]]]
        ranges = strmap(int, "1,2-4[0-4,2],6[1-3],9-10")
        assert lstranges(ranges) == [
            1,
            [2, [0, 1, 2, 3, 4, 2]],
            [3, [0, 1, 2, 3, 4, 2]],
            [4, [0, 1, 2, 3, 4, 2]],
            [6, [1, 2, 3]],
            9,
            10,
        ]
        ranges = strmap(int, "1,[2,4][0-4,2],6[1-3],9-10")
        assert lstranges(ranges) == [
            1,
            [2, [0, 1, 2, 3, 4, 2]],
            [4, [0, 1, 2, 3, 4, 2]],
            [6, [1, 2, 3]],
            9,
            10,
        ]
        ranges = strmap(int, "1,[2,4,6-7][0,3-4,2],6[1-3],9-10")
        assert lstranges(ranges) == [
            1,
            [2, [0, 3, 4, 2]],
            [4, [0, 3, 4, 2]],
            [6, [0, 3, 4, 2]],
            [7, [0, 3, 4, 2]],
            [6, [1, 2, 3]],
            9,
            10,
        ]
        ranges = strmap(int, "[82,83][10-13]")
        assert lstranges(ranges) == [[82, [10, 11, 12, 13]], [83, [10, 11, 12, 13]]]
        ranges = strmap(int, " [82,85][3]")
        assert lstranges(ranges) == [[82, [3]], [85, [3]]]
        ranges = strmap(int, "81,[82,85][3]")
        assert lstranges(ranges) == [81, [82, [3]], [85, [3]]]

    def test_lstranges2(self):
        ranges = strmap(int, "1:2:5")
        assert lstranges(ranges) == [1, 3, 5]
        ranges = strmap(int, "1-2-5")
        assert lstranges(ranges) == [1, 3, 5]

    def test_fileindex1(self):
        fname = "hello[1]"
        assert fileindex("hehlo")[1] is None
        assert fileindex(fname)[1] == 1
        assert fileindex("hehlo[1,2]")[1] == [1, 2]
        assert fileindex("hehlo[1-2]")[1] == [1, 2]
        assert fileindex("hehlo[1[1],2]")[1] == [[1, [1]], 2]

    def test_list2str(self):
        a = list2str([2, 4, 5, 6])
        assert a == "2, 4-6"
        a = list2str([2, 4, 5, 6, 8, 9])
        assert a == "2, 4-6, 8-9"
