#include "Behaviors/StateMachine.h"

$nodeclass HW4_Parse_TicTacToe : VisualRoutinesStateNode {
private:
  enum TTTBoardPiece {
    noPiece=0,
    xPiece,
    oPiece
  };

public:
  $nodeclass MakeRequest : MapBuilderNode($,MapBuilderRequest::groundMap) : doStart {
    // info about lines
    mapreq.addObjectColor(lineDataType, "yellow");
    mapreq.addOccluderColor(lineDataType, "blue");
    mapreq.addOccluderColor(lineDataType, "pink");

    // info about easter eggs
    mapreq.addObjectColor(ellipseDataType, "blue");
    mapreq.addObjectColor(ellipseDataType, "pink");
  }

  static char tttPieceToChar(TTTBoardPiece piece) {
    switch (piece) {
    case xPiece:  return 'X';
    case oPiece:  return 'O';
    case noPiece: // fall through
    default:      return ' ';
    }
  }

  class ParallelTo : public UnaryShapePred<LineData> {
  private:
    Shape<LineData> &line1;
  public:
    ParallelTo(Shape<LineData> &line) : UnaryShapePred<LineData>(), line1(line) {}
    bool operator() (const Shape<LineData> &line2) const {
      return LineData::ParallelTest() (line1, line2);
    }
  };

  $nodeclass ParseBoard : VisualRoutinesStateNode : doStart {
    // Separate lines and blobs
    NEW_SHAPEVEC(lines, LineData, select_type<LineData>(camShS));
    NEW_SHAPEVEC(ellipses, EllipseData, select_type<EllipseData>(camShS));

    // Select 2 longest horizontal lines
    NEW_SHAPEVEC(hlines, LineData, subset(lines, LineData::IsHorizontal()));
    NEW_SHAPEVEC(long_hlines, LineData, stable_sort(hlines, not2(LineData::LengthLessThan())));
    if (long_hlines.size()<2) {
      cout << "Only found " << long_hlines.size() << " horizontal lines" << endl;
      return;
    }
    NEW_SHAPE(hlineT, LineData, (long_hlines[0]->leftPtShape()->isAbove(long_hlines[1]->leftPtShape()))?long_hlines[0]:long_hlines[1]);
    NEW_SHAPE(hlineB, LineData, (long_hlines[0]->leftPtShape()->isAbove(long_hlines[1]->leftPtShape()))?long_hlines[1]:long_hlines[0]);
    
    // Select 2 longest vertical lines
    NEW_SHAPEVEC(vlines, LineData, subset(lines, not1(OrPred(ParallelTo(hlineT), ParallelTo(hlineB)))));
    NEW_SHAPEVEC(long_vlines, LineData, stable_sort(vlines, not2(LineData::LengthLessThan())));
    if (long_vlines.size()<2) {
      cout << "Only found " << long_vlines.size() << " vertical lines" << endl;
      return;
    }
    NEW_SHAPE(vlineL, LineData, (long_vlines[0]->topPtShape()->isLeftOf(long_vlines[1]->topPtShape()))?long_vlines[0]:long_vlines[1]);
    NEW_SHAPE(vlineR, LineData, (long_vlines[0]->topPtShape()->isLeftOf(long_vlines[1]->topPtShape()))?long_vlines[1]:long_vlines[0]);

    // Get extreme points
    NEW_SHAPE(leftPt,   PointData, (hlineT->leftPtShape()->isLeftOf(hlineB->leftPtShape()))    ? hlineT->leftPtShape()   : hlineB->leftPtShape());
    NEW_SHAPE(rightPt,  PointData, (hlineT->rightPtShape()->isRightOf(hlineB->rightPtShape())) ? hlineT->rightPtShape()  : hlineB->rightPtShape());
    NEW_SHAPE(topPt,    PointData, (vlineL->topPtShape()->isAbove(vlineR->topPtShape()))       ? vlineL->topPtShape()    : vlineR->topPtShape());
    NEW_SHAPE(bottomPt, PointData, (vlineL->bottomPtShape()->isBelow(vlineR->bottomPtShape())) ? vlineL->bottomPtShape() : vlineR->bottomPtShape());

    // Create boundries
    NEW_SHAPE(leftLine, LineData, new LineData(camShS, leftPt, vlineL->getOrientation()));
    NEW_SHAPE(rightLine, LineData, new LineData(camShS, rightPt, vlineR->getOrientation()));
    NEW_SHAPE(topLine, LineData, new LineData(camShS, topPt, hlineT->getOrientation()));
    NEW_SHAPE(bottomLine, LineData, new LineData(camShS, bottomPt, hlineB->getOrientation()));

    // Generate masks for each grid cell
    NEW_SKETCH(boardArea, bool,
      visops::rightHalfPlane(leftLine) & visops::leftHalfPlane(rightLine) &
      visops::bottomHalfPlane(topLine) & visops::topHalfPlane(bottomLine));
    NEW_SKETCH(row1, bool, boardArea & visops::topHalfPlane(hlineT));
    NEW_SKETCH(row2, bool, boardArea & visops::bottomHalfPlane(hlineT) & visops::topHalfPlane(hlineB));
    NEW_SKETCH(row3, bool, boardArea & visops::bottomHalfPlane(hlineB));
    NEW_SKETCH(col1, bool, boardArea & visops::leftHalfPlane(vlineL));
    NEW_SKETCH(col2, bool, boardArea & visops::rightHalfPlane(vlineL) & visops::leftHalfPlane(vlineR));
    NEW_SKETCH(col3, bool, boardArea & visops::rightHalfPlane(vlineR));

    Sketch<bool> rows[3] = {row1, row2, row3};
    Sketch<bool> cols[3] = {col1, col2, col3};

    // Fill in the tic-tac-toe board
    TTTBoardPiece tttboard[3][3] = {
      {noPiece,noPiece,noPiece},
      {noPiece,noPiece,noPiece},
      {noPiece,noPiece,noPiece}};

    SHAPEVEC_ITERATE(ellipses, EllipseData, ellipse)
      NEW_SKETCH(ellipseSk, bool, ellipse->getRendering());
      NEW_SKETCH(botPixels, bool, ellipseSk & ! ellipseSk[*camSkS.idxS]);
      for (int r=0; r<3; r++) {
        for (int c=0; c<3; c++) {
          NEW_SKETCH_N(intersect, uint, botPixels & rows[r] & cols[c]);
          if (intersect->sum() > 5) {
            if (IsColor("blue")(ellipse)) {
              tttboard[r][c] = oPiece;
            } else {
              tttboard[r][c] = xPiece;
            }
          }
        }
      }
    END_ITERATE

    // Display the tic-tac-toe board
    for (int r=0; r<3; r++) {
      cout << " " << tttPieceToChar(tttboard[r][0]) << " | "
        << tttPieceToChar(tttboard[r][1]) << " | "
        << tttPieceToChar(tttboard[r][2]) << " " << endl;
      if (r<2) cout << "---+---+---" << endl;
    }

    // Create a sketch with the parsed tic-tac-toe board
    NEW_SKETCH(parsedBoard, uchar, visops::zeros(camSkS));
    for (int r=0; r<3; r++) {
      for (int c=0; c<3; c++) {
        NEW_SKETCH_N(cell, bool, rows[r] & cols[c]);
        if (tttboard[r][c]==xPiece) {
          parsedBoard += cell * ProjectInterface::getColorIndex("pink");
        } else if (tttboard[r][c]==oPiece) {
          parsedBoard += cell * ProjectInterface::getColorIndex("blue");
        }
      }
    }
  }

  $setupmachine {
    startnode: MakeRequest =MAP=> ParseBoard
  }
}

REGISTER_BEHAVIOR(HW4_Parse_TicTacToe);