struct Wmap {
  enum Entity: Equatable {
    case vertex(Vertex)
    case edge(Edge)
    case note(Note)
    case stage(Stage)
    case group(Group)
    case inertia(Inertia)
    case evolution(Evolution)

    var position: SourcePosition {
      switch self {
      case .vertex(let vertex): return vertex.position
      case .edge(let edge): return edge.position
      case .note(let note): return note.position
      case .stage(let stage): return stage.position
      case .group(let group): return group.position
      case .inertia(let inertia): return inertia.position
      case .evolution(let evolution): return evolution.position
      }
    }

    var range: SourceRange {
      switch self {
      case .vertex(let vertex): return vertex.range
      case .edge(let edge): return edge.range
      case .note(let note): return note.range
      case .stage(let stage): return stage.range
      case .group(let group): return group.range
      case .inertia(let inertia): return inertia.range
      case .evolution(let evolution): return evolution.range
      }
    }

    var syntaxElements: [SyntaxElement] {
      switch self {
      case .vertex(let vertex): return vertex.syntaxElements
      case .edge(let edge): return edge.syntaxElements
      case .note(let note): return note.syntaxElements
      case .stage(let stage): return stage.syntaxElements
      case .group(let group): return group.syntaxElements
      case .inertia(let inertia): return inertia.syntaxElements
      case .evolution(let evolution): return evolution.syntaxElements
      }
    }
  }

  struct Vertex: Equatable {
    let label: String
    let position: SourcePosition
    let coordinates: (Double, Double)
    let shape: String?
    let range: SourceRange

    // Syntax element ranges (populated during parsing)
    let labelRange: SourceRange
    let leftParenRange: SourceRange
    let xCoordRange: SourceRange
    let commaRange: SourceRange
    let yCoordRange: SourceRange
    let rightParenRange: SourceRange
    let shapeRange: SourceRange?

    static func == (lhs: Vertex, rhs: Vertex) -> Bool {
      return lhs.label == rhs.label && lhs.position == rhs.position
        && lhs.coordinates.0 == rhs.coordinates.0 && lhs.coordinates.1 == rhs.coordinates.1
        && lhs.shape == rhs.shape && lhs.range == rhs.range
    }

    var syntaxElements: [SyntaxElement] {
      var elements: [SyntaxElement] = []

      // Vertex label
      elements.append(
        SyntaxElement(
          type: .vertexLabel,
          range: labelRange,
          value: label
        ))

      // Left parenthesis
      elements.append(
        SyntaxElement(
          type: .punctuation,
          range: leftParenRange,
          value: "("
        ))

      // X coordinate
      elements.append(
        SyntaxElement(
          type: .number,
          range: xCoordRange,
          value: String(coordinates.0)
        ))

      // Comma
      elements.append(
        SyntaxElement(
          type: .punctuation,
          range: commaRange,
          value: ","
        ))

      // Y coordinate
      elements.append(
        SyntaxElement(
          type: .number,
          range: yCoordRange,
          value: String(coordinates.1)
        ))

      // Right parenthesis
      elements.append(
        SyntaxElement(
          type: .punctuation,
          range: rightParenRange,
          value: ")"
        ))

      // Shape if present
      if let shape = shape, let shapeRange = shapeRange {
        elements.append(
          SyntaxElement(
            type: .shape,
            range: shapeRange,
            value: "[\(shape)]"
          ))
      }

      return elements
    }
  }

  struct Edge: Equatable {
    let from: String
    let to: String
    let isDirected: Bool
    let position: SourcePosition
    let range: SourceRange

    // Syntax element ranges
    let fromLabelRange: SourceRange
    let edgeTypeRange: SourceRange
    let toLabelRange: SourceRange

    var syntaxElements: [SyntaxElement] {
      var elements: [SyntaxElement] = []

      // From vertex label
      elements.append(
        SyntaxElement(
          type: .vertexLabel,
          range: fromLabelRange,
          value: from
        ))

      // Edge symbol
      elements.append(
        SyntaxElement(
          type: .symbol,
          range: edgeTypeRange,
          value: isDirected ? "->" : "--"
        ))

      // To vertex label
      elements.append(
        SyntaxElement(
          type: .vertexLabel,
          range: toLabelRange,
          value: to
        ))

      return elements
    }
  }

  struct Note: Equatable {
    let coordinates: (Double, Double)
    let text: String
    let position: SourcePosition
    let range: SourceRange

    // Syntax element ranges
    let keywordRange: SourceRange
    let leftParenRange: SourceRange
    let xCoordRange: SourceRange
    let commaRange: SourceRange
    let yCoordRange: SourceRange
    let rightParenRange: SourceRange
    let textRange: SourceRange

    static func == (lhs: Note, rhs: Note) -> Bool {
      return lhs.coordinates.0 == rhs.coordinates.0 && lhs.coordinates.1 == rhs.coordinates.1
        && lhs.text == rhs.text && lhs.position == rhs.position && lhs.range == rhs.range
    }

