#import <stdlib.h>

#import "c4model.h"

@interface Connect4Model (Private)
/**
 * Determines the game-theoretic value of a state in which the given mark
 * has won and it is the given player's turn.  The value is from the
 * point of view of the current player.  The marks corresponnding to
 * players are assumed to be 'X' for P0 and 'O' for P1.
 *
 * @param mark 'X' or 'O'
 * @param num 0 or 1
 * @return the value of a state with that winner on that player's turn
 */
 
+(int) determineValue:(char)mark onTurn:(int)num;

/**
 * Determines if either player has 4 in a row starting at the given
 * position and going in the given direction.
 *
 * @param r a row index on this board
 * @param c a column index on this board
 * @param dr -1, 0, or 1
 * @param dc -1, 0, or 1
 * @return true iff a player has a win at that location in that direction
 */
 
-(BOOL) checkWinAtRow:(int)r andCol:(int)c inDirR:(int)dr andC:(int)dc;
@end

@implementation Connect4Model

-(id) init
{
  return [self initWithWidth:7 andHeight:6];
}

-(id) initWithWidth:(int)w andHeight:(int)h
{
  [super init];

  if (w <= 0 || h <=0)
    {
      [self error:"Dimensions must be positive"];
    }

  width = w;
  height = h;
  turn = 0;

  board = malloc(sizeof(char *) * height);
  for (int r = 0; r < height; r++)
    {
      board[r] = malloc(sizeof(char) * width);
      for (int c = 0; c < width; c++)
	{
	  board[r][c] = '.';
	}
    }
  
  return self;
}

-(void) dealloc
{
  for (int r = 0; r < height; r++)
    {
      free(board[r]);
    }
  free(board);

  [super dealloc];
}

-(id) copyWithZone:(NSZone*)zone
{
  Connect4Model* clone = [Connect4Model allocWithZone:zone];
  [clone init];

  clone->width = width;
  clone->height = height;
  clone->turn = turn;

  clone->board = malloc(sizeof(char *) * height);
  for (int r = 0; r < height; r++)
    {
      clone->board[r] = malloc(sizeof(char) * width);
      for (int c = 0; c < width; c++)
	{
	  clone->board[r][c] = board[r][c];
	}
    }

  return clone;
}

-(BOOL) isLegalMove:(int)c
{
  // check for game over

  if ([self getWinner] != 0)
    {
      return NO;
    }

  // check for valid column

  if (![self isValidCol:c])
    {
      return NO;
    }

  // check for room left in col

  if ([self isFull:c]) 
    {
      return NO;
    }

  // nothing else could go wrong; move is legal

  return YES;
}

-(void) makeMoveAtCol:(int)c
{
  // find top of stack in column c
  
  int r = 0;

  while (r < height && board[r][c] == '.')
    {
      r++;
    }

  if (turn == 0)
    {
      board[r - 1][c] = 'X';
      turn = 1;
    }
  else
    {
      board[r - 1][c] = 'O';
      turn = 0;
    }
}

-(void) undoMoveAtCol:(int)c
{
  // find top of stack in column c
  
  int r = 0;

  while (r < height && board[r][c] == '.')
    {
      r++;
    }

  board[r][c] = '.';

  if (turn == 0)
    {
      turn = 1;
    }
  else
    {
      turn = 0;
    }
}

-(BOOL) isValidCol:(int)c
{
  return (c >= 0 && c < width);
}

-(BOOL) isFull:(int)c
{
  return (board[0][c] != '.');
}

-(NSArray*) getLegalMoves
{
  NSMutableArray* moves = [[[NSMutableArray alloc] init] autorelease];

  for (int c = 0 ; c < width; c++)
    {
      if (![self isFull:c])
	{
	  [moves addObject:[NSNumber numberWithInt:c]];
	}
    }

  return moves;
}

-(BOOL) isGameOver
{
  // first, check for a draw

  if ([self allFull])
    {
      return YES;
    }
  
  int winner = [self getWinner];

  return (winner != 0);
}

-(BOOL) allFull
{
  int c = 0;

  while (c < width && [self isFull:c])
    {
      c++;
    }

  return (c == width);
}

