import pytest
import numpy as np
import matplotlib as mpl
import scipy.spatial.distance as scipy_distance

from whatlies import Embedding, EmbeddingSet
from .common import validate_plot_general_properties


"""
*Guide*

Here are the plot's propertites which could be checked (some of them may not be applicable
for a particular plot):
    - type: the class type of collection in matplotlib to ensure the right kind of plot
        has been created.
    - data: the position of points, arrows or texts in the plot, depending on the plot's type.
    - x_label: label of x-axis.
    - y_label: label of y-axis.
    - tilte: title of the plot.
    - aspect: aspect ratio of plot, usually 'auto' unless `axis_option` argument is set.
    - color: color of points (in scatter plot) or arrows (in arrow plot). It should be rgba values.
    - label: label of points (in scatter plot) or arrows (in arrow plot).
"""


@pytest.fixture
def embset():
    names = ["red", "blue", "green", "yellow", "white"]
    vectors = np.random.rand(5, 3)
    embeddings = [Embedding(name, vector) for name, vector in zip(names, vectors)]
    return EmbeddingSet(*embeddings)


def test_embeddingset_plot_scatter_str_axis(embset):
    fig, ax = mpl.pyplot.subplots()
    embset.plot(kind="scatter", x_axis="blue", y_axis="red")
    vectors = []
    for emb in embset.embeddings.values():
        vec = []
        vec.append(emb > embset["blue"])
        vec.append(emb > embset["red"])
        vectors.append(vec)
    vectors = np.array(vectors)
    props = {
        "type": mpl.collections.PathCollection,
        "data": vectors,
        "x_label": "blue",
        "y_label": "red",
        "title": "",
        "label": list(embset.embeddings.keys()),
        "color": mpl.colors.to_rgba_array("steelblue"),
        "aspect": "auto",
    }
    assert isinstance(ax.collections[0], props["type"])
    assert np.array_equal(ax.collections[0].get_offsets(), props["data"])
    assert [t.get_text() for t in ax.texts] == props["label"]
    assert np.array_equal(ax.collections[0].get_facecolors(), props["color"])
    validate_plot_general_properties(ax, props)
    mpl.pyplot.close(fig)


def test_embeddingset_plot_arrow_str_axis(embset):
    fig, ax = mpl.pyplot.subplots()
    embset.plot(
        kind="arrow",
        x_axis="blue",
        y_axis="red",
        x_label="blue vec",
        y_label="red vec",
        title="test plot",
        color="yellow",
    )
    vectors = []
    for emb in embset.embeddings.values():
        vec = []
        vec.append(emb > embset["blue"])
        vec.append(emb > embset["red"])
        vectors.append(vec)
    vectors = np.array(vectors)
    props = {
        "type": mpl.collections.PolyCollection,
        "data": vectors,
        "x_label": "blue vec",
        "y_label": "red vec",
        "title": "test plot",
        "label": list(embset.embeddings.keys()),
        "color": mpl.colors.to_rgba_array("yellow"),
        "aspect": "auto",
    }
    UV = np.concatenate(
        (ax.collections[1].U[:, None], ax.collections[1].V[:, None]), axis=-1
    )
    assert isinstance(ax.collections[1], props["type"])
    assert np.array_equal(UV, props["data"])
    assert [t.get_text() for t in ax.texts] == props["label"]
    assert np.array_equal(ax.collections[1].get_facecolors(), props["color"])
    validate_plot_general_properties(ax, props)
    mpl.pyplot.close(fig)


def test_embeddingset_plot_scatter_integer_axis(embset):
    fig, ax = mpl.pyplot.subplots()
    embset.plot(kind="scatter", x_axis=0, y_axis=1, annot=False, color="black")
    vectors = np.concatenate((embset.to_X()[:, 0:1], embset.to_X()[:, 1:2]), axis=-1)
    props = {
        "type": mpl.collections.PathCollection,
        "data": vectors,
        "x_label": "Dimension 0",
        "y_label": "Dimension 1",
        "title": "",
        "color": mpl.colors.to_rgba_array("black"),
        "aspect": "auto",
        # Not applicable: label
    }
    assert isinstance(ax.collections[0], props["type"])
    assert np.array_equal(ax.collections[0].get_offsets(), props["data"])
    assert ax.texts == []
    assert np.array_equal(ax.collections[0].get_facecolors(), props["color"])
    validate_plot_general_properties(ax, props)
    mpl.pyplot.close(fig)