    var syntaxElements: [SyntaxElement] {
      var elements: [SyntaxElement] = []

      // [Note] keyword
      elements.append(
        SyntaxElement(
          type: .keyword,
          range: keywordRange,
          value: "[Note]"
        ))

      // Left parenthesis
      elements.append(
        SyntaxElement(
          type: .punctuation,
          range: leftParenRange,
          value: "("
        ))

      // X coordinate
      elements.append(
        SyntaxElement(
          type: .number,
          range: xCoordRange,
          value: String(coordinates.0)
        ))

      // Comma
      elements.append(
        SyntaxElement(
          type: .punctuation,
          range: commaRange,
          value: ","
        ))

      // Y coordinate
      elements.append(
        SyntaxElement(
          type: .number,
          range: yCoordRange,
          value: String(coordinates.1)
        ))

      // Right parenthesis
      elements.append(
        SyntaxElement(
          type: .punctuation,
          range: rightParenRange,
          value: ")"
        ))

      // Text content
      elements.append(
        SyntaxElement(
          type: .text,
          range: textRange,
          value: text
        ))

      return elements
    }
  }

  struct Stage: Equatable {
    let number: String  // i, ii, iii, iv
    let value: Double
    let position: SourcePosition
    let range: SourceRange

    // Syntax element ranges
    let leftBracketRange: SourceRange
    let stageNumberRange: SourceRange
    let rightBracketRange: SourceRange
    let valueRange: SourceRange

    var syntaxElements: [SyntaxElement] {
      var elements: [SyntaxElement] = []

      // Left bracket
      elements.append(
        SyntaxElement(
          type: .punctuation,
          range: leftBracketRange,
          value: "["
        ))

      // Stage number
      elements.append(
        SyntaxElement(
          type: .stageNumber,
          range: stageNumberRange,
          value: number
        ))

      // Right bracket
      elements.append(
        SyntaxElement(
          type: .punctuation,
          range: rightBracketRange,
          value: "]"
        ))

      // Value
      elements.append(
        SyntaxElement(
          type: .number,
          range: valueRange,
          value: String(value)
        ))

      return elements
    }
  }

  struct Group: Equatable {
    let vertices: [String]
    let position: SourcePosition
    let range: SourceRange

    // Syntax element ranges
    let keywordRange: SourceRange
    let vertexRanges: [SourceRange]
    let commaRanges: [SourceRange]

    var syntaxElements: [SyntaxElement] {
      var elements: [SyntaxElement] = []

      // [Group] keyword
      elements.append(
        SyntaxElement(
          type: .keyword,
          range: keywordRange,
          value: "[Group]"
        ))

      // Vertex labels and commas
      for (index, vertex) in vertices.enumerated() {
        elements.append(
          SyntaxElement(
            type: .vertexLabel,
            range: vertexRanges[index],
            value: vertex
          ))

        // Add comma if not the last vertex
        if index < vertices.count - 1 {
          elements.append(
            SyntaxElement(
              type: .punctuation,
              range: commaRanges[index],
              value: ","
            ))
        }
      }

      return elements
    }
  }

  struct Inertia: Equatable {
    let vertex: String
    let position: SourcePosition
    let range: SourceRange

    // Syntax element ranges
    let keywordRange: SourceRange
    let vertexRange: SourceRange

    var syntaxElements: [SyntaxElement] {
      var elements: [SyntaxElement] = []

      // [Inertia] keyword
      elements.append(
        SyntaxElement(
          type: .keyword,
          range: keywordRange,
          value: "[Inertia]"
        ))

      // Vertex label
      elements.append(
        SyntaxElement(
          type: .vertexLabel,
          range: vertexRange,
          value: vertex
        ))

      return elements
    }
  }

  struct Evolution: Equatable {
    let vertex: String
    let isPositive: Bool
    let value: Double
    let position: SourcePosition
    let range: SourceRange

    // Syntax element ranges
    let keywordRange: SourceRange
    let vertexRange: SourceRange
    let signRange: SourceRange
    let valueRange: SourceRange

    var syntaxElements: [SyntaxElement] {
      var elements: [SyntaxElement] = []

      // [Evolution] keyword
      elements.append(
        SyntaxElement(
          type: .keyword,
          range: keywordRange,
          value: "[Evolution]"
        ))

      // Vertex label
      elements.append(
        SyntaxElement(
          type: .vertexLabel,
          range: vertexRange,
          value: vertex
        ))

      // Sign
      elements.append(
        SyntaxElement(
          type: .symbol,
          range: signRange,
          value: isPositive ? "+" : "-"
        ))

      // Value
      elements.append(
        SyntaxElement(
          type: .number,
          range: valueRange,
          value: String(value)
        ))

      return elements
    }
  }
}