# chaotic system of equations visualizer
# code by : nando takeshiro
# email: nandotakeshiro@hotmail.com
# please keep the above as part of the code



import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import colorsys
import numpy as np
from random import random
import time
import winsound


# main function for saving multiple runs of frames
def main():

    # set up figures and axes
    fig = plt.figure(figsize=(12,12))
    ax = fig.add_subplot(111, projection = '3d', aspect='equal')
    ax.axis('off')
    ax.set_axis_bgcolor((0,  0,  0))

    # lists of x, y and z coordinates
    xl, yl, zl = [], [], []

    # b sets up axis limits
    b = 3.5

    # time step
    dt = np.pi/32.

    # number of lines to draw
    sc = 250.
    sci = int(sc)

    # scalar size
    u = random() * 5


    # random scalars
    p = (-(u/2) + (u * random())) * np.pi
    q = -(u/2) + (u * random())
    r = -(u/2) + (u * random())
    s = -(u/2) + (u * random())
    v = -(u/2) + (u * random())
    w = -(u/2) + (u * random())


    # initial radius
    f = 4 * random()

    # index to choose x, y or z
    def rnd():
        i = int(3 * random())
        return i

    # indices
    a1, a2, a3 = rnd(), rnd(), rnd()

    # vectorfield - main formula
    def vectorfield(x,y,z):
        xyz = [x,y,z]
        i = q * np.cos(xyz[a3] / (xyz[a2] + p * (xyz[a1] + q) * np.pi + z) * np.pi)
        j = r * np.sin(xyz[a2] * np.cos(np.pi * x) / (xyz[a2] + r * -(x + s)) * np.pi + x)
        k = w * np.sin(xyz[a1] / (xyz[a2] + v / (xyz[a3] - xyz[a2] + w)) * np.pi + xyz[a1])

        return i, j, k

    # viewing angle
    ax.view_init(-90, 30)

    print "Calculating lines..."
    for pts in range(0,(sci*2)):

        hue = random()
        # set up initial coordinates
        x = f * np.cos((pts/sc)*np.pi)
        y = f * np.sin((pts/sc)*np.pi)
        z = f * np.sin((pts/sc)*np.pi)

        xl, yl, zl = [], [], []

        # length of curves
        for t in np.arange(0, 2000, dt):
            i, j, k = vectorfield(x, y, z)
            x = x + (i * dt)
            y = y + (j * dt)
            z = z + (k * dt)
            xl.append(x)
            yl.append(y)
            zl.append(z)

        # show curves
        ax.plot(xl,yl,zl, '-', zdir='z',alpha=0.4,
            color = colorsys.hsv_to_rgb(hue,1.0,1.0))



    ax.set_xlim((-b,b))
    ax.set_ylim((-b,b))
    ax.set_zlim((-b,b))





    save_set = 1
    if save_set == 1:


        # save individual frames
        # comment this out if you only want to see the plot
        unique_id = str(random()*3) + '_'
        n_frames = 1
        zoom_depth = 3.

        # create log file of all parameters
        set_wr = open('FS%s log.txt' % unique_id, 'w+')
        set_wr.write('ID: FS %s' % unique_id + '\n')
        set_wr.write('b = %s (initial x,y,z lim)' % str(b) + '\n')
        set_wr.write('u = %s' % str(u) + '\n')
        set_wr.write('p = %s' % str(p) + '\n')
        set_wr.write('q = %s' % str(q) + '\n')
        set_wr.write('r = %s' % str(r) + '\n')
        set_wr.write('s = %s' % str(s) + '\n')
        set_wr.write('v = %s' % str(v) + '\n')
        set_wr.write('w = %s' % str(w) + '\n')
        set_wr.write('f = %s' % str(f) + '\n')
        set_wr.write('for a1,a2,a3: 0 = x, 1 = y, 2 = z' + '\n')
        set_wr.write('a1 = %s' % str(a1) + '\n')
        set_wr.write('a1 = %s' % str(a2) + '\n')
        set_wr.write('a1 = %s' % str(a3) + '\n')
        set_wr.write('number of drawn lines: %.0f' % sc)
        set_wr.close()



        for frames in range(0,n_frames):

            # zooms from b to b-zoom_depth and back

            fstep = (frames/(1.*n_frames)) * np.pi
            fsin = zoom_depth * (np.sin(fstep))**2
            lim = b - fsin



            ax.set_xlim((-lim,lim))
            ax.set_ylim((-lim,lim))
            ax.set_zlim((-lim,lim))


            ax.view_init(-90+frames,30+frames)
            remaining_time = (2880 - (8 * frames))/60.


            fname = 'FS' + unique_id + str(frames) + '.png'
            print "Saving: %s" % fname
            plt.savefig(fname, format='png',dpi=300,
                        bbox_inches='tight', pad_inches=0)

        winsound.Beep(2000,500)
    else:
        plt.show()



#main function, generates 50 images
for images in range(0,50):
    print "Image: %d" % (images + 1)
    main()