def test_embeddingset_plot_arrow_integer_axis(embset):
    fig, ax = mpl.pyplot.subplots()
    embset.plot(
        kind="arrow",
        x_axis=1,
        y_axis=2,
        x_label="1",
        y_label="2",
        color="yellow",
        axis_option="scaled",
    )
    vectors = np.concatenate((embset.to_X()[:, 1:2], embset.to_X()[:, 2:3]), axis=-1)
    props = {
        "type": mpl.collections.PolyCollection,
        "data": vectors,
        "x_label": "1",
        "y_label": "2",
        "title": "",
        "label": list(embset.embeddings.keys()),
        "color": mpl.colors.to_rgba_array("yellow"),
        "aspect": 1.0,
    }
    UV = np.concatenate(
        (ax.collections[1].U[:, None], ax.collections[1].V[:, None]), axis=-1
    )
    assert isinstance(ax.collections[1], props["type"])
    assert np.array_equal(UV, props["data"])
    assert [t.get_text() for t in ax.texts] == props["label"]
    assert np.array_equal(ax.collections[1].get_facecolors(), props["color"])
    validate_plot_general_properties(ax, props)
    mpl.pyplot.close(fig)


def test_embeddingset_plot_scatter_emb_axis(embset):
    fig, ax = mpl.pyplot.subplots()
    embset.plot(kind="scatter", x_axis=embset["green"], y_axis=embset["white"])
    vectors = []
    for emb in embset.embeddings.values():
        vec = []
        vec.append(emb > embset["green"])
        vec.append(emb > embset["white"])
        vectors.append(vec)
    vectors = np.array(vectors)
    props = {
        "type": mpl.collections.PathCollection,
        "data": vectors,
        "x_label": "green",
        "y_label": "white",
        "title": "",
        "label": list(embset.embeddings.keys()),
        "color": mpl.colors.to_rgba_array("steelblue"),
        "aspect": "auto",
    }
    assert isinstance(ax.collections[0], props["type"])
    assert np.array_equal(ax.collections[0].get_offsets(), props["data"])
    assert [t.get_text() for t in ax.texts] == props["label"]
    assert np.array_equal(ax.collections[0].get_facecolors(), props["color"])
    validate_plot_general_properties(ax, props)
    mpl.pyplot.close(fig)


def test_embeddingset_plot_arrow_emb_axis(embset):
    fig, ax = mpl.pyplot.subplots()
    embset.plot(
        kind="arrow",
        x_axis=embset["blue"],
        y_axis=embset["red"],
        x_label="xx",
        color="magenta",
    )
    vectors = []
    for emb in embset.embeddings.values():
        vec = []
        vec.append(emb > embset["blue"])
        vec.append(emb > embset["red"])
        vectors.append(vec)
    vectors = np.array(vectors)
    props = {
        "type": mpl.collections.PolyCollection,
        "data": vectors,
        "x_label": "xx",
        "y_label": "red",
        "title": "",
        "label": list(embset.embeddings.keys()),
        "color": mpl.colors.to_rgba_array("magenta"),
        "aspect": "auto",
    }
    UV = np.concatenate(
        (ax.collections[1].U[:, None], ax.collections[1].V[:, None]), axis=-1
    )
    assert isinstance(ax.collections[1], props["type"])
    assert np.array_equal(UV, props["data"])
    assert [t.get_text() for t in ax.texts] == props["label"]
    assert np.array_equal(ax.collections[1].get_facecolors(), props["color"])
    validate_plot_general_properties(ax, props)
    mpl.pyplot.close(fig)


def test_embeddingset_plot_arrow_integer_axis_with_str_axis_metric(embset):
    fig, ax = mpl.pyplot.subplots()
    embset.plot(
        kind="arrow",
        x_axis=1,
        y_axis=2,
        axis_metric="cosine_distance",
        x_label="1",
        y_label="2",
        color="yellow",
        axis_option="scaled",
    )
    vectors = np.concatenate((embset.to_X()[:, 1:2], embset.to_X()[:, 2:3]), axis=-1)
    props = {
        "type": mpl.collections.PolyCollection,
        "data": vectors,
        "x_label": "1",
        "y_label": "2",
        "title": "",
        "label": list(embset.embeddings.keys()),
        "color": mpl.colors.to_rgba_array("yellow"),
        "aspect": 1.0,
    }
    UV = np.concatenate(
        (ax.collections[1].U[:, None], ax.collections[1].V[:, None]), axis=-1
    )
    assert isinstance(ax.collections[1], props["type"])
    assert np.array_equal(UV, props["data"])
    assert [t.get_text() for t in ax.texts] == props["label"]
    assert np.array_equal(ax.collections[1].get_facecolors(), props["color"])
    validate_plot_general_properties(ax, props)
    mpl.pyplot.close(fig)


