#!/usr/bin/env python3
import numpy as np
import sys
import os
import time
import colorsys

Nx = int(sys.argv[1])
Ny = int(sys.argv[2])
N = 100
iterations = 2
try:
    N = int(sys.argv[3])
except:
    pass

def rk4(f, x, dt):
    k1 = f(x)
    k2 = f(x+k1*dt/2)
    k3 = f(x+k2*dt/2)
    k4 = f(x+k3*dt)
    return x+dt*(k1+2*k2+2*k3+k4)/6.0
dt = 0.08 
m = 1
l = 1
g = 1

def sq(x): return x*x

def rhs(x):
    phi1 = x[0]
    p1 = x[1]
    phi2 = x[2]
    p2 = x[3]

    dphi1 = 6/(m*l*l)*(2*p1-3*np.cos(phi1-phi2)*p2)/(16-9*sq(np.cos(phi1-phi2))) 
    dphi2 = 6/(m*l*l)*(8*p2-3*np.cos(phi1-phi2)*p1)/(16-9*sq(np.cos(phi1-phi2)))

    dp1 = -0.5*m*l*l*(dphi1*dphi2*np.sin(phi1-phi2)+3*g/l*np.sin(phi1))
    dp2 = -0.5*m*l*l*(-dphi1*dphi2*np.sin(phi1-phi2)+g/l*np.sin(phi2))
    return np.array([dphi1, dp1, dphi2, dp2])

s = np.array([np.pi+0.1, 0, np.pi-0.1, 0])

angles = np.linspace(np.pi, np.pi/3, N)+np.random.uniform(-np.pi/6/N, np.pi/6/N, N)
angles = np.random.uniform(-np.pi/6, np.pi/3, N)+np.pi
states = [np.array([np.pi, 0, a, 0]) for a in angles]

def clamp(x, min_, max_):
    return  max(min_, min(max_, x))

buffer = bytearray(b"\x00"*Nx*Ny*3)

def setpixel(x, y, r, g, b):
    xi = int(Nx*(x+1.2)/2.4)
    yi = int(Ny*(y+1.2)/2.4)
    if xi < 0 or xi >= Nx or yi < 0 or yi >= Ny:
        return
    idx = xi+Nx*yi
    buffer[3*idx+0] = r
    buffer[3*idx+1] = g
    buffer[3*idx+2] = b
timestep = 0
timefac = 0.03
while True:
    timestep += 1
    for s, i in zip(states, range(len(states))):
        phi1 = s[0]
        phi2 = s[2]
        x1 = np.sin(phi1)*l
        y1 = np.cos(phi1)*l
        x2 = np.sin(phi2)*l+x1
        y2 = np.cos(phi2)*l+y1
        h = 0.2*i/N+0*timestep*timefac/iterations
        r, g, b = colorsys.hsv_to_rgb(h%1, 1, 1)
        setpixel(x2, y2, int(r*255), int(g*255), int(b*255))
        states[i] = rk4(rhs, s, dt/iterations)
    if timestep > 10*iterations and timestep % iterations == 0:
        #time.sleep(0.01)
        os.write(1, buffer)