---- types and terminology
-- side: white | black
-- row: 1..8  -- counts in graphics order, from black's side down to white's
-- file: 1..8
-- coord: row, file
-- square: light | dark
-- piecekind: k | q | r | b | n | p                     -- just a char below
-- piece: side, piecekind                               -- just a char, using uppercase for white
-- side = 8
-- maybe_piece = piece | _                              -- just a char below
-- pos: maybe_piece[side][side]
-- moved: k:bool, 1:bool, 8:bool  -- integers for starting files of rooks
-- simple_move: s:coord, e:coord
-- move:
--   simple_move,
--   also_move: simple_move,  -- only for castling
--   also_clear: square,  -- only for en passant
--
-- board:
--   turn: side
--   pos
--   moved: {k=moved, K=moved}
--   pawn_starting_move_just_happened: simple_move

run_screen('functional')

local utils = require 'utils'  -- undocumented so far

function side(maybe_piece)
  if maybe_piece == '_' then
    return nil
  elseif maybe_piece:match('%l') ~= nil then
    return 'black'
  else
    return 'white'
  end end

function pos_side(pos, row, file)
  return side(pos[row][file])
end

function piecekind(maybe_piece)
  if maybe_piece == '_' then
    return nil
  else
    return maybe_piece:lower()
  end end

function pos_piecekind(pos, row, file)
  return piecekind(pos[row][file])
end

function starting_pos()
  local K,Q,R,B,N,P = 'K','Q','R','B','N','P'
  local k,q,r,b,n,p = 'k','q','r','b','n','p'
  local _ = '_'
  return {
    {r,n,b,q,k,b,n,r},
    {p,p,p,p,p,p,p,p},
    {_,_,_,_,_,_,_,_},
    {_,_,_,_,_,_,_,_},
    {_,_,_,_,_,_,_,_},
    {_,_,_,_,_,_,_,_},
    {P,P,P,P,P,P,P,P},
    {R,N,B,Q,K,B,N,R},
  }
end

function empty_pos()
  local K,Q,R,B,N,P = 'K','Q','R','B','N','P'
  local k,q,r,b,n,p = 'k','q','r','b','n','p'
  local _ = '_'
  return {
    {_,_,_,_,_,_,_,_},
    {_,_,_,_,_,_,_,_},
    {_,_,_,_,_,_,_,_},
    {_,_,_,_,_,_,_,_},
    {_,_,_,_,_,_,_,_},
    {_,_,_,_,_,_,_,_},
    {_,_,_,_,_,_,_,_},
    {_,_,_,_,_,_,_,_},
  }
end

---- drawing

local d = 96
local left, top = 30+10, Menu_bottom+10+d

-- squares

local dark_color = {0.6,0.4,0.2}
local light_color = {1, 0.8, 0.6}
local square_color = {light=light_color, dark=dark_color}

function draw_square(square, x,y, d)
  color(unpack(square_color[square]))
  rect('fill', x,y, d,d)
end

function opposite_square(square)
  return (square == 'dark') and 'light' or 'dark'
end

function draw_empty_board()
  local square = 'light'
  local x, y = left, top
  for _ = 1,8 do
    for _ = 1,8 do
      draw_square(square, x,y, d)
      square = opposite_square(square)
      x = x+d
    end
    square = opposite_square(square)
    y = y+d
    x = x - 8*d
  end
end

-- pieces

local black_color = {0, 0, 0}
local white_color = {1, 1, 1}
local side_color = {black=black_color, white=white_color}

function opposite_side(side)
  return (side == 'black') and 'white' or 'black'
end

local draw = {}

function draw.k(x,y, d)
  rect('fill', x+d/6, y+d/3, d*2/3, d/3)
  rect('fill', x+d/3, y+d/6, d/3, d*2/3)
end

function draw.q(x,y, d)
  circle('fill', x+d/2, y+d/2, d/3)
end

function draw.r(x,y, d)
  rect('fill', x+d/4, y+d/4, d/2, d/2)
