import Mathlib

inductive Color
| Red
| Black
deriving Inhabited, BEq

inductive Node
| Empty : Node
| Tree  (color : Color) (val : Int) (left : Node) (right : Node) : Node
deriving Inhabited

def view (t : Node) : Set Int :=
  match t with
  | Node.Empty => ∅
  | Node.Tree _ v l r => view l ∪ view r ∪ {v}

def is_red (t : Node) : Bool :=
  match t with
  | Node.Tree Color.Red _ _ _ => true
  | _ => false

def get_color (t : Node) : Color :=
  match t with
  | Node.Tree c _ _ _ => c
  | Node.Empty => Color.Black -- Empty nodes are black

def is_bst (t : Node) : Prop :=
  match t with
  | Node.Empty => True
  | Node.Tree _ v l r =>
    (∀ x ∈ view l, x < v) ∧
    is_bst l ∧
    (∀ x ∈ view r, x > v) ∧
    is_bst r

-- Black Height: Returns -1 if invalid
def black_height (t : Node) : Int :=
  match t with
  | Node.Empty => 1
  | Node.Tree c _ l r =>
    let lh := black_height l
    let rh := black_height r
    if lh ≠ -1 ∧ rh ≠ -1 ∧ lh = rh then
      if c == Color.Black then lh + 1 else lh
    else
      -1

def satisfies_red_props (t : Node) : Prop :=
  match t with
  | Node.Empty => True
  | Node.Tree _ _ l r =>
    satisfies_red_props l ∧
    satisfies_red_props r ∧
    -- Rule: No Right-Leaning Red
    (¬ is_red r) ∧
    -- Rule: No Double Red on Left
    (match l with
     | Node.Tree Color.Red _ ll _ => ¬ is_red ll
     | _ => True)

-- Main Invariant
def is_llrbt (t : Node) : Prop :=
  is_bst t ∧
  black_height t ≠ -1 ∧
  satisfies_red_props t

def get_left (t : Node) : Option Node :=
  match t with | Node.Tree _ _ l _ => some l | _ => none

def get_right (t : Node) : Option Node :=
  match t with | Node.Tree _ _ _ r => some r | _ => none

def get_val (t : Node) : Option Int :=
  match t with | Node.Tree _ v _ _ => some v | _ => none

def is_not_empty (t : Node) : Prop :=
  match t with | Node.Empty => False | _ => True

-- Precondition definitions
@[reducible, simp]
def flip_colors_precond (t : Node) : Prop :=
  -- !benchmark @start precond
  -- t is not empty
  is_not_empty t ∧
  match t with
  | Node.Tree c _ l r =>
    is_not_empty l ∧ is_not_empty r ∧
    (get_color l = get_color r) ∧
    (c ≠ get_color l)
  | _ => False
  -- !benchmark @end precond

-- !benchmark @start auxcode
-- !benchmark @end auxcode

-- Main function definition
def flip_colors (t : Node) 
    (h_precond : flip_colors_precond t) : Node :=
  -- !benchmark @start code
  sorry
  -- !benchmark @end code

-- Postcondition definitions
@[reducible, simp]
def flip_colors_postcond (t : Node) (result : Node)
    (_ : flip_colors_precond t) : Prop :=
  -- !benchmark @start postcond
  view result = view t ∧
  is_bst result ∧ -- Requirement 1: BST structure preserved
  black_height result = black_height t ∧ -- Requirement 2: Balance preserved (Critical)
  match result, t with
  | Node.Tree c' _ l' r', Node.Tree c _ l r =>
      c' ≠ c ∧
      get_color l' ≠ get_color l ∧
      get_color r' ≠ get_color r
  | _, _ => False
  -- !benchmark @end postcond

-- !benchmark @start lemma
-- !benchmark @end lemma

-- Proof content
theorem flip_colors_postcond_satisfied (t : Node)
    (h_precond : flip_colors_precond t) :
    flip_colors_postcond t (flip_colors t h_precond) h_precond := by
  -- !benchmark @start proof
  sorry
  -- !benchmark @end proof