import numpy as np, pandas as pd, datashader as ds
from datashader import transfer_functions as tf
from datashader.colors import inferno, viridis
from numba import jit
from math import sin, cos, sqrt, fabs

@jit(nopython=True)
def Clifford(x, y, a, b, c, d, *o):
    return sin(a * y) + c * cos(a * x), \
           sin(b * x) + d * cos(b * y)

n=10000000

@jit(nopython=True)
def trajectory_coords(fn, x0, y0, a, b=0, c=0, d=0, e=0, f=0, n=n):
    x, y = np.zeros(n), np.zeros(n)
    x[0], y[0] = x0, y0
    for i in np.arange(n-1):
        x[i+1], y[i+1] = fn(x[i], y[i], a, b, c, d, e, f)
    return x,y

def trajectory(fn, x0, y0, a, b=0, c=0, d=0, e=0, f=0, n=n):
    x, y = trajectory_coords(fn, x0, y0, a, b, c, d, e, f, n)
    return pd.DataFrame(dict(x=x,y=y))

%%time
df = trajectory(Clifford, 0, 0, -1.3, -1.3, -1.8, -1.9)

df.tail()

%%time

cvs = ds.Canvas(plot_width = 700, plot_height = 700)
agg = cvs.points(df, 'x', 'y')
print(agg.values[190:195,190:195],"\n")

ds.transfer_functions.Image.border=0

tf.shade(agg, cmap = ["white", "black"])

def dsplot(fn, vals, n=n, cmap=viridis, label=True):
    """Return a Datashader image by collecting `n` trajectory points for the given attractor `fn`"""
    lab = ("{}, "*(len(vals)-1)+" {}").format(*vals) if label else None
    df  = trajectory(fn, *vals, n=n)
    cvs = ds.Canvas(plot_width = 300, plot_height = 300)
    agg = cvs.points(df, 'x', 'y')
    img = tf.shade(agg, cmap=cmap, name=lab)
    return img

from colorcet import palette
palette["viridis"]=viridis
palette["inferno"]=inferno

import yaml
vals = yaml.load(open("strange_attractors.yml","r"), Loader=yaml.FullLoader)

def args(name):
    """Return a list of available argument lists for the given type of attractor"""
    return [v[1:] for v in vals if v[0]==name]  

def plot(fn, vals=None, **kw):
    """Plot the given attractor `fn` once per provided set of arguments."""
    vargs=args(fn.__name__) if vals is None else vals
    return tf.Images(*[dsplot(fn, v[1:], cmap=palette[v[0]][::-1], **kw) for v in vargs]).cols(4)

plot(Clifford)

import numpy.random
numpy.random.seed(21)
num = 4

rvals=np.c_[np.zeros((num,2)), numpy.random.random((num,4))*4-2]
plot(Clifford, vals=[["kbc"]+list(rvals[i]) for i in range(len(rvals))], label=True)

@jit(nopython=True)
def De_Jong(x, y, a, b, c, d, *o):
    return sin(a * y) - cos(b * x), \
           sin(c * x) - cos(d * y)

plot(De_Jong)

@jit(nopython=True)
def Svensson(x, y, a, b, c, d, *o):
    return d * sin(a * x) - sin(b * y), \
           c * cos(a * x) + cos(b * y)

plot(Svensson)

@jit(nopython=True)
def Bedhead(x, y, a, b, *o):
    return sin(x*y/b)*y + cos(a*x-y), \
           x + sin(y)/b

plot(Bedhead)

@jit(nopython=True)
def Fractal_Dream(x, y, a, b, c, d, *o):
    return sin(y*b)+c*sin(x*b), \
           sin(x*a)+d*sin(y*a)

plot(Fractal_Dream)

@jit(nopython=True)
def Hopalong1(x, y, a, b, c, *o):
    return y - sqrt(fabs(b * x - c)) * np.sign(x), \
           a - x
@jit(nopython=True)
def Hopalong2(x, y, a, b, c, *o):
    return y - 1.0 - sqrt(fabs(b * x - 1.0 - c)) * np.sign(x - 1.0), \
           a - x - 1.0

plot(Hopalong1)

plot(Hopalong2)

@jit(nopython=True)
def G(x, mu):
    return mu * x + 2 * (1 - mu) * x**2 / (1.0 + x**2)

@jit(nopython=True)
def Gumowski_Mira(x, y, a, b, mu, *o):
    xn = y + a*(1 - b*y**2)*y  +  G(x, mu)
    yn = -x + G(xn, mu)
    return xn, yn

plot(Gumowski_Mira)

@jit(nopython=True)
def Symmetric_Icon(x, y, a, b, g, om, l, d, *o):
    zzbar = x*x + y*y
    p = a*zzbar + l
    zreal, zimag = x, y
    
    for i in range(1, d-1):
        za, zb = zreal * x - zimag * y, zimag * x + zreal * y
        zreal, zimag = za, zb
    
    zn = x*zreal - y*zimag
    p += b*zn
    
    return p*x + g*zreal - om*y, \
           p*y - g*zimag + om*x

plot(Symmetric_Icon)

import hvplot.pandas # noqa

df = trajectory(Clifford, *(args("Clifford")[5][1:]))
opts = dict(rasterize=True, dynspread=True, cnorm='eq_hist', width=400, height=400, colorbar=False)

df.hvplot.scatter(x='x', y='y', **opts)

df.hvplot.line(x='x', y='y', **opts)

def hv_clif(a,b,c,d,x0=0,y0=0,n=n):
    df = trajectory(Clifford, x0, y0, a, b, c, d, n)
    return df.hvplot.scatter(x='x', y='y', **opts)

import panel as pn
init = zip(['a', 'b', 'c', 'd'], args("Clifford")[6][-4:])
widgets = [pn.widgets.FloatSlider(name=n, start=-2.0, end=2.0, value=v) for n,v in init]
pn.Row(pn.bind(hv_clif, *widgets), pn.Column(*widgets))