end

function draw.b(x,y, d)
  ellipse('fill', x+d/2, y+d/2, d/6, d/4)
end

function draw.n(x,y, d)
  poly('fill', x+d/4, y+2*d/3, x+3*d/4, y+d/4, x+3*d/4, y+3*d/4)
end

function draw.p(x,y, d)
  circle('fill', x+d/2, y+d/2, d/8)
end

function draw_pieces(pos)
  local x, y = left, top
  for r,row in ipairs(pos) do
    for f,p in ipairs(row) do
      local side = pos_side(pos, r, f)
      if side then
        color(unpack(side_color[pos_side(pos, r, f)]))
        draw[p:lower()](x,y, d)
      end
      x = x+d
    end
    y = y+d
    x = x - 8*d
  end
end

---- moves

function is_valid_square(row, file)
  return file >= 1 and file <= 8 and row >= 1 and row <= 8
end

function legal_simple_moves(board, move_start)
  local result = {}
  local sr, sf = move_start[1], move_start[2]
  local move_ends = legal_move_ends(board, sr,sf)
  for _,m in ipairs(move_ends) do
    table.insert(result, {s=move_start, e=m})
  end
  return result
end

function is_legal_simple_move(board, start_coord, end_coord)
  local simple_moves = legal_simple_moves(board, start_coord)
  local sr, sf = start_coord[1], start_coord[2]
  local r1, f1 = end_coord[1], end_coord[2]
  for _,sm in ipairs(simple_moves) do
    assert(sr == sm.s[1] and sf == sm.s[2])
    local r2, f2 = sm.e[1], sm.e[2]
    if r1 == r2 and f1 == f2 then return true end
  end
end

function legal_move_ends(board, rank, file)
  local result = legal_move_ends_nochecks(board, rank,file)
  -- remove moves that would leave the king in check
  result = remove(result,
    function(e)
      local nboard = make_simple_move(board, {rank, file}, e)
      local kr, kf = find_piece(nboard, board.turn == 'white' and 'K' or 'k')
      return under_attack(nboard, kr, kf, board.turn)
    end)
  return result
end

local moves = {}  -- simple move destinations; no castling or en passant

function moves.k(pos, row, file)
  return array_join(
    available_squares(pos, row, file, unfold{0, 1}),
    available_squares(pos, row, file, unfold{1, 1}))
end

function moves.q(pos, row, file)
  return array_join(
    project(pos, row, file, unfold{0, 1}),
    project(pos, row, file, unfold{1, 1}))
end

function moves.r(pos, row, file)
  return project(pos, row, file, unfold{0, 1})
end

function moves.b(pos, row, file)
  return project(pos, row, file, unfold{1, 1})
end

function moves.n(pos, row, file)
  return available_squares(pos, row, file, unfold{1, 2})
end

function moves.p(pos, row, file)
  local side = pos_side(pos, row, file)
  assert(row > 1 and row < 8)
  local result = {}
  local newr = (side == 'black') and row+1 or row-1
  if pos[newr][file] == '_' then
    table.insert(result, {newr, file})
    -- first move
    if side == 'black' then
      if row == 2 and pos[row+2][file] == '_' then
        table.insert(result, {row+2, file})
      end
    else
      if row == 7 and pos[row-2][file] == '_' then
        table.insert(result, {row-2, file})
      end
    end
  end
  -- captures
  if file > 1 then
    if pos_side(pos, newr, file-1) == opposite_side(side) then
      table.insert(result, {newr, file-1})
    end
  end
  if file < 8 then
    if pos_side(pos, newr, file+1) == opposite_side(side) then
      table.insert(result, {newr, file+1})
    end
  end
  return result
end

-- possible ends of legal simple moves without considering checks
function legal_move_ends_nochecks(board, row, file, skip_castling)
  local piecekind = pos_piecekind(board.pos, row, file)
  assert(piecekind)
  local result = moves[piecekind](board.pos, row, file)
  if piecekind == 'k' then
    if not skip_castling then
      maybe_add_castling(board, row, file, result)
    end
  elseif piecekind == 'p' then
    maybe_add_en_passant(board, row, file, result)
  end
  return result
