# -*- coding: utf-8 -*-
# Copyright 2016-2020 The pyXem developers
#
# This file is part of pyXem.
#
# pyXem 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.
#
# pyXem 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 pyXem.  If not, see <http://www.gnu.org/licenses/>.

import numpy as np
import pytest
from pyxem.signals.crystallographic_map import CrystallographicMap
from pyxem.signals.crystallographic_map import load_mtex_map
from pyxem.signals.crystallographic_map import _distance_from_fixed_angle
from transforms3d.euler import euler2quat, quat2axangle
from transforms3d.quaternions import qmult, qinverse
import os


def worker_for_test_CrystallographicMap_io(mapp):
    mapp.save_mtex_map("file_01.txt")
    lmap = load_mtex_map("file_01.txt")
    os.remove("file_01.txt")
    return mapp, lmap


def get_distance_between_two_angles_longform(angle_1, angle_2):
    """
    Using the long form to find the distance between two angles in euler form
    """
    q1 = euler2quat(*np.deg2rad(angle_1), axes="rzxz")
    q2 = euler2quat(*np.deg2rad(angle_2), axes="rzxz")
    # now assume transform of the form MODAL then Something = TOTAL
    # so we want to calculate MODAL^{-1} TOTAL

    q_from_mode = qmult(qinverse(q2), q1)
    axis, angle = quat2axangle(q_from_mode)
    return np.rad2deg(angle)


@pytest.fixture()
def sp_cryst_map():
    """
    Generates a single phase Crystallographic Map
    """
    base = np.empty((4, 3), dtype="object")
    base[0] = [
        0,
        np.array([5, 17, 6]),
        {"correlation": 3e-17, "orientation_reliability": 0.5},
    ]
    base[1] = [
        0,
        np.array([6, 17, 6]),
        {"correlation": 2e-17, "orientation_reliability": 0.4},
    ]
    base[2] = [
        0,
        np.array([12, 3, 6]),
        {"correlation": 4e-17, "orientation_reliability": 0.3},
    ]
    base[3] = [
        0,
        np.array([12, 3, 5]),
        {"correlation": 8e-16, "orientation_reliability": 0.2},
    ]
    crystal_map = CrystallographicMap(base.reshape((2, 2, 3)))
    crystal_map.method = "template_matching"
    return crystal_map


@pytest.fixture()
def dp_cryst_map():
    """
    Generates a Crystallographic Map with two phases
    """
    base = np.empty((4, 3), dtype="object")
    base[0] = [
        0,
        np.array([5, 17, 6]),
        {
            "correlation": 3e-17,
            "orientation_reliability": 0.5,
            "phase_reliability": 0.6,
        },
    ]
    base[1] = [
        1,
        np.array([6, 17, 6]),
        {
            "correlation": 2e-17,
            "orientation_reliability": 0.4,
            "phase_reliability": 0.7,
        },
    ]
    base[2] = [
        0,
        np.array([12, 3, 6]),
        {
            "correlation": 4e-17,
            "orientation_reliability": 0.3,
            "phase_reliability": 0.1,
        },
    ]
    base[3] = [
        0,
        np.array([12, 3, 5]),
        {
            "correlation": 8e-16,
            "orientation_reliability": 0.2,
            "phase_reliability": 0.8,
        },
    ]
    crystal_map = CrystallographicMap(base.reshape((2, 2, 3)))
    crystal_map.method = "template_matching"
    return crystal_map


@pytest.fixture()
def mod_cryst_map():
    """
    Generates a Crystallographic Map with (5,17,6) as the modal angle
    """
    base = np.empty((6, 3), dtype="object")
    base[0] = [
        0,
        np.array([5, 17, 6]),
        {"correlation": 5e-17, "orientation_reliability": 0.5},
    ]
    base[1] = [
        0,
        np.array([5, 17, 6]),
        {"correlation": 5e-17, "orientation_reliability": 0.5},
    ]
    base[2] = [
        0,
        np.array([6, 19, 6]),
        {"correlation": 5e-17, "orientation_reliability": 0.5},
    ]
    base[3] = [
        0,
        np.array([7, 19, 6]),
        {"correlation": 5e-17, "orientation_reliability": 0.5},
    ]
    base[4] = [
        0,
        np.array([8, 19, 6]),
        {"correlation": 5e-17, "orientation_reliability": 0.5},
    ]
    base[5] = [
        0,
        np.array([9, 19, 6]),
        {"correlation": 5e-17, "orientation_reliability": 0.5},
    ]
    crystal_map = CrystallographicMap(base.reshape((3, 2, 3)))
    crystal_map.method = "template_matching"
    return crystal_map