-(int) getWinner
{
  // now check for a winner at each top or left or bottom/left or top/left
  // position

  for (int startR = 0; startR < height; startR++)
    {
      for (int startC = 0; startC < width; startC++)
	{
	  if (board[startR][startC] != '.')
	    {
	      // check for vertical win
	      
	      if (startR + 4 <= height
		  && [self checkWinAtRow:startR andCol:startC inDirR:1 andC:0])
		{
		  return [Connect4Model determineValue:board[startR][startC] onTurn:turn];
		}
	  
	      // check for horizontal win
	      
	      if (startC + 4 <= width
		  && [self checkWinAtRow:startR andCol:startC inDirR:0 andC:1]) 
		{
		  return [Connect4Model determineValue:board[startR][startC] onTurn:turn];
		}
	      
	      // check for diag top/left -> bot/right win
	      
	      if (startC + 4 <= width && startR + 4 <= height
		  && [self checkWinAtRow:startR andCol:startC inDirR:1 andC:1]) 
		{
		  return [Connect4Model determineValue:board[startR][startC] onTurn:turn];
		}
	      
	      // check for diag bot/left -> top/right win
	      
	      if (startC + 4 <= width && startR - 4 >= -1
		  && [self checkWinAtRow:startR andCol:startC inDirR:-1 andC:1]) 
		{
		  return [Connect4Model determineValue:board[startR][startC] onTurn:turn];
		}
	    }
	}
    }

  return 0;
}

-(BOOL) checkWinAtRow:(int)r andCol:(int)c inDirR:(int)dr andC:(int)dc
{
  int step = 1;

  while (step < 4 && board[r + step * dr][c + step * dc] == board[r][c])
    {
      step++;
    }

  return (step == 4);
}

-(BOOL) scorePossibleWinAtRow:(int)r andCol:(int)c inDirR:(int)dr andC:(int)dc
{
  int step = 0;
  char mark = '.';
  int stepR = r;
  int stepC = c;
  int count = 0;

  while (step < 4 && (mark == '.' || board[stepR][stepC] == mark))
    {
      if (board[stepR][stepC] != '.')
	{
	  mark = board[stepR][stepC];
	  count++;
	}
      step++;
      stepR += dr;
      stepC += dc;
    }

  if (step == 4)
    {
      if (mark == '.')
	{
	  return 0;
	}
      else if (mark == turn)
	{
	  return count;
	}
      else
	{
	  return -count;
	}
    }
  else
    {
      return 0;
    }
}

-(double) getHeuristicValue
{
  // TODO: make this more efficient: scorePossible win will return +/- 4
  // if someone has won, so we don't have to go over all the possible
  // 4-in-a-row locations twice as we do below (once in getWinner and
  // then again in the loops)

  // return actual value if game is over...

  int winner = [self getWinner];
  
  if (winner != 0)
    {
      return winner;
    }

  if ([self allFull])
    {
      return 0;
    }

  // ...otherwise, see who has control of all the possible 4-in-a-row
  // locations and how close they are to them (1 of 4, 2 of 4, 3 of 4)

  int onMoveScore;
  int offMoveScore;

  for (int startR = 0; startR < height; startR++)
    {
      for (int startC = 0; startC < width; startC++)
	{
	  // get score for vertical run
	  
	  if (startR + 4 <= height)
	    {
	      int score = [self scorePossibleWinAtRow:startR andCol:startC inDirR:1 andC:0];

	      if (score < 0)
		{
		  offMoveScore += -score;
		}
	      else
		{
		  onMoveScore += score;
		}
	    }
	  
	  // check for horizontal win
	      
	  if (startC + 4 <= width)
	    {
	      int score = [self scorePossibleWinAtRow:startR andCol:startC inDirR:0 andC:1];
	      
	      if (score < 0)
		{
		  offMoveScore += -score;
		}
	      else
		{
		  onMoveScore += score;
		}
	    }
	  
	  // check for diag top/left -> bot/right win
	  
	  if (startC + 4 <= width && startR + 4 <= height)
	    {
	      int score = [self scorePossibleWinAtRow:startR andCol:startC inDirR:1 andC:1];
	      
	      if (score < 0)
		{
		  offMoveScore += -score;
		}
	      else
		{
		  onMoveScore += score;
		}
	    }
	  
	  // check for diag bot/left -> top/right win
	  
	  if (startC + 4 <= width && startR - 4 >= -1)
	    {
	      int score = [self scorePossibleWinAtRow:startR andCol:startC inDirR:-1 andC:1];
	      
	      if (score < 0)
		{
		  offMoveScore += -score;
		}
	      else
		{
		  onMoveScore += score;
		}
	    }
	}
    }
  
  if (onMoveScore == 0 && onMoveScore == 0)
    {
      return 0.0;
    }
  else
    {
      return (double)onMoveScore / (onMoveScore + offMoveScore) - 0.5;
    }
}

