import math
import random

import numpy as np
import pyglet
from fishpy.utility.debug import profile
from numba import cuda

SIMULATION_RATE = 50

SLIMES = 6000
SLIME_SIZE = 2.5
SLIME_COLOR = (173,216,230)
SLIME_SPEED = 3

GPU_BLOCKS = 96
GPU_BLOCKS_SHAPE = (8,12)
GPU_THREADS_PER_BLOCK = 64
GPU_THREADS_SHAPE = (8,8)
GPU_THREADS = GPU_THREADS_PER_BLOCK * GPU_BLOCKS
GPU_THREADS_X = (GPU_BLOCKS_SHAPE[0]*GPU_THREADS_SHAPE[0])
GPU_THREADS_Y = (GPU_BLOCKS_SHAPE[1]*GPU_THREADS_SHAPE[1])

INITIAL_DIRECTION_RANDOM = False
INITIAL_POSITION_RANDOM = False

WALL_BOUNCE_RANDOM_ANGLE = True
RANDOM_DIRECTION_CHANGE_CHANCE = 0.1

DISPERSAL_ENABLED = True
DEGRADATION_STEEPNESS = 1.44
DEGRADATION_OFFSET = 1.07

SENSOR_ANGLE = math.pi/4
SENSOR_OFFSET = 30
ROTATION_ANGLE = math.pi/4