@pytest.fixture()
def dp_cryst_map_vector():
    """
    Generates a Crystallographic Map with two phases from vector matching
    """
    base = np.empty((4, 3), dtype="object")
    base[0] = [
        0,
        np.array([5, 17, 6]),
        {
            "match_rate": 0.5,
            "ehkls": np.array([0.1, 0.05, 0.2]),
            "total_error": 0.1,
            "orientation_reliability": 13.2,
            "phase_reliability": 42.0,
        },
    ]
    base[1] = [
        1,
        np.array([6, 17, 6]),
        {
            "match_rate": 0.5,
            "ehkls": np.array([0.1, 0.05, 0.2]),
            "total_error": 0.1,
            "orientation_reliability": 13.2,
            "phase_reliability": 42.0,
        },
    ]
    base[2] = [
        0,
        np.array([12, 3, 6]),
        {
            "match_rate": 0.5,
            "ehkls": np.array([0.1, 0.05, 0.2]),
            "total_error": 0.1,
            "orientation_reliability": 13.2,
            "phase_reliability": 42.0,
        },
    ]
    base[3] = [
        0,
        np.array([12, 3, 5]),
        {
            "match_rate": 0.5,
            "ehkls": np.array([0.1, 0.05, 0.2]),
            "total_error": 0.1,
            "orientation_reliability": 13.2,
            "phase_reliability": 42.0,
        },
    ]
    crystal_map = CrystallographicMap(base.reshape((2, 2, 3)))
    crystal_map.method = "vector_matching"
    return crystal_map


class TestMapCreation:
    def test_get_phase_map(self, sp_cryst_map):
        phasemap = sp_cryst_map.get_phase_map()
        assert phasemap.isig[0, 0] == 0

    def test_get_orientation_map(self, sp_cryst_map):
        orimap = sp_cryst_map.get_orientation_map()
        assert orimap.isig[0, 0] == 0

    @pytest.mark.parametrize(
        "metric, value",
        [
            ("correlation", 3e-17),
            ("orientation_reliability", 0.5),
            ("phase_reliability", 0.6),
        ],
    )
    def test_get_metric_map_template_match(self, dp_cryst_map, metric, value):
        metric_map = dp_cryst_map.get_metric_map(metric)
        assert metric_map.isig[0, 0] == value

    @pytest.mark.parametrize(
        "metric, value",
        [
            ("match_rate", 0.5),
            ("ehkls", np.array([0.1, 0.05, 0.2])),
            ("total_error", 0.1),
            ("orientation_reliability", 13.2),
            ("phase_reliability", 42.0),
        ],
    )
    def test_get_metric_map_vector_match(self, dp_cryst_map_vector, metric, value):
        metric_map = dp_cryst_map_vector.get_metric_map(metric)
        assert np.allclose(metric_map.isig[0, 0], value)

    def test_get_metric_map_template_match_bad_metric(self, sp_cryst_map):
        # Note - uses regex via re.search()
        with pytest.raises(
            ValueError, match=r"metric .* is not valid for template matching"
        ):
            metric_map = sp_cryst_map.get_metric_map("no metric")

    def test_get_metric_map_vector_match_bad_metric(self, dp_cryst_map_vector):
        # Note - uses regex via re.search()
        with pytest.raises(
            ValueError, match=r"metric .* is not valid for vector matching"
        ):
            metric_map = dp_cryst_map_vector.get_metric_map("no metric")

    def test_get_metric_map_no_method(self):
        crystal_map = CrystallographicMap(np.array([[1]]))

        with pytest.raises(
            ValueError, match="crystallographic mapping method must be specified"
        ):
            metric_map = crystal_map.get_metric_map("no metric")


class TestMTEXIO:
    def test_Crystallographic_Map_io_sp(self, sp_cryst_map):
        saved, loaded = worker_for_test_CrystallographicMap_io(sp_cryst_map)
        saved.method = "template_matching"
        loaded.method = "template_matching"
        np.testing.assert_allclose(
            saved.isig[0].data.astype(np.int), loaded.isig[0].data.astype(np.int)
        )
        np.testing.assert_allclose(
            np.array(saved.isig[1].data.tolist()),
            np.array(loaded.isig[1].data.tolist()),
        )
        np.testing.assert_allclose(
            saved.get_metric_map("correlation").data,
            loaded.get_metric_map("correlation").data,
        )

    def test_Crystallographic_Map_io_dp(self, dp_cryst_map):
        saved, loaded = worker_for_test_CrystallographicMap_io(dp_cryst_map)
        saved.method = "template_matching"
        loaded.method = "template_matching"
        np.testing.assert_allclose(
            saved.isig[0].data.astype(np.int), loaded.isig[0].data.astype(np.int)
        )
        np.testing.assert_allclose(
            np.array(saved.isig[1].data.tolist()),
            np.array(loaded.isig[1].data.tolist()),
        )
        np.testing.assert_allclose(
            saved.get_metric_map("correlation").data,
            loaded.get_metric_map("correlation").data,
        )


class TestModalAngularFunctionality:
    def test_get_distance_from_modal(self, mod_cryst_map):
        # function runs without error
        formal = mod_cryst_map.get_distance_from_modal_angle()
        assert True

    def test_get_modal_angles(self, mod_cryst_map):
        # modal angle is found correctly
        out = mod_cryst_map.get_modal_angles()
        assert np.allclose(out[0], [5, 17, 6])
        assert np.allclose(out[1], (2 / 6))

    def test_get_distance_from_fixed_angle(self):
        # distance between two angles is found correctly
        angle_1 = [1, 1, 3]
        angle_2 = [1, 1, 4]
        implemented = _distance_from_fixed_angle([angle_1], angle_2)
        testing = get_distance_between_two_angles_longform(angle_1, angle_2)
        assert np.allclose(implemented, testing)
        assert np.allclose(implemented, 1)