end

-- offsets is a 2-D sparse table
-- order is not guaranteed
function available_squares(pos, row, file, offsets)
  local piece_side = pos_side(pos, row, file)
  local result = {}
  for r,h in pairs(offsets) do
    for f in pairs(h) do
      local row, file = row+r, file+f
      if is_valid_square(row, file)
          and pos_side(pos, row, file) ~= piece_side
      then
        table.insert(result, {row, file})
      end end end
  return result
end

-- like available squares, but now the offsets are directions
-- that you can repeatedly move along.
function project(pos, row, file, dirs)
  local piece_side = pos_side(pos, row, file)
  local result = {}
  for r,h in pairs(dirs) do
    for f in pairs(h) do
      local row, file = row, file
      while true do
        row, file = row+r, file+f
        if is_valid_square(row, file) then
          if pos[row][file] == '_' then
            table.insert(result, {row, file})
          elseif pos_side(pos, row, file) == piece_side then
            break
          elseif pos_side(pos, row, file) ~= piece_side then
            -- can capture, but not move past
            table.insert(result, {row, file})
            break
         end
        else
          break
        end end end end
  return result
end

-- 8-fold symmetry on a displacement
function unfold(t)
  local m = {}
  local function add(r, f)
    m[r] = m[r] or {}
    m[r][f] = true
  end
  local function add2(r, f)
    add(r, f) add(r, -f) add(-r, f) add(-r, -f)
  end
  add2(t[1], t[2])
  add2(t[2], t[1])
  return m
end

function maybe_add_castling(board, row, file, result)
  if can_castle(board, row, file, 1) then
    table.insert(result, {row, 3})
  end
  if can_castle(board, row, file, 8) then
    table.insert(result, {row, 7})
  end
end

function can_castle(board, row, file, rook_file)
  local p = board.pos[row][file]
  if p == 'k' and row ~= 1 then return end
  if p == 'K' and row ~= 8 then return end
  if board.moved[p][p] then return end
  if board.moved[p][rook_file] then return end
  assert(file == 5)
  -- can't pass through check
  -- no need to move the king
  if rook_file == 1 then
    for f=2,5 do
      if f ~= 5 and board.pos[row][f] ~= '_' then return end
      if under_attack(board, row, f, side(p)) then return end
    end
  else assert(rook_file == 8)
    for f=5,7 do
      if f ~= 5 and board.pos[row][f] ~= '_' then return end
      if under_attack(board, row, f, side(p)) then return end
    end
  end
  return true
end

function maybe_add_en_passant(board, row, file, result)
  local p = board.pos[row][file]
  if p == 'p' and row ~= 5 then return end
  if p == 'P' and row ~= 4 then return end
--?   print('E c', board.pawn_starting_move_just_happened)
  if not board.pawn_starting_move_just_happened then
    return
  end
  local q = board.pawn_starting_move_just_happened
  if q[1] ~= row then return end
  assert(0 < q[2] and q[2] <= 8)
--?   print('E d', q[2], file)
  if q[2] == file+1 or q[2] == file-1 then
    local side = pos_side(board.pos, row, file)
    local newr = (side == 'black') and row+1 or row-1
    assert(board.pos[newr][q[2]] == '_')
    table.insert(result, {newr, q[2]})
  end
end

function make_simple_move(board, s, e)
  local nboard = utils.deepcopy(board)
  nboard.pos[e[1]][e[2]] = board.pos[s[1]][s[2]]
  nboard.pos[s[1]][s[2]] = '_'
  return nboard
end

function find_piece(board, piece)
  for r=1,8 do
    for f=1,8 do
      if board.pos[r][f] == piece then
        return r, f
      end end end end

