'''

playfield:
box = 63*63 + 4
20x20 boxes
2d array filled or not
shape class

methods:
shift direction
is row full


actions:
mini frame to detect a new spawned piece in the middle of the board
direction decision
identify shape
find posititon to store on y axis and move n fields
drop piece down


extras:
pre run calc with next piece and bypass calc


'''

import cv2 as cv
import keyboard
import numpy as np
from utils import mse
from game_base_class import GameBase
import random
from pynput.keyboard import Key, Controller


BLOCK_FULL = [[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 1, 0], [0, 0, 0, 0]]
BLOCK_SMALL = [[1, 1], [1, 1]]
BLOCK_COL = 9

BL3_FULL = [[0, 0, 0, 0], [0, 1, 1, 0], [0, 0, 1, 0], [0, 0, 0, 0]]
Bl3_COL = 9

L1_FULL = [[0, 0, 0, 0], [0, 0, 1, 0], [1, 1, 1, 0], [0, 0, 0, 0]]
L1_COL = 8

L2_FULL = [[0, 0, 0, 0], [1, 1, 1, 0], [0, 0, 1, 0], [0, 0, 0, 0]]
L2_COL = 8

LINE_FULL = [[0, 0, 0, 0], [1, 1, 1, 1], [0, 0, 0, 0], [0, 0, 0, 0]]
LINE_COL = 8

DOT_FULL = [[0, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 0], [0, 0, 0, 0]]
DOT_COL = 10

DDOT_FULL = [[0, 0, 0, 0], [0, 1, 1, 0], [0, 0, 0, 0], [0, 0, 0, 0]]
DDOT_COL = 9

DDDOT_FULL = [[0, 0, 0, 0], [0, 1, 1, 1], [0, 0, 0, 0], [0, 0, 0, 0]]
DDDOT_COL = 9

Z1_FULL = [[0, 0, 0, 0], [0, 0, 1, 1], [0, 1, 1, 0], [0, 0, 0, 0]]
Z1_COL = 9

Z2_FULL = [[0, 0, 0, 0], [0, 1, 1, 0], [0, 0, 1, 1], [0, 0, 0, 0]]
Z2_COL = 9

T1_FULL = [[0, 0, 0, 0], [0, 1, 1, 1], [0, 0, 1, 0], [0, 0, 0, 0]]
T1_COL = 9

