#!/usr/bin/ruby

require "io/console"

# Grid dimensions.  Standard Puyo Puyo is 6x12.
WIDTH = 6
HEIGHT = 12

# Delay between each frame.
FRAME_DELAY = 0.03

# Tile sets and palette.
TILE =        ["<>",         "()",         "{}",         "[]"]
COLOR =       ["\e[97;101m", "\e[97;102m", "\e[97;103m", "\e[97;104m"]
HINT_COLOR =  ["\e[91;40m",  "\e[92;40m",  "\e[93;40m",  "\e[94;40m"]
FLASH_COLOR = ["\e[31;107m", "\e[32;107m", "\e[33;107m", "\e[34;107m"]

# Text positions, relative to upper left corner of grid.
TEXT_X = WIDTH * 2 + 4
TEXT_NEXT_Y = 0
TEXT_BLOCK_INDEX_Y = 4
TEXT_SCORE_Y = 7
TEXT_CHAIN_Y = HEIGHT - 1

# Empty tile index.
EMPTY = nil

# Offsets for the second tile.  0 is the initial orientation, with the
# second tile just beyond top of screen.  Subsequent rotations are
# clockwise such that first rotation would make the block flat and centered.
OFFSET_X = [0, 1, 0, -1]
OFFSET_Y = [1, 0, -1, 0]


# Block data.
class Block
   attr_accessor :x, :y, :orientation, :color
end

# Initialize PRNG seed to first command line argument if it's defined,
# otherwise use current time.
$seed = (ARGV.length > 0 and ARGV[0] or Time.now).to_i

# If additional command line arguments are available, use those as the
# input stream.  Otherwise, determine whether to read from keyboard or
# pipe based on whether stdin is a TTY.
if ARGV.length > 1
   $input_bytes = ARGV[1, ARGV.length - 1] * ""
else
   if STDIN.isatty
      $input_bytes = nil
   else
      $input_bytes = STDIN.read
   end
end

# Render score label.
def render_score(score)
   # Score
   print "\e[#{TEXT_X}C\e[#{TEXT_SCORE_Y + 1}B"
   label = " #{score} "
   print "\e[97;40m#{label}\e[0m"
   print "\e[#{TEXT_X + label.length}D\e[#{TEXT_SCORE_Y + 1}A"
end

# Generate grid image with overlays.
#  grid = grid to draw.
#  block = block data.
#  flash = set of cells that should be highlighted.
def grid_image(grid, block, flash)
   # Convert column-major bottom-up grid to row-major top-down cells.
   # grid uses [x][y] indices where y=0 is bottommost row.
   # cells uses [y][x] indices where y=0 is top of screen.
   cells = []
   (HEIGHT - 1).downto(0) {|y|
      cells.push([])
      WIDTH.times{|x|
         t = grid[x][y]
         if t
            if flash[[x, y]]
               cells[-1].push(FLASH_COLOR[t] + TILE[t])
            else
               cells[-1].push(COLOR[t] + TILE[t])
            end
         else
            cells[-1].push(nil)
         end
      }
   }

   # Draw pending tile.
   if block
      # Make sure we draw the bottom cell first.
      t = [
         [block.x, block.y, block.color[0]],
         [block.x + OFFSET_X[block.orientation],
          block.y + OFFSET_Y[block.orientation],
          block.color[1]]
      ]
      if t[0][1] > t[1][1]
         t[0], t[1] = t[1], t[0]
      end

      2.times{|i|
         x = t[i][0]
         y = t[i][1]
         next if x < 0 or x >= WIDTH

         # Draw in-flight cell.
         if y >= 0 and y < HEIGHT
            cells[HEIGHT - y - 1][x] = COLOR[t[i][2]] + TILE[t[i][2]]
         end

         # Draw drop position.
         y.times{|dy|
            if dy >= 0 and dy < HEIGHT and not cells[HEIGHT - dy - 1][x]
               cells[HEIGHT - dy - 1][x] = HINT_COLOR[t[i][2]] + TILE[t[i][2]]
               break
            end
         }
      }
   end

   # Flatten cells, joining rows with escape code to move cursor to next line.
   text = cells.map{|row|
      # Flatten each row, replacing empty cells with two spaces.
      row.map{|c| c or "\e[40m  "} * ""
   } * ("\e[1B\e[#{WIDTH * 2}D")

   # Move cursor back to original position
   text += "\e[#{HEIGHT - 1}A\e[#{WIDTH * 2}D"
end

