import numpy as np
from gt.pipeline_forward import forward_transform
from gt.pipeline_inverse import inverse_transform
from gt.visualization.fourier_animation import animate_fourier_orbits
import matplotlib.pyplot as plt
from matplotlib import animation
from IPython.display import Video
from IPython.display import Image
import matplotlib as mpl
from matplotlib.animation import FFMpegWriter
from IPython.display import HTML

%matplotlib inline

def encode_text(text: str):
    return np.array([ord(c) for c in text], dtype=float)

text = "ab"
signal = encode_text(text)

positions, meta = forward_transform(signal, steps=2000, dt=0.01)
positions.shape, meta

mpl.rcParams['animation.embed_limit'] = 20  # MB

ani = animate_fourier_orbits(
    positions,
    meta,
    message="ab",
    interval=20
)

writer = FFMpegWriter(
    fps=40,
    metadata=dict(artist='Gravitational Transform'),
    bitrate=1800
)

ani.save("ab.gif", writer="pillow", dpi=90)

text = "aa"
signal = encode_text(text)

positions, meta = forward_transform(signal, steps=2000, dt=0.01)
positions.shape, meta

mpl.rcParams['animation.embed_limit'] = 20  # MB

ani = animate_fourier_orbits(
    positions,
    meta,
    message=text,
    interval=20
)

writer = FFMpegWriter(
    fps=40,
    metadata=dict(artist='Gravitational Transform'),
    bitrate=1800
)

ani.save(f"{text}.gif", writer="pillow", dpi=90)

text = "abcdefghijklmnopqrstuvwxyz"
signal = encode_text(text)

positions, meta = forward_transform(signal, steps=2000, dt=0.01, ignore_dc=True)
positions.shape, meta

mpl.rcParams['animation.embed_limit'] = 20  # MB

ani = animate_fourier_orbits(
    positions,
    meta,
    message=text,
    interval=20
)

writer = FFMpegWriter(
    fps=40,
    metadata=dict(artist='Gravitational Transform'),
    bitrate=1800
)

ani.save(f"{text}.gif", writer="pillow", dpi=90)

text = "3.14159"
signal = encode_text(text)

positions, meta = forward_transform(signal, steps=2000, dt=0.01, ignore_dc=False)
positions.shape, meta

mpl.rcParams['animation.embed_limit'] = 20  # MB

ani = animate_fourier_orbits(
    positions,
    meta,
    message=text,
    interval=20
)

writer = FFMpegWriter(
    fps=40,
    metadata=dict(artist='Gravitational Transform'),
    bitrate=1800
)

ani.save(f"{text}.gif", writer="pillow", dpi=90)

decoded_t0 = inverse_transform(positions, meta, t_index=100)
decoded_t0

'364241.'

def plot_frame(positions, t):
    pos = positions[t]
    x = pos[:,0]
    y = pos[:,1]

    plt.figure(figsize=(6,6))
    plt.scatter(x[1:], y[1:], s=30)
    plt.scatter([0],[0], s=80, color='orange')
    plt.title(f"Orbital state at t = {t}")
    plt.axis('equal')
    plt.show()

plot_frame(positions, 1000)