+(int) determineValue:(char)mark onTurn:(int)num
{
  if (mark == 'X' && num == 0)
    {
      return 1;
    }
  else if (mark == 'O' && num == 1)
    {
      return 1;
    }
  else if (mark != '.')
    {
      return -1;
    }
  else
    {
      return 0;
    }
}

-(int) alphaBetaMoveWithMaxDepth:(int)d
{
  double max; // lower than any possible value
  int bestMove = -1;

  NSEnumerator* e = [[self getLegalMoves] objectEnumerator];
  NSNumber* object;
  while ((object = [e nextObject]) != nil)
    {
      int col = [object intValue];

      [self makeMoveAtCol:col];

      double value = -[self minimaxValueWithMaxDepth:d];

      printf("%d:%f ", col, value);

      [self undoMoveAtCol:col];

      if (bestMove == -1 || value > max)
	{
	  max = value;
	  bestMove = col;
	}
    }

  printf("\n");

  return bestMove;
}
 
-(double) minimaxValueWithMaxDepth:(int)d
{
  return [self alphaBetaValue:-1 :1 :d];
}

-(double) alphaBetaValue:(double)alpha :(double)beta :(int)d
{
  NSAutoreleasePool *pool = [[NSAutoreleasePool alloc] init];

  if ([self getWinner] != 0 || [self allFull])
    {
      return [self getWinner];
    }
  
  if (d == 0)
    {
      return [self getHeuristicValue];
    }

  NSEnumerator* e = [[self getLegalMoves] objectEnumerator];
  NSNumber* move;
  while (alpha < beta && (move = [e nextObject]) != nil)
    {
      // make the move

      int moveCol = [move intValue];

      [self makeMoveAtCol: moveCol];

      // get alpha-beta value of new state

      double a = -[self alphaBetaValue:-beta :-alpha :d - 1];  

      // undo move

      [self undoMoveAtCol: moveCol];

      if (a > alpha)
	{
	  alpha = a;
	}

      if (alpha >= beta)
	{
	  [pool release];
	  return alpha;
	}
    }

  [pool release];
  return alpha;
}

-(double) monteCarloValueWithVariations:(int)variations andMaxTurns:(int)max
{
  double totValue = 0.0;

  for (int v = 0; v < variations; v++)
    {
      NSAutoreleasePool *pool = [[NSAutoreleasePool alloc] init];

      totValue += [self monteCarloValueWithMaxTurns:max];

      [pool release];
    }

  return totValue / variations;
}

-(double) monteCarloValueWithMaxTurns:(int)max
{
  // check for game over

  int winner = [self getWinner];

  if (winner != 0)
    {
      return winner;
    }

  if ([self allFull])
    {
      return 0.0;
    }

  // apply heuristic if at depth bound

  if (max == 0)
    {
      return [self getHeuristicValue];
    }

  // make a random move and find Monte Carlo value of the resulting state

  NSArray* moves = [self getLegalMoves];
  
  int i = random() % [moves count];
  NSNumber* move = [moves objectAtIndex:i];
  int col = [move intValue];

  [self makeMoveAtCol:col];

  double value = -[self monteCarloValueWithMaxTurns:max - 1];

  [self undoMoveAtCol:col];

  return value;
}

-(int) getPlayerOnMove
{
  return turn;
}

-(int) getWidth
{
  return width;
}

-(int) getHeight
{
  return height;
}

-(int)getPieceAtRow:(int)r andCol:(int)c
{
  switch (board[r][c])
    {
    case 'X':
      return 0;
      
    case 'O':
      return 1;

    default:
      return -1;
    }
}

-(int) monteCarloMoveWithVariations:(int)variations andMaxTurns:(int)max
{
  double bestValue; // lower than any possible value
  int bestMove = -1;

  NSEnumerator* e = [[self getLegalMoves] objectEnumerator];
  NSNumber* object;
  while ((object = [e nextObject]) != nil)
    {
      int col = [object intValue];

      [self makeMoveAtCol:col];

      double value = -[self monteCarloValueWithVariations:variations andMaxTurns:max - 1];

      [self undoMoveAtCol:col];

      if (bestMove == -1 || value > bestValue)
	{
	  bestValue = value;
	  bestMove = col;
	}
    }

  return bestMove;
}

@end