import penzai
from penzai import pz
import IPython
import jax.numpy as np
from chalk import *
import chalk
from colour import Color
from chalk.core import Primitive, ApplyStyle, ApplyName

def draw_fn(value, path):
    x = rectangle(1, 1).align_bl().line_color(Color("red")).fill_opacity(0)
    if isinstance(value, Primitive):
        chalk.set_svg_height(200)
        return pz.ts.IPythonVisualization(
            value.show_origin(), True
        )
    
    if isinstance(value, Path):
        chalk.set_svg_height(100)
        return pz.ts.IPythonVisualization(
            value.stroke(), True
        )
    
    if isinstance(value, ApplyName) and value.dname == "group":
        chalk.set_svg_height(100)
        return pz.ts.IPythonVisualization(
            value.diagram, True
        )
    
    if isinstance(value, Envelope):
        chalk.set_svg_height(100)
        return pz.ts.IPythonVisualization(
            empty().compose(value).show_envelope(False, 10), True
        )
    if isinstance(value, Style):
        chalk.set_svg_height(75)
        return pz.ts.IPythonVisualization(
            rectangle(7, 1).apply_style(value), True
        )
    if isinstance(value, Affine):
       chalk.set_svg_height(100)

       y = x.apply_transform(value).fill_opacity(0).dashing([20, 2, 2], 0.1).line_color(Color("green"))
       return pz.ts.IPythonVisualization(
           (x + y), True
       )

pz.ts.register_as_default()

with pz.ts.active_autovisualizer.set_scoped(draw_fn):
    d = rectangle(2, 2) | circle(0.75) +  rectangle(0.1, 3).fill_color(Color("blue"))                            
    pz.ts.display(d)

(Loading...)