PIXEL_SIZE = (4,4)
WINDOW_SIZE = (500,500)
PIXEL_COUNTS = (WINDOW_SIZE[0]//PIXEL_SIZE[0],WINDOW_SIZE[1]//PIXEL_SIZE[1])

window = pyglet.window.Window(width=WINDOW_SIZE[0],height=WINDOW_SIZE[1],resizable=False)
batch = pyglet.graphics.Batch()
opacity_buffer = np.zeros(PIXEL_COUNTS[0]*PIXEL_COUNTS[1])
pixels = []
slimes = []

for y in range(PIXEL_COUNTS[1]):
    row = []
    for x in range(PIXEL_COUNTS[0]):
        row.append(pyglet.shapes.Rectangle(x*PIXEL_SIZE[0],y*PIXEL_SIZE[1],PIXEL_SIZE[0],PIXEL_SIZE[1],SLIME_COLOR,batch=batch))
        row[-1].opacity = opacity_buffer[y*PIXEL_COUNTS[0]+x]
    pixels.append(row)

for s in range(SLIMES):
    x = random.randint(0,WINDOW_SIZE[0]-1) if INITIAL_POSITION_RANDOM else WINDOW_SIZE[0]//2
    y = random.randint(0,WINDOW_SIZE[1]-1) if INITIAL_POSITION_RANDOM else WINDOW_SIZE[1]//2
    theta = random.uniform(0,2*math.pi) if INITIAL_DIRECTION_RANDOM else s/SLIMES*2*math.pi
    slimes.append((x,y,theta))

def assess_square(x,y,theta):
    dx,dy = math.cos(theta),math.sin(theta)
    window_x,window_y = x+dx*SENSOR_OFFSET,y+dy*SENSOR_OFFSET
    pixel_x,pixel_y = int(window_x//PIXEL_SIZE[0]),int(window_y//PIXEL_SIZE[1])

    pixel_bounded_x = 0 if pixel_x < 0 else PIXEL_COUNTS[0]-1 if pixel_x > PIXEL_COUNTS[0]-1 else pixel_x
    pixel_bounded_y = 0 if pixel_y < 0 else PIXEL_COUNTS[1]-1 if pixel_y > PIXEL_COUNTS[1]-1 else pixel_y
    return opacity_buffer[pixel_bounded_y*PIXEL_COUNTS[0]+pixel_bounded_x]

def follow_trail(slime):
    x,y,theta = slime[0],slime[1],slime[2]

    forward = assess_square(x,y,theta)
    left = assess_square(x,y,theta+SENSOR_ANGLE)
    right = assess_square(x,y,theta-SENSOR_ANGLE)

    if forward > left and forward > right:
        return 0
    elif forward < left and forward < right and (left != 0 and right != 0):
        return (random.randint(0,1)*2-1) * ROTATION_ANGLE
    elif left > right:
        return ROTATION_ANGLE
    elif right > left:
        return -ROTATION_ANGLE
    else:
        return 0

def update_slime(slime):
    x,y,theta = slime[0],slime[1],slime[2]
    dx,dy = SLIME_SPEED*math.cos(theta),SLIME_SPEED*math.sin(theta)
    reflect_x,reflect_y = False,False

    if x+dx > WINDOW_SIZE[0] and dx > 0:
        dx = -dx
        reflect_x = True
    elif x+dx < 0 and dx < 0:
        dx = -dx
        reflect_x = True
    if y+dy > WINDOW_SIZE[1] and dy > 0:
        dy = -dy
        reflect_y = True
    elif y+dy < 0 and dy < 0:
        dy = -dy
        reflect_y = True

    if reflect_x:
        if WALL_BOUNCE_RANDOM_ANGLE:
            theta = random.uniform(0,2*math.pi)
        else:
            theta = (3*math.pi - theta) % (2*math.pi)
    elif reflect_y:
        if WALL_BOUNCE_RANDOM_ANGLE:
            theta = random.uniform(0,2*math.pi)
        else:
            theta = math.pi*2 - theta
    else:
        if random.random() < RANDOM_DIRECTION_CHANGE_CHANCE:
            theta += random.uniform(-math.pi/4,math.pi/4)
        else:
            theta += follow_trail(slime)
        theta %= 2*math.pi

    x,y = x+dx,y+dy
    return (x,y,theta)

def update(_):
    slimes_pos = np.array(slimes).flatten()
    opacity_out = np.zeros(len(opacity_buffer))
    draw[GPU_BLOCKS_SHAPE,GPU_THREADS_SHAPE](opacity_buffer, slimes_pos, opacity_out)

    for i in range(len(slimes)):
        slimes[i] = update_slime(slimes[i])

    for x in range(PIXEL_COUNTS[0]):
        for y in range(PIXEL_COUNTS[1]):
            i = y*PIXEL_COUNTS[0]+x
            out = opacity_out[i]
            if opacity_buffer[i] != out:
                opacity_buffer[i] = out
                pixels[y][x].opacity = out
    window.clear()
    batch.draw()

@cuda.jit
def draw(opacity_in, slimes_pos, opacity_out):
    thread_id = cuda.grid(2)

    min_pixel_x = int(thread_id[0]*PIXEL_COUNTS[0]/GPU_THREADS_X)
    max_pixel_x = int((thread_id[0]+1)*PIXEL_COUNTS[0]/GPU_THREADS_X)
    min_pixel_y = int(thread_id[1]*PIXEL_COUNTS[1]/GPU_THREADS_Y)
    max_pixel_y = int((thread_id[1]+1)*PIXEL_COUNTS[1]/GPU_THREADS_Y)

    min_pos_x = min_pixel_x*PIXEL_SIZE[0]
    max_pos_x = max_pixel_x*PIXEL_SIZE[0]
    min_pos_y = min_pixel_y*PIXEL_SIZE[1]
    max_pos_y = max_pixel_y*PIXEL_SIZE[1]

    for y in range(min_pixel_y,max_pixel_y):
        for x in range(min_pixel_x,max_pixel_x):
            if DISPERSAL_ENABLED:
                opacity,count = 0,0
                for dis_x in range(x-1 if x > 0 else 0, x+2 if x < PIXEL_COUNTS[0]-1 else PIXEL_COUNTS[0]-1):
                    for dis_y in range(y-1 if y > 0 else 0, y+2 if y < PIXEL_COUNTS[1]-1 else PIXEL_COUNTS[1]-1):
                        opacity += opacity_in[dis_y*PIXEL_COUNTS[0]+dis_x]
                        count += 1
                opacity /= count
            else:
                opacity = opacity_in[y*PIXEL_COUNTS[0]+x]

            if opacity < 5:
                opacity_out[y*PIXEL_COUNTS[0]+x] = 0
            else:
                opacity = opacity/255
                opacity = (opacity ** DEGRADATION_STEEPNESS / (opacity ** DEGRADATION_STEEPNESS + (DEGRADATION_OFFSET * (1-opacity)) ** DEGRADATION_STEEPNESS))
                opacity_out[y*PIXEL_COUNTS[0]+x] = int(255 * 0.9 * opacity)

    for i in range(0,len(slimes_pos),3):
        x,y = slimes_pos[i],slimes_pos[i+1]
        if min_pos_x <= x < max_pos_x and min_pos_y <= y < max_pos_y:
            pixel_x,pixel_y = int(x/PIXEL_SIZE[0]),int(y/PIXEL_SIZE[1])
            opacity_out[pixel_y*PIXEL_COUNTS[0]+pixel_x] = 255

# @window.event
# def on_mouse_drag(x,y,dx,dy,buttons,modifiers):
#     pixel_x,pixel_y = int(x//PIXEL_SIZE[0]),int(y//PIXEL_SIZE[1])
#     opacity_buffer[pixel_y*PIXEL_COUNTS[0]+pixel_x] = 255

if __name__ == '__main__':
    pyglet.clock.schedule_interval(update, 1/SIMULATION_RATE)
    pyglet.app.run()


import statistics

import numpy as np
import cv2
from matplotlib import pyplot as plt
import math
import modeling

from extraction.FormatModel import UtilDebug

dx = [-1, 1, 0, 0]
dy = [0, 0, -1, 1]

def expandOnlyIntersections(BinaryOriginal, globalMask):
    V = np.zeros(BinaryOriginal.shape, dtype=bool)
    intersectionOfMask = cv2.bitwise_and(BinaryOriginal, globalMask)
    Q = [(1, 1)] * (BinaryOriginal.shape[0] * BinaryOriginal.shape[0])
    indxHead = 0
    indxTail = 0
    for i in range(BinaryOriginal.shape[0]):
        for j in range(BinaryOriginal.shape[1]):
            if intersectionOfMask[i, j] > 0:
                Q[indxTail] = (i, j)
                indxTail += 1
                V[i, j] = True
                # intersectionOfMask[i,j] = 0
    intersectionOfMask[intersectionOfMask >= 0] = 0  # OFF , 0 es ON
    while indxHead < indxTail:
        (i, j) = Q[indxHead]
        indxHead += 1
        intersectionOfMask[i, j] = 255
        for k in range(0, 4):
            ni = i + dx[k]
            nj = j + dy[k]
            if 0 <= ni < BinaryOriginal.shape[0] and 0 <= nj < BinaryOriginal.shape[1]:
                if BinaryOriginal[ni, nj] > 0 and not V[ni, nj]:
                    Q[indxTail] = (ni, nj)
                    indxTail += 1
                    V[ni, nj] = True
    return intersectionOfMask

def getPointProportion(A, B, a, b):
    px = (A[0] * b + B[0] * a) / (a + b)
    py = (A[1] * b + B[1] * a) / (a + b)
    return (int(px), int(py))

def closestNonZero(img, p, maxSize=21):
    if img[p[0], p[1]] > 0:
        return p
    if img[p[0], p[1]] > 0:
        return p
    dx_step = [-1, 0, 1, 0]
    dy_step = [0, 1, 0, -1]
    currentK = 0
    copyP = p
    for k in range(1, maxSize):
        for twice in range(0, 2):
            for times in range(0, k):
                p = (p[0] + dx_step[currentK]), (p[1] + dy_step[currentK])
                if p[0] >= 0 and p[1] >= 0 and p[0] < img.shape[0] and p[1] < img.shape[1] and img[p[0], p[1]] > 0:
                    # print('found: ', p)
                    return p
            currentK = (currentK + 1) % 4
    return copyP

class MyClass:
    count_letters = 290

def filterSingleCharacter_new(letter_original_and_thersh):
    letter_original = letter_original_and_thersh[0]

    threshold_border = letter_original_and_thersh[1]
    #

    blur = cv2.GaussianBlur(letter_original, (1, 1), 0)
    ret5, to_extract_letters = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    # plt.subplot(1, 2, 1), plt.imshow(letter_original), plt.title('letter_original')
    # plt.subplot(1, 2, 2), plt.imshow(to_extract_letters), plt.title('to_extract_letters')
    # plt.show()

    #

    to_extract_letters = cv2.bitwise_not(to_extract_letters)

    If = cv2.GaussianBlur(letter_original, (3, 3), 0)
    If = cv2.adaptiveThreshold(If, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 5, 2)
    If = cv2.bitwise_not(If)

    top_left_L, bottom_right_L = getBestRectangle(letter_original)

    # letter_original_copy = letter_original.copy()
    # letter_original = letter_original[top_left_L[1]:bottom_right_L[1], top_left_L[0]:bottom_right_L[0]]

    delta_inside = 3
    centro = to_extract_letters[top_left_L[1] + delta_inside:bottom_right_L[1] - delta_inside,
             top_left_L[0] + delta_inside:bottom_right_L[0] - delta_inside]

    th_macha_grande = 25

    stats_centro = cv2.connectedComponentsWithStats(centro, connectivity=8)
    num_labels_centro = stats_centro[0]
    labelStats_centro = stats_centro[2]
    empty_square = True
    for label in range(num_labels_centro):
        if label == 0:  # background
            continue
        area = labelStats_centro[label, cv2.CC_STAT_AREA]
        if area > th_macha_grande:
            empty_square = False

    #
    # if MyClass.count_letters < 0:
    #     plt.subplot(1, 3, 1), plt.imshow(letter_original, 'gray'), plt.title('letter_original')
    #     plt.subplot(1, 3, 2), plt.imshow(centro, 'gray'), plt.title('centro')
    #     # plt.subplot(1, 3, 3), plt.imshow(letter_no_borders, 'gray'), plt.title('letter_no_borders')
    #     plt.show()
    # MyClass.count_letters -= 1

    if empty_square:
        imgResult = None
    else:

        gb = 3  # grosor del borde

        borde = np.zeros(letter_original.shape, np.uint8)
        borde[top_left_L[1] - gb // 2:top_left_L[1] + (gb - gb // 2), :] = 255
        borde[bottom_right_L[1] - (gb // 2):bottom_right_L[1] + (gb - (gb // 2)), :] = 255
        borde[:, top_left_L[0] - gb // 2:top_left_L[0] + (gb - gb // 2)] = 255
        borde[:, bottom_right_L[0] - gb // 2:bottom_right_L[0] + (gb - (gb // 2))] = 255

        borde[letter_original < threshold_border] = 0
        letter_no_borders = to_extract_letters.copy()
        letter_no_borders[borde > 125] = 0

        stats_sin_bordes = cv2.connectedComponentsWithStats(letter_no_borders, connectivity=8)
        num_labels_sin_bordes = stats_sin_bordes[0]
        labelStats_sin_bordes = stats_sin_bordes[2]
        centroides_sin_bordes = stats_sin_bordes[3]

        foreground = []
        background = []
        area_th = 5
        centro_vertical = (bottom_right_L[1] + top_left_L[1]) // 2

        for label in range(num_labels_sin_bordes):
            if label == 0:  # background
                continue
            area = labelStats_sin_bordes[label, cv2.CC_STAT_AREA]
            cy = centroides_sin_bordes[label][0]
            cx = centroides_sin_bordes[label][1]
            outsideOfY = cy <= top_left_L[0] + (gb - (gb // 2)) or cy >= bottom_right_L[0] - (gb - (gb // 2))
            outsideOfX = cx <= top_left_L[1] + (gb - (gb // 2)) or cx >= bottom_right_L[1] - (gb - (gb // 2))
            deleteGroup = False
            if outsideOfX or outsideOfY:
                deleteGroup = True

            Left = labelStats_sin_bordes[label, cv2.CC_STAT_LEFT]
            Top = labelStats_sin_bordes[label, cv2.CC_STAT_TOP]
            Right = Left + labelStats_sin_bordes[label, cv2.CC_STAT_WIDTH]
            Bottom = Top + labelStats_sin_bordes[label, cv2.CC_STAT_HEIGHT]

            if centro_vertical - 2 < cx < centro_vertical + 2:
                if area < 10 and (cy <= top_left_L[0] + gb // 2 or cy >= bottom_right_L[0] - gb // 2):
                    deleteGroup = True

            if deleteGroup:
                letter_no_borders[Top:Bottom, Left:Right] = 0
            else:
                if area > area_th:
                    foreground.append([(Left, Top), (Right, Bottom)])
                else:
                    background.append([(Left, Top), (Right, Bottom)])

        def isNearForegraound(tl_br):

            def pointIsNearForeGround(xy):
                x = xy[0]
                y = xy[1]
                dist_th = 5
                minDistance_global = dist_th + 1
                for sq in foreground:
                    sq_x1 = sq[0][0]
                    sq_x2 = sq[1][0]

                    sq_y1 = sq[0][1]
                    sq_y2 = sq[1][1]
                    if sq_x1 <= x <= sq_x2:
                        if sq_y1 <= y <= sq_y2:
                            return True
                        else:
                            minDist_local = min(abs(sq_y1 - y), abs(y - sq_y2))

                    else:
                        if sq_y1 <= y <= sq_y2:
                            minDist_local = min(abs(sq_x1 - x), abs(x - sq_x2))
                        else:
                            minDist_local_1 = min(abs(sq_x1 - x), abs(x - sq_x2))
                            minDist_local_2 = min(abs(sq_y1 - y), abs(y - sq_y2))
                            minDist_local = minDist_local_1 + minDist_local_2

                    if minDist_local <= dist_th:
                        return True

                return False

            x1 = tl_br[0][0]
            x2 = tl_br[1][0]

            y1 = tl_br[0][1]
            y2 = tl_br[1][1]

            a = pointIsNearForeGround((x1, y1))
            b = pointIsNearForeGround((x1, y2))
            c = pointIsNearForeGround((x2, y1))
            d = pointIsNearForeGround((x2, y2))

            return a or b or c or d

            # letter_original_no_borders[top_left_L[1] -gb//2:top_left_L[1] +gb//2, :]

        background = sorted(background, key=lambda p: p[1][0] - p[0][0] + p[1][1] - p[0][1])

        for bk in background:
            if isNearForegraound(bk):
                foreground.append(bk)
            else:
                Left = bk[0][0]
                Top = bk[0][1]
                Right = bk[1][0]
                Bottom = bk[1][1]
                letter_no_borders[Top:Bottom, Left:Right] = 0

        if len(foreground) == 0:
            imgResult = None
        else:

            img_copy = letter_original.copy()

            img_copy[letter_no_borders < 125] = 255  # lo que es negro en only match, ahora sera blanco en img_copy
            img_copy[letter_no_borders >= 125] = 0  # to_do lo que no es absolutamente blanco, pasa a ser negro

            if cv2.countNonZero(letter_no_borders) == 0:
                img_copy[letter_no_borders >= 0] = 255
            debugThisCharacter = False

            try:

                # imgResult = GenerateData.myImResize_20x20_32x32(img_copy)
                imgResult = modeling.GenerateTrainDataAZ.myImResize_forDataTraining(img_copy, None)

                # for i in range(0,32):
                #     for j in range(0, 32):
                #         if shouldFill(imgResult, (i,j)):
                #             imgResult[i, j] = 0
                #         if shouldClear(imgResult, (i, j)):
                #             imgResult[i, j] = 255

                imgResult = (imgResult -
                             255.0 / 2) / 255.0

            except Exception as e:
                print('error filtering: ', e)
                imgResult = None

            if debugThisCharacter:
                plt.subplot(1, 4, 1), plt.imshow(letter_original, 'gray'), plt.title('letter_original_copy')
                plt.subplot(1, 4, 2), plt.imshow(to_extract_letters, 'gray'), plt.title('letter_original')
                plt.subplot(1, 4, 3), plt.imshow(letter_no_borders, 'gray'), plt.title('letter_no_borders')

                if imgResult is not None:
                    plt.subplot(1, 4, 4), plt.imshow(imgResult, 'gray'), plt.title('imgResult resized to 32x32')
                plt.show()

    characterDebugger = UtilDebug.CharacterDebugger()
    characterDebugger.add((letter_original, imgResult))
    return imgResult

def filterSingleCharacter(letter_original_and_mask):
    # todavia puede tener los bordes
    letter_original = letter_original_and_mask[0]
    mask = letter_original_and_mask[1]
    img = letter_original.copy()
    # se va a crear una nueva mascara, pero mask, el parametro deberia ser util

    ret3, th3 = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    rows, cols = th3.shape

    # th3 tiene el character correcto, pero tambien tiene basura, solo se extendera los que hagan match
    #
    # SEh = cv2.getStructuringElement(cv2.MORPH_RECT, (30, 1))
    # SEv = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 30))
    # opHorizontal = cv2.morphologyEx(th3, cv2.MORPH_OPEN, SEh)
    # opVertical = cv2.morphologyEx(th3, cv2.MORPH_OPEN, SEv)

    # Ifp = cv2.bitwise_xor(th3, cv2.bitwise_or(opHorizontal, opVertical))

    # Dk = getDk()
    # Ifp_union = Ifp.copy()
    # Ifp_union[:,:]=0
    # for D in Dk:
    #     closeD = cv2.morphologyEx(Ifp, cv2.MORPH_CLOSE, D)
    #     Ifp_union = cv2.bitwise_or(Ifp_union, closeD)
    # L = cv2.bitwise_or(opHorizontal, opVertical)
    # for i in L.shape[0]:
    #    for j in L.shape[1]:
    #        if L[i,j]>0:
    #            Ifp[i,j] = Ifp_union[i,j]
    # Ifp[L>0]=Ifp_union[L>0]
    onlyMatch = expandOnlyIntersections(th3, mask)

    # newMask = th3.copy()
    # for i in range(0,rows):
    #    for j in range(0, cols):
    #        if shouldFill(th3, (i,j)):
    #            th3[i,j]=255
    #        if shouldClear(th3, (i, j)):
    #            th3[i, j] = 0

    # plt.subplot(1,2,1), plt.imshow(newMask,'gray'), plt.title('newMask (th original)')
    # plt.subplot(1, 2, 2), plt.imshow(th3,'gray'), plt.title('th3, modified')
    # plt.show()

    newMask = cv2.medianBlur(th3, 3)

    img_copy = img.copy()

    img_copy[onlyMatch < 125] = 255
    img_copy[img_copy < 255] = 0
    img_copy[img_copy >= 255] = 255

    if cv2.countNonZero(onlyMatch) == 0:
        img_copy[mask >= 0] = 255

    try:

        # imgResult = GenerateData.myImResize_20x20_32x32(img_copy)
        imgResult = modeling.GenerateTrainDataAZ.myImResize_20x20_32x32(img_copy)

        # for i in range(0,32):
        #     for j in range(0, 32):
        #         if shouldFill(imgResult, (i,j)):
        #             imgResult[i, j] = 0
        #         if shouldClear(imgResult, (i, j)):
        #             imgResult[i, j] = 255

        imgResult = (imgResult -
                     255.0 / 2) / 255.0
    except Exception as e:
        # print('error filtering: ', e)
        imgResult = None

    # plt.subplot(1, 5, 1), plt.imshow(img, 'gray'), plt.title('img original')
    # plt.subplot(1, 5, 2), plt.imshow(mask, 'gray'), plt.title('mask from input')
    # plt.subplot(1, 5, 3), plt.imshow(onlyMatch, 'gray'), plt.title('mask, mezclado')
    # plt.subplot(1, 5, 4), plt.imshow(img_copy, 'gray'), plt.title('to pass resize32x32')
    # if( imgResult is not None):
    #    plt.subplot(1, 5, 5), plt.imshow(imgResult, 'gray'), plt.title('imgResult resized to 32x32')
    # plt.show()

    return imgResult

def findApropiateTemplate(ratio):
    current_image = None
    bestRatio = 100
    for k in range(1, 5):
        img = cv2.imread('extraction/FormatModel/cuadro_template_' + str(k) + '.png', 0)
        current_ratio = img.shape[0] / img.shape[1]

        if abs(current_ratio - ratio) < abs(bestRatio - ratio):
            bestRatio = current_ratio
            current_image = img
            # print('best ratio: '+str(k))
    return current_image

def plotImagesWithPrediction(preditectArray, images):
    cols = len(images)
    for k in range(0, cols):
        if images[k] is not None:
            plt.subplot(1, cols, k + 1), plt.imshow(images[k], 'gray'), plt.title(preditectArray[k]), plt.axis('off')
    plt.show()

def findMaxElement(A):
    currentValue = -1.0
    currentI = 0
    currentJ = 0
    for i in range(A.shape[0]):
        for j in range(A.shape[1]):
            if A[i, j] > currentValue:
                currentValue = A[i, j]
                currentI = i
                currentJ = j
    return (currentI, currentJ)

def countNonZerosRows(sumRows, I, j1, j2):
    if j1 <= 0:
        return sumRows[I, j2]
    else:
        return sumRows[I, j2] - sumRows[I, (j1 - 1)]

def countNonZerosCols(sumCols, J, i1, i2):
    if i1 <= 0:
        return sumCols[i2, J]
    else:
        return sumCols[i2, J] - sumCols[(i1 - 1), J]

def countNonZeros(sumRows, sumCols, pi, pf):
    top = countNonZerosRows(sumRows, pi[0], pi[1], pf[1])
    bottom = countNonZerosRows(sumRows, pf[0], pi[1], pf[1])
    left = countNonZerosCols(sumCols, pi[1], pi[0] + 1, pf[0] - 1)
    right = countNonZerosCols(sumCols, pf[1], pi[0] + 1, pf[0] - 1)
    return top + bottom + left + right

def getFirstGroupLargerThan(a, L):
    b = np.zeros(a.shape)

    for i in range(1, len(b)):
        if a[i] > 0:
            b[i] = b[i - 1] + 1
        else:
            if b[i - 1] >= L:
                # print(b,' - ', i)
                return int(round(i - b[i - 1])), int(i)
    return -1, -1

def getBestRectangle_big(region, ratio_cols_over_rows):
    B = range(max(0, region.shape[1] - 30), region.shape[1])
    copia = region.copy()

    copia[copia >= 0] = 0
    rows, cols = region.shape
    print(region.shape)
    bestValue = -1.0
    bestA = 0
    bestB = 0
    bestPos = (-1, -1)

    sumCols = np.asarray(np.sum(region, 0) / 255.0)
    diffSumCols = np.ediff1d(sumCols)
    sortedCols = np.sort(sumCols)
    indexToCutCols = (cols - 2 * 22) // 2
    toCut = 2 * sortedCols[indexToCutCols] + 2
    sumCols[sumCols < toCut] = 0
    left, right = getFirstGroupLargerThan(sumCols, 15)
    if left + right < 0:
        return (0, 0), (20, 20)
    sumCols[:left] = 0
    sumCols[(right + 1):] = 0
    print(left, right)
    importanColum = region[:, left:right]
    sumRows = np.asarray(np.sum(importanColum, 1) / 255.0)
    diffSumRows = np.ediff1d(sumRows)
    sortedRows = np.sort(sumRows)
    indexToCutRows = (rows - 7 * 16) // 2
    valueToCut = 2 * sortedRows[indexToCutRows] + 2
    sumRows[sumRows < valueToCut] = 0

    for i in range(0, 7):
        left, right = getFirstGroupLargerThan(sumRows, 10)
        if right - left > 0:
            marca = sumRows[left:right]
            sortedMarca = np.sort(marca)
            print(i, ' : ', sortedMarca[len(sortedMarca) // 2])
            sumRows[left:right] = 0

    width = 1

    #
    #
    # max_indx_row, max_value_row = max(enumerate(sumRows), key=lambda p: p[1])
    # max_indx_col, max_value_col = max(enumerate(sumCols), key=lambda p: p[1])
    #
    # for b in B:
    #     minA = int(round(b/(ratio_cols_over_rows+0.1)))
    #     maxA = int(round(b/(ratio_cols_over_rows-0.1)))
    #     for a in range(minA,maxA):
    #         cum = np.zeros((rows, cols))
    #         for i in range(rows):
    #             if i + a >= rows:
    #                 break
    #             for j in range(cols):
    #                 if j + b >= cols:
    #                     break
    #                 #copia[copia >= 0] = 0
    #                 # pi = (j, i)
    #                 # pf = (j + b, i + a)
    #                 # cv2.rectangle(copia, pi, pf, 255, thickness=1)
    #                 # copia = cv2.bitwise_and(copia, region)
    #                 # cantMatch = cv2.countNonZero(copia)
    #                 # print('cant match: ', cantMatch)
    #                 myCantMatch = countNonZeros(sumRows, sumCols, (i, j), (i+a, j+b))
    #                 # print('cant mAtch: ', myCantMatch)
    #                 cum[i, j] = myCantMatch / (2 * (a + b))
    #                 # (I, J) = findMaxElement(cum)
    #                 # print(I,J)
    #                 # print('inicio', i, j)
    #                 # print('longitudes', a, b)
    #                 # print(cum)
    #                 # plt.subplot(1, 2, 1), plt.imshow(region, 'gray'), plt.title('region')
    #                 # plt.subplot(1, 2, 2), plt.imshow(copia, 'gray'), plt.title('copia con rect de 255')
    #                 # plt.show()
    #                 # cv2.rectangle(copia, pi, pf, 0, thickness=1)
    #
    #
    #         (I, J) = findMaxElement(cum)
    #         if cum[I, J] > bestValue:
    #             bestValue = cum[I, J]
    #             bestA = a
    #             bestB = b
    #             bestPos = (I, J)
    #
    # copia[copia >= 0] = 0
    # pi = (bestPos[1], bestPos[0])
    # pf = (bestPos[1] + bestB, bestPos[0] + bestA)
    # cv2.rectangle(copia, pi, pf, 255, thickness=1)
    # section = region[pi[1]:pf[1], pi[0]:pf[0]]
    # plt.subplot(1, 3, 1), plt.imshow(region, 'gray'), plt.title('region')
    # plt.subplot(1, 3, 2), plt.imshow(copia, 'gray'), plt.title('copia con rect de 255')
    # plt.subplot(1, 3, 3), plt.imshow(section, 'gray'), plt.title('best mark')
    plt.subplot(4, 2, 1), plt.imshow(region, 'gray')
    plt.subplot(4, 2, 2), plt.imshow(importanColum, 'gray')

    plt.subplot(4, 2, 3), plt.bar(range(len(sumRows)), sumRows, width, color="blue")
    plt.subplot(4, 2, 4), plt.bar(range(len(sumCols)), sumCols, width, color="blue")
    plt.subplot(4, 2, 5), plt.bar(range(len(diffSumRows)), diffSumRows, width, color="blue")
    plt.subplot(4, 2, 6), plt.bar(range(len(diffSumCols)), diffSumCols, width, color="blue")
    plt.subplot(4, 2, 7), plt.bar(range(len(sortedRows)), sortedRows, width, color="blue")
    plt.subplot(4, 2, 8), plt.bar(range(len(sortedCols)), sortedCols, width, color="blue")

    plt.show()
    return (0, 0), (20, 20)

# recibe como parametro una region en escala de grises
def getBestRectangle(region, default_th=0.5):
    # B = range(max(0, region.shape[1] - 20), region.shape[1])
    # copia = region.copy()
    #
    # copia[copia >= 0] = 0
    # rows, cols = region.shape
    #
    # bestValue = -1.0
    # bestA = 0
    # bestB = 0
    # bestPos = (-1, -1)
    #
    # sumRows = np.zeros((rows, cols))
    # sumCols = np.zeros((rows, cols))
    # for i in range(rows):
    #     for j in range(cols):
    #         if j == 0:
    #             sumRows[i, j] = (1 if region[i, j] > 0 else 0)
    #         else:
    #             sumRows[i, j] = (1 if region[i, j] > 0 else 0) + sumRows[i, j - 1]
    #
    #         if i == 0:
    #             sumCols[i, j] = (1 if region[i, j] > 0 else 0)
    #         else:
    #             sumCols[i, j] = (1 if region[i, j] > 0 else 0) + sumCols[i - 1, j]
    # # print(region)
    # # print(sumRows)
    # # print(sumCols)
    #
    #
    # for b in B:
    #     minA = int(round(b / (ratio_cols_over_rows + 0.1)))
    #     maxA = int(round(b / (ratio_cols_over_rows - 0.1)))
    #     for a in range(minA, maxA):
    #         cum = np.zeros((rows, cols))
    #         for i in range(rows):
    #             if i + a >= rows:
    #                 break
    #             for j in range(cols):
    #                 if j + b >= cols:
    #                     break
    #                 # copia[copia >= 0] = 0
    #                 # pi = (j, i)
    #                 # pf = (j + b, i + a)
    #                 # cv2.rectangle(copia, pi, pf, 255, thickness=1)
    #                 # copia = cv2.bitwise_and(copia, region)
    #                 # cantMatch = cv2.countNonZero(copia)
    #                 # print('cant match: ', cantMatch)
    #                 myCantMatch = countNonZeros(sumRows, sumCols, (i, j), (i + a, j + b))
    #                 # print('cant mAtch: ', myCantMatch)
    #                 cum[i, j] = myCantMatch / (2 * (a + b))
    #                 # (I, J) = findMaxElement(cum)
    #                 # print(I,J)
    #                 # print('inicio', i, j)
    #                 # print('longitudes', a, b)
    #                 # print(cum)
    #                 # plt.subplot(1, 2, 1), plt.imshow(region, 'gray'), plt.title('region')
    #                 # plt.subplot(1, 2, 2), plt.imshow(copia, 'gray'), plt.title('copia con rect de 255')
    #                 # plt.show()
    #                 # cv2.rectangle(copia, pi, pf, 0, thickness=1)
    #
    #         (I, J) = findMaxElement(cum)
    #         if cum[I, J] > bestValue:
    #             bestValue = cum[I, J]
    #             bestA = a
    #             bestB = b
    #             bestPos = (I, J)
    #
    # copia[copia >= 0] = 0
    # pi = (bestPos[1], bestPos[0])
    # pf = (bestPos[1] + bestB, bestPos[0] + bestA)
    # # cv2.rectangle(copia, pi, pf, 255, thickness=1)
    # # section = region[pi[1]:pf[1], pi[0]:pf[0]]
    # # plt.subplot(1, 3, 1), plt.imshow(region, 'gray'), plt.title('region')
    # # plt.subplot(1, 3, 2), plt.imshow(copia, 'gray'), plt.title('copia con rect de 255')
    # # plt.subplot(1, 3, 3), plt.imshow(section, 'gray'), plt.title('best mark')
    # # plt.show()
    # ratioBuffer = UtilDebug.RatiosBuffer()
    # ratioBuffer.append((bestA, bestB))
    # retFirstAlgorithm = pi, pf

    img = cv2.medianBlur(region, 1)
    to_find_border = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 5, 1)

    # plt.subplot(1, 3, 1), plt.imshow(region,'gray'), plt.title('region')
    # plt.subplot(1, 3, 2), plt.imshow(img,'gray'), plt.title('img')
    # plt.subplot(1, 3, 3), plt.imshow(to_find_border,'gray'), plt.title('to_find_border')
    # plt.show()

    #
    to_find_border = cv2.bitwise_not(to_find_border)

    rows, cols = region.shape

    bestValue = -1.0
    bestA = 0
    bestB = 0
    bestPos = (-1, -1)

    acumSumRows = np.cumsum(to_find_border, 1)
    acumSumCols = np.cumsum(to_find_border, 0)

    totSumRows = np.sum(to_find_border, 1)
    totSumCols = np.sum(to_find_border, 0)
    delta = 1

    cb = CuadroBuffer()

    if not cb.calculated:
        threshold_minPercent = default_th
    else:
        threshold_minPercent = 0.3

    maxCols = filter_and_getMaxElements(totSumCols, 0, 100, minPercent=threshold_minPercent)
    maxRows = filter_and_getMaxElements(totSumRows, 0, 100, minPercent=threshold_minPercent)
    # print('SS:', len(maxCols), len(maxRows))
    bestValue = 0
    bestA = None
    bestB = None

    minB = region.shape[1] - 20

    if not cb.calculated:
        for p in range(len(maxRows) - 1):
            for q in range(p + 1, len(maxRows)):
                for m in range(len(maxCols) - 1):
                    for n in range(m + 1, len(maxCols)):
                        i = min(maxRows[p], maxRows[q])
                        a = max(maxRows[p], maxRows[q]) - i
                        j = min(maxCols[m], maxCols[n])
                        b = max(maxCols[m], maxCols[n]) - j

                        if 0.75 < b / a < 0.95 and b >= minB:

                            myCantMatch = countNonZeros(acumSumRows, acumSumCols, (i, j), (i + a, j + b))
                            # print('cant mAtch: ', myCantMatch)
                            cum = myCantMatch / (2 * (a + b))
                            if cum > bestValue:
                                bestValue = cum
                                bestA = a
                                bestB = b
                                bestPos = (i, j)
    else:
        maxCols = sorted(maxCols)
        maxRows = sorted(maxRows)

        bestA_calc = cb.A_predicted
        bestB_calc = cb.B_predicted

        exists_cols = [False] * (max(maxCols) + 1)
        for k in maxCols:
            exists_cols[k] = True

        exists_rows = [False] * (max(maxRows) + 1)
        for k in maxRows:
            exists_rows[k] = True

        for i in maxRows:
            for a in (range(bestA_calc - 4, bestA_calc + 4)):
                end_row = i + a
                if end_row >= 0 and end_row < len(exists_rows) and exists_rows[end_row]:

                    for j in maxCols:
                        for b in (range(bestB_calc - 5, bestB_calc + 5)):
                            end_col = j + b
                            if end_col >= 0 and end_col < len(exists_cols) and exists_cols[end_col]:

                                myCantMatch = countNonZeros(acumSumRows, acumSumCols, (i, j), (i + a, j + b))
                                # print('cant mAtch: ', myCantMatch)
                                cum = myCantMatch / (2 * (a + b))
                                if cum > bestValue:
                                    bestValue = cum
                                    bestA = a
                                    bestB = b
                                    bestPos = (i, j)

                                    # print(maxCols)
                                    # print(exists_cols)
                                    # print(maxRows)
                                    # print(exists_rows)
                                    # print(bestA,bestB,bestPos)
                                    # plt.imshow(region)
                                    # plt.show()

    pi = (bestPos[1], bestPos[0])

    if pi[0] == -1:
        return getBestRectangle(region, default_th / 2.0)

        # print('SS:', len(maxCols), len(maxRows), ' calculated: ', cb.calculated)
        #
        # print(maxCols)
        # print(maxRows)
        #
        # plt.imshow(region)
        # plt.show()

    pf = (bestPos[1] + bestB, bestPos[0] + bestA)
    retSecondtAlgorithm = (pi, pf)
    # print(retFirstAlgorithm, retSecondtAlgorithm)

    if not cb.calculated:
        cb.add((bestA, bestB))
    return retSecondtAlgorithm

def filter_and_getMaxElements(A, delta, cantMax, minPercent=0.5):
    filtered = np.zeros(len(A))
    for i in range(delta, len(A) - delta):
        k = max(A[i - delta:i + delta + 1])
        if k == A[i]:
            filtered[i] = k

    b = enumerate(filtered)
    c = sorted(b, key=lambda p: -p[1])

    ret = []
    count = 0
    for k in c:
        if count < cantMax and k[1] > minPercent * c[0][1]:
            ret.append(k[0])
            count += 1
    if 0 not in ret:
        ret.append(0)
    if len(A) - 1 not in ret:
        ret.append(len(A) - 1)

    # plt.bar(range(len(filtered)), filtered, 1)
    # print(ret)
    # plt.show()
    return ret

def predictCategoric_column_labels_inside(column, labels):
    sumRows = np.asarray(np.sum(column, 1) // 255)
    resp = extractLabelsBySquares(column, sumRows, labels)
    return resp

def predictCategoric_column_labels_sex(column, labels):
    # plt.imshow(column,'gray')
    # plt.show()
    sumRows = np.asarray(np.sum(column, 1) // 255)
    resp = extractLabelsBySquaresSex(column, sumRows, labels)
    return resp

def predictCategoric_column_labels_documento(column, labels):
    sumRows = np.asarray(np.sum(column, 1) // 255)
    resp = extractLabelsBySquaresDocument(column, sumRows, labels)
    return resp

def predictCategoric_column_labels_SingleButton(column, labels):
    sumRows = np.asarray(np.sum(column, 1) // 255)
    if len(labels) != 1:
        raise Exception("expected len(labels) = 1 on: " + str(labels))
    width = column.shape[1]
    button_height = column.shape[0]
    i = column.shape[0] // 2
    rth = getRatioOn_number(i, img=None, width=width, buttonHeight=button_height, sumRows=sumRows)
    if rth > 0.5:
        resp = labels[0]
    else:
        resp = '?'
    # print(resp)
    # plt.subplot(1,2,1),plt.bar(range(len(sumRows)),sumRows, 1)
    # plt.subplot(1, 2, 2), plt.imshow(column)
    # plt.show()
    return resp

def predictCategoric_column_labels_left(column, labels):
    cantLabels = len(labels)
    onlyGlobles = column[:, :]
    sumRowsGlobes = np.asarray(np.sum(onlyGlobles, 1) // 255)
    results = ''
    bHeight = 7
    if cantLabels > 1:
        space_betweenButtons = (column.shape[0] - cantLabels * bHeight) // (cantLabels - 1)
    else:
        space_betweenButtons = 0
    for k in range(cantLabels):
        j = 4 + k * (bHeight + space_betweenButtons)
        if isOn(j, width=column.shape[1], buttonHeight=6, sumRows=sumRowsGlobes):
            if len(results) > 0:
                results = results + ';' + labels[k]
            else:
                results = labels[k]

    if results is None or len(results) == 0:
        results = '?'

    # print(max_i_2)
    # print(labels)
    # print(results)
    #
    # plt.subplot(1, 3, 1), plt.imshow(column)
    # plt.subplot(1, 3, 2), plt.imshow(onlyGlobles,'gray')
    # plt.subplot(1, 3, 3), plt.bar(range(len(sumRowsGlobes)), sumRowsGlobes, 1)
    # plt.show()
    # print(results)
    return results
    # return results

def predictValuesCategory(arrayOfImages, array_labels, func):
    if len(arrayOfImages) != len(array_labels):
        print('error: sz arrayOfImages != sz arrayLabels')
        return ['None'] * len(array_labels)

    result = []
    for indx, img in enumerate(arrayOfImages):
        if img is not None:
            labels = array_labels[indx]
            predictions = func(img, labels)
            if predictions is not None and result is not None:
                result.append(predictions)
            else:
                result = None
        else:
            print('wtf is none')

    if result is None:
        result = ['Unknow'] * len(arrayOfImages)
    print('Returning:', result)
    return result

def predictValuesCategory_labelsLeft(arrayOfImages, array_labels):
    return predictValuesCategory(arrayOfImages, array_labels, func=predictCategoric_column_labels_left)

def predictValuesCategory_labelsInside(arrayOfImages, array_labels):
    return predictValuesCategory(arrayOfImages, array_labels, func=predictCategoric_column_labels_inside)

def predictValuesCategory_labelsSex(arrayOfImages, array_labels):
    return predictValuesCategory(arrayOfImages, array_labels, func=predictCategoric_column_labels_sex)

def predictValuesCategory_labelsDocumento(arrayOfImages, array_labels):
    return predictValuesCategory(arrayOfImages, array_labels, func=predictCategoric_column_labels_documento)

def predictValuesCategory_labelsSingleButtons(arrayOfImages, array_labels):
    return predictValuesCategory(arrayOfImages, array_labels, func=predictCategoric_column_labels_SingleButton)

def countBlocks(a, sz):
    count = 0
    while True:
        left, right = getFirstGroupLargerThan(a, sz)
        if left + right < 0:
            return count
        count += 1
        print(left, right)
        a[left:right] = 0
        if left == right:
            a[left] = 0

def calcMeans(a, limit, iterations=1):
    n = len(a)
    medias = np.zeros(n)
    count_items = 0
    acum = []

    for i in range(n):
        acum.append(a[i])
        count_items += 1
        if count_items > limit:
            acum = acum[1:]
            count_items -= 1
        if count_items == limit:
            medias[i - limit // 2] = statistics.median(acum)
    if iterations > 1:
        return calcMeans(medias, limit, iterations - 1)
    return medias

def dropMinsTo0(a, limit):
    n = len(a)
    medias = np.zeros(n, dtype=np.int32)
    count_items = 0
    acum = []

    for i in range(n):
        acum.append(a[i])
        count_items += 1
        if count_items > limit:
            acum = acum[1:]
            count_items -= 1
        if count_items == limit:
            j = i - limit // 2
            val = a[j]
            if abs(val - min(acum)) < 2 and val * 3 < max(acum[0:2]) and val * 3 < max(acum[-2:]):
                medias[j] = 0
            else:
                medias[j] = val

    return medias

def isOn(row_i, img=None, width=None, buttonHeight=None, sumRows=None):
    if img is None:
        if width is None or sumRows is None or buttonHeight is None:
            raise Exception('img is None but widht or sumRows or buttonHeight are not defined')

        r = sum(sumRows[row_i - buttonHeight // 2:row_i + buttonHeight // 2])
        print(UtilDebug.bcolors.OKBLUE + "Recieving data w: " + str(width) + " buttonHeight:" + str(buttonHeight) +
              " r: " + str(r) + " " + str(r * 100.0 / (buttonHeight * width)) + "% " + UtilDebug.bcolors.ENDC)
        return 0.25 * buttonHeight * width < r
    else:
        width = img.shape[1]
        sumRows = np.asarray(np.sum(img, 1) / 255.0)
        print(UtilDebug.bcolors.OKBLUE + "Recieving data img.shape " + str(img.shape) + UtilDebug.bcolors.ENDC)
        return isOn(row_i, width=width, sumRows=sumRows, img=None, buttonHeight=buttonHeight)

def getRatioOn_number(row_i, img=None, width=None, buttonHeight=None, sumRows=None):
    if img is None:
        if width is None or sumRows is None or buttonHeight is None:
            raise Exception('img is None but widht or sumRows or buttonHeight are not defined')

        r = sum(sumRows[row_i - buttonHeight // 2:row_i + buttonHeight // 2])
        print(UtilDebug.bcolors.OKBLUE + "Recieving data w: " + str(width) + " buttonHeight:" + str(buttonHeight) +
              " r: " + str(r) + " " + str(r * 100.0 / (buttonHeight * width)) + "% " + UtilDebug.bcolors.ENDC)

        return  r/(buttonHeight * width)
    else:
        width = img.shape[1]
        sumRows = np.asarray(np.sum(img, 1) / 255.0)
        print(UtilDebug.bcolors.OKBLUE + "Recieving data img.shape " + str(img.shape) + UtilDebug.bcolors.ENDC)
        return getRatioOn_number(row_i, width=width, sumRows=sumRows, img=None, buttonHeight=buttonHeight)

def extractLabelsBySquares(column, sumRows, labels):
    cantRows = len(labels)

    sumRows = sumRows.copy()
    originalRows = sumRows.copy()
    i = 8

    results = ''
    widthColumn = column.shape[1]

    ratios = []
    for k in range(cantRows):
        j = i + int(round(19.6 * k))
        percent = 0.65
        if k >= 9 or cantRows == 2:
            if cantRows == 2 and k == 1:
                percent = 0.9
            else:
                percent = 0.8
        ratios.append(getRatioOn_number(j, img=None, width=widthColumn, sumRows=originalRows, buttonHeight=8))
        #
        # if isOn_number(j, img=None, width=widthColumn, sumRows=originalRows, buttonHeight=8, ratioTh=percent):
        #     if len(results) > 0:
        #         results = results + ';' + labels[k]
        #     else:
        #         results = labels[k]
    max_ratio = max(ratios)
    min_ratio = min(ratios)
    th = (max_ratio + min_ratio) / 2.0
    if max_ratio - min_ratio < 0.15:
        th = 0.8
    if cantRows == 1:
        th = 0.6

    # print(max_ratio, min_ratio, th, 'max min th')
    for k in range(cantRows):
        if ratios[k] > th:
            if len(results) > 0:
                results = results + ';' + labels[k]
            else:
                results = labels[k]

    # print(labels)
    # print(results)
    # plt.subplot(1, 2, 1), plt.imshow(column, 'gray')
    # plt.subplot(1, 2, 2), plt.bar(range(len(originalRows)), originalRows, 1)
    # plt.show()
    if len(results) == 0:
        results = '?'
    return results

def extractLabelsBySquaresDocument(column, sumRows, labels):
    cantRows = len(labels)

    sumRows = sumRows.copy()
    originalRows = sumRows.copy()

    results = ''
    widthColumn = column.shape[1]

    # plt.subplot(1,2,1), plt.imshow(column,'gray')
    # plt.subplot(1,2,2), plt.bar(range(len(originalRows)),originalRows, 1)
    # plt.show()
    ratios = []

    for k, j in enumerate([8, 23]):
        ratios.append(getRatioOn_number(j, img=None, width=widthColumn, sumRows=originalRows, buttonHeight=6))
        #
        # if isOn_number(j, img=None, width=widthColumn, sumRows=originalRows, buttonHeight=6, ratioTh=0.5):
        #     if len(results) > 0:
        #         results = results + ';' + labels[k]
        #     else:
        #         results = labels[k]
    max_ratio = max(ratios)
    min_ratio = min(ratios)
    th = (max_ratio + min_ratio) / 2.0
    if max_ratio - min_ratio < 0.15:
        th = 0.9
    if cantRows == 1:
        th = 0.6
    for k in range(cantRows):
        if ratios[k] > th:
            if len(results) > 0:
                results = results + ';' + labels[k]
            else:
                results = labels[k]

    if len(results) == 0:
        results = '?'

    # print(labels)
    # print(results)

    return results

def extractLabelsBySquaresSex(column, sumRows, labels):
    cantRows = len(labels)

    sumRows = sumRows.copy()
    originalRows = sumRows.copy()

    results = ''
    widthColumn = column.shape[1]

    ratios = []
    for k, addJ in enumerate([10, 30, 108, 128]):
        j = addJ
        percent = 0.7
        if k >= 2:
            percent = 0.85
        print('debugging sex')
        ratios.append(getRatioOn_number(j, img=None, width=widthColumn, sumRows=originalRows, buttonHeight=8))
        # if isOn_number(j, img=None, width=widthColumn, sumRows=originalRows, buttonHeight=8, ratioTh=percent):
        #     if len(results) > 0:
        #         results = results + ';' + labels[k]
        #     else:
        #         results = labels[k]

    max_ratio = max(ratios)
    min_ratio = min(ratios)
    th = (max_ratio + min_ratio) / 2.0
    if max_ratio - min_ratio < 0.2:
        th = 0.9
    if cantRows == 1:
        th = 0.6
    for k in range(cantRows):
        if ratios[k] > th:
            if len(results) > 0:
                results = results + ';' + labels[k]
            else:
                results = labels[k]

    if len(results) == 0:
        results = '?'
    return results
    # print(labels)
    # print(results)
    # plt.subplot(1,2,1), plt.imshow(column,'gray')
    # plt.subplot(1,2,2), plt.bar(range(len(originalRows)),originalRows, 1)
    # plt.show()

def extractColumnsBySquares(If, cantColumns):
    If_original = If.copy()
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2))
    If = cv2.dilate(If, kernel)

    sumRows = np.asarray(np.sum(If, 1) // 255)
    maxRows = filter_and_getMaxElements(sumRows, 0, 100, minPercent=0.8)
    if len(maxRows) > 2:
        maxRows = maxRows[:-2]
    maxDistance_rows = 0
    for i in maxRows:
        for j in maxRows:
            if abs(i - j) > maxDistance_rows:
                maxDistance_rows = abs(i - j)

    baseDistanceRows = maxDistance_rows / 2
    row_ini = 0
    row_fin = len(sumRows) - 1
    for i in maxRows:
        for j in maxRows:
            if abs(i - j) > baseDistanceRows:
                if abs(i - j) < row_fin - row_ini:
                    row_ini = min(i, j)
                    row_fin = max(i, j)

    # GG
    row_ini += 9
    sumCols = np.asarray(np.sum(If, 0) // 255)
    sumCols = sumCols.copy()

    sumCols = sumCols - int(min(sumCols))
    sumCols = dropMinsTo0(sumCols, 3)
    left = 0
    right = max(sumCols)
    if cantColumns == 2:
        delta = 60
    else:
        delta = 45

    sumCols[sumCols < 125] = 0
    sc = sumCols.copy()
    blocks = countBlocks(sumCols, 1)
    arrayResult = []
    if blocks == 2:

        sumCols = sc.copy()
        returnNone = False

        center = []
        for k in range(0, 2):
            left, right = getFirstGroupLargerThan(sumCols, 1)

            if left + right < 0:
                print('this shouldn be happening')
                returnNone = True
                continue
            sumCols[left:right] = 0
            if left == right:
                sumCols[left] = 0

            center.append((left + right) // 2)

            # importanColum = If[:, left:right]
            # arrayResult.append(importanColum)

        if returnNone or len(center) < 2:
            return [None] * cantColumns

        # print('maxRows:', maxRows)
        # print('row ini:', row_ini)
        # print('row fin:', row_fin)
        if cantColumns == 2:
            i = getPointProportion((center[0], 0), (center[1], 0), 20, 47)[0]
            j = getPointProportion((center[0], 0), (center[1], 0), 47, 20)[0]
            importanColum_I = If_original[row_ini:row_fin, (i - 11):(i + 11)]
            importanColum_J = If_original[row_ini:row_fin, (j - 11):(j + 11)]
            arrayResult.append(importanColum_I)
            arrayResult.append(importanColum_J)
            # plt.subplot(1, 3, 1), plt.imshow(If), plt.title('If')
            # plt.subplot(1, 3, 2), plt.imshow(importanColum_I), plt.title('FIrst Column')
            # plt.subplot(1, 3, 3), plt.imshow(importanColum_J), plt.title('Second Column')
            # plt.show()

        else:
            i = (center[0] + center[1]) // 2
            importanColum_I = If_original[row_ini:row_fin, (i - 11):(i + 11)]
            # plt.subplot(1, 3, 1), plt.imshow(If), plt.title('If')
            # plt.subplot(1, 3, 2), plt.imshow(importanColum_I), plt.title('Column')
            # plt.show()
            arrayResult.append(importanColum_I)

        return arrayResult

    else:
        return [None] * cantColumns

def extractSimpleButton(img):
    ret, If = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    If = cv2.bitwise_not(If)
    sumRows = np.asarray(np.sum(If, 1) // 255)
    sumCols = np.asarray(np.sum(If, 0) // 255)
    acumBy12 = sumRows.copy()
    for i in range(len(acumBy12)):
        if i > 0:
            acumBy12[i] += acumBy12[i - 1]
        if i >= 12:
            acumBy12[i] -= sumRows[i - 12]

    acumBy19 = sumCols.copy()
    for i in range(len(acumBy19)):
        if i > 0:
            acumBy19[i] += acumBy19[i - 1]
        if i >= 19:
            acumBy19[i] -= sumCols[i - 19]

    max_indx_12, max_value_12 = max(enumerate(acumBy12), key=lambda p: p[1])
    max_indx_19, max_value_19 = max(enumerate(acumBy19), key=lambda p: p[1])

    return If[max_indx_12 - 11:max_indx_12 + 1, max_indx_19 - 18:max_indx_19 + 1]

def extractCategory_extractColumnLabelsTipoSiNo(img, TL, BR, cantColumns):
    deltaAmpliacion = 5
    ROI_base = img[TL[1] - deltaAmpliacion:BR[1] + deltaAmpliacion, TL[0] - deltaAmpliacion:BR[0] + deltaAmpliacion]
    cols = ROI_base.shape[1]
    left = extractSimpleButton(ROI_base[:, :cols // 2])
    right = extractSimpleButton(ROI_base[:, cols // 2:])
    # plt.subplot(3, 1, 1), plt.imshow(ROI_base, 'gray'), plt.title('ROIS')
    # plt.subplot(3, 1, 2), plt.imshow(left, 'gray'), plt.title('left')
    # plt.subplot(3, 1, 3), plt.imshow(right, 'gray'), plt.title('right')

    # plt.show()
    arrayResult = []
    arrayResult.append(left)
    arrayResult.append(right)
    return arrayResult

def extractCategory_extractColumnLabelsDocumento(img, TL, BR, cantColumns):
    deltaAmpliacion = 5
    ROI_base = img[TL[1] - deltaAmpliacion:BR[1] + deltaAmpliacion, TL[0] - deltaAmpliacion:BR[0] + deltaAmpliacion]
    ROI = ROI_base.copy()
    #

    #GG_GG
    ret, If = cv2.threshold(ROI, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    If = cv2.bitwise_not(If)

    If_marks = cv2.adaptiveThreshold(ROI, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 21, 2)
    If_marks = cv2.bitwise_not(If_marks)
    # PART_NAC = If[26:33, 25:77]
    # cv2.imwrite('resources/PART_NAC.png', PART_NAC)

    PART_NAC = cv2.imread('resources/PART_NAC.png', 0)
    neg_PART_NAC = cv2.bitwise_not(PART_NAC)
    rows, cols = PART_NAC.shape
    ROWS, COLS = If.shape
    I_J = None
    val = 0
    for i in range(ROWS):
        for j in range(COLS):
            i_2 = i + rows
            j_2 = j + cols
            if i_2 >= ROWS or j_2 >= COLS:
                break
            pos_no_tiene = If[i:i_2, j:j_2]

            c1 = cv2.countNonZero(cv2.bitwise_and(PART_NAC, pos_no_tiene))
            c2 = cv2.countNonZero(cv2.bitwise_and(neg_PART_NAC, pos_no_tiene))
            v = c1 - c2
            if v > val:
                val = v
                I_J = (i, j)

    if I_J is None or I_J[1] <= 12 or I_J[1] > If.shape[1] - 72 or I_J[0] <= 11 or I_J[0] > If.shape[0] - 5:
        return [None] * 2
    arrayResult = []
    print('I_J', I_J)

    for k in [-12, 72]:
        left = I_J[1] + k - 10
        right = left + 17

        top = I_J[0] - 10 - 9
        bot = I_J[0] + 4 + 7
        # print('left:right', left,right)
        # print('top:bot', top, bot)
        globe = If_marks[top:bot, left:right]
        arrayResult.append(globe)

    # plt.subplot(2,2,1),plt.imshow(PART_NAC,'gray'), plt.title('PART_NAC')
    # plt.subplot(2, 2, 2), plt.imshow(If, 'gray'), plt.title('If')
    # plt.subplot(2, 2, 3), plt.imshow(arrayResult[0], 'gray'), plt.title('left')
    # plt.subplot(2, 2, 4), plt.imshow(arrayResult[1], 'gray'), plt.title('right')
    # plt.show()

    return arrayResult

def extractCategory_extractColumnLabelsTipoSuministro(img, TL, BR, cantColumns):
    deltaAmpliacion = 5
    ROI_base = img[TL[1] - deltaAmpliacion:BR[1] + deltaAmpliacion, TL[0] - deltaAmpliacion:BR[0] + deltaAmpliacion]
    ROI = ROI_base.copy()
    #

    ret, If = cv2.threshold(ROI, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    If = cv2.bitwise_not(If)
    # NO_TIENE = If[14:25, 222:277]
    # cv2.imwrite('NO_TIENE.png',NO_TIENE)

    sumRows = np.asarray(np.sum(If, 1) // 255)
    sumCols = np.asarray(np.sum(If, 0) // 255)
    NO_TIENE = cv2.imread('resources/NO_TIENE.png', 0)
    neg_NO_TIENE = cv2.bitwise_not(NO_TIENE)
    rows, cols = NO_TIENE.shape
    ROWS, COLS = If.shape
    I_J = None
    val = 0
    for i in range(ROWS):
        for j in range(COLS):
            i_2 = i + rows
            j_2 = j + cols
            if i_2 >= ROWS or j_2 >= COLS:
                break
            pos_no_tiene = If[i:i_2, j:j_2]

            c1 = cv2.countNonZero(cv2.bitwise_and(NO_TIENE, pos_no_tiene))
            c2 = cv2.countNonZero(cv2.bitwise_and(neg_NO_TIENE, pos_no_tiene))
            v = c1 - c2
            if v > val:
                val = v
                I_J = (i, j)

    if I_J is None or I_J[1] <= 199:
        return [None] * 3
    arrayResult = []
    d_j = 8
    d_i = 3
    for k in [-199, -105, -13]:
        j = I_J[1] + k
        i = I_J[0] + 5
        globe = If[(i - d_i):(i + d_i), (j - d_j):(j + d_j)]
        arrayResult.append(globe)

    # print(I_J)
    # plt.subplot(2, 2, 1), plt.imshow(ROI_base, 'gray'), plt.title('img')
    # plt.subplot(2, 2, 2), plt.imshow(If, 'gray'), plt.title('Solo columna ancha importante')
    # plt.subplot(2, 2, 3), plt.bar(range(len(sumRows)), sumRows, width=1), plt.title('sumRows')
    # plt.subplot(2, 2, 4), plt.bar(range(len(sumCols)), sumCols, width=1), plt.title('sumCols')
    # # plt.subplot(2, 2, 4), plt.imshow(right, 'gray'), plt.title('bot')
    # plt.show()

    return arrayResult

def extractCategory_extractColumnLabelsTipoVia(img, TL, BR, cantColumns):
    deltaAmpliacion = 5
    ROI_base = img[TL[1] - deltaAmpliacion:BR[1] + deltaAmpliacion, TL[0] - deltaAmpliacion:BR[0] + deltaAmpliacion]
    ROI = ROI_base.copy()
    #

    ret, If = cv2.threshold(ROI, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    If = cv2.bitwise_not(If)

    # CARRETERA = If[8:19, 523:589]
    # cv2.imwrite('CARRETERA.png',CARRETERA)
    #
    sumRows = np.asarray(np.sum(If, 1) // 255)
    sumCols = np.asarray(np.sum(If, 0) // 255)
    CARRETERA = cv2.imread('resources/CARRETERA.png', 0)
    neg_CARRETERA = cv2.bitwise_not(CARRETERA)
    rows, cols = CARRETERA.shape
    ROWS, COLS = If.shape
    I_J = None
    val = 0
    for i in range(ROWS):
        for j in range(COLS):
            i_2 = i + rows
            j_2 = j + cols
            if i_2 >= ROWS or j_2 >= COLS:
                break
            pos_no_tiene = If[i:i_2, j:j_2]

            c1 = cv2.countNonZero(cv2.bitwise_and(CARRETERA, pos_no_tiene))
            c2 = cv2.countNonZero(cv2.bitwise_and(neg_CARRETERA, pos_no_tiene))
            v = c1 - c2
            if v > val:
                val = v
                I_J = (i, j)

    if I_J is None or I_J[1] <= 500:
        return [None] * 6
    arrayResult = []
    d_j = 8
    d_i = 3
    for k in range(6):
        j = I_J[1] - 500 + 122 * k
        i = I_J[0] + 4
        globe = If[(i - d_i):(i + d_i), (j - d_j):(j + d_j)]
        arrayResult.append(globe)

    #
    # print(I_J)
    # plt.subplot(2, 2, 1), plt.imshow(ROI_base, 'gray'), plt.title('img')
    # plt.subplot(2, 2, 2), plt.imshow(If, 'gray'), plt.title('Solo columna ancha importante')
    # # plt.subplot(2, 2, 3), plt.bar(range(len(sumRows)), sumRows, width=1), plt.title('sumRows')
    # # plt.subplot(2, 2, 4), plt.bar(range(len(sumCols)), sumCols, width=1), plt.title('sumCols')
    # # plt.subplot(2, 2, 4), plt.imshow(right, 'gray'), plt.title('bot')
    # plt.show()
    return arrayResult

def extractCategory_extractColumnLabelsInside(img, TL, BR, cantColumns):
    deltaAmpliacion = 10
    ROI = img[TL[1] - deltaAmpliacion:BR[1] + deltaAmpliacion, TL[0] - deltaAmpliacion:BR[0] + deltaAmpliacion]
    # If = cv2.GaussianBlur(ROI, (3, 3), 0)

    # If = cv2.adaptiveThreshold(ROI, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 21, 2)

    #
    ret, If = cv2.threshold(ROI, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    If = cv2.bitwise_not(If)

    rows, cols = If.shape

    # plt.subplot(1, 2, 1), plt.imshow(ROI), plt.title('ROI')
    # plt.subplot(1, 2, 2), plt.imshow(If), plt.title('If')
    # plt.show()
    resp = extractColumnsBySquares(If, cantColumns)
    # GG2
    #
    # sumCols = np.asarray(np.sum(If, 0) // 255)
    # sumCols = dropMinsTo0(sumCols, 11)
    #
    # left = 0
    # right = max(sumCols)
    #
    # toCut = -1
    # for i in range(left, right):
    #     sc = sumCols.copy()
    #     sc[sc < i] = 0
    #     print('testing with ', i)
    #     numBlocks = countBlocks(sc, 20)
    #
    #     if numBlocks == cantColumns:
    #         toCut = i
    #         break
    # #
    # # media_21 = calcMeans(sumCols, 3,3)
    # # with0 = dropMinsTo0(sumCols, 11)
    # # sum_medias = np.zeros(len(media_21))
    # # for i in range(len(sum_medias)):
    # #     sum_medias[i] = with0[i]+media_21[i]
    # # plt.subplot(2, 2, 1), plt.imshow(If, 'gray')
    # # plt.subplot(2, 2, 2), plt.bar(range(len(sumCols)), sumCols, 1, color="blue")
    # # plt.subplot(2, 2, 3), plt.bar(range(len(media_21)), media_21, 1, color="blue")
    # # plt.subplot(2, 2, 4), plt.bar(range(len(with0)), with0, 1, color="blue")
    # # plt.show()
    # sumCols[sumCols < toCut] = 0
    #
    # arrayResult = []
    # arResultOk = True
    # for k in range(0, cantColumns):
    #     left, right = getFirstGroupLargerThan(sumCols, 15)
    #     if left + right < 0:
    #         arrayResult.append(None)
    #         arResultOk = False
    #         continue
    #
    #     sumCols[left:right] = 0
    #     importanColum = If[:, left:right]
    #     arrayResult.append(importanColum)
    # print('Returning images A')
    # toPrint = False
    # for I in arrayResult:
    #     if I is None:
    #         toPrint = False
    # # if arResultOk == False:
    # #     arrayResult = resp
    # if toPrint:
    #
    #     sumCols = np.asarray(np.sum(If, 0) // 255)
    #     plt.subplot(3, 2, 1), plt.imshow(If, 'gray')
    #     plt.subplot(3, 2, 2), plt.bar(range(len(sumCols)), sumCols, 1, color="blue")
    #     if resp is not None and len(resp) > 0 and resp[0] is not None:
    #         plt.subplot(3, 2, 3), plt.imshow(resp[0], 'gray')
    #         sumRows = np.asarray(np.sum(resp[0], 1) // 255)
    #         plt.subplot(3, 2, 4), plt.bar(range(len(sumRows)), sumRows, 1, color="blue")
    #     if resp is not None and len(resp) > 1 and resp[1] is not None:
    #         plt.subplot(3, 2, 5), plt.imshow(resp[1], 'gray')
    #         sumRows = np.asarray(np.sum(resp[1], 1) // 255)
    #         plt.subplot(3, 2, 6), plt.bar(range(len(sumRows)), sumRows, 1, color="blue")
    #     plt.show()
    return resp

def extractCategory_extractColumnLabelsLeft(img, TL, BR, cantColumns, cantRows):
    deltaAmpliacion = 10
    ROI = img[TL[1] - deltaAmpliacion:BR[1] + deltaAmpliacion, TL[0] - deltaAmpliacion - 5:BR[0] + deltaAmpliacion + 5]
    # If = cv2.GaussianBlur(ROI, (3, 3), 0)
    ret, If = cv2.threshold(ROI, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    If = cv2.bitwise_not(If)

    sumCols = np.asarray(np.sum(If, 0) // 255)
    cumSumCols = np.cumsum(sumCols)

    compCols = (sumCols * (-1)) + If.shape[0]
    max_j = 0
    max_value_j = 0
    for j in range(0, len(sumCols) - 44):
        val = cumSumCols[j + 6] - cumSumCols[j] + cumSumCols[j + 43] - cumSumCols[j + 37]

        if val > max_value_j:
            max_value_j = val
            max_j = j

    leftPart = If[:, 0:max_j + 5]
    rightPart = If[:, max_j + 37:]

    sumRows_Left = np.asarray(np.sum(leftPart, 1) // 255)
    sumRows_Right = np.asarray(np.sum(rightPart, 1) // 255)

    sumRows = np.zeros(len(sumRows_Left))
    for indx in range(0, len(sumRows_Left)):
        sumRows[indx] = sumRows_Left[indx] + sumRows_Right[indx]
    mean_rows = 4

    important_rows = sumRows[:]

    test_filter = important_rows < mean_rows

    important_rows[important_rows < mean_rows] = 0
    iniGroups = []
    endGroups = []
    for indx in range(0, len(important_rows)):
        k = important_rows[indx]
        if k > 0:
            if len(iniGroups) == 0 or indx - iniGroups[-1] > 13:
                iniGroups.append(indx)
                endGroups.append(indx)
            else:
                endGroups[-1] = indx

    # print('ini Groups:', iniGroups)
    # print('fin Groups:', endGroups)
    arrayResult = []

    if len(iniGroups) < cantRows:
        cumSumRows = np.cumsum(sumRows)

        If2 = If[:, max_j + 14:max_j + 32]

        arrayResult.append(If2)

    else:
        while len(iniGroups) > cantRows:
            indx_min_length = 322
            min_length = 322
            for indx in range(0, len(iniGroups)):
                if min_length > endGroups[indx] - iniGroups[indx]:
                    min_length = endGroups[indx] - iniGroups[indx]
                    indx_min_length = indx

            del iniGroups[indx_min_length]
            del endGroups[indx_min_length]

        If2 = If[iniGroups[0]:endGroups[-1] + 1, max_j + 14:max_j + 32]

        arrayResult.append(If2)

    # print(max_j)
    # plt.subplot(2, 2, 1), plt.imshow(If, 'gray')
    # plt.subplot(2, 2, 2), plt.bar(range(len(sumCols)), sumCols, width=1)
    # plt.subplot(2, 2, 3), plt.imshow(If2, 'gray')
    # plt.subplot(2, 2, 4), plt.bar(range(len(important_rows)), important_rows, width=1)
    # plt.show()
    return arrayResult

def extractCategory_extractColumnLabelsSex(img, TL, BR, cantColumns):
    return extractCategory_extractColumnLabelsInside(img, TL, BR, cantColumns)

def getPixels(img, P, Q):
    minX = min(P[0], Q[0]) - 2
    maxX = max(P[0], Q[0]) + 2
    minY = min(P[1], Q[1]) - 2
    maxY = max(P[1], Q[1]) + 2

    roi = img[minX:maxX, minY:maxY]
    # print('getting:',P,Q)
    # plt.imshow(img,'gray')
    # plt.show()
    dx = Q[0] - P[0]
    dy = Q[1] - P[1]
    D = math.sqrt(dx * dx + dy * dy)
    res = []
    for d in np.arange(0, D, step=1.0):
        a = d
        b = D - a
        q = getPointProportion(P, Q, a, b)
        res.append(img[q])
    return res

def predictCuadros(img, TL, BR, count):
    numRows = (BR[0] - TL[0]) / count
    numCols = BR[1] - TL[1]
    # print('finding ratio nr/nc : ' + str(numRows)+' / ' + str(numCols)+'  divided by '+ str(count))
    template = findApropiateTemplate(numRows / numCols)

    deltaAmpliacion = 5

    ROI = img[TL[1] - deltaAmpliacion:BR[1] + deltaAmpliacion, TL[0] - deltaAmpliacion:BR[0] + deltaAmpliacion]

    leftPart = ROI[:, 0:(template.shape[1] + (deltaAmpliacion * 2 - 1))]
    rightPart = ROI[:, -(template.shape[1] + (deltaAmpliacion * 2 - 1)):]

    top_left_L, bottom_right_L = getBestRectangle(leftPart)
    delta_L = (bottom_right_L[0] - top_left_L[0], bottom_right_L[1] - top_left_L[1])

    top_left_R, bottom_right_R = getBestRectangle(rightPart)
    delta_R = (bottom_right_R[0] - top_left_R[0], bottom_right_R[1] - top_left_R[1])

    top_left_R = (top_left_R[0] + ROI.shape[1] - (template.shape[1] + ((deltaAmpliacion * 2 - 1))), top_left_R[1])
    bottom_right_R = (top_left_R[0] + delta_R[0], top_left_R[1] + delta_R[1])

    pointA = (top_left_L[1], top_left_L[0])
    pointY = (bottom_right_R[1], bottom_right_R[0])

    pointB = (pointY[0], pointA[1])
    pointX = (pointA[0], pointY[1])

    res = []
    TL = pointA
    BR = pointY
    BL = (bottom_right_L[1], top_left_L[0])
    TR = (top_left_R[1], bottom_right_R[0])
    res.extend(getPixels(ROI, TL, TR))
    res.extend(getPixels(ROI, BL, BR))
    for k in range(0, count):
        s = getPointProportion(TL, TR, k, count - k)
        t = getPointProportion(BL, BR, k, count - k)
        res.extend(getPixels(ROI, s, t))
    res.extend(getPixels(ROI, TR, BR))

    for k in range(0, count):
        upperLeft = getPointProportion(pointA, pointX, k, count - k)
        bottomLeft = getPointProportion(pointB, pointY, k, count - k)
        upperRight = getPointProportion(pointA, pointX, k + 1, count - (k + 1))
        bottomRight = getPointProportion(pointB, pointY, k + 1, count - (k + 1))
        delta = 2
        minX = max(0, min(upperLeft[0], bottomLeft[0]) - delta)
        maxX = min(ROI.shape[0], max(upperRight[0], bottomRight[0]) + delta)

        minY = max(0, min(bottomLeft[1], bottomRight[1]) - delta)
        maxY = min(ROI.shape[1], max(upperLeft[1], upperRight[1]) + delta)

        getBestRectangle(ROI[minX:maxX, minY:maxY])

def extractCharacters(img, TL, BR, count):
    numRows = (BR[0] - TL[0]) / count
    numCols = BR[1] - TL[1]
    # print('finding ratio nr/nc : ' + str(numRows)+' / ' + str(numCols)+'  divided by '+ str(count))
    template = findApropiateTemplate(numRows / numCols)

    deltaAmpliacion = 5

    ROI = img[TL[1] - deltaAmpliacion:BR[1] + deltaAmpliacion, TL[0] - deltaAmpliacion:BR[0] + deltaAmpliacion]

    #
    # plt.imshow(ROI,'gray')
    # plt.show()

    # TODO revisar ampliacion de imagen si es necesario

    leftPart = ROI[:, 0:(template.shape[1] + (deltaAmpliacion * 2 - 1))]
    rightPart = ROI[:, -(template.shape[1] + (deltaAmpliacion * 2 - 1)):]

    # plt.subplot(2, 2, 1), plt.imshow(ROI, 'gray')
    # plt.subplot(2, 2, 3), plt.imshow(leftPart, 'gray')
    # plt.subplot(2, 2, 4), plt.imshow(rightPart, 'gray')
    # plt.show()

    top_left_L, bottom_right_L = getBestRectangle(leftPart)
    delta_L = (bottom_right_L[0] - top_left_L[0], bottom_right_L[1] - top_left_L[1])

    top_left_R, bottom_right_R = getBestRectangle(rightPart)
    delta_R = (bottom_right_R[0] - top_left_R[0], bottom_right_R[1] - top_left_R[1])

    bestLeft = leftPart[top_left_L[1]:bottom_right_L[1], top_left_L[0]:bottom_right_L[0]]
    bestRight = rightPart[top_left_R[1]:bottom_right_R[1], top_left_R[0]:bottom_right_R[0]]

    # print('If shape: ', If.shape)
    # print('template shape: ', template.shape)
    # print('current top_left_R', top_left_R)
    top_left_R = (top_left_R[0] + ROI.shape[1] - (template.shape[1] + ((deltaAmpliacion * 2 - 1))), top_left_R[1])
    bottom_right_R = (top_left_R[0] + delta_R[0], top_left_R[1] + delta_R[1])
    # print('after top_left_R', top_left_R)
    possibleBestLeft = ROI[top_left_L[1]:bottom_right_L[1], top_left_L[0]:bottom_right_L[0]]
    possibleBestRight = ROI[top_left_R[1]:bottom_right_R[1], top_left_R[0]:bottom_right_R[0]]

    # plt.subplot(1,8,1), plt.imshow(If), plt.title('If')
    # plt.subplot(1, 8, 2), plt.imshow(template), plt.title('template')
    # plt.subplot(1, 8, 3), plt.imshow(leftPart), plt.title('leftPart')
    # plt.subplot(1, 8, 4), plt.imshow(rightPart), plt.title('rightPart')
    # plt.subplot(1, 8, 5), plt.imshow(bestLeft), plt.title('bestLeft')
    # plt.subplot(1, 8, 6), plt.imshow(possibleBestLeft), plt.title('possibleBestLeft')
    # plt.subplot(1, 8, 7), plt.imshow(bestRight), plt.title('bestRight')
    # plt.subplot(1, 8, 8), plt.imshow(possibleBestRight), plt.title('possbielBestRight')
    # plt.show()

    pointA = (top_left_L[1], top_left_L[0])
    pointY = (bottom_right_R[1], bottom_right_R[0])

    pointB = (pointY[0], pointA[1])
    pointX = (pointA[0], pointY[1])

    res = []
    TL = pointA
    BR = pointY
    BL = (bottom_right_L[1], top_left_L[0])
    TR = (top_left_R[1], bottom_right_R[0])
    res.extend(getPixels(ROI, TL, TR))
    res.extend(getPixels(ROI, BL, BR))
    for k in range(0, count):
        s = getPointProportion(TL, TR, k, count - k)
        t = getPointProportion(BL, BR, k, count - k)
        res.extend(getPixels(ROI, s, t))
    res.extend(getPixels(ROI, TR, BR))
    np_mean = int(np.mean(res))
    np_stdv = int(np.std(res))
    # print(pointA,pointB,pointX,pointY,ROI.shape)
    ROI_2 = ROI[pointA[0]:pointY[0], pointA[1]:pointY[1]]

    # plt.subplot(2,1,1), plt.imshow(ROI)
    # plt.subplot(2, 1, 2), plt.imshow(ROI_2)
    # plt.show()

    letters = []
    for k in range(0, count):
        upperLeft = getPointProportion(pointA, pointX, k, count - k)
        bottomLeft = getPointProportion(pointB, pointY, k, count - k)
        upperRight = getPointProportion(pointA, pointX, k + 1, count - (k + 1))
        bottomRight = getPointProportion(pointB, pointY, k + 1, count - (k + 1))
        delta = 2
        minX = max(0, min(upperLeft[0], bottomLeft[0]) - delta)
        maxX = min(ROI.shape[0], max(upperRight[0], bottomRight[0]) + delta)

        minY = max(0, min(bottomLeft[1], bottomRight[1]) - delta*2)
        maxY = min(ROI.shape[1], max(upperLeft[1], upperRight[1]) + delta*2)

        singleCharacter = (ROI[minX:maxX, minY:maxY], (np_mean - 2 * np_stdv))
        letters.append(singleCharacter)

    filteredLetters = []

    for letter in letters:
        singleLetterFiltered = filterSingleCharacter_new(letter)
        filteredLetters.append(singleLetterFiltered)

        # if singleLetterFiltered != None:
        #    plt.imshow(singleLetterFiltered)
        #    plt.show()

    return filteredLetters

def extractCharacters_old(img, onlyUserMarks, TL, BR, count):
    letters = []

    ROI = img[TL[1] - 3:BR[1] + 3, TL[0] - 3:BR[0] + 3]
    ROI_onlyUserMarks = onlyUserMarks[TL[1] - 3:BR[1] + 3, TL[0] - 3:BR[0] + 3]
    # ROI = cv2.medianBlur(ROI, 3)
    If = cv2.GaussianBlur(ROI, (3, 3), 0)
    # If = cv2.Canny(ROI,50,200)

    If = cv2.adaptiveThreshold(If, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 5, 2)
    If = cv2.bitwise_not(If)

    dst = cv2.cornerHarris(If, 2, 3, 0.04)
    # dst = cv2.dilate(dst, None)

    ret, dst = cv2.threshold(dst, 0.01 * dst.max(), 255, 0)
    dst = np.uint8(dst)
    # find centroids
    ret, labels, stats, centroids = cv2.connectedComponentsWithStats(dst)

    # define the criteria to stop and refine the corners
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.001)
    corners = cv2.cornerSubPix(If, np.float32(centroids), (5, 5), (-1, -1), criteria)

    # Now draw them
    res = np.hstack((centroids, corners))
    res = np.int0(res)
    If_copy = If.copy()
    If_copy[If_copy >= 125] = 125
    If[If >= 0] = 0

    list2 = res[:, 2]
    for indx, x in enumerate(res[:, 3]):
        y = list2[indx]
        if 0 <= x < If.shape[0] and 0 <= y < If.shape[1]:
            If[x, y] = 255
            If_copy[x, y] = 255
    # print('looking for A,B,X,Y')
    pointA = closestNonZero(If, (3, 3), 12)
    pointB = closestNonZero(If, (ROI.shape[0] - 3, 3), 12)
    pointX = closestNonZero(If, (3, ROI.shape[1] - 3), 12)
    pointY = closestNonZero(If, (ROI.shape[0] - 3, ROI.shape[1] - 3), 12)

    #
    # #ret3, If = cv2.threshold(ROI, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    # #If = cv2.dilate(If, cv2.getStructuringElement(cv2.MORPH_RECT, (3,1)), iterations=1)
    # D9 = cv2.getStructuringElement(cv2.MORPH_CROSS, (15,15))
    # D9[1:,1:] = 0
    #
    # #D9[0:8, :] = 0
    # print(D9)
    # ROI_Open_D9 = cv2.morphologyEx(If, cv2.MORPH_ERODE, D9)

    # plt.subplot(3,1,1), plt.imshow(ROI), plt.title('ROI')
    # plt.subplot(3,1,2), plt.imshow(If), plt.title('If')
    # plt.subplot(3, 1, 3), plt.imshow(If_copy), plt.title('If_copy')
    # plt.show()

    for k in range(0, count):
        upperLeft = getPointProportion(pointA, pointX, k, count - k)

        # if k ==0:
        #    upperLeft = getCross(edges, 9, upperLeft)
        # else:
        #    upperLeft = getCross(edges, 13, upperLeft)

        bottomLeft = getPointProportion(pointB, pointY, k, count - k)

        # if k == 0:
        #    bottomLeft = getCross(edges, 3, bottomLeft)
        # else:
        #    bottomLeft = getCross(edges, 7, bottomLeft)

        upperRight = getPointProportion(pointA, pointX, k + 1, count - (k + 1))

        # if k == cant-1:
        #    upperRight = getCross(edges, 12, upperRight)
        # else:
        #    upperRight = getCross(edges, 13, upperRight)

        bottomRight = getPointProportion(pointB, pointY, k + 1, count - (k + 1))

        # if k == cant-1:
        #    bottomRight = getCross(edges, 6, bottomRight)
        # else:
        #    bottomRight = getCross(edges, 7, bottomRight)

        minX = min(upperLeft[0], bottomLeft[0]) + 2
        maxX = max(upperRight[0], bottomRight[0]) - 2

        minY = min(bottomLeft[1], bottomRight[1]) + 2
        maxY = max(upperLeft[1], upperRight[1]) - 2

        singleCharacter = (ROI[minX:maxX, minY:maxY], ROI_onlyUserMarks[minX:maxX, minY:maxY])
        letters.append(singleCharacter)

    filteredLetters = []

    for letter in letters:
        singleLetterFiltered = filterSingleCharacter(letter)
        filteredLetters.append(singleLetterFiltered)
    return filteredLetters

class CuadroCharacter(object):
    def __init__(self):
        self.A = []
        self.B = []
        self.A_predicted = 0
        self.B_predicted = 0
        self.calculated = False

    def calc(self):
        print('A')
        print(self.A)
        print('B')
        print(self.B)
        self.A_predicted = int(np.mean(self.A))
        self.B_predicted = int(np.mean(self.B))
        self.calculated = True

    def add(self, A_B):
        # print('appended')
        if A_B[0] is None or A_B[1] is None:
            print('A_B is None ', A_B)
        self.A.append(A_B[0])
        self.B.append(A_B[1])

class CuadroBuffer(CuadroCharacter):
    instance = None

    def __new__(cls):
        if not CuadroBuffer.instance:
            CuadroBuffer.instance = CuadroCharacter()
        return CuadroBuffer.instance


def latexDeskNode(name, x, y):
    return f"\\node[desk] at ({x}, {y}) {{{name}}};"

def latexDeskPositionsSixByFourStandard(num_students):
    num_rows = (num_students - 1) // 6 + 1
    xs = [0, 2, 5, 7, 10, 12] * num_rows
    ys = []
    for row in range(num_rows):
        ys += [row*4]*6
    xs = xs[:num_students]
    ys = ys[:num_students]
    return xs, ys

def seatingListToLatexDesks(students):
    xs, ys = latexDeskPositionsSixByFourStandard(len(students))
    latex_desks = [latexDeskNode(s, x, y) for s, x, y in zip(students, xs, ys)]
    return '\n'.join(latex_desks)


from typing import Optional
import hashlib

class Solution:
    key = "04"

    def __init__(self, input: str):
        self.secret = input

    def part_one(self) -> Optional[int]:
        for i in range(10**9):
            h = self.hash(i)
            if h.startswith("0" * 5):
                return i
        return None

    def part_two(self) -> Optional[int]:
        for i in range(10**9):
            h = self.hash(i)
            if h.startswith("0" * 6):
                return i
        return None

    def hash(self, num: int) -> str:
        h = hashlib.new("md5")
        h.update(f"{self.secret}{num}".encode())
        return h.hexdigest()

if __name__ == "__main__":
    with open(f"./inputs/{Solution.key}.txt", "r") as fh:
        input = fh.read().strip()
    print("Part 1:", Solution(input).part_one())
    print("Part 2:", Solution(input).part_two())


import sys
import pandas as pd
from openpyxl import Workbook
from openpyxl.utils.dataframe import dataframe_to_rows
from openpyxl.styles import Alignment
import io

def generar_excel(paises):
    """
    Esta función genera un archivo de Excel con datos de una lista de países.

    :param paises: Una lista de objetos que representan países, con atributos como nombre, capital,
    moneda, población, actividades, idiomas y continente
    """
    # Creo un dataframe con los datos de los paises
    data = []
    for pais in paises:
        actividades = ", ".join([actividad.nombre for actividad in pais.actividades])
        idiomas = ", ".join([idioma.nombre for idioma in pais.idiomas])
        continentes = ", ".join([continente.nombre for continente in pais.continente])

        pais_data = {
            "Nombre": pais.nombre,
            "Capital": pais.capital,
            "Moneda": pais.moneda,
            "Población": pais.poblacion,
            "Actividades": actividades,
            "Idiomas": idiomas,
            "Continente": continentes,
        }
        data.append(pais_data)

    df = pd.DataFrame(data)

    # Creo un excel
    wb = Workbook()
    ws = wb.active

    # cargo los datos del dataframe en el excel
    for row in dataframe_to_rows(df, index=False, header=True):
        ws.append(row)

    # Un poco de formato y blink blink
    for column in ws.columns:
        max_length = 0
        column = [cell for cell in column]
        for cell in column:
            try:
                if len(str(cell.value)) > max_length:
                    max_length = len(cell.value)
            except:
                pass
        adjusted_width = max_length + 2
        ws.column_dimensions[column[0].column_letter].width = adjusted_width

    for row in ws.iter_rows(min_row=2):
        for cell in row:
            cell.alignment = Alignment(horizontal="center", vertical="center")

    # guardar excel en memoria para enviar por correo
    buffer = io.BytesIO()
    try:
        wb.save(buffer)
    except:
        print("El archivo no se pudo guardar en memoria.")

    return buffer


from typing import Callable, Tuple, Any, List, Optional, Union

import numpy as np
import functools

from scipy import optimize
from sklearn.base import BaseEstimator, RegressorMixin
from sklearn.utils.validation import check_X_y, check_array, check_is_fitted

class CurvefitEstimator(BaseEstimator, RegressorMixin):

    def __init__(self,
        model_func: Callable[[], np.array]=None,
        p0: Optional[List[float]]=None,
        bounds: Union[ Tuple[np.dtype, np.dtype],
            List[Tuple[np.dtype, np.dtype]],
            Callable[ [], List[Tuple[np.dtype, np.dtype]]] ]=(-np.inf, np.inf),
        method: str='trf',
        jac: Union[str, Callable[[np.array, Any], np.array], None ]=None,
        lsq_kwargs: dict=None
        ) -> None:
        """Wraps the scipy.optimize.curve_fit function for non-linear least squares. The curve_fit function is itself a wrapper around
        scipy.optimize.leastsq and/or scipy.optimize.least_squares that aims to simplfy some of the calling mechanisms. An entrypoint
        to kwargs for these lower level method is provided by the lsq_kwargs dictionary.

        On success, the curve_fit function will return a tuple of the optimized parameters to the function (popt) as well as the estimated
        covariance of these parameters. These values are used in the predict method and can be accessed after the model has been fit
            as ``model.popt_`` and ``model.pcov_``.

        It is best to refer to these docs to understand the methods being wrapped:
        https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.curve_fit.html
        https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.leastsq.html
        https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.least_squares.html

        Args:
            model_func (Callable[[], np.array], optional): The function you wish to model. Defaults to None.
            p0 (Optional[List[float]], optional): The intial guess. Defaults to None.
            bounds (Union[ Tuple[np.dtype, np.dtype], List[Tuple[np.dtype, np.dtype]],
                Callable[ [], List[Tuple[np.dtype, np.dtype]]] ], optional): Bounds for trf. Defaults to (-np.inf, np.inf).
            method (str, optional): The curve fit method. Defaults to 'trf'.
            jac (Union[str, Callable[[np.array, Any], np.array], None ], optional): The jacobian matrix.
                If one is not provided then curve_fit will calculate it. Defaults to None.
            lsq_kwargs (dict, optional): Extra arguments for underlying lsq implementation. See `scipy.optimize.least_squares`. Defaults to None.
        """
        self.model_func = model_func
        self.p0 = p0
        self.bounds = bounds
        self.method = method
        self.jac = jac
        self.lsq_kwargs = lsq_kwargs if lsq_kwargs is not None else {}

    def fit(self,
        X: np.array,
        y: np.array=None,
        sigma: Optional[np.array]=None,
        absolute_sigma: bool=False) -> 'CurvefitEstimator':
        """ Fit X features to target y.

        Refer to scipy.optimize.curve_fit docs for details on sigma values.

        Args:
            X (np.array): The feature matrix we are using to fit.
            y (np.array): The target array.
            sigma (Optional[np.array], optional): Determines uncertainty in the ydata. Defaults to None.
            absolute_sigma (bool, optional): Uses sigma in an absolute sense and reflects this in the pcov. Defaults to True.
            squeeze_1d: (bool, optional): Squeeze X into a 1 dimensional array for curve fitting. This is useful if you are fitting
                a function with an X array and do not want to squeeze before it enters curve_fit. Defaults to True.

        Returns:
            GeneralizedCurveFitEstimator: self
        """
        # NOTE the user defined function should handle the neccesary array manipulation (squeeze, reshape etc.)
        X, y = check_X_y(X, y)  # pass the sklearn estimator dimensionality check

        if callable(self.bounds):  # we allow bounds to be a callable
            bounds = self.bounds(X)
        else:
            bounds = self.bounds

        self.X_ = X
        self.y_ = y

        popt, pcov = optimize.curve_fit(f=self.model_func,
            xdata=X,
            ydata=y,
            p0=self.p0,
            method=self.method,
            sigma=sigma,
            absolute_sigma=absolute_sigma,
            bounds=bounds,
            jac=self.jac,
            **self.lsq_kwargs
            )

        self.popt_ = popt # set optimized parameters on the instance
        self.pcov_ = pcov # set optimzed covariances on the instance
        self.name_ = self.model_func.__name__ # name of func in case we are trying to fit multiple funcs in a Pipeline

        return self

    def predict(self, X: np.array) -> np.array:
        """ Predict the target y values given X features using the best fit
        model (model_func) and best fit model parameters (popt)

        Args:
            X (np.array): The X matrix

        Returns:
            np.array: The predicted y values
        """
        check_is_fitted(self, ["popt_", "pcov_", "name_"])
        X = check_array(X)

        return self.model_func(X, *self.popt_)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Jan  6 10:24:40 2023

@author: junqi
"""
import os
import argparse
import random
import numpy as np

# PyTorch
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.optim.lr_scheduler import CosineAnnealingLR

# Torch Quantum
import torchquantum as tq
from torchquantum.datasets import MNIST

# Tensor Train Network
from tc.tc_fc import TTLinear

# Variational Quantum Circuit
from vqc import VQC

from utils import add_normal_noise, add_laplace_noise

class TTN_VQC(nn.Module):
    def __init__(self,
                 input_dims: int,
                 n_qubits: int = 8,
                 n_class: int = 3,
                 n_qlayers: int = 1,
                 use_tt: bool = True,
                 tt_shape = [[7, 16, 7], [2, 2, 2]],
                 tt_rank = [1, 2, 2, 1]):
        super(TTN_VQC, self).__init__()
        self.dev = tq.QuantumDevice(n_wires=n_qubits)
        if use_tt:
            # Using a tensor-train layer
            assert np.prod(tt_shape[1]) == n_qubits
            assert len(tt_shape[0]) == len(tt_rank) - 1
            self.pre_net = TTLinear(inp_modes=tt_shape[0],
                                    out_modes=tt_shape[1],
                                    tt_rank=tt_rank)
        else:
            # Using a classical feed-forward layer
            self.pre_net = nn.Linear(input_dims, n_qubits)
        self.vqc = VQC(n_wires=n_qubits, n_qlayers=n_qlayers)
        self.post_net = nn.Linear(n_qubits, n_class)

    def forward(self, input_features):
        pre_out = self.pre_net(input_features)
        q_in = torch.sigmoid(pre_out) * np.pi / 2.0
        q_out = self.vqc(q_in, self.dev)
        q_class = self.post_net(q_out)

        return F.log_softmax(q_class, dim=1)

# The console of training a TTN-VQC model
def train(data_flow: dict,
          model: TTN_VQC,
          device: tq.QuantumDevice,
          optimizer: torch.optim):
    batch_idx = 1
    for feed_dct in dataflow['train']:
        inputs = feed_dct['image'].reshape((-1, args.input_dims)).to(device)
        targets = feed_dct['digit'].to(device)

        outputs = model(inputs)
        loss = F.nll_loss(outputs, targets)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        print(f"batch: {batch_idx}, loss: {loss.item()}\n", end='\r')
        batch_idx += 1

def valid_test(dataflow: dict,
               split,
               model: TTN_VQC,
               device: tq.QuantumDevice,
               add_noise: bool = False):
    target_all = []
    output_all = []
    with torch.no_grad():
        for feed_dict in dataflow[split]:
            if not add_noise:
                inputs = feed_dict['image'].reshape((-1, args.input_dims)).to(device)
            else:
                inputs = feed_dict['image'].reshape((-1, args.input_dims))
                for idx, _input in enumerate(inputs):
                    if args.noise_kind == 'normal':
                        inputs[idx, :] = add_normal_noise(_input, args.noise_snr)
                    else:
                        inputs[idx, :] = add_laplace_noise(_input, args.noise_snr)
                inputs = inputs.to(device)
            targets = feed_dict['digit'].to(device)
            outputs = model(inputs)

            target_all.append(targets)
            output_all.append(outputs)
        target_all = torch.cat(target_all, dim=0)
        output_all = torch.cat(output_all, dim=0)

    _, indices = output_all.topk(1, dim=1)
    masks = indices.eq(target_all.view(-1, 1).expand_as(indices))
    size = target_all.shape[0]
    corrects = masks.sum().item()
    accuracy = corrects / size
    loss = F.nll_loss(output_all, target_all).item()

    print(f"{split} set accuracy: {accuracy}")
    print(f"{split} set loss: {loss}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--epochs', type=int, default=12,
                    help='number of training epochs')
    parser.add_argument('--static', action='store_true',
                    help='compute with static mode')
    parser.add_argument('--wires-per-block', type=int, default=2,
                    help='wires per block int static mode')
    parser.add_argument('--input_dims', type=int, default=784,
                    help='input dimensions')
    parser.add_argument('--digits_of_interest', type=list[int], default=[1, 3],
                    help='digits of interest')
    parser.add_argument('--n_class', type=int, default=2,
                    help='number of classes')
    parser.add_argument('--n_qubits', type=int, default=8,
                    help='number of qubits')
    parser.add_argument('--n_qlayers', type=int, default=1,
                    help='number of PQC layers')
    parser.add_argument('--use_tt', type=bool, default=True,
                    help='whether we use tensor-train network?')
    parser.add_argument('--tt_shape', type=list[list[int]], default=[[7, 16, 7], [2, 2, 2]],
                    help='the tensor-train shape')
    parser.add_argument('--tt_ranks', type=list[int], default=[1, 2, 2, 1],
                    help='the ranks of tensor-train')
    parser.add_argument('--add_test_noise', type=bool, default=True,
                    help='whether we add noise to test data')
    parser.add_argument('--noise_snr', type=int, default=6,
                    help='noise snr level')
    parser.add_argument('--noise_kind', metavar='DIR', default='laplace',
                    help='noise kind')

    args = parser.parse_args()

    seed = 4321
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

    dataset = MNIST(
        root='./mnist_data',
        train_valid_split_ratio=[0.9, 0.1],
        digits_of_interest=args.digits_of_interest,
        n_test_samples=1000,
    )
    dataflow = dict()

    for split in dataset:
        sampler = torch.utils.data.RandomSampler(dataset[split])
        dataflow[split] = torch.utils.data.DataLoader(
            dataset[split],
            batch_size=256,
            sampler=sampler,
            num_workers=8,
            pin_memory=True)

    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")

    model = TTN_VQC(args.input_dims,
                    n_qubits=args.n_qubits,
                    n_class=args.n_class,
                    n_qlayers=args.n_qlayers,
                    use_tt=args.use_tt,
                    tt_shape=args.tt_shape,
                    tt_rank=args.tt_ranks).to(device)

    n_epochs = args.epochs
    optimizer = optim.Adam(model.parameters(), lr=5e-2, weight_decay=5e-4)
    scheduler = CosineAnnealingLR(optimizer, T_max=n_epochs)

    if args.static:
    # optionally to switch to the static mode, which can bring speedup on training
        model.VQC.static_on(wires_per_block=args.wires_per_block)

    for epoch in range(1, n_epochs + 1):
        # train
        print(f"Epoch {epoch}:")
        train(dataflow, model, device, optimizer)
        print(f"learning rate: {optimizer.param_groups[0]['lr']}")

        # valid
        valid_test(dataflow, 'valid', model, device, args.add_test_noise)
        scheduler.step()

    # test
    valid_test(dataflow, 'test', model, device, args.add_test_noise)

    # Saving the TTN-VQC model as a whole
    path = os.path.join(os.path.abspath(''), 'models/model_ttn-vqc.pth')
    torch.save(model.state_dict(), path)


from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt

import argparse
from tqdm import tqdm

from utils.IO import read_attr, get_space_mat, normalize_space_mat
from utils.draw import plot_attr_map

def main(MODEL, DATE, TIME, CHANNEL, W,
         attribution_root, figure_log_root,
         channel_of_interest: int = 11, quantile: float = 0.99,
         visualize: bool = True,
         save_visualize: bool = True,
         save_pickle: bool = True
         ):
    ## Read Attribution
    # path
    attribution_root = Path(attribution_root)
    figure_log_root = Path(figure_log_root)

    attr, figure_log_path = read_attr(DATE, TIME, CHANNEL, W, MODEL, attribution_root, figure_log_root, method_dir="simpleThreshold")

    # print(figure_log_path)

    if attr is None:
        # print("Not exists")
        return

    filename = f"simpleThreshold_{DATE}_{TIME}_C{CHANNEL}-W{W}" if channel_of_interest == 11 else f"simpleThreshold_{DATE}_{TIME}_C{CHANNEL}-W{W}-feat{channel_of_interest}"

    space_mat = get_space_mat(attr)
    # of shape [19,512,448]
    # 19 = 12 timeepochs + 7 static features

    # Normalize space mat such that in each channel, the attribution lies in the range [0,255]
    norm_space_mat = normalize_space_mat(space_mat)  # of shape [512,448,19]

    # Find the threshold value defined by the quantile
    flat_channel = norm_space_mat[..., channel_of_interest].flatten()
    flat_channel = np.sort(flat_channel)
    threshold = flat_channel[np.round(quantile * flat_channel.shape[0]).astype(int)]
    # print(f"Threshold {threshold}")

    if visualize:
        import matplotlib as mpl

        # Mask out values below threshold
        channel = norm_space_mat[..., channel_of_interest]
        channel = np.ma.masked_where(channel < threshold, channel)

        cmap = mpl.cm.get_cmap("OrRd").copy()
        cmap.set_bad(color="white")

        # Plot
        fig, ax = plt.subplots(1, 1, figsize=(10, 10))
        plot_attr_map(channel, ax, cmap)

        # Save plot
        if save_visualize:
            plt.savefig(figure_log_path / f"{filename}.png", bbox_inches='tight')
            plt.close(fig)
        del fig, ax

    if save_pickle:
        import pickle
        channel = norm_space_mat[..., channel_of_interest]
        channel[channel < threshold] = 0
        np.save(figure_log_path / filename, channel)

if __name__ == "__main__":
    a = argparse.ArgumentParser(description='conduct simple threshold')
    a.add_argument('--attribution_root',
                   default='/cluster/scratch/jingyli/interpretable/attribution_result/', type=str,
                   help='root path to attribution pickle')
    a.add_argument('--figure_log_root',
                   default='/cluster/scratch/jingyli/interpretable/data_intermediate/visualization', type=str,
                   help='root path to save figures')
    a.add_argument('--quantile', default=0.99, type=float,
                   help='quantile as threshold to filter out hotspots')
    a.add_argument('--channel_of_interest', default=11, type=int,
                   help='which feature to be conducted watershed [0,18]')
    a.add_argument('--save_pickle', default=True, type=bool,
                   help='whether to save watershed and attr map as .npy')
    a.add_argument('--visualize', default=True, type=bool,
                   help='whether to visualize the attribution map')
    a.add_argument('--save_visualize', default=True, type=bool,
                   help='whether to save the visualization of the attribution map')
    args = a.parse_args()

    print(args)

    MODEL = "resUnet"
    DATE = "2019-07-01"
    TIME = 144
    CHANNEL = 0

    for i in tqdm(range(22)):
        for j in range(22):
            W = f"{i}_{j}"
            main(MODEL, DATE, TIME, CHANNEL, W, **vars(args))

    # main(MODEL, DATE, TIME, CHANNEL, W="16_18", **vars(args))