import gym
import math
import numpy as np

class MyWrapper(gym.Wrapper):
    def __init__(self, env):
        super().__init__(env)
        self.prevBoardState = np.zeros((self.env.height, self.env.width))

    def check_cleared_rows(self, board):
        to_delete = []
        for i, row in enumerate(board[::-1]):
            if 0 not in row:
                to_delete.append(len(board) - 1 - i)
        if len(to_delete) > 0:
            board = self.remove_row(board, to_delete)
        return len(to_delete)


    def get_bumpiness_and_height(self, board):
        board = np.array(board)
        mask = board != 0
        invert_heights = np.where(mask.any(axis=0), np.argmax(mask, axis=0), self.height)
        heights = self.height - invert_heights
        total_height = np.sum(heights)
        currs = heights[:-1]
        nexts = heights[1:]
        diffs = np.abs(currs - nexts)
        total_bumpiness = np.sum(diffs)
        return total_bumpiness, total_height, np.max(heights)

    def get_holes(self, board):
        num_holes = 0
        for col in zip(*board):
            row = 0
            while row < self.height and col[row] == 0:
                row += 1
            num_holes += len([x for x in col[row + 1:] if x == 0])
        return num_holes


    def step(self, action):
        
        ##Eredeti  step fv meghívása

        obs, reward, done, info = self.env.step(action)
        ##Itt tudom felül írni a reward számítást.

        currBoard = obs['board'].copy()
        cleared_lines = self.check_cleared_rows(currBoard)
        bumpiness, total_height, max_height = self.get_bumpiness_and_height(currBoard)
        num_holes = self.get_holes(currBoard)

        reward = 1 + math.pow(cleared_lines, 4)  * obs['board'].shape[1] - 0.28 * bumpiness - 0.36 * num_holes

        if done:
            reward -= 1000

        self.prevBoardState = obs['board'].copy()

        return obs, reward, done, info