# 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/.
import math as m
from numbers import Integral, Real

import numpy as np
import pytest

from sisl import Atom, Geometry, Lattice, Sphere


@pytest.fixture
def setup():
    class t():
        def __init__(self):
            bond = 1.42
            sq3h = 3.**.5 * 0.5
            self.lattice = Lattice(np.array([[1.5, sq3h, 0.],
                                             [1.5, -sq3h, 0.],
                                             [0., 0., 10.]], np.float64) * bond, nsc=[3, 3, 1])
            C = Atom(Z=6, R=bond * 1.01)
            self.g = Geometry(np.array([[0., 0., 0.],
                                        [1., 0., 0.]], np.float64) * bond,
                              atoms=C, lattice=self.lattice)

            C = Atom(Z=6, R=[bond * 1.01] * 2)
            self.g2 = Geometry(np.array([[0., 0., 0.],
                                         [1., 0., 0.]], np.float64) * bond,
                               atoms=C, lattice=self.lattice)
    return t()


@pytest.mark.geometry
class TestGeometryReturn:

    def test_fl_o(self, setup):
        # first o is always one element longer than the number of atoms
        assert np.all(setup.g.firsto == [0, 1, 2])
        assert np.all(setup.g.lasto == [0, 1])

        # first o is always one element longer than the number of atoms
        assert np.all(setup.g2.firsto == [0, 2, 4])
        assert np.all(setup.g2.lasto == [1, 3])

    def test_rij1(self, setup):
        d = setup.g.rij(0, 1)
        assert isinstance(d, Real)
        d = setup.g.rij(0, [0, 1])
        assert len(d) == 2
        d = setup.g.rij([0], [0, 1])
        assert len(d) == 1

        # also test orij
        d = setup.g.orij(0, 1)
        assert isinstance(d, Real)
        d = setup.g.orij(0, [0, 1])
        assert len(d) == 2
        d = setup.g.orij([0], [0, 1])
        assert len(d) == 1

    def test_rij2(self, setup):
        d = setup.g.rij([0, 1], [0, 1])
        assert np.allclose(d, [0., 0.])

    def test_osc2uc(self, setup):
        # single value
        d = setup.g.osc2uc(0)
        assert d == 0
        # more values
        d = setup.g.osc2uc([0, 1, 2])
        assert len(d) == 3
        assert np.allclose(d, [0, 1, 0])

    def test_slice1(self, setup):
        d = setup.g[1]
        assert len(d) == 3
        d = setup.g[[1, 2]]
        assert d.shape == (2, 3)
        d = setup.g[1:3]
        assert d.shape == (2, 3)
        d = setup.g[2:4]
        assert d.shape == (2, 3)

        d = setup.g[1, 2]
        assert d == 0.
        d = setup.g[1, 1:3]
        assert d.shape == (2,)
        d = setup.g[1, :]
        assert d.shape == (3,)
        d = setup.g[2:4, :]
        assert d.shape == (2, 3)
        d = setup.g[2:4, 1]
        assert d.shape == (2,)
        d = setup.g[2:4, 1:3]
        assert d.shape == (2, 2)
        d = setup.g[2:10:2, 1:3]
        assert d.shape == (4, 2)
        d = setup.g[2:10:2, 2]
        assert d.shape == (4,)

        d = setup.g[None, 2]
        assert d.shape == (len(setup.g),)

    def test_a2o(self, setup):
        d = setup.g.a2o(1)
        assert isinstance(d, Integral)
        d = setup.g.a2o([0, 1])
        assert len(d) == 2
        d = setup.g2.a2o(1)
        assert isinstance(d, Integral)
        d = setup.g2.a2o([0, 1])
        assert len(d) == 2

        d = setup.g.a2o(1, True)
        assert len(d) == 1
        d = setup.g.a2o([0, 1], True)
        assert len(d) == 2
        d = setup.g2.a2o(1, True)
        assert len(d) == 2
        d = setup.g2.a2o([0, 1], True)
        assert len(d) == 4

    def test_o2a(self, setup):
        d = setup.g.o2a(1)
        assert isinstance(d, Integral)
        d = setup.g.o2a([0, 1])
        assert len(d) == 2

    def test_axyz(self, setup):
        d = setup.g.axyz(1)
        assert len(d) == 3
        d = setup.g.axyz([0, 1])
        assert d.shape == (2, 3)
