#import "@preview/cetz:0.4.2"

#set page(width: auto, height: auto, margin: 1cm)

// Example: Neural Network Loss Visualization
#align(center)[
  #text(size: 14pt, weight: "bold")[Inner Contrastive Loss Flow]
]

#v(0.5cm)

#cetz.canvas({
  import cetz.draw: *

  // Set up style
  set-style(
    stroke: (thickness: 1.5pt),
    mark: (end: ">", fill: black),
  )

  // Input layer
  rect((-2, 6), (2, 8), fill: rgb("#BBDEFB"), stroke: black, name: "input")
  content((0, 7), [Hidden States])
  content((0, 5.5), [$h_s$])

  // Positive and negative branches
  rect((-6, 3), (-2, 5), fill: rgb("#C8E6C9"), stroke: black, name: "pos")
  content((-4, 4), [Positive])
  content((-4, 2.5), [$h_"pos"$])

  rect((2, 3), (6, 5), fill: rgb("#FFCDD2"), stroke: black, name: "neg")
  content((4, 4), [Negative])
  content((4, 2.5), [$h_"neg"$])

  // Projection
  rect((-3, -1), (3, 1), fill: rgb("#FFF9C4"), stroke: black, name: "proj")
  content((0, 0), [Projection])
  content((0, -1.5), [$bold(V)$])

  // Energy computation
  rect((-7, -5), (-3, -3), fill: rgb("#E1BEE7"), stroke: black, name: "in")
  content((-5, -4), [In-Space])

  rect((3, -5), (7, -3), fill: rgb("#FFCCBC"), stroke: black, name: "null")
  content((5, -4), [Null-Space])

  // Final loss
  rect((-2, -9), (2, -7), fill: rgb("#E0E0E0"), stroke: black, name: "loss")
  content((0, -8), [Loss])
  content((0, -9.5), [$cal(L)$])

  // Arrows
  line((0, 6), (-4, 5), mark: (end: ">"))
  line((0, 6), (4, 5), mark: (end: ">"))
  line((-4, 3), (0, 1), mark: (end: ">"))
  line((4, 3), (0, 1), mark: (end: ">"))
  line((-1.5, -1), (-5, -3), mark: (end: ">"))
  line((1.5, -1), (5, -3), mark: (end: ">"))
  line((-5, -5), (0, -7), mark: (end: ">"))
  line((5, -5), (0, -7), mark: (end: ">"))
})