# -*- coding: utf-8 -*-
# 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


def test_init_rodrigues_plot():
    fig = plt.figure()
    ax = fig.add_subplot(projection="rodrigues", auto_add_to_figure=False)
    assert isinstance(ax, RodriguesPlot)


def test_init_axangle_plot():
    fig = plt.figure()
    ax = fig.add_subplot(projection="axangle", auto_add_to_figure=False)
    assert isinstance(ax, AxAnglePlot)


def test_RotationPlot_methods():
    """This code is lifted from demo-3-v0.1."""
    misori = Misorientation([1, 1, 1, 1])  # any will do
    ori = Orientation.random()
    fig = plt.figure()
    ax = fig.add_subplot(
        projection="axangle", proj_type="ortho", auto_add_to_figure=False
    )
    ax.scatter(misori)
    ax.scatter(ori)
    ax.plot(misori)
    ax.plot(ori)
    ax.plot_wireframe(OrientationRegion.from_symmetry(D6, D6))
    plt.close("all")

    # Clear the edge case
    ax.transform(np.asarray([1, 1, 1]))


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


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


def test_RotationPlot_map_into_symmetry_reduced_zone():
    # orientations are (in, out) of D6 fundamental zone
    ori = Orientation(((1, 0, 0, 0), (0.5, 0.5, 0.5, 0.5)))
    ori.symmetry = D6
    fz = OrientationRegion.from_symmetry(ori.symmetry)
    assert np.allclose(ori < fz, (True, False))
    # test map_into_symmetry_reduced_zone in RotationPlot.transform
    fig = ori.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
    ori.symmetry = C1
    fig2 = ori.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)


def test_correct_aspect_ratio():
    # Set up figure the "old" way
    fig = plt.figure()
    ax = fig.add_subplot(
        projection="axangle", proj_type="ortho", auto_add_to_figure=False
    )

    # 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)

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