-- return true if square s, e is under attack by the opponent of 'side'
function under_attack(board, s, e, side)
  local opposite_side = opposite_side(side)
  local attacked_squares = all_legal_move_ends_nochecks_nocastling(board, opposite_side)
  return find(attacked_squares,
    function(sq) return sq[1] == s and sq[2] == e end) ~= nil
end

function all_legal_move_ends_nochecks_nocastling(board, for_side)
  local result = {}
  for r,row in ipairs(board.pos) do
    for f,p in ipairs(row) do
      local piece_side = pos_side(board.pos, r,f)
      if piece_side == for_side then
        result = array_join(result, legal_move_ends_nochecks(board, r,f, --[[skip castling]] true))
      end end end
  return result
end

---- handlers

function sy(row) return top+(row-1)*d end
function sx(file) return left+(file-1)*d end

function to_square(x, y)
  local row = 1+floor((y-top)/d)
  local file = 1+floor((x-left)/d)
  if is_valid_square(row, file) then
    return row, file
  end
end

function custom_pos()
  -- modify to taste
  local K,Q,R,B,N,P = 'K','Q','R','B','N','P'
  local k,q,r,b,n,p = 'k','q','r','b','n','p'
  local _ = '_'
  local result = {
    {r,n,b,q,k,_,n,r},
    {p,p,p,p,P,P,p,p},
--?     {_,_,_,_,_,_,_,_},
    {_,_,_,_,_,_,_,_},
    {n,_,_,_,_,_,_,_},
    {_,_,_,_,p,_,r,_},
    {_,_,_,_,_,_,_,_},
--?     {_,_,_,_,_,_,_,_},
    {P,P,P,P,P,P,P,P},
    {R,N,B,Q,K,_,_,R},
  }
  return result
end

function initial_board()
  -- modify to taste
  local pos = custom_pos()
  return {
    turn='white',
    pos=pos,
    moved={
      k={k=nil, [1]=nil, [8]=nil},
      K={k=nil, [1]=nil, [8]=nil},
    },
    pawn_starting_move_just_happened=nil,  -- for en passant
  }
end

-- UI state while a move is in the process of being made
local moving_piece_coord = nil
local promoting = nil

local board = initial_board()
local history = {board}

function car.keychord_press(chord)
  if chord == 'C-z' then
    if #history > 1 then
      board = table.remove(history)
      moving_piece_coord, promoting = nil
    end
  end
end

function draw_board(board)
  draw_empty_board()
  if moving_piece_coord then
    local coords = legal_move_ends(board, moving_piece_coord[1], moving_piece_coord[2])
    highlight_square(moving_piece_coord)
    for _,coord in ipairs(coords) do
      highlight_move({s=moving_piece_coord, e=coord})
    end
  end
  draw_pieces(board.pos)
end

function highlight_move(m)
  highlight_square(m.e)
  -- draw 'arrow'
  g.setLineWidth(3)
  color(0,0,1)
  local x2, y2 = sx(m.e[2])+d/2, sy(m.e[1])+d/2
  line(sx(m.s[2])+d/2, sy(m.s[1])+d/2, x2, y2)
  circle('fill', x2, y2, 5)
end

function highlight_square(s)
  color(0.4,1,0.4, 0.1)
  rect('fill', sx(s[2]), sy(s[1]), d, d)
end

function car.draw()
  draw_board(board)
  if promoting then
    draw_promotion_options(board)
  end
  color(0.2,0.2,0.2)
  local y = top+d*2
  if board.turn == 'white' then
    y = top+d*6
  end
  g.print(board.turn, left+d*8+20, y)
end

function car.mouse_press(x,y)
  if promoting then
    return select_piece_to_promote(board, x, y)
  end
  local row, file = to_square(x, y)
  if moving_piece_coord == nil then
    -- pick up piece
    if row and pos_side(board.pos, row, file) == board.turn then
      moving_piece_coord = {row, file}
    end
  else
    -- attempt to put down piece
    board = maybe_make_move(board, moving_piece_coord, {row, file})
    moving_piece_coord = nil
  end