# Test for collision with grid, returns nil if there is no collision.
def collision(block, grid)
   if block.x < 0 or block.x >= grid.length or grid[block.x][block.y]
      return 1
   end
   x1 = block.x + OFFSET_X[block.orientation]
   y1 = block.y + OFFSET_Y[block.orientation]
   if x1 < 0 or x1 >= grid.length or grid[x1][y1]
      return 1
   end
   return nil
end

# Save initial seed.  This will be printed at the end.
seed = $seed

# Hide cursor, disable echo, draw screen decorations, and leave cursor at
# upper left corner of grid.
STDIN.echo=false
print "\e[?25l"
(HEIGHT + 2).times{|y|
   row = "##" + ("##" * WIDTH) + ("#" * 18) + "\n"
   if y - 1 == TEXT_NEXT_Y
      row[TEXT_X + 2, 7] = " Next: "
   elsif y - 1 == TEXT_BLOCK_INDEX_Y
      row[TEXT_X + 2, 8] = " Block: "
   elsif y - 1 == TEXT_SCORE_Y
      row[TEXT_X + 2, 8] = " Score: "
   end
   print row
}
print "\e[#{HEIGHT + 1}A\e[2C"

# Initialize game state.
grid = []
WIDTH.times{ grid.push([]) }

# Generate blocks with equal distribution.
blocks = []
6.times{
   TILE.length.times{|i|
      TILE.length.times{|j|
         blocks.push([i, j])
      }
   }
}

# Add a few more single-color ones so that we get exactly 100 blocks.
TILE.length.times{|i|
   blocks.push([i, i])
}

# Fisher-Yates shuffle.
(blocks.length - 1).downto 1 do |i|
   # Generate a 31bit random number.  This is almost identical to the
   # algorithm described here:
   # https://en.wikipedia.org/wiki/Xorshift
   #
   # Except it's 31bits because we drop the sign bit, and we drop the
   # sign bit because right shift may or may not do sign extension
   # depending on language, and we want to have a portable random number
   # sequence that we can use elsewhere.
   $seed ^= $seed << 13
   $seed &= 0x7fffffff
   $seed ^= $seed >> 17
   $seed ^= $seed << 5
   $seed &= 0x7fffffff

   j = $seed % (i + 1)
   blocks[i], blocks[j] = blocks[j], blocks[i]
end

chain_history = []
total_tiles_cleared = 0
score = 0
input_history = ""

