from math import sqrt
from collections import OrderedDict

import numpy as np
import matplotlib.pyplot as plt
from prettytable import PrettyTable

if __name__ == '__main__':
    solutions_table = PrettyTable('x1 x2 x3 x4 p1 p2 p3 p4'.split())
    solutions_table.align = 'l'

    verification_table = PrettyTable('x1 dx1/dt dx2/dt dx3/dt dx4/dt Jacobian'.split())
    verification_table.align = 'l'

    plt.title('Bifurcation Diagram')
    plt.xlabel('p3')
    plt.ylabel('p2')
    plt.grid(True)

    points = OrderedDict()
    p1 = 2

    for x1 in (val * 0.1 for val in range(1, 40)):
        x3 = 2 * p1 - x1

        if x1 == p1:
            solutions_table.add_row('-' * 8)
            continue

        alpha1 = 20 * (x1 ** 2 + x3 ** 2)
        beta1 = 2 * x1 ** 2 * x3 ** 2 + 10 * (x1 ** 2 + x3 ** 2)
        gamma1 = -20 * x1 * x3
        delta1 = x1 ** 2 * x3 ** 2

        alpha2 = alpha1 - (80 * p1 * x3 * (x1 - p1)) / x1
        beta2 = beta1 - 4 * x3 * (x1 - p1) * (x1 ** 2 + (10 * p1) / x1)
        gamma2 = 10 * (x1 ** 2 + x3 ** 2 - 4 * p1 * x3)
        delta2 = delta1 - 2 * x1 ** 2 * x3 * (x1 - p1)

        alpha = alpha1 * gamma2 - alpha2 * gamma1
        beta = beta1 * gamma2 - beta2 * gamma1
        delta = delta1 * gamma2 - delta2 * gamma1

        p3 = (-beta + sqrt(beta ** 2 - 4 * alpha * delta)) / (2 * alpha)

        if not p3:
            solutions_table.add_row('-' * 8)
            continue

        p4 = 10 * p3
        p2 = -(alpha2 * p3 + beta2 + delta2 / p3) / gamma2
        x2 = (p2 * x1 + (2 * p3 + 1) * (x1 - p1)) / (x1 ** 2)
        x4 = x2 + (2 * p3 + 1) * (x1 - p1) / p4

        points[p3] = p2
        solutions_table.add_row((x1, x2, x3, x4, p1, p2, p3, p4))

        jacobian_matrix = np.array((2 * x1 * x2 - p2 - p3 - 1, x1 ** 2, p3, 0,
                                    -2 * x1 * x2 + p2, -x1 ** 2 - p4, 0, p4, p3, 0,
                                    2 * x3 * x4 - p2 - p3 - 1, x3 ** 2,
                                    0, p4, -2 * x3 * x4 + p2, -x3 ** 2 - p4))
        jacobian_matrix = jacobian_matrix.reshape(4, 4)
        fmt = '{:.0e}'
        dx1 = fmt.format(p1 - (p2 + 1) * x1 + x1 ** 2 * x2 + p3 * (x3 - x1))
        dx2 = fmt.format(p2 * x1 - x1 ** 2 * x2 + p4 * (x4 - x2))
        dx3 = fmt.format(p1 - (p2 + 1) * x3 + x3 ** 2 * x4 - p3 * (x3 - x1))
        dx4 = fmt.format(p2 * x3 - x3 ** 2 * x4 - p4 * (x4 - x2))
        det = fmt.format(np.linalg.det(jacobian_matrix))
        verification_table.add_row((x1, dx1, dx2, dx3, dx4, det))

    print 'Solution:\n', solutions_table
    print '\n Verification:\n', verification_table

    plt.plot(points.keys(), points.values())[0].axes.set_xscale('log')
    points.clear()

    for p3 in (val * 0.1 for val in range(-10, 5) if val != 0):
        points[p3] = (2 * p3 + 1) * (p1 ** 2 + 20 * p3) / (20 * p3)

    plt.plot(points.keys(), points.values())
    plt.show()