# this is by Simone Conradi but the orbit is generated in Apple
# https://github.com/profConradi/Python_Simulations/blob/main/Nice_orbits.ipynb

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap, LinearSegmentedColormap, to_rgb
from matplotlib import rcParams

#Latex is needed to create plot titles
rcParams['text.usetex'] = True #set this to False if Latex distribution is not installed
rcParams["text.latex.preamble"] = r"\usepackage{amsmath}\usepackage{amsfonts}\usepackage{systeme}"
%config InlineBackend.figure_format = 'retina'

def create_colormap(hex_color1, hex_color2, reverse=False):
    """
    Create a linearly interpolated colormap from two hex colors.

    -hex_color1: String, the first hex color (e.g., "#FF0000" for red)
    -hex_color2: String, the second hex color (e.g., "#0000FF" for blue)
    -reverse: Boolean, invert colors.
    
    Return: LinearSegmentedColormap
    """
    color1 = to_rgb(hex_color1)
    color2 = to_rgb(hex_color2)
    if reverse:
        colors = [color2, color1]
    else:
        colors = [color1, color2]
    cmap = LinearSegmentedColormap.from_list("custom_colormap", colors)

    return cmap

# a custom colormap
c_labels = "#000000"
yellow_map = create_colormap("#000000", "#ffe100", True)

import apple

orb=apple.jit('''
λa.λb.
{
  unzip ← [{a⟜x;((->1)`{0} a, (->2)`{0} a)}];
  g ← λ(x,y). (sin. (x^2-y^2+a), cos. (2*x*y+b));
  unzip ((frange _1 1 500) [(gen. (x,y) g 200)]⊗ (frange _1 1 500))
}
''')

a,b=5.46,4.55

l_c=orb(a,b)

%timeit orb(a,b)

842 ms ± 3.66 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

l_cx=l_c[0].flatten()
l_cy=l_c[1].flatten()

area = [[-1,1],[-1,1]]
h,_,_ = np.histogram2d(l_cx,l_cy,bins=3000,range=area)

fig, ax = plt.subplots(figsize=(10,10))
fig.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=None, hspace=None)
ax.imshow(np.log(h+1),vmin=0,vmax=5,cmap=yellow_map)
plt.xticks([]), plt.yticks([])
ax.set_title(r'$\begin{array}{lr} x_{t+1} = \sin(x_t^2-y_t^2 + %s) \\ y_{t+1} = \cos(2x_t y_t + %s) \end{array}$'%(a,b), 
             y=0.05, fontsize=15, color=c_labels )
plt.show()