def test_embeddingset_plot_scatter_emb_axis_with_common_str_axis_metric(embset):
    fig, ax = mpl.pyplot.subplots()
    embset.plot(
        kind="scatter",
        x_axis=embset["green"],
        y_axis=embset["white"],
        axis_metric="euclidean",
    )
    vectors = []
    for emb in embset.embeddings.values():
        vectors.append(
            [
                scipy_distance.euclidean(emb.vector, embset["green"].vector),
                scipy_distance.euclidean(emb.vector, embset["white"].vector),
            ]
        )
    vectors = np.array(vectors)
    props = {
        "type": mpl.collections.PathCollection,
        "data": vectors,
        "x_label": "green",
        "y_label": "white",
        "title": "",
        "label": list(embset.embeddings.keys()),
        "color": mpl.colors.to_rgba_array("steelblue"),
        "aspect": "auto",
    }
    assert isinstance(ax.collections[0], props["type"])
    assert np.array_equal(ax.collections[0].get_offsets(), props["data"])
    assert [t.get_text() for t in ax.texts] == props["label"]
    assert np.array_equal(ax.collections[0].get_facecolors(), props["color"])
    validate_plot_general_properties(ax, props)
    mpl.pyplot.close(fig)


def test_embeddingset_plot_arrow_emb_axis_with_different_axis_metric(embset):
    fig, ax = mpl.pyplot.subplots()
    embset.plot(
        kind="arrow",
        x_axis=embset["blue"],
        y_axis="red",
        axis_metric=[scipy_distance.correlation, "cosine_similarity"],
        x_label="xx",
        color="magenta",
    )
    vectors = []
    for emb in embset.embeddings.values():
        vectors.append(
            [
                scipy_distance.correlation(emb.vector, embset["blue"].vector),
                1.0 - scipy_distance.cosine(emb.vector, embset["red"].vector),
            ]
        )
    vectors = np.array(vectors)
    props = {
        "type": mpl.collections.PolyCollection,
        "data": vectors,
        "x_label": "xx",
        "y_label": "red",
        "title": "",
        "label": list(embset.embeddings.keys()),
        "color": mpl.colors.to_rgba_array("magenta"),
        "aspect": "auto",
    }
    UV = np.concatenate(
        (ax.collections[1].U[:, None], ax.collections[1].V[:, None]), axis=-1
    )
    assert isinstance(ax.collections[1], props["type"])
    assert np.array_equal(UV, props["data"])
    assert [t.get_text() for t in ax.texts] == props["label"]
    assert np.array_equal(ax.collections[1].get_facecolors(), props["color"])
    validate_plot_general_properties(ax, props)
    mpl.pyplot.close(fig)


def test_embeddingset_plot_scatter_mixed_axis(embset):
    fig, ax = mpl.pyplot.subplots()
    embset.plot(kind="scatter", x_axis=embset["green"], y_axis=2)
    vectors = []
    for emb in embset.embeddings.values():
        vec = []
        vec.append(emb > embset["green"])
        vec.append(emb.vector[2])
        vectors.append(vec)
    vectors = np.array(vectors)
    props = {
        "type": mpl.collections.PathCollection,
        "data": vectors,
        "x_label": "green",
        "y_label": "Dimension 2",
        "title": "",
        "label": list(embset.embeddings.keys()),
        "color": mpl.colors.to_rgba_array("steelblue"),
        "aspect": "auto",
    }
    assert isinstance(ax.collections[0], props["type"])
    assert np.array_equal(ax.collections[0].get_offsets(), props["data"])
    assert [t.get_text() for t in ax.texts] == props["label"]
    assert np.array_equal(ax.collections[0].get_facecolors(), props["color"])
    validate_plot_general_properties(ax, props)
    mpl.pyplot.close(fig)


def test_embeddingset_plot_arrow_mixed_axis(embset):
    fig, ax = mpl.pyplot.subplots()
    embset.plot(kind="arrow", x_axis=1, y_axis="red", x_label="xx", color="magenta")
    vectors = []
    for emb in embset.embeddings.values():
        vec = []
        vec.append(emb.vector[1])
        vec.append(emb > embset["red"])
        vectors.append(vec)
    vectors = np.array(vectors)
    props = {
        "type": mpl.collections.PolyCollection,
        "data": vectors,
        "x_label": "xx",
        "y_label": "red",
        "title": "",
        "label": list(embset.embeddings.keys()),
        "color": mpl.colors.to_rgba_array("magenta"),
        "aspect": "auto",
    }
    UV = np.concatenate(
        (ax.collections[1].U[:, None], ax.collections[1].V[:, None]), axis=-1
    )
    assert isinstance(ax.collections[1], props["type"])
    assert np.array_equal(UV, props["data"])
    assert [t.get_text() for t in ax.texts] == props["label"]
    assert np.array_equal(ax.collections[1].get_facecolors(), props["color"])
    validate_plot_general_properties(ax, props)
    mpl.pyplot.close(fig)


def test_embeddingset_plot_raises_error_when_str_axis_not_exists(embset):
    with pytest.raises(KeyError):
        embset.plot(x_axis="bnb", y_axis="blue")
    with pytest.raises(KeyError):
        embset.plot(x_axis="red", y_axis="clk")
