datatype BYTree = BlueLeaf | YellowLeaf | Node(left: BYTree, right: BYTree)

function ReverseColors(t: BYTree): BYTree {
    match t
    case BlueLeaf => YellowLeaf
    case YellowLeaf => BlueLeaf 
    case Node(left, right) => BYTree.Node(ReverseColors(left), ReverseColors(right))
}

lemma {:induction false} ReverseColorsInvolution(t: BYTree)
    ensures ReverseColors(ReverseColors(t)) == t
{     
    match t
    case BlueLeaf =>
    case YellowLeaf =>
    case Node(left, right) =>
        calc {
            ReverseColors(ReverseColors(Node(left,right)));
        ==
            ReverseColors(Node(ReverseColors(left), ReverseColors(right)));
        ==  
            Node(ReverseColors(ReverseColors(left)), ReverseColors(ReverseColors(right)));
        ==  {ReverseColorsInvolution(left); ReverseColorsInvolution(right);}
            Node(left, right);
        }
}

function LeafCount(t: BYTree): nat 
{
    match t
    case BlueLeaf | YellowLeaf => 1
    case Node(left, right) => LeafCount(left) + LeafCount(right)
}

lemma {:induction false} ReversePreserveLeafCount(t: BYTree) 
    ensures LeafCount(ReverseColors(t)) == LeafCount(t)
    decreases t
{
    match t
    case BlueLeaf =>
    case YellowLeaf =>
    case Node(left, right) => 
        calc {
            LeafCount(ReverseColors(Node(left, right)));
        ==
            LeafCount(Node(ReverseColors(left), ReverseColors(right)));
        ==
            LeafCount(ReverseColors(left)) + LeafCount(ReverseColors(right));
        ==  {ReversePreserveLeafCount(left); ReversePreserveLeafCount(right);}
            LeafCount(left) + LeafCount(right);
        ==
            LeafCount(Node(left, right));
        }
}

function Oceanize(t: BYTree): BYTree {
    match t
        case BlueLeaf | YellowLeaf => BlueLeaf
        case Node(left, right) => Node(Oceanize(left), Oceanize(right))
}

lemma {:induction false} IdempotentOceanize(t: BYTree) 
    ensures Oceanize(t) == Oceanize(Oceanize(t))
    decreases t
{
    match t
    case BlueLeaf => 
        calc {
            Oceanize(Oceanize(BlueLeaf));
        ==
            Oceanize(BlueLeaf);
        ==
            BlueLeaf;
        }
    case YellowLeaf =>
        calc {
            Oceanize(Oceanize(YellowLeaf));
        ==
            Oceanize(BlueLeaf);
        ==
            BlueLeaf;
        ==
            Oceanize(YellowLeaf);
        }
        
    case Node(left, right) =>
        calc {
            Oceanize(Node(left, right));
        ==
            Node(Oceanize(left), Oceanize(right));
        ==  {IdempotentOceanize(left); IdempotentOceanize(right);}
            Node(Oceanize(Oceanize(left)), Oceanize(Oceanize(right)));
        ==
            Oceanize(Node(Oceanize(left), Oceanize(right)));
        ==
            Oceanize(Oceanize(Node(left,right)));
        }
}

function BlueCount(t: BYTree): nat {
    match t
    case BlueLeaf => 1
    case YellowLeaf => 0
    case Node(left, right) =>
        BlueCount(left) + BlueCount(right)
}

lemma {:induction false} BlueCountBound(t: BYTree) 
    ensures BlueCount(t) <= BlueCount(Oceanize(t))
{
    match t
    case BlueLeaf => 
        calc {
            BlueCount(BlueLeaf);
        ==
            BlueCount(Oceanize(BlueLeaf));
        }
    case YellowLeaf =>
        calc {
            BlueCount(YellowLeaf);
        <
            BlueCount(BlueLeaf);
        ==
            BlueCount(Oceanize(YellowLeaf));
        }
    case Node(left, right) =>
        calc {
            BlueCount(Node(left, right));
        ==
            BlueCount(left) + BlueCount(right);
        <= {BlueCountBound(left); BlueCountBound(right);}
            BlueCount(Oceanize(left)) + BlueCount(Oceanize(right));
        }
}