#!/usr/bin/env python
import numpy as np
import matplotlib.pyplot as plt


def vectors_uniform(k):
    """Uniformly generates k vectors."""
    vectors = []
    for a in np.linspace(0, 2 * np.pi, k, endpoint=False):
        vectors.append(2 * np.array([np.sin(a), np.cos(a)]))
    return vectors


def visualize_transformation(A, vectors):
    """Plots original and transformed vectors for a given 2x2 transformation matrix A and a list of 2D vectors."""
    for i, v in enumerate(vectors):
        # Plot original vector.
        plt.quiver(0.0, 0.0, v[0], v[1], width=0.008, color="blue", scale_units='xy', angles='xy', scale=1,
                   zorder=4)
        plt.text(v[0]/2 + 0.25, v[1]/2, "v{0}".format(i), color="blue")

        # Plot transformed vector.
        tv = A.dot(v)
        plt.quiver(0.0, 0.0, tv[0], tv[1], width=0.005, color="magenta", scale_units='xy', angles='xy', scale=1,
                   zorder=4)
        plt.text(tv[0] / 2 + 0.25, tv[1] / 2, "v{0}'".format(i), color="magenta")
    plt.xlim([-6, 6])
    plt.ylim([-6, 6])
    plt.margins(0.05)
    # Plot eigenvectors
    plot_eigenvectors(A)
    plt.show()


def visualize_vectors(vectors, color="green"):
    """Plots all vectors in the list."""
    for i, v in enumerate(vectors):
        plt.quiver(0.0, 0.0, v[0], v[1], width=0.006, color=color, scale_units='xy', angles='xy', scale=1,
                   zorder=4)
        plt.text(v[0] / 2 + 0.25, v[1] / 2, "eigv{0}".format(i), color=color)


def plot_eigenvectors(A):
    """Plots all eigenvectors of the given 2x2 matrix A."""
    # TODO: Zad. 4.1. Oblicz wektory własne A. Możesz wykorzystać funkcję np.linalg.eig
    # eigvec = []
    (eigvalues, eigvec) = np.linalg.eig(A)
    # TODO: Zad. 4.1. Upewnij się poprzez analizę wykresów, że rysowane są poprawne wektory własne (łatwo tu o pomyłkę).
    visualize_vectors(eigvec.T)


def EVD_decomposition(A):
    # TODO: Zad. 4.2. Uzupełnij funkcję tak by obliczała rozkład EVD zgodnie z zadaniem.
    (eigvalues, K) = np.linalg.eig(A)
    K_inv = np.linalg.inv(K)
    L = np.dot(np.dot(K_inv, A), K)

    print("Eigvalues")
    print(eigvalues)
    print("K")
    print(K)
    print("L")
    print(L)

    print("A")
    print(A)
    print("KLK^-1")
    print(np.dot(np.dot(K, L), K_inv))

def calculate_vector_length(vector):
    result = 0
    for v in vector:
        result += v**2
    return np.sqrt(result)


def plot_attractors(A):

    vectors = vectors_uniform(k=20)

    (eigvalues, K) = np.linalg.eig(A)

    normalized_eigen_vectors = []

    for k in K.T:
        k_length = calculate_vector_length(k)
        normalized_eigen_vectors.append(k/k_length)
        normalized_eigen_vectors.append(-k/k_length)


    colors = []
    for i in range(int(len(normalized_eigen_vectors)/2)):
        colors.append(((0.1/(i+0.1))**2, (i-(i/(1+i)))/(i+1), (0.5/(i+0.5)), 0.8))
        colors.append((1-(1/(i+1))**2, (i/(1+i)), 1-(1/(i+1)), 0.4))

    for i, eigen_vectors in enumerate(normalized_eigen_vectors):
        plt.quiver(0.0, 0.0, eigen_vectors[0], eigen_vectors[1], width=0.02, color=colors[i], scale_units='xy', angles='xy', scale=1,
                   zorder=4)

    # TODO: Zad. 4.3. Uzupełnij funkcję tak by generowała wykres z atraktorami.
    for i, v in enumerate(vectors):
        # Plot original vector.
        new_v = v/calculate_vector_length(v)
        for _ in range(10):
            new_v = np.dot(A, new_v)
            new_v = new_v/calculate_vector_length(new_v)

        min_distance = float("inf")
        min_index = None
        for index, normalized_eigen_vector in enumerate(normalized_eigen_vectors):
            dist = np.linalg.norm(new_v - normalized_eigen_vector)
            if dist < min_distance:
                min_distance = dist
                min_index = index

        v /= calculate_vector_length(v)
        # plt.quiver(0.0, 0.0, new_v[0], new_v[1], width=0.008, color="yellow", scale_units='xy', angles='xy', scale=1,
        #            zorder=4)

        plt.quiver(0.0, 0.0, v[0], v[1], width=0.008, color=colors[min_index], scale_units='xy', angles='xy', scale=1,
                   zorder=4)
        # plt.text(v[0] / 2 + 0.25, v[1] / 2, "v{0}".format(i), color="blue")
    plt.xlim([-1.1, 1.1])
    plt.ylim([-1.1, 1.1])
    plt.margins(0.05)
    plt.show()


def test_A1(vectors):
    """Standard scaling transformation."""
    A = np.array([[2, 0],
                  [0, 2]])
    EVD_decomposition(A)
    visualize_transformation(A, vectors)
    plot_attractors(A)


def test_A2(vectors):
    A = np.array([[-1, 2],
                  [2, 1]])
    EVD_decomposition(A)
    visualize_transformation(A, vectors)
    plot_attractors(A)


def test_A3(vectors):
    A = np.array([[3, 1],
                  [0, 2]])
    EVD_decomposition(A)
    visualize_transformation(A, vectors)
    plot_attractors(A)



if __name__ == "__main__":
    vectors = vectors_uniform(k=8)


    test_A1(vectors)
    test_A2(vectors)
    test_A3(vectors)