#!/usr/bin/env python
# coding: utf-8

# In[1]:


from IPython.display import HTML
from matplotlib.animation import FuncAnimation
import matplotlib.pyplot as plt
import numpy as np


# In[ ]:


class Circles:
    def __init__(self, frames):
        self.frames = frames
        self.ds = 2.1
        self.data = []
        self.edge_data = []
        self.xlims = [-0.5 * self.ds, 7.5 * self.ds]
        self.ylims = [-7.5 * self.ds, 0.5 * self.ds]
        t = np.linspace(np.pi / 2, -3 * np.pi / 2, frames)
        for row in range(1, 8):
            self.edge_data.append((np.cos(row * t) + self.ds * row, np.sin(row * t), f"C{row}"))
            self.edge_data.append((np.cos(row * t), np.sin(row * t) - self.ds * row, f"C{row}"))

        for row in range(1, 8):
            for col in range(1, 8):
                self.data.append((np.cos(row * t) + self.ds * row, 
                                  np.sin(col * t) - self.ds * col, f"C{(row - col + 1) % 8}"))
        self.fig, self.ax = plt.subplots(figsize=(10, 10), constrained_layout=True)

                
    def init(self):
        self.axlines = []
        for idx, (xdata, ydata, c) in enumerate(self.edge_data):
            self.ax.plot(xdata, ydata, color=c, lw=3)
            if idx % 2 == 0:
                self.axlines.append(self.ax.plot([], [], lw=1, color='w', alpha=0.2)[0])
            else:
                self.axlines.append(self.ax.plot([], [], lw=1, color='w', alpha=0.2)[0])

        self.edge_circs = [self.ax.plot([], [], 'wo', lw=0, mfc='w')[0] for _ in self.edge_data]
        self.lines = [self.ax.plot([], [], color=c, lw=3)[0] for _, _, c in self.data]
        self.circs = [self.ax.plot([], [], 'wo', lw=0, mfc='w')[0] for _ in self.data]
        
        self.ax.set_facecolor('black')
        self.ax.yaxis.set_visible(False)
        self.ax.xaxis.set_visible(False)
        self.ax.set_xlim(*self.xlims)
        self.ax.set_ylim(*self.ylims)

        self.artists = self.edge_circs + self.lines + self.circs + self.axlines
        return self.artists

    def update(self, idx):
        for line, circ, (xdata, ydata, _) in zip(self.lines, self.circs, self.data):
            line.set_data(xdata[:idx], ydata[:idx])
            circ.set_data(xdata[idx], ydata[idx])
        
        for j, ((xdata, ydata, _), circ, line) in enumerate(zip(self.edge_data, self.edge_circs, self.axlines)):
            circ.set_data(xdata[idx], ydata[idx])
            if j % 2 == 0:
                line.set_data([xdata[idx], xdata[idx]], self.ylims)
            else:
                line.set_data(self.xlims, [ydata[idx], ydata[idx]])

        return self.artists

circ = Circles(200)

def init():
    return circ.init()
    

anim = FuncAnimation(circ.fig, circ.update, frames=circ.frames,
                     init_func=circ.init, blit=True, interval=50)
HTML(anim.to_html5_video())