class Litris(GameBase):

    def __init__(self, overlay):
        super().__init__(overlay)

        self.keyboard = Controller()

        self.data_coordinates = np.zeros((20, 20), dtype=object)
        self.stone_coordinates = np.zeros((4, 4), dtype=object)
        self.observation = np.zeros((20, 20), dtype=int)

        self.colors = [1, 2, 3, 4, 5, 6, 7, 8, 9]

        self.offset_left = 610
        self.offset_down = 40

        self.fill_data_coordinates()

        #self.sd_reset_board = cv.imread("control_elements/sodoku_reset_button.jpg", cv.IMREAD_COLOR)

        self.needles = {1: cv.imread("litris/blue_needle.jpg", cv.IMREAD_COLOR)
        #                2: cv.imread("sodoku/2.jpg", cv.IMREAD_COLOR),
                        }

        self.full_stones_dic = {1: BLOCK_FULL,
                                2: L1_FULL,
                                3: LINE_FULL,
                                4: DOT_FULL
                                }

        self.col_stones_dic = {1: BLOCK_COL,
                               2: L1_COL,
                               3: LINE_COL,
                               4: DOT_COL
                               }

    def fill_data_coordinates(self):
        # 610 to 1950 = 1340 - 76 / 20 = 63
        # 40 to 1380 = 1340 - 76 / 20 = 63
        # spacing 19 * 4
        dim = 63
        e_spacing = 4
        i_spacing = 4
        for e in range(0, 20, 1):
            for i in range(0, 20, 1):
                self.data_coordinates[e][i] = [(i * dim) + (i * i_spacing), (e * dim) + (e * e_spacing), dim, dim]
        for e in range(0, 4, 1):
            for i in range(0, 4, 1):
                self.stone_coordinates[e][i] = [(i * dim) + (i * i_spacing), (e * dim) + (e * e_spacing), dim, dim]


    def assess_playfield_and_make_move(self):

        #if self.check_for_button_and_execute(self.capture_window.get_screenshot(), self.sd_reset_board):
        #    cv.waitKey(2000)
        current_stone = self.new_stone_detection_and_identification()
        new_observation, new_screenshot = self.get_current_board_state()
        col = self.find_place_for_stone(current_stone, new_observation)
        self.move_stone(col)


        self.observation = new_observation
        return new_observation

    def get_current_board_state(self):
        # get an updated image of the game
        screenshot = self.capture_window.get_screenshot()
        #screenshot = cv.imread("litris/main_playfield.jpg")

        # gray = cv.cvtColor(screenshot, cv.COLOR_BGR2GRAY)
        # thresh = cv.threshold(gray, 0, 255, cv.THRESH_BINARY_INV + cv.THRESH_OTSU)[1]
        #if self.check_for_button_and_execute(screenshot, self.ok_button):
        #    cv.waitKey(500)

        #screenshot = self.capture_window.get_screenshot()
        screenshot = screenshot[40:1380, 610:1950]

        #cv.imshow("screenshot", screenshot)
        #cv.waitKey(150)
        #continue
        data_coords = np.zeros((20, 20), dtype=object)
        # field = Pickaxe_Field()
        for needle_key in self.needles.keys():
            # gray_needle = cv.cvtColor(self.needles[needle_key], cv.COLOR_BGR2GRAY)
            # thresh_needle = cv.threshold(gray_needle, 0, 255, cv.THRESH_BINARY_INV + cv.THRESH_OTSU)[1]
            rectangles = self.vision_stun.find(screenshot, self.needles[needle_key], 0.85, 56)
            if len(rectangles) == 0:
                continue
            points = self.vision_stun.get_click_points(rectangles)

            for point in points:
                x, y = self.point_in_rect(point)
                if x is not None and y is not None:
                    data_coords[x][y] = int(needle_key)
                    # self.change_value(x, y, int(needle_key))
                    # print(field.data_value_grid)
                    #cv.circle(screenshot, points[0], 7, (0, 255, 0), -1)
                    #output_image = self.vision_stun.draw_rectangles(screenshot, rectangles)
                    #cv.imshow("output_image", output_image)
                    #cv.waitKey(150)
        return data_coords, screenshot

    def new_stone_detection_and_identification(self):
        stone_coords = np.zeros((4, 4), dtype=object)
        while True:
            screenshot = self.capture_window.get_screenshot()
            #screenshot = cv.imread("litris/main_playfield.jpg")
            # 1148 1412 580 845
            screenshot = screenshot[580:845, 1148:1412]
            #cv.imshow("screenshot", screenshot)
            #cv.waitKey(150)

            # gray_needle = cv.cvtColor(self.needles[needle_key], cv.COLOR_BGR2GRAY)
            # thresh_needle = cv.threshold(gray_needle, 0, 255, cv.THRESH_BINARY_INV + cv.THRESH_OTSU)[1]
            rectangles = self.vision_stun.find(screenshot, self.needles[1], 0.85, 16)
            if len(rectangles) == 0:
                cv.waitKey(100)
                continue
            points = self.vision_stun.get_click_points(rectangles)

            for point in points:
                x, y = self.point_in_smal_rect(point)
                if x is not None and y is not None:
                    stone_coords[x][y] = 1
                    # self.change_value(x, y, int(needle_key))
                    # print(field.data_value_grid)
                    #cv.circle(screenshot, points[0], 7, (0, 255, 0), -1)
                    #output_image = self.vision_stun.draw_rectangles(screenshot, rectangles)
                    #cv.imshow("output_image", output_image)
                    #cv.waitKey(150)
            return stone_coords

    def find_place_for_stone(self, stone, current_board):

        if np.array_equal(stone, BLOCK_FULL):
            # block
            for e in range(19, 17, - 1):
                for i in range(0, 19, 1):
                    if current_board[e][i] == 0 and current_board[e - 1][i] == 0 and current_board[e][i + 1] == 0 and current_board[e - 1][i + 1] == 0:
                        if e == 19:
                            return i - BLOCK_COL
                        elif e < 19 and current_board[e + 1][i] == 1 and current_board[e + 1][i + 1] == 1:
                            return i - BLOCK_COL
        if np.array_equal(stone, BL3_FULL):
            # block
            for e in range(19, 17, - 1):
                for i in range(0, 19, 1):
                    if current_board[e][i] == 1 and current_board[e - 1][i] == 0 and current_board[e][i + 1] == 0 and current_board[e - 1][i + 1] == 0:
                        return i - Bl3_COL
        elif np.array_equal(stone, L1_FULL):
            # L1
            for e in range(19, 17, - 1):
                for i in range(0, 18, 1):
                    if current_board[e][i] == 0 and current_board[e][i + 1] == 0 and current_board[e][i + 2] == 0 and \
                            current_board[e - 1][i + 2] == 0:
                        return i - L1_COL
        elif np.array_equal(stone, L2_FULL):
            # L1
            for e in range(19, 17, - 1):
                for i in range(0, 18, 1):
                    if current_board[e - 1][i] == 0 and current_board[e - 1][i + 1] == 0 and current_board[e - 1][i + 2] == 0 and \
                            current_board[e][i + 2] == 0 and current_board[e][i] == 1 and current_board[e][i + 1] == 1:
                        return i - L2_COL
        elif np.array_equal(stone, LINE_FULL):
            # Line
            for e in range(19, 18, - 1):
                for i in range(0, 17, 1):
                    if current_board[e][i] == 0 and current_board[e][i + 1] == 0 and current_board[e][i + 2] == 0 and \
                            current_board[e][i + 3] == 0:
                        if e == 19:
                            return i - LINE_COL
                        elif e < 19 and current_board[e + 1][i] == 1 and current_board[e + 1][i + 1] == 1 and current_board[e + 1][i + 2] and current_board[e + 1][i + 3]:
                            return i - LINE_COL
        elif np.array_equal(stone, DOT_FULL):
            # Dot
            for e in range(19, 18, - 1):
                for i in range(0, 20, 1):
                    if current_board[e][i] == 0:
                        if e == 19:
                            return i - DOT_COL
                        elif e < 19 and current_board[e + 1][i] == 1:
                            return i - DOT_COL

        elif np.array_equal(stone, DDOT_FULL) :
            # DDot
            for e in range(19, 18, - 1):
                for i in range(0, 19, 1):
                    if current_board[e][i] == 0 and current_board[e][i + 1] == 0:
                        if e == 19:
                            return i - DDOT_COL
                        elif e < 19 and current_board[e + 1][i] == 1 and current_board[e + 1][i + 1]:
                            return i - DDOT_COL
        elif np.array_equal(stone, DDDOT_FULL) :
            # DDDot
            for e in range(19, 18, - 1):
                for i in range(0, 18, 1):
                    if current_board[e][i] == 0 and current_board[e][i + 1] == 0 and current_board[e][i + 2] == 0:
                        if e == 19:
                            return i - DDDOT_COL
                        elif e < 19 and current_board[e + 1][i] == 1 and current_board[e + 1][i + 1] == 1 and current_board[e + 1][i + 2]:
                            return i - DDDOT_COL
                        return i - DDDOT_COL
        elif np.array_equal(stone, Z1_FULL):
            # Z1
            for e in range(19, 17, - 1):
                for i in range(0, 18, 1):
                    if current_board[e][i] == 0 and current_board[e][i + 1] == 0 and current_board[e][i + 2] == 1 and current_board[e - 1][i + 1] == 0 and current_board[e - 1][i + 2] == 0:
                        return i - Z1_COL
        elif np.array_equal(stone, Z2_FULL):
            # Z2
            for e in range(19, 17, - 1):
                for i in range(0, 18, 1):
                    if current_board[e][i] == 1 and current_board[e][i + 1] == 0 and current_board[e][i + 2] == 0 and current_board[e - 1][i] == 0 and current_board[e - 1][i + 1] == 0:
                        return i - Z2_COL
        elif np.array_equal(stone, T1_FULL):
            # T1
            for e in range(19, 17, - 1):
                for i in range(0, 18, 1):
                    if current_board[e][i] == 1 and current_board[e][i + 1] == 0 and current_board[e][i + 2] == 1 and current_board[e - 1][i] == 0 and current_board[e - 1][i + 1] == 0 and current_board[e - 1][i + 2] == 0:
                        return i - T1_COL
    def move_stone(self, col_movement):
        if col_movement is None:
            return
        # Press and release space
        self.keyboard.press(Key.down)
        self.keyboard.release(Key.down)
        cv.waitKey(250)
        if col_movement < 0:
            for i in range(0, col_movement, - 1):
                self.keyboard.press(Key.left)
                self.keyboard.release(Key.left)
                cv.waitKey(250)
        else:
            for i in range(0, col_movement, 1):
                self.keyboard.press(Key.right)
                self.keyboard.release(Key.right)
                cv.waitKey(250)

    def point_in_rect(self, point):
        for e in range(0, 20, 1):
            for i in range(0, 20, 1):
                x1, y1, w, h = self.data_coordinates[e][i]
                x2, y2 = x1 + w, y1 + h
                x, y = point
                if x1 < x and x < x2:
                    if y1 < y and y < y2:
                        return e, i
        return None, None

    def point_in_smal_rect(self, point):
        for e in range(0, 4, 1):
            for i in range(0, 4, 1):
                x1, y1, w, h = self.stone_coordinates[e][i]
                x2, y2 = x1 + w, y1 + h
                x, y = point
                if x1 < x and x < x2:
                    if y1 < y and y < y2:
                        return e, i
        return None, None