diff --git a/.gitignore b/.gitignore
index 1998e51..967cb0b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,2 +1,5 @@
 # vscode
 .history
+
+# Gem
+Gemfile.lock
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..9ad197b
--- /dev/null
+++ b/README.md
@@ -0,0 +1,18 @@
+# Setup
+- Install ruby2d
+- Clone the repo
+- Run the app using this command ```ruby main.rb```
+  
+# Missing Features
+- Timer
+- Score Board
+# Preview
+
+[Video Preview](https://drive.google.com/file/d/1zZ987ZvxyP2RvDqjjQ8MjmZ4AaKLwjy1/view?usp=sharing)
+
+---
+#### Screenshots
+![Screenshot from 2024-08-27 18-31-23](https://github.com/user-attachments/assets/00da1f90-d12f-42a3-b42b-bdf7a754b732)
+![Screenshot from 2024-08-27 18-31-27](https://github.com/user-attachments/assets/44428fa9-126f-456b-81f1-9c107f7dfbd9)
+![Screenshot from 2024-08-27 18-31-32](https://github.com/user-attachments/assets/7bfdc413-62d8-44c2-ac94-53779beac9a5)
+![Screenshot from 2024-08-27 18-32-10](https://github.com/user-attachments/assets/4b435bde-f9be-4299-8476-a3ee966c3740)
diff --git a/lib/concerns/game_display.rb b/lib/concerns/game_display.rb
new file mode 100644
index 0000000..89c5e7e
--- /dev/null
+++ b/lib/concerns/game_display.rb
@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module GameDisplay
+  FOUR_BY_FOUR = "4 x 4"
+  SIX_BY_SIX = "6 x 6"
+
+  def is_four_by_four? dimension
+    dimension.equal? FOUR_BY_FOUR
+  end
+
+  def is_six_by_six? dimension
+    dimension.equal? SIX_BY_SIX
+  end
+
+  def set_timeout seconds, &block
+    Thread.new do
+      sleep seconds
+      block.call
+    end
+  end
+end
diff --git a/lib/custom_classes/button.rb b/lib/custom_classes/button.rb
index 1d8528f..1cd5db0 100644
--- a/lib/custom_classes/button.rb
+++ b/lib/custom_classes/button.rb
@@ -35,14 +35,19 @@ def is_clicked?
     end
   end
 
+  def hide
+    @button.hide
+  end
+
   private
 
   attr_reader :label, :btn_x, :btn_y, :btn_width, :btn_height, :color
 
   def draw_button
-    CRectangle.new(
+    @button = CRectangle.new(
       x: btn_x, y: btn_y, width: btn_width, height: btn_height, color: color
-    ).draw
+    )
+    @button.draw
   end
 
   def draw_text
diff --git a/lib/game.rb b/lib/game.rb
index d36c7f0..90cb142 100644
--- a/lib/game.rb
+++ b/lib/game.rb
@@ -2,7 +2,9 @@
 
 require "ruby2d"
 require "aasm"
-require_relative "game_state"
+require_relative "states/play_button"
+require_relative "states/dimension_selection"
+require_relative "states/image_selection"
 require_relative "custom_classes/button"
 
 class Game
@@ -10,95 +12,82 @@ class Game
 
   attr_accessor :mouse_location
 
-  def initialize screen_width, screen_height
-    @game_state = State::GameState.new screen_width, screen_height
-    @mouse_location = nil
-    @selected_dimension = nil
-    super()
-  end
-
   aasm do
     state :init, initial: true
     state :play_button
     state :dimension_selection
     state :image_selection
-    state :ended
+    state :game_over
 
     event :start_game do
       transitions from: :init, to: :play_button
     end
 
-    event :dimension_selection do
+    event :transition_to_dimension_selection do
       transitions from: :play_button, to: :dimension_selection
     end
 
-    event :show_images do
-      transitions from: :dimension_selection, to: :show_images
-    end
-
-    event :image_selection do
+    event :transition_to_image_selection do
       transitions from: :dimension_selection, to: :image_selection
     end
 
-    event :ended do
-      transitions from: [:play_button, :dimension_selection, :image_selection], to: :ended
+    event :transition_to_game_over do
+      transitions from: [:play_button, :dimension_selection, :image_selection], to: :game_over
     end
   end
 
+  def start
+    start_game!
+    handle_state
+  end
+
   def running?
     !aasm.current_state.eql?(:ended)
   end
 
   def update
+  end
+
+  def mouse_click click_type
+    case click_type
+    when :mouse_down
+      @current_state.mouse_down mouse_location
+    when :mouse_up
+      @current_state.mouse_up mouse_location
+    end
+  end
+
+  private
+
+  attr_reader :game_state
+
+  def handle_state
     case aasm.current_state
     when :init
       start_game!
     when :play_button
-      @current_game_state ||= State::PlayButton.new @game_state
+      @current_state = PlayButton.new next_event
     when :dimension_selection
-      @current_game_state ||= State::DimensionSelection.new @game_state
-    when :show_images
-      @current_game_state ||= State::ShowImages.new @game_state
+      @current_state = DimensionSelection.new next_event
     when :image_selection
-      @current_game_state ||= State::ImageSelection.new @game_state, @selected_dimension
-    when :ended
+      @current_state = ImageSelection.new next_event, @current_state.dimension
+    when :game_over
     end
+  end
 
-    @current_game_state&.update
+  def next_event
+    -> { transition_event; handle_state; }
   end
 
-  # Handle click will be called in the main.rb when there is click events.
-  # Each state can have different process.
-  def mouse_click click_type # :mouse_down, :mouse_up
+  def transition_event
     case aasm.current_state
     when :play_button
-      # Do something if user click play button
-      callback = -> { dimension_selection!;
-                      reset_screen;
-                    }
-      game_state_mouse_click callback: callback
+      transition_to_dimension_selection!
     when :dimension_selection
-      @selected_dimension = game_state_mouse_click
-
-      unless @selected_dimension.nil?
-        image_selection!
-        reset_screen
-      end
+      transition_to_image_selection!
     when :image_selection
-      @current_game_state.mouse_click mouse_location, click_type
+      transition_to_game_over!
+    when :game_over
     end
   end
-
-  private
-
-  def game_state_mouse_click callback: nil
-    return @current_game_state.mouse_click mouse_location if callback.nil?
-
-    @current_game_state.mouse_click mouse_location, callback: callback
-  end
-
-  def reset_screen
-    Screen.clear
-    @current_game_state = nil # Assign nil to change screen properly
-  end
 end
diff --git a/lib/game_state.rb b/lib/game_state.rb
deleted file mode 100644
index 44c9f0c..0000000
--- a/lib/game_state.rb
+++ /dev/null
@@ -1,287 +0,0 @@
-# frozen_string_literal: true
-
-require "ruby2d"
-require "debug"
-require_relative "custom_classes/screen"
-require_relative "custom_classes/button"
-require_relative "custom_classes/c_rectangle"
-require_relative "asset"
-
-FOUR_BY_FOUR = "4 x 4"
-SIX_BY_SIX = "6 x 6"
-
-module State
-  class GameState
-    def initialize screen_width, screen_height
-      @screen_width = screen_width
-      @screen_height = screen_height
-      @update_callback = nil
-    end
-
-    def update
-      @update_callback.call unless @update_callback.nil?
-    end
-
-    def mouse_click mouse_location, callback: nil
-    end
-
-    protected
-
-    attr_reader :screen_width, :screen_height
-
-    def draw_button label, x, y, width, height, image_path = ""
-      Button.new(
-        label: label, x: x, y: y, width: width,
-        height: height, image_path: image_path
-      )
-    end
-
-    def set_timeout seconds, &block
-      Thread.new do
-        sleep seconds
-        block.call
-      end
-    end
-  end
-
-  class PlayButton < GameState
-    def initialize game_state
-      @game_state = game_state
-
-      show
-    end
-
-    def mouse_click mouse_location, callback: nil
-      @button&.mouse_location = mouse_location
-
-      if @button&.is_clicked? && callback
-        callback.call
-      end
-    end
-
-    private
-
-    attr_reader :game_state
-
-    def show
-      show_play_button
-    end
-
-    def center_x width
-      (game_state.screen_width - width) / 2
-    end
-
-    def center_y height
-      (game_state.screen_height - height) / 2
-    end
-
-    def show_play_button
-      width = 200
-      height = 60
-
-      @button = draw_button "Play", center_x(width), center_y(height), width, height
-      @button.draw
-    end
-  end
-
-  class DimensionSelection < GameState
-    def initialize game_state
-      @game_state = game_state
-      show
-    end
-
-    def mouse_click mouse_location
-      @btn_1.mouse_location = mouse_location
-      @btn_2.mouse_location = mouse_location
-
-      if @btn_1&.is_clicked?
-        FOUR_BY_FOUR
-      elsif @btn_2&.is_clicked?
-        SIX_BY_SIX
-      else
-        return
-      end
-    end
-
-    private
-
-    attr_reader :game_state
-
-    def screen_width
-      game_state.screen_width
-    end
-
-    def show
-      draw_selection_button
-    end
-
-    def draw_selection_button
-      width = 150
-      height = 60
-
-      new_y = screen_y_center height
-
-      @btn_1 = draw_button FOUR_BY_FOUR, position_x(1), new_y, width, height
-      @btn_1.draw
-      @btn_2 = draw_button SIX_BY_SIX, position_x(4), new_y, width, height
-      @btn_2.draw
-    end
-
-    def position_x column_position
-      # Divide width by 6 to have 6 columns
-      if column_position <= 6
-        column_size = screen_width / 6
-        position = column_size * column_position
-
-        return position
-      end
-    end
-
-    def screen_y_center btn_height
-      (game_state.screen_height / 2) - btn_height
-    end
-  end
-
-  class ImageSelection < GameState
-    def initialize game_state, dimension
-      @game_state = game_state
-      @dimension = dimension
-      @box_positions = [] # image_box, front_box, is_show ( default: false )
-      @enable_click = true
-      @click_count = 0
-      show
-    end
-
-    def mouse_click mouse_location, click_type, callback: nil
-      if click_type == :mouse_down && @enable_click
-        @box_positions.each_with_index do |box, index|
-          box[:image].mouse_location = mouse_location
-
-          if box[:image]&.is_clicked?
-            @click_count += 1
-            @current_index = index
-
-            box[:front].hide
-
-            show_hide_box
-
-            break
-          end
-        end
-      end
-
-      if click_type == :mouse_up
-        @box_positions.each_with_index do |box, index|
-          box[:image].mouse_location = mouse_location
-
-          if box[:image]&.is_clicked? && @enable_click
-            if !@second_selection.nil?
-              @enable_click = false
-              set_timeout(0.3) do
-                @first_selection[:front].show unless box[:show_image]
-                @second_selection[:front].show unless box[:show_image]
-                @first_selection = nil
-                @second_selection = nil
-                @enable_click = true
-              end
-            end
-
-            break
-          end
-        end
-      end
-    end
-
-    private
-
-    attr_reader :game_state, :dimension, :assets
-
-    def show
-      create_image_boxes
-    end
-
-    def create_dimension box_dimension_size, assets_count, dimension
-      counter = 0
-      # Assets count should be half of the total boxes
-      load_assets assets_count
-      box_width = box_dimension_size
-      box_height = box_dimension_size
-      margin = 10
-      total_width = (dimension * box_dimension_size) + ((dimension - 1) * margin)
-      total_height = (dimension * box_dimension_size) + ((dimension - 1) * margin)
-      start_x = (game_state.screen_width - total_width) / 2
-      start_y = (game_state.screen_height - total_height) / 2
-
-      dimension.times.each do |row|
-        dimension.times.each do |col|
-          x = start_x + ( col * (box_dimension_size + margin))
-          y = start_y + ( row * (box_dimension_size + margin))
-          new_box = draw_button(nil, x, y, box_width, box_height, @assets[counter])
-          new_box.draw_with_image
-          front_box = CRectangle.new(x: x,
-                                     y: y,
-                                     width: box_width,
-                                     height: box_height,
-                                     color: "black")
-          front_box.draw
-
-          @box_positions << {
-            image: new_box,
-            front: front_box,
-            show_image: false
-          }
-
-          counter += 1
-        end
-      end
-    end
-
-    def load_assets count
-      assets = Asset.new(count).random_assets
-      new_assets = []
-      # We create shuffled duplicates by assets count
-      new_assets << assets.shuffle
-      new_assets << assets.shuffle
-      @assets = new_assets.flatten
-    end
-
-    def create_image_boxes
-      if dimension.equal? FOUR_BY_FOUR
-        create_four_dimension
-      elsif dimension.equal? SIX_BY_SIX
-        create_six_dimension
-      end
-    end
-
-    def create_four_dimension
-      create_dimension 100, 8, 4 # 4 x 4
-    end
-
-    def create_six_dimension
-      create_dimension 80, 18, 6 # 6 x 6
-    end
-
-    def show_hide_box
-      if @first_selection.nil?
-        return if @box_positions[@current_index][:show_image] # Skip if image already shown
-        @first_selection = @box_positions[@current_index]
-      end
-
-      if !@first_selection.nil? && @click_count == 2
-        return if @box_positions[@current_index][:show_image] # Skip if image already shown
-        @second_selection = @box_positions[@current_index]
-        @click_count = 0
-        if is_same_selection?
-          @first_selection[:show_image] = true
-          @second_selection[:show_image] = true
-        end
-      end
-    end
-
-    def is_same_selection?
-      return false if @first_selection.nil? || @second_selection.nil?
-
-      @first_selection[:image].image_path == @second_selection[:image].image_path
-    end
-  end
-end
diff --git a/lib/states/dimension_selection.rb b/lib/states/dimension_selection.rb
new file mode 100644
index 0000000..2b453bc
--- /dev/null
+++ b/lib/states/dimension_selection.rb
@@ -0,0 +1,78 @@
+# frozen_string_literal: true
+
+require "ruby2d"
+require_relative "../custom_classes/button"
+require_relative "../concerns/game_display"
+require_relative "../asset"
+
+class DimensionSelection
+  include GameDisplay
+
+  attr_accessor :dimension
+
+  def initialize next_event
+    @next_event = next_event
+    @dimension = nil
+
+    show_button
+  end
+
+  def mouse_down mouse_location
+    unless @btn_1.nil? && @btn_2.nil?
+      @btn_1.mouse_location = mouse_location
+      @btn_2.mouse_location = mouse_location
+      handle_click
+    end
+  end
+
+  def mouse_up mouse_location
+  end
+
+  private
+
+  attr_reader :next_event
+
+  def show_button
+    width = 150
+    @height = 60
+
+    new_y = screen_y_center
+
+    @btn_1 = Button.new label: FOUR_BY_FOUR, x: position_x(1), y: new_y, width: width, height: height
+    @btn_1.draw
+    @btn_2 = Button.new label: SIX_BY_SIX, x: position_x(4), y: new_y, width: width, height: height
+    @btn_2.draw
+  end
+
+
+  def handle_click
+    if @btn_1&.is_clicked?
+      hide_buttons
+      @dimension = FOUR_BY_FOUR
+      next_event.call
+    elsif @btn_2&.is_clicked?
+      hide_buttons
+      @dimension = SIX_BY_SIX
+      next_event.call
+    end
+  end
+
+  def position_x column_position
+    # Divide width by 6 to have 6 columns
+    if column_position <= 6
+      column_size = Window.width / 6
+      position = column_size * column_position
+
+      return position
+    end
+  end
+
+  def screen_y_center
+    (Window.height / 2) - @height
+  end
+
+  def hide_buttons
+    @btn_1.hide
+    @btn_2.hide
+  end
+end
diff --git a/lib/states/image_selection.rb b/lib/states/image_selection.rb
new file mode 100644
index 0000000..6f79762
--- /dev/null
+++ b/lib/states/image_selection.rb
@@ -0,0 +1,170 @@
+# frozen_string_literal: true
+
+require "ruby2d"
+require_relative "../custom_classes/button"
+require_relative "../concerns/game_display"
+
+class ImageSelection
+  include GameDisplay
+
+  def initialize next_event, dimension
+    @next_event = next_event
+    @dimension = dimension
+    @box_positions = [] # image_box, front_box, is_show ( default: false )
+    @enable_click = true
+    @click_count = 0
+    create_image_boxes
+  end
+
+  def mouse_down mouse_location
+    if @enable_click
+      @box_positions.each_with_index do |box, index|
+        box[:image].mouse_location = mouse_location
+
+        if box[:image]&.is_clicked?
+          @click_count += 1
+          @current_index = index
+
+          box[:front].hide
+
+          show_hide_box
+
+          break
+        end
+      end
+    end
+  end
+
+  def mouse_up mouse_location
+    if @enable_click
+      @box_positions.each_with_index do |box, index|
+        box[:image].mouse_location = mouse_location
+
+        if box[:image]&.is_clicked? && @enable_click
+          if !@second_selection.nil?
+            @enable_click = false
+            set_timeout(0.3) do
+              @first_selection[:front].show unless box[:show_image]
+              @second_selection[:front].show unless box[:show_image]
+              @first_selection = nil
+              @second_selection = nil
+              @enable_click = true
+            end
+          end
+
+          break
+        end
+      end
+    end
+  end
+
+  private
+
+  attr_reader :next_event, :dimension, :assets
+
+  def create_image_boxes
+    if dimension.equal? FOUR_BY_FOUR
+      create_four_dimension
+    elsif dimension.equal? SIX_BY_SIX
+      create_six_dimension
+    end
+  end
+
+  def create_four_dimension
+    create_dimension 100, 8, 4 # 4 x 4
+  end
+
+  def create_six_dimension
+    create_dimension 80, 18, 6 # 6 x 6
+  end
+
+  def create_dimension box_dimension_size, assets_count, dimension
+    counter = 0
+    # Assets count should be half of the total boxes
+    load_assets assets_count
+    box_width = box_dimension_size
+    box_height = box_dimension_size
+    margin = 10
+    total_width = (dimension * box_dimension_size) + ((dimension - 1) * margin)
+    total_height = (dimension * box_dimension_size) + ((dimension - 1) * margin)
+    start_x = (Window.width - total_width) / 2
+    start_y = (Window.height - total_height) / 2
+
+    dimension.times.each do |row|
+      dimension.times.each do |col|
+        x = start_x + ( col * (box_dimension_size + margin))
+        y = start_y + ( row * (box_dimension_size + margin))
+        new_box = Button.new label: nil, x: x, y: y, width: box_width, height: box_height, image_path: @assets[counter]
+        new_box.draw_with_image
+        front_box = CRectangle.new(x: x,
+                                    y: y,
+                                    width: box_width,
+                                    height: box_height,
+                                    color: "black")
+        front_box.draw
+
+        @box_positions << {
+          image: new_box,
+          front: front_box,
+          show_image: false
+        }
+
+        counter += 1
+      end
+    end
+  end
+
+  def load_assets count
+    assets = Asset.new(count).random_assets
+    new_assets = []
+    # We create shuffled duplicates by assets count
+    new_assets << assets.shuffle
+    new_assets << assets.shuffle
+    @assets = new_assets.flatten
+  end
+
+  def create_image_boxes
+    if is_four_by_four? dimension
+      create_four_dimension
+    elsif is_six_by_six? dimension
+      create_six_dimension
+    end
+
+    show_boxes
+  end
+
+  def show_boxes
+    @enable_click = false
+    @box_positions.each{ |bp| bp[:front].hide }
+
+    time_delay = is_six_by_six?(dimension) ? 3 : 2
+
+    set_timeout(time_delay) do
+      @box_positions.each{ |bp| bp[:front].show }
+      @enable_click = true
+    end
+  end
+
+  def show_hide_box
+    if @first_selection.nil?
+      return if @box_positions[@current_index][:show_image] # Skip if image already shown
+      @first_selection = @box_positions[@current_index]
+    end
+
+    if !@first_selection.nil? && @click_count == 2
+      return if @box_positions[@current_index][:show_image] # Skip if image already shown
+      @second_selection = @box_positions[@current_index]
+      @click_count = 0
+      if is_same_selection?
+        @first_selection[:show_image] = true
+        @second_selection[:show_image] = true
+      end
+    end
+  end
+
+  def is_same_selection?
+    return false if @first_selection.nil? || @second_selection.nil?
+
+    @first_selection[:image].image_path == @second_selection[:image].image_path
+  end
+end
\ No newline at end of file
diff --git a/lib/states/play_button.rb b/lib/states/play_button.rb
new file mode 100644
index 0000000..49fbf9b
--- /dev/null
+++ b/lib/states/play_button.rb
@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+require "ruby2d"
+require_relative "../custom_classes/button"
+require_relative "../concerns/game_display"
+
+class PlayButton
+  include GameDisplay
+
+  def initialize next_event
+    @next_event = next_event
+    @width = 200
+    @height = 60
+
+    show_button
+  end
+
+  def mouse_down mouse_location
+    unless @button.nil?
+      @button.mouse_location = mouse_location
+      handle_click
+    end
+  end
+
+  def mouse_up mouse_location
+  end
+
+  private
+
+  attr_reader :next_event, :width, :height
+
+  def show_button
+    @button = Button.new(
+      label: "Play", x: center_x, y: center_y, width: width, height: height
+    )
+    @button.draw
+  end
+
+  def center_x
+    (Window.width - width) / 2
+  end
+
+  def center_y
+    (Window.height - height) / 2
+  end
+
+  def handle_click
+    if @button.is_clicked?
+      @button.hide
+      next_event.call
+    end
+  end
+end
diff --git a/main.rb b/main.rb
index e9072c5..8848d20 100644
--- a/main.rb
+++ b/main.rb
@@ -12,16 +12,14 @@
 def run
   Screen.set width: SCREEN_WIDTH, height: SCREEN_HEIGHT, title: TITLE, background: BG_COLOR
 
-  game = Game.new SCREEN_WIDTH, SCREEN_HEIGHT
+  game = Game.new
+  game.start
 
   if game.running?
     update do
       game.update
     end
 
-    render do
-    end
-
     on :mouse_down do |event|
       game.mouse_location = [event.x, event.y]
       game.mouse_click :mouse_down