import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
from random import randint
from matplotlib import animation

#x - prey
#y - predator

def get_chart(x0, y0, time, h, alpha, beta, gamma, delta):
    all_x = []
    all_y = []
    x_prior = x0
    y_prior = y0
    for i in range(0, len(time)):
        x = x_prior + h*(alpha*x_prior - beta*x_prior*y_prior)
        y = y_prior + h*(delta*x*y_prior - gamma*y_prior)
        all_x.append(x_prior)
        all_y.append(y_prior)
        x_prior = x
        y_prior = y
    return all_x, all_y

def printPoint(x,y):
    point = '(' + str(x) + ',' + str(y) + ')'
    return point

def rand_colour_set(charts):
    colour = []
    for i in range(charts):
        colour.append(np.random.random(3))
    return colour

def rand_chart_set(charts, h, alpha, beta, gamma, delta, time, min, max):
    X_set = []
    Y_set = []
    for i in range(charts):
        prey = randint(min, max)
        predator = randint(min, max)
        line1, line2 = get_chart(prey, predator, time, h, alpha, beta, gamma, delta)
        X_set.append(line1)
        Y_set.append(line2)
    return (X_set, Y_set)

def chart_3D(X_set, Y_set, time, ax3D, colour_set):
    for i in range(0, len(X_set)):
        line1 = X_set[i]
        line2 = Y_set[i]
        ch = str(i + 1) + '(ofiara, drapieżnik): ' + printPoint(line1[0], line2[0])
        color = colour_set[i]
        ax3D.plot(line1, line2, time, color = color, label = ch)
        if(i == 0):
            ax3D.scatter(line1[0], line2[0], time[0], c='g', label='start')
            ax3D.scatter(line1[len(line1) - 1], line2[len(line1) - 1], time[len(line1) - 1], c='r', label='end')
        else:
            ax3D.scatter(line1[0], line2[0], time[0], c='g')
            ax3D.scatter(line1[len(line1) - 1], line2[len(line1) - 1], time[len(line1) - 1], c='r')

T, t0 = 50, 0
h = 0.01
alpha, beta, gamma, delta = 2, 1.1 , 1 ,0.9
min, max = 1, 30
charts = 7
time = np.arange(t0, T, h).tolist()
X_set, Y_set = rand_chart_set(charts, h, alpha, beta, gamma, delta, time, min, max)

fig2D = plt.figure()
ax2D = fig2D.add_subplot(111)
colours = rand_colour_set(charts)
charts_animate = [ax2D.plot([],[], c = colours[i], label = str(i + 1))[0] for i in range(charts)]


def animate(frame):
    for i,ch in enumerate(charts_animate):
        ch.set_data(X_set[i][:frame], Y_set[i][:frame])
    return charts_animate

anim = animation.FuncAnimation(fig2D, animate, interval=0.01, blit=True)

ax2D.set_xlabel('ofiary')
ax2D.set_ylabel('drapieżnicy')
ax2D.set_title("model ofirara-drapieżnik 2D")
ax2D.axis([-5,max + 20,-5,max + 10])

fig3D = plt.figure()
ax3D = fig3D.add_subplot(111, projection = '3d')
ax3D.set_xlabel("ofiary")
ax3D.set_ylabel("drapieżnicy")
ax3D.set_zlabel("czas")
ax3D.set_title("model ofirara-drapieżnik 3D")

chart_3D(X_set, Y_set, time, ax3D, colours)
ax2D.legend(loc = 'best')
ax3D.legend(loc = 'best')

ax2D.grid()
ax3D.grid()

plt.show()