end

function maybe_make_move(board, start_coord, end_coord)
  if is_legal_simple_move(board, start_coord, end_coord) then
    table.insert(history, board)
    board = make_move(board, start_coord, end_coord)
    if not promoting then
      board.turn = opposite_side(board.turn)
    end
  end
  return board
end

-- precondition: move is valid
function make_move(board, start_coord, end_coord)
  local s, e = start_coord, end_coord
  local nboard = make_simple_move(board, s, e)
  -- additional mutations
  maybe_finish_castling(board, nboard, s, e)
  maybe_finish_en_passant(board, nboard, s, e)
  -- some state for future castling and en passant
  nboard.pawn_starting_move_just_happened = nil
  local piece = board.pos[s[1]][s[2]]
  local piece_kind = pos_piecekind(board.pos, s[1], s[2])
  if piece_kind == 'k' then
    nboard.moved[piece][piece] = true
  elseif piece == 'r' then
    if s[2] == 1 or s[2] == 8 then
      nboard.moved.k[s[2]] = true
    end
  elseif piece == 'R' then
    if s[2] == 1 or s[2] == 8 then
      nboard.moved.K[s[2]] = true
    end
  elseif piece_kind == 'p' then
    if piece == 'p' and e[1] == 8 then promoting = e end
    if piece == 'P' and e[1] == 1 then promoting = e end
    if math.abs(s[1]-e[1]) == 2 then
--?     print('set', s[1], s[2], e[1], e[2])
      nboard.pawn_starting_move_just_happened = e
    end
  end
  -- commit
  return nboard
end

-- special moves

function maybe_finish_castling(board, nboard, s, e)
  local piece = board.pos[s[1]][s[2]]
  if piece ~= 'k' and piece ~= 'K' then return end
  local r = e[1]
  if e[2] == s[2]+2 then
    assert(pos_piecekind(board.pos, r, 8) == 'r')
    nboard.pos[r][6] = board.pos[r][8]
    nboard.pos[r][8] = '_'
    nboard.moved[piece][8] = true
  elseif e[2] == s[2]-2 then
    assert(pos_piecekind(board.pos, r, 1) == 'r')
    nboard.pos[r][4] = board.pos[r][1]
    nboard.pos[r][1] = '_'
    nboard.moved[piece][1] = true
  end
end

function maybe_finish_en_passant(board, nboard, s, e)
  local piece = board.pos[s[1]][s[2]]
--?   print('e a', piece)
  if piece ~= 'p' and piece ~= 'P' then return end
--?   print('e b', s[2], e[2])
  if e[2] == s[2] then return end
--?   print('e c', e[1], e[2])
  if pos_piecekind(board.pos, e[1], e[2]) then return end
  nboard.pos[s[1]][e[2]] = '_'
end

function draw_promotion_options(board)
  local row = (board.turn == 'white') and 0 or 9
  for f,piece in ipairs{'q','r','b','n'} do
    local x, y = sx(f), sy(row)
    draw_square('light', x,y, d-2)
    color(unpack(side_color[board.turn]))
    draw[piece](x,y, d)
  end
end

function select_piece_to_promote(board, x,y)
  local watch_row = (board.turn == 'white') and 0 or 9
  local r = 1+floor((y-top)/d)
  local f = 1+floor((x-left)/d)
  if r ~= watch_row then return end
  local piece = ({'q', 'r', 'b', 'n'})[f]
  if piece == nil then return end
  if board.turn == 'white' then
    assert(promoting[1] == 1)
    assert(board.pos[1][promoting[2]] == 'P')
    board.pos[1][promoting[2]] = piece:upper()
  else
    assert(promoting[1] == 8)
    assert(board.pos[8][promoting[2]] == 'p')
    board.pos[8][promoting[2]] = piece
  end
  promoting = nil
  board.turn = (board.turn == 'white') and 'black' or 'white'
end