# Copyright 2018-2024 the orix developers
#
# This file is part of orix.
#
# orix is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# orix is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with orix.  If not, see <http://www.gnu.org/licenses/>.

from matplotlib import pyplot as plt
import numpy as np
import pytest

from orix.plot import AxAnglePlot, RodriguesPlot, RotationPlot
from orix.quaternion import Misorientation, Orientation, OrientationRegion
from orix.quaternion.symmetry import C1, D6


class TestRodriguesPlot:
    def test_creation(self):
        fig = plt.figure()
        ax = fig.add_subplot(projection="rodrigues")
        assert isinstance(ax, RodriguesPlot)


class TestAxisAnglePlot:
    def test_creation(self):
        fig = plt.figure()
        ax = fig.add_subplot(projection="axangle")
        assert isinstance(ax, AxAnglePlot)

        plt.close("all")

    def test_rotation_plot(self):
        M = Misorientation.random()
        O = Orientation.random()
        fig = plt.figure()
        ax = fig.add_subplot(projection="axangle", proj_type="ortho")
        ax.scatter(M)
        ax.scatter(O)
        ax.plot(M)
        ax.plot(O)
        ax.plot_wireframe(OrientationRegion.from_symmetry(D6, D6))

        ax.transform(np.asarray([1, 1, 1]))  # Edge case

        plt.close("all")

    def test_get_plot_data(self):
        empty = OrientationRegion.from_symmetry(C1, C1)
        _ = empty.get_plot_data()

    def test_rotation_plot_transform_fundamental_zone_raises(self):
        fig = plt.figure()
        ax = RotationPlot(fig)
        fig.add_axes(ax)
        with pytest.raises(TypeError, match="fundamental_zone is not an "):
            ax.transform(Orientation.random(), fundamental_zone=1)

    def test_rotation_plot_map_into_symmetry_reduced_zone(self):
        # Orientations are (in, out) of D6 fundamental zone
        O = Orientation(((1, 0, 0, 0), (0.5, 0.5, 0.5, 0.5)))
        O.symmetry = D6
        fz = OrientationRegion.from_symmetry(O.symmetry)
        assert np.allclose(O < fz, (True, False))

        # test map_into_symmetry_reduced_zone in RotationPlot.transform
        fig = O.scatter(return_figure=True)
        xyz_symmetry = fig.axes[0].collections[1]._offsets3d

        # compute same plot again but with C1 symmetry where both orientations are in C1 FZ
        O.symmetry = C1
        fig2 = O.scatter(return_figure=True)
        xyz = fig2.axes[0].collections[1]._offsets3d

        # test that the plotted points are not the same
        assert not np.allclose(xyz_symmetry, xyz)

        plt.close("all")

    def test_correct_aspect_ratio(self):
        fig = plt.figure()
        ax = fig.add_subplot(projection="axangle", proj_type="ortho")

        # Check aspect ratio
        x_old, _, z_old = ax.get_box_aspect()
        assert np.allclose(x_old / z_old, 1.334, atol=1e-3)

        fr = OrientationRegion.from_symmetry(D6)
        ax._correct_aspect_ratio(fr, set_limits=False)

        x_new, _, z_new = ax.get_box_aspect()
        assert np.allclose(x_new / z_new, 3, atol=1e-3)

        assert np.allclose(ax.get_xlim(), [0, 1], atol=0.1)
        ax._correct_aspect_ratio(fr)  # set_limits=True is default
        assert np.allclose(ax.get_xlim(), [-np.pi / 2, np.pi / 2])

        plt.close("all")
