diff --git a/demo3.cpp b/demo3.cpp
index f9c0d49..493a3a8 100644
--- a/demo3.cpp
+++ b/demo3.cpp
@@ -32,7 +32,7 @@ class Change
 public:
   Vector2D mCoords;
   Unit mHeight;
-  uint8_t mColor;
+  int8_t mColor;
 };
 
 #define MAX_CHANGES 16
@@ -44,6 +44,9 @@ bool editing = false;
 uint16_t changeIndex = 0;
 int16_t editCounter = 0;
 
+#define SQUARE_COLORS 4
+uint8_t squareColors[SQUARE_COLORS];
+
 #define HEIGHT_PROFILE_LENGTH 256
 const int8_t heightProfile[] = {
 9,9,9,10,10,10,11,11,12,13,13,14,14,14,15,15,15,16,16,16,16,16,15,15,15,14,14,
@@ -149,6 +152,11 @@ const unsigned char imageBackground[] =
 
 Unit floorHeightAt(int16_t x, int16_t y)
 {
+  /*
+    This for loop may become a bottleneck, since this function is called
+    very often - if more changes are to be kept, optimizations are needed -
+    probably a hash table, sorting by coordinates etc.
+  */
   for (uint16_t i = 0; i < MAX_CHANGES; ++i)
     if (changes[i].mCoords.x == x && changes[i].mCoords.y == y)
       return changes[i].mHeight;
@@ -158,6 +166,16 @@ Unit floorHeightAt(int16_t x, int16_t y)
          UNITS_PER_SQUARE;
 }
 
+Unit colorAt(int16_t x, int16_t y)
+{
+  for (uint16_t i = 0; i < MAX_CHANGES; ++i)
+    if (changes[i].mCoords.x == x && changes[i].mCoords.y == y)
+      return changes[i].mColor;
+
+  return min((heightProfile[absVal(x * 2) % HEIGHT_PROFILE_LENGTH] +
+         heightProfile[absVal(y) % HEIGHT_PROFILE_LENGTH]) / 10,3);
+}
+
 uint16_t previousColumn = 255;
 uint16_t backgroundColumn = 0;
 
@@ -172,7 +190,8 @@ inline void pixelFunc(PixelInfo pixel)
 
   if (pixel.isWall)
   {
-    c = pixel.hit.square.x != selectedSquare.x || pixel.hit.square.y != selectedSquare.y || (editing && pokitto.frameCount % 2) == 0 ? 50 : 30;
+    c = pixel.hit.square.x != selectedSquare.x || pixel.hit.square.y != selectedSquare.y || (editing && pokitto.frameCount % 2) == 0 ?
+        squareColors[pixel.hit.type] : 30;
     
     intensity = pixel.depth / (UNITS_PER_SQUARE * 3);
     intensity += pixel.hit.direction % 2 == 0 ? 2 : 0;
@@ -219,7 +238,7 @@ void draw()
   c.maxSteps = 20;
   c.computeTextureCoords = 0;
 
-  render(player.mCamera,floorHeightAt,0,0,pixelFunc,c);
+  render(player.mCamera,floorHeightAt,0,colorAt,pixelFunc,c);
 /*
 pokitto.display.setColor(rgbToIndex(7,7,3));
 pokitto.display.setCursor(1,1);
@@ -237,6 +256,11 @@ int main()
 
   floorColor = rgbToIndex(4,2,0);
 
+  squareColors[0] = rgbToIndex(0,0,3);
+  squareColors[1] = rgbToIndex(7,0,0);
+  squareColors[2] = rgbToIndex(0,7,0);
+  squareColors[3] = rgbToIndex(4,4,0);
+
   player.setPositionSquare(4,5);
 
   for (uint16_t i = 0; i < MAX_CHANGES; ++i)
@@ -320,6 +344,7 @@ int main()
             }
 
           changes[changeIndex].mHeight = floorHeightAt(selectedSquare.x,selectedSquare.y);
+          changes[changeIndex].mColor = colorAt(selectedSquare.x,selectedSquare.y);
           changes[changeIndex].mCoords.x = selectedSquare.x;
           changes[changeIndex].mCoords.y = selectedSquare.y;
           editReleased = false;
@@ -338,8 +363,6 @@ int main()
             changes[changeIndex].mHeight += UNITS_PER_SQUARE / 4;
             editCounter = 4;
           }
-          else
-            editCounter--;
         }
         else if (aButton)
           player.mCamera.shear = min(player.mCamera.shear + 10,60);
@@ -355,8 +378,6 @@ int main()
             changes[changeIndex].mHeight -= UNITS_PER_SQUARE / 4;
             editCounter = 4;
           }
-          else
-            editCounter--;
         }
         else if (aButton)
           player.mCamera.shear = max(player.mCamera.shear - 10,-60);
@@ -369,33 +390,42 @@ int main()
       else
       {
         player.mCamera.shear /= 2;
-        editCounter = 0;
       }
-      
+ 
       if (!aButton)  
         player.mCamera.shear /= 2;
 
       int addition = 0;
 
-      if (!editing)
-      {
-        if (pokitto.rightBtn())
-          addition = 1;      
-        else if (pokitto.leftBtn())
-          addition = -1;
-      }
+      if (pokitto.rightBtn())
+        addition = 1;      
+      else if (pokitto.leftBtn())
+        addition = -1;
 
-      if (aButton)
+      if (editing)
       {
-        d = angleToDirection(player.mCamera.direction + UNITS_PER_SQUARE / 4);
-        d.x = (d.x * horizontalStep * addition) / UNITS_PER_SQUARE;
-        d.y = (d.y * horizontalStep * addition) / UNITS_PER_SQUARE;
-
-        moveOffset = d;
+        if (editCounter == 0)
+        {
+          changes[changeIndex].mColor = wrap(changes[changeIndex].mColor + addition,SQUARE_COLORS);
+          editCounter = 4;
+        }
       }
       else
-        player.mCamera.direction = wrap(player.mCamera.direction + addition * rotationStep,UNITS_PER_SQUARE);
-    
+      {
+        if (aButton)
+        {
+          d = angleToDirection(player.mCamera.direction + UNITS_PER_SQUARE / 4);
+          d.x = (d.x * horizontalStep * addition) / UNITS_PER_SQUARE;
+          d.y = (d.y * horizontalStep * addition) / UNITS_PER_SQUARE;
+
+          moveOffset = d;
+        }
+        else
+          player.mCamera.direction = wrap(player.mCamera.direction + addition * rotationStep,UNITS_PER_SQUARE);
+      }    
+
+      editCounter = max(0, editCounter - 1);
+
       Unit prevHeight = player.mCamera.height;
 
       moveCameraWithCollision(&player.mCamera,moveOffset,player.mVericalSpeed,