# Run game for each block.
blocks.length.times{|block_index|
   # Display next two blocks.
   print "\e[#{TEXT_NEXT_Y + 1}B\e[#{TEXT_X}C"
   2.times{|i|
      b = blocks[block_index + i + 1]
      if b
         t = ["\e[40m #{COLOR[b[0]]}#{TILE[b[0]]}\e[40m \e[0m",
              "\e[40m #{COLOR[b[1]]}#{TILE[b[1]]}\e[40m \e[0m"]
      else
         t = ["\e[0m####"] * 2
      end

      # Note that because the initial orientation has tile[1] on top
      # of tile[0], we will render the tiles upside down here as well.
      print "#{t[1]}\e[4D\e[1B#{t[0]}\e[1A"
   }
   print "\e[#{TEXT_X + 8}D\e[#{TEXT_NEXT_Y + 1}A"

   # Display block index.
   print "\e[#{TEXT_X}C\e[#{TEXT_BLOCK_INDEX_Y + 1}B"
   label = " #{block_index + 1} / #{blocks.length} "
   print "\e[97;40m#{label}\e[0m"
   print "\e[#{TEXT_X + label.length}D\e[#{TEXT_BLOCK_INDEX_Y + 1}A"

   # Display score.
   render_score(score)

   # Get next block to be placed.
   pending = Block.new
   pending.x = (WIDTH - 1) / 2
   pending.y = HEIGHT - 1
   pending.orientation = 0
   pending.color = blocks[block_index]

   # If we got a collision immediately after generating a new
   # block, it means grid is completely filled, i.e. game over.
   break if collision(pending, grid)

   # If we didn't see a collision right after the block is spawned,
   # we will not see another collision for the lifetime of this
   # block, since the various movement functions will enforce no
   # collisions.  That said, it may very well be the case that the
   # block can't be moved at all, and can only be dropped.

   while true
      print grid_image(grid, pending, {})

      # Read next input command from either keyboard or $input_bytes.
      key = nil

      # Repeat until we got a valid key.  This also skips over all the
      # extra bytes from arrow keys encoded as escape sequences.
      while not key
         if $input_bytes
            break if $input_bytes.length == 0
            tmp_key = $input_bytes[0]
            $input_bytes = $input_bytes[1, $input_bytes.length - 1]
         else
            tmp_key = STDIN.getch
         end

         # Convert arrow keys to VI keys.
         case tmp_key
         when 'A', 'k'; key = 'k'  # Up
         when 'B', 'j'; key = 'j'  # Down
         when 'C', 'l'; key = 'l'  # Right
         when 'D', 'h'; key = 'h'  # Left
         when 'q', 'x'
            # Quit
            key = nil
            break
         end
      end

      case key
      when 'k'  # Up = rotate block
         updated_block = pending.clone
         updated_block.orientation = (updated_block.orientation + 1) % 4

         # Check if the second tile collides with the grid after rotation.
         if collision(updated_block, grid)
            if OFFSET_X[updated_block.orientation] == 0
               # Colliding vertically, try shift the block up just
               # above the grid.
               updated_block.y += 1
            else
               # Colliding horizontally.  Try shift the block in the
               # opposite direction to resolve the collision, i.e.
               # translate the rotation to a "kick" off the wall.
               updated_block.x -= OFFSET_X[updated_block.orientation]
               updated_block = nil if collision(updated_block, grid)
            end
         end
         if updated_block
            input_history += key
            pending = updated_block
         end
      when 'j'  # Down = drop block
         # Mark grid location with pending block.
         grid[pending.x][pending.y] = pending.color[0]

         x1 = pending.x + OFFSET_X[pending.orientation]
         y1 = pending.y + OFFSET_Y[pending.orientation]
         grid[x1][y1] = pending.color[1]

         input_history += key
         break
      when 'h', 'l'  # Left or right = lateral movement.
         updated_block = pending.clone
         if key == 'h'
            updated_block.x -= 1
         else
            updated_block.x += 1
         end
         if not collision(updated_block, grid)
            # Try bring the block down to usual height if possible.
            if updated_block.y != HEIGHT - 1
               d = updated_block.clone
               d.y = HEIGHT - 1
               updated_block = d if not collision(d, grid)
            end
            input_history += key
            pending = updated_block
         end
      else
         key = nil
         break
      end
   end

   # Check if user wants to quit.
   break if not key

   # One point for each block dropped.
   #
   # Puyo Puyo has some opaque formula on how drop bonus is computed and
   # I can't find any documentation on it.  It's probably something based
   # on time, but since this version doesn't have any time pressure for
   # dropping blocks, we will just give a constant score all the time.
   score += 1

   # Repeat until there are no blocks to be cleared.
   chain_length = 0
   while true
      # Apply gravity until blocks fell to the bottom.
      print grid_image(grid, nil, {})
      while true
         updated_grid = []
         updated = false
         grid.length.times{|x|
            y = grid[x].index(EMPTY)
            if y
               updated_grid.push(grid[x][0, y] +
                                 grid[x][y + 1, grid[x].length - 1])
               updated = true
            else
               updated_grid.push(grid[x])
            end
         }

         break if not updated
         grid = updated_grid
         sleep FRAME_DELAY
         print grid_image(grid, nil, {})
      end

      # Mark clusters that are suitable for removal in grid.
      clusters = []

      # Set of tiles that have already been marked from previous depth-first
      # search expansions.  This is to prevent double-counting tiles.
      visited = {}
      grid.length.times{|i|
         grid[i].length.times{|j|
            # Perform depth-first search starting at current cell.
            color = grid[i][j]
            sub_cluster = []
            stack = [[i, j]]
            while stack.length > 0
               cursor = stack.pop
               x, y = cursor[0, 2]
               next if x < 0 or y < 0 or
                       x >= grid.length or y >= grid[x].length
               next if visited[[x, y]] or grid[x][y] != color

               visited[[x, y]] = 1
               sub_cluster.push([x, y])
               stack.push([x + 1, y])
               stack.push([x - 1, y])
               stack.push([x, y + 1])
               stack.push([x, y - 1])
            end

            # Add sub-cluster to list of clusters if it's sufficiently large.
            if sub_cluster.length >= 4
               clusters.push(sub_cluster)
            end
         }
      }

      # Remove clusters.
      break if clusters.length == 0
      chain_length += 1
      label = " #{chain_length} chain "
      print "\e[#{TEXT_CHAIN_Y}B\e[#{TEXT_X}C\e[97;40m#{label}"
      print "\e[0m\e[#{TEXT_X + label.length}D\e[#{TEXT_CHAIN_Y}A"

      remove_tiles = {}
      clusters.each{|group|
         group.each{|tile|
            remove_tiles[[tile[0], tile[1]]] = 1
         }
      }
      normal_image = grid_image(grid, nil, {})
      flash_image = grid_image(grid, nil, remove_tiles)
      10.times{
         print normal_image
         sleep FRAME_DELAY
         print flash_image
         sleep FRAME_DELAY
      }

      # Compute score based on list of clusters cleared.
      #
      # Rules taken from this page:
      # https://puyonexus.com/wiki/Scoring
      tiles_cleared = 0
      group_bonus = 0
      clear_colors = {}
      clusters.each{|group|
         clear_colors[grid[group[0][0]][group[0][1]]] = 1
         tiles_cleared += group.length
         if group.length < 11
            group_bonus += [nil, nil, nil, nil,
                            0, 2, 3, 4, 5, 6, 7][group.length]
         else
            group_bonus += 10
         end
      }

      if chain_length < 9
         # Classic Puyo Puyo (1992) table.
         # https://puyonexus.com/wiki/List_of_attack_powers#Classic_Rules
         chain_power = [nil, 0, 8, 16, 32, 64, 128, 256, 512][chain_length]
      else
         chain_power = 999
      end

      color_bonus = [nil, 0, 3, 6, 12][clear_colors.keys.length]

      bonus = chain_power + color_bonus + group_bonus
      bonus = [[1, bonus].max, 999].min
      score += 10 * tiles_cleared * bonus

      render_score(score)

      # Remove tiles.  Do this after computing score since we need
      # the tiles to compute color bonus.
      remove_tiles.each_key{|tile|
         grid[tile[0]][tile[1]] = nil
      }
      total_tiles_cleared += remove_tiles.keys.length
   end
   chain_history.push(chain_length)

   # Overwrite chain label text.
   print "\e[", HEIGHT - 1, "B\e[", WIDTH * 2 + 4, "C\e[0m"
   print "#" * 12
   print "\e[", HEIGHT - 1, "A\e[", WIDTH * 2 + 16, "D\e[0m"
}

# Move cursor back to bottom of screen and reset cursor + colors.
STDIN.echo=true
print "\e[0m\n" * (HEIGHT + 2), "\e[?25h"

# Output termination label text.
if chain_history.length < blocks.length
      rank = ""
   print "\e[91mGAME OVER\e[0m\n"
else
   # Give a ranking based on score.
   #
   # To get the "minimalist" ranking (inspired by Ikaruga's "Dot Eater"
   # ranking), you will need to score the minimum number of points
   # possible:
   #
   #  (total number of tiles - tiles that can be held in the field) * 10
   #  + (points for dropping all the tiles)
   #
   # It's actually possible to get slightly lower score by making use of
   # the overflow rows, but they only help a little bit.  The main
   # challenge is to make sure to not make accidental chains, which is
   # counter intuitive to how people usually play Puyo Puyo.
   rank = "MINIMALIST"
   [[(blocks.length * 2 - WIDTH * HEIGHT) * 10 + blocks.length, "NOVICE"],
    [9999, "HOBBYIST"],
    [19999, "ENTHUSIAST"],
    [49999, "PROFESSIONAL"],
    [99999, "MASTER"],
    [499999, "GRANDMASTER"]].each{|s|
      rank = s[1] if score > s[0]
   }
   rank = "rank = \e[4;1m#{rank}\e[0m\n"
   print "\e[92mYOU WIN :)\e[0m\n"
end

# Output general stats.
print "\nscore = #{score}\n"
print "blocks = #{chain_history.length} / #{blocks.length}\n"
print "tiles cleared = #{total_tiles_cleared}\n", rank

# Output chain histogram for each block.
if chain_history.length > 0 and chain_history.max > 0
   print "\nChain history:\n"
   y = chain_history.max + 2
   y += 1 if y % 2 == 1
   until y == 0
      y -= 2
      x = 0
      until x >= chain_history.length
         v = chain_history[x, 2].max
         if v > y + 1
            print "8"
         elsif v == y + 1
            print "o"
         elsif y == 0
            print "_"
         else
            print " "
         end
         x += 2
      end
      print "\n"
   end

   # Mark block index to first occurrence of maximum chain.
   max_chain_index = chain_history.index(chain_history.max)
   print " " * (max_chain_index / 2), "^\n"
   print " " * (max_chain_index / 2), "block #{max_chain_index + 1} = "
   print "#{chain_history.max} chain\n"
end

print "\nReplay:\nruby #{$0} #{seed} #{input_history}\n"