import pygame
import numpy
import numba

class Player:
	def __init__(self, position, fov, view_distance):
		self.position = position
		self.angle = 0
		self.fov = fov
		self.view_distance = view_distance
		self.offset_y = 0.0

@numba.jit(nopython=True, nogil=True, cache=True, fastmath=True)
def check_intersection(wall_1, wall_2):
	x1, y1 = wall_1[0]
	x2, y2 = wall_1[1]
	x3, y3 = wall_2[0]
	x4, y4 = wall_2[1]

	denom = (y4-y3)*(x2-x1) - (x4-x3)*(y2-y1)

	if denom == 0:
		return (0, 0)
	ua = ((x4-x3)*(y1-y3) - (y4-y3)*(x1-x3)) / denom
	if ua < 0 or ua > 1: # out of range
		return (0, 0)
	ub = ((x2-x1)*(y1-y3) - (y2-y1)*(x1-x3)) / denom
	if ub < 0 or ub > 1: # out of range
		return (0, 0)
	x = x1 + ua * (x2-x1)
	y = y1 + ua * (y2-y1)
	return (x,y)

@numba.jit(nopython=True, nogil=True, cache=True, fastmath=True)
def clamp(value, minimum, maximum):
	return max(minimum, min(value, maximum))

@numba.jit(nopython=True, nogil=True, cache=True, fastmath=True)
def get_closest_wall(position, translated_angle, angle, translated_point, level):
	closest_wall = level[1]
	intersection_position = (0, 0)
	closest_distance = 0

	initial_request = True

	all_walls_intersected = []

	for wall in level:
		checked_intersection = check_intersection(translated_point, wall)

		if checked_intersection != (0, 0):
			closest_distance = numpy.sqrt(
				numpy.power(position[0] - checked_intersection[0], 2) +
				numpy.power(position[1] - checked_intersection[1], 2))

			all_walls_intersected.append((closest_distance, checked_intersection, wall))

	for i in range(len(all_walls_intersected)):
		for j in range(0, len(all_walls_intersected) - i - 1):
			if all_walls_intersected[j][0] > all_walls_intersected[j + 1][0]:
				all_walls_intersected[j], all_walls_intersected[j + 1] = all_walls_intersected[j + 1], all_walls_intersected[j]

	return all_walls_intersected

@numba.jit(nopython=True, nogil=True, cache=True, fastmath=True)
def normalize(direction):
	magnitude = numpy.sqrt(direction[0] * direction[0] + direction[1] * direction[1])
	if magnitude > 0:
		return direction[0] / magnitude, direction[1] / magnitude

	return (0, 0)

@numba.jit(nopython=True, nogil=True, cache=True, fastmath=True)
def collision_check(old_position, new_position, velocity, level):
	for wall in level:
		move_difference_segment = ((old_position[0], old_position[1]), (new_position[0], new_position[1]))
		checked_intersection = check_intersection(move_difference_segment, wall)

		if checked_intersection != (0, 0):
			return old_position

	return new_position

@numba.jit(nopython=True, nogil=True, cache=True, fastmath=True)
def lerp(a, b, t):
	return a + (b - a) * t

@numba.jit(nopython=True, nogil=True, cache=True, fastmath=True)
def mix(rgb_1, rgb_2):
	mixed_r = int(rgb_1[0] * rgb_2[0]) << 16
	mixed_g = int(rgb_1[1] * rgb_2[1]) << 8
	mixed_b = int(rgb_1[2] * rgb_2[2])

	return mixed_r + mixed_g + mixed_b

@numba.jit(nopython=True, nogil=True, cache=True, fastmath=True)
def convert_int_rgb(code):
	converted_r = (code >> 16) & 0xff
	converted_g = (code >> 8) & 0xff
	converted_b = code & 0xff

	return converted_r, converted_g, converted_b

@numba.jit(nopython=True, nogil=True, cache=True, fastmath=True, debug=True)
def scan_line(position, angle, fov, view_distance, offset_y, level, floor, ceiling, buffer, offset):
	#Get The Interval Angle To Loop Through Every X-Coordinate Correctly
	interval_angle = fov / buffer.shape[0]
	half_height = int(buffer.shape[1] / 2)

	for x in range(buffer.shape[0]):
		translated_angle = numpy.radians((angle - fov / 2) + interval_angle * x)

		translated_point = (
			(position[0], position[1]),
			(position[0] + view_distance * numpy.cos(translated_angle), position[1] + view_distance * numpy.sin(translated_angle)))

		intersected_walls = get_closest_wall(position, translated_angle, angle, translated_point, level)

		previous_top_wall_height = 0
		previous_bottom_wall_height = 0
		previous_wall_height = 0

		previous_top_wall_ceil_height = 0
		previous_bottom_wall_ceil_height = 0
		previous_wall_ceil_height = 0

		previous_sector = 0
		previous_before_sector = 0
		previous_wall_offset = 0

		for wall in range(len(intersected_walls)):
			distance = intersected_walls[wall][0]
			cos_distance = intersected_walls[wall][0] * numpy.cos(translated_angle - numpy.radians(angle))

			intersection_position = intersected_walls[wall][1]
			wall_reference = intersected_walls[wall][2]
			sector = wall_reference[5]

			if distance != 0:
				wall_height = (half_height / cos_distance) * (buffer.shape[0] / buffer.shape[1])

				texture_distance = numpy.sqrt(
					numpy.power(wall_reference[0][0] - intersection_position[0], 2) +
					numpy.power(wall_reference[0][1] - intersection_position[1], 2)) % 1

				darkness = clamp(lerp(0, .5, 1 / distance), 0, 1)

				#Floor Casting Time!
				bottom_height = (half_height + wall_height) - 2 * wall_height * offset_y
				top_height = (half_height + wall_height) - 2 * wall_height * (wall_reference[3] + offset_y)

				#Ceiling Casting Time!
				top_height_ceiling = (half_height - wall_height)
				bottom_height_ceiling = ((half_height - wall_height) + 2 * wall_height * (wall_reference[4]))

				#We Don't Overdraw!
				if wall > 0:
					bottom_height = clamp(bottom_height, 0, previous_top_wall_height)
					top_height = clamp(top_height, 0, previous_top_wall_height)

					bottom_height_ceiling = clamp(bottom_height_ceiling, 0, previous_bottom_wall_ceil_height)
					top_height_ceiling = clamp(top_height_ceiling, 0, previous_top_wall_ceil_height)

				#FLOOR

				#If We Are In The Same Sector, Don't Draw The Other Walls!
				if sector == previous_sector and wall > 0:
					for y in range(clamp(top_height, 0, buffer.shape[1]), clamp(previous_top_wall_height, 0, buffer.shape[1])):
						#print("Before Divide")

						if (2 * y - buffer.shape[1]) > 0:
							floor_distance = buffer.shape[1] / (2 * y - buffer.shape[1])

							#print("After Divide")
							floor_distance /= numpy.cos(translated_angle - numpy.radians(angle))

							darkness = clamp(lerp(0, 1, 1 / floor_distance), 0, 1)

							#We Multiply The Final Result To Offset The Floor Correctly
							translated_floor_point = (
								position[0] + (floor_distance) * numpy.cos(translated_angle) * (1 - (wall_reference[3] + offset_y) * 2),
								position[1] + (floor_distance) * numpy.sin(translated_angle) * (1 - (wall_reference[3] + offset_y) * 2))

							final_point = (
								int((translated_floor_point[0] * 64) % 64),
								int((translated_floor_point[1] * 64) % 64))
							floor_rgb = convert_int_rgb(floor[final_point[0], final_point[1]])
							buffer[x, y] = mix(floor_rgb, (1, 1, 1))

				#Here The Sector Has Changed Most Likely
				else:
					#We Are Going To Check If The Future Wall That Will Be Drawn Is In The Same Sector, If It Isn't, Don't Draw The Wall
					if wall + 1 < len(intersected_walls):
						#Get The Sector & Compare
						if sector == intersected_walls[wall + 1][2][5]:

							#If It Is The Same Sector, Draw It
							for y in range(clamp(top_height, 0, buffer.shape[1]), clamp(bottom_height, 0, buffer.shape[1])):
								texture_rgb = convert_int_rgb(wall_reference[2][int(texture_distance * 64), int((y - ((half_height + wall_height) - 2 * wall_height * offset_y)) / wall_height * 32)])
								buffer[x, y] = mix(texture_rgb, (1, 1, 1))
						else:
							for y in range(clamp(top_height, 0, buffer.shape[1]), buffer.shape[1]):
								if (2 * y - buffer.shape[1]) > 0:
									floor_distance = buffer.shape[1] / (2 * y - buffer.shape[1])
									floor_distance /= numpy.cos(translated_angle - numpy.radians(angle))

									darkness = clamp(lerp(0, 1, 1 / floor_distance), 0, 1)

									#We Multiply The Final Result To Offset The Floor Correctly
									translated_floor_point = (
										position[0] + (floor_distance) * numpy.cos(translated_angle) * (1 - (wall_reference[3] + offset_y) * 2),
										position[1] + (floor_distance) * numpy.sin(translated_angle) * (1 - (wall_reference[3] + offset_y) * 2))

									final_point = (
										int((translated_floor_point[0] * 64) % 64),
										int((translated_floor_point[1] * 64) % 64))
									floor_rgb = convert_int_rgb(floor[final_point[0], final_point[1]])
									buffer[x, y] = mix(floor_rgb, (1, 1, 1))

					if wall + 1 >= len(intersected_walls):
						for y in range(clamp(top_height, 0, buffer.shape[1]), buffer.shape[1]):
							if (2 * y - buffer.shape[1]) > 0:
								floor_distance = buffer.shape[1] / (2 * y - buffer.shape[1])
								floor_distance /= numpy.cos(translated_angle - numpy.radians(angle))

								darkness = clamp(lerp(0, 1, 1 / floor_distance), 0, 1)

								#We Multiply The Final Result To Offset The Floor Correctly
								translated_floor_point = (
									position[0] + (floor_distance) * numpy.cos(translated_angle) * (1 - (wall_reference[3] + offset_y) * 2),
									position[1] + (floor_distance) * numpy.sin(translated_angle) * (1 - (wall_reference[3] + offset_y) * 2))

								final_point = (
									int((translated_floor_point[0] * 64) % 64),
									int((translated_floor_point[1] * 64) % 64))
								floor_rgb = convert_int_rgb(floor[final_point[0], final_point[1]])
								buffer[x, y] = mix(floor_rgb, (1, 1, 1))

				#CEILING

				#If We Are In The Same Sector, Don't Draw The Other Walls!
				if sector == previous_sector and wall > 0:
					for y in range(clamp(previous_bottom_wall_ceil_height, 0, buffer.shape[1]), clamp(bottom_height_ceiling, 0, buffer.shape[1])):
						#print("Before Divide")

						if (2 * y - buffer.shape[1]) < 0:
							floor_distance = buffer.shape[1] / (2 * y - buffer.shape[1])

							#print("After Divide")
							floor_distance /= numpy.cos(translated_angle - numpy.radians(angle))

							darkness = clamp(lerp(0, 1, 1 / floor_distance), 0, 1)

							#We Multiply The Final Result To Offset The Floor Correctly
							translated_floor_point = (
								position[0] + (floor_distance) * numpy.cos(translated_angle) * (1 - (wall_reference[4]) * 2),
								position[1] + (floor_distance) * numpy.sin(translated_angle) * (1 - (wall_reference[4]) * 2))

							final_point = (
								int((translated_floor_point[0] * 64) % 64),
								int((translated_floor_point[1] * 64) % 64))
							floor_rgb = convert_int_rgb(floor[final_point[0], final_point[1]])
							buffer[x, y] = mix(floor_rgb, (1, 1, 1))

				#Here The Sector Has Changed Most Likely
				else:
					#We Are Going To Check If The Future Wall That Will Be Drawn Is In The Same Sector, If It Isn't, Don't Draw The Wall
					if wall + 1 < len(intersected_walls):
						#Get The Sector & Compare
						if sector == intersected_walls[wall + 1][2][5]:

							#If It Is The Same Sector, Draw It
							for y in range(clamp(top_height_ceiling, 0, buffer.shape[1]), clamp(bottom_height_ceiling, 0, buffer.shape[1])):
								texture_rgb = convert_int_rgb(wall_reference[2][int(texture_distance * 64), int((y - ((half_height - wall_height))) / wall_height * 32)])
								buffer[x, y] = mix(texture_rgb, (1, 1, 1))
						else:
							for y in range(0, top_height_ceiling):
								if (2 * y - buffer.shape[1]) < 0:
									floor_distance = buffer.shape[1] / (2 * y - buffer.shape[1])
									floor_distance /= numpy.cos(translated_angle - numpy.radians(angle))

									darkness = clamp(lerp(0, 1, 1 / floor_distance), 0, 1)

									#We Multiply The Final Result To Offset The Floor Correctly
									translated_floor_point = (
										position[0] + (floor_distance) * numpy.cos(translated_angle) * (1 - (wall_reference[4]) * 2),
										position[1] + (floor_distance) * numpy.sin(translated_angle) * (1 - (wall_reference[4]) * 2))

									final_point = (
										int((translated_floor_point[0] * 64) % 64),
										int((translated_floor_point[1] * 64) % 64))
									floor_rgb = convert_int_rgb(floor[final_point[0], final_point[1]])
									buffer[x, y] = mix(floor_rgb, (1, 1, 1))

					if wall + 1 >= len(intersected_walls):
						for y in range(0, top_height_ceiling):
							if (2 * y - buffer.shape[1]) < 0:
								floor_distance = buffer.shape[1] / (2 * y - buffer.shape[1])
								floor_distance /= numpy.cos(translated_angle - numpy.radians(angle))

								darkness = clamp(lerp(0, 1, 1 / floor_distance), 0, 1)

								#We Multiply The Final Result To Offset The Floor Correctly
								translated_floor_point = (
									position[0] + (floor_distance) * numpy.cos(translated_angle) * (1 - (wall_reference[4]) * 2),
									position[1] + (floor_distance) * numpy.sin(translated_angle) * (1 - (wall_reference[4]) * 2))

								final_point = (
									int((translated_floor_point[0] * 64) % 64),
									int((translated_floor_point[1] * 64) % 64))
								floor_rgb = convert_int_rgb(floor[final_point[0], final_point[1]])
								buffer[x, y] = mix(floor_rgb, (1, 1, 1))

				previous_top_wall_height = top_height
				previous_bottom_wall_height = bottom_height
				previous_wall_height = wall_height
				previous_wall_offset = wall_reference[3]

				previous_top_wall_ceil_height = top_height_ceiling
				previous_bottom_wall_ceil_height = bottom_height_ceiling
				previous_wall_ceil_height = wall_height
				previous_wall_ceil_offset = wall_reference[4]

				previous_before_sector = previous_sector
				previous_sector = sector

pygame.init()

screen_surface = pygame.display.set_mode((256, 256), pygame.SCALED, vsync=True)
pygame.mouse.set_visible(False)
pygame.event.set_grab(True)

running = True

#Create Buffer
buffer = numpy.zeros((screen_surface.get_width(), screen_surface.get_height()), dtype=numpy.int32)

#Create Textures
basic_wall_1 = pygame.surfarray.array2d(pygame.image.load("texture.png").convert())
basic_wall_2 = pygame.surfarray.array2d(pygame.image.load("texture2.png").convert())
basic_wall_3 = pygame.surfarray.array2d(pygame.image.load("texture3.png").convert())

#Create Player
player = Player((63, 63), 60, 128)

#Create Frame Counter
clock = pygame.time.Clock()
font = pygame.font.SysFont("Monospace" , 16 , bold = False)

#correct offset:
# 0.1 -> 0.8
# 0.2 -> 0.6
# 0.3 -> 0.4
# 0.4 -> 0.2
# 0.5 -> 0.0

#Therefore
#Multiply The Offset By Two
#Then Calculate The Difference From 1?

#0.1 -> 0.1 * 2 = 0.2
#1 - 0.2 = 0.8

#0.2 -> 0.2 * 2 = 0.4
#1 - 0.4 = 0.6

second_offset = 7

level = (
	((64, 64), (70, 64), basic_wall_2, 0.8, 0.2, 1),
	((70, 64), (70, 70), basic_wall_2, 0.8, 0.2, 1),
	((64, 64), (70, 70), basic_wall_2, 0.8, 0.2, 1),

	((64, 64 + second_offset), (70, 64 + second_offset), basic_wall_1, 0.4, 0.2, 0),
	((70, 64 + second_offset), (70, 70 + second_offset), basic_wall_1, 0.4, 0.2, 0),
	((64, 64 + second_offset), (70, 70 + second_offset), basic_wall_1, 0.4, 0.2, 0),

	((64 + second_offset, 64 + second_offset), (70 + second_offset, 64 + second_offset), basic_wall_1, 0.0, 0.2, 2),
	((70 + second_offset, 64 + second_offset), (70 + second_offset, 70 + second_offset), basic_wall_1, 0.0, 0.2, 2),
	((64 + second_offset, 64 + second_offset), (70 + second_offset, 70 + second_offset), basic_wall_1, 0.0, 0.2, 2),
)

floor = basic_wall_3
ceiling = basic_wall_3

first_part_wall = None
second_part_wall = None

dt = 0

offset = 1

while running:
	keys = pygame.key.get_pressed()

	for event in pygame.event.get():
		if event.type == pygame.QUIT:
			running = False

		if event.type == pygame.MOUSEMOTION:
			player.angle += event.rel[0] * .1

		if event.type == pygame.KEYDOWN:
			down_keys = pygame.key.get_pressed()
			if down_keys[pygame.K_DOWN] or down_keys[pygame.K_UP]:
				pass
				#player.fov += (down_keys[pygame.K_UP] - down_keys[pygame.K_DOWN])

				#player.offset_y = .09
				#print("Changing Offset To:", player.offset_y)

	#Player Movement
	move_position = (
		player.position[0] + (numpy.cos(numpy.radians(player.angle)) * (keys[pygame.K_w] - keys[pygame.K_s]) * 4 * dt),
		player.position[1] + (numpy.sin(numpy.radians(player.angle)) * (keys[pygame.K_w] - keys[pygame.K_s]) * 4 * dt)
	)

	move_position = (
		move_position[0] + (numpy.cos(numpy.radians(player.angle + 90)) * (keys[pygame.K_d] - keys[pygame.K_a]) * 4 * dt),
		move_position[1] + (numpy.sin(numpy.radians(player.angle + 90)) * (keys[pygame.K_d] - keys[pygame.K_a]) * 4 * dt),
	)
	player.angle += (keys[pygame.K_RIGHT] - keys[pygame.K_LEFT]) * 128 * dt

	#player.position = collision_check(player.position, move_position, 8 * dt, level)
	player.offset_y += (keys[pygame.K_UP] - keys[pygame.K_DOWN]) * dt
	player.position = move_position

	scan_line(player.position, player.angle, player.fov, player.view_distance, player.offset_y, level, floor, ceiling, buffer, offset)

	pygame.surfarray.blit_array(screen_surface, buffer)
	screen_surface.blit(font.render("FPS: " + str(int(clock.get_fps())), False, (255, 255, 255)), (0, 0))

	#This Is Unecessary In Closed Areas
	pygame.display.flip()
	buffer.fill(0)

	dt = clock.tick() / 1000


import collections
import pickle
import numpy as np
import matplotlib.pyplot as plt

def mean_elo(game):
    return((game['white']['rating']+game['black']['rating'])/2)

games = pickle.load( open( "martinovichev.p", "rb" ))
elos = []

for game in (games):
    elos.append(mean_elo(game))

pickle.dump(games, open( "martinovichev.p", "wb" ) )

print(len(games))

array = np.array(elos)
bins = np.arange(0,3500,100)
hist = np.histogram(array, bins=bins)
print(hist)
plt.hist(array, bins=bins)
plt.show()


import argparse
import pickle
import ssl
import os

import torch.optim as optim
import torch.nn as nn
import torchvision
from torch.utils.data import DataLoader, Dataset
import matplotlib.pyplot as plt
import torch
import torchvision.transforms as transforms
import time as t

import vanilla_model as vanilla
import moded_model as moded

def get_output_accuracy(pred, confirmed, top_res=(1,)):
    max_num_results = max(top_res) #get highest result looking for
    data_size = confirmed.size(0)

    pred_values, pred_indx = pred.topk(k=max_num_results, dim=1)
    pred_indx = pred_indx.t()
#   Reshape confirmed indxs
    confirmed_indx_reshape = confirmed.view(1, -1).expand_as(pred_indx)
    answers = pred_indx == confirmed_indx_reshape

    top_res_acc = []
    for k in top_res:
        temp_answers = answers[:k] #limit each image answer to k values
        temp_answers = temp_answers.reshape(-1).float() #flatten to see whats right form ALL images in batch
        temp_answers = temp_answers.float().sum(dim=0, keepdim=True)
        curr_acc = temp_answers/data_size
        top_res_acc.append(curr_acc)
    return top_res_acc
def evaluate_epoch_top1_5(model, data, device, test_loader=None):
    model.eval() #Set to evaluate
    tot_top1_accuracy = 0
    tot_top5_accuracy = 0
    test_tot_top1_accuracy = 0
    test_tot_top5_accuracy = 0
    for imgs, labels in data:
        imgs = imgs.to(device)
        labels = labels.to(device)

        with torch.no_grad():
            output = model(imgs)
        curr_top1, curr_top5 = get_output_accuracy(output, labels, (1, 5))
        tot_top1_accuracy += curr_top1
        tot_top5_accuracy += curr_top5
    avg_test_top5_epoch_acc = 0
    avg_test_top1_epoch_acc = 0
    if test_loader is not None:
        for imgs, labels in test_loader:
            imgs = imgs.to(device)
            labels = labels.to(device)

            with torch.no_grad():
                output = model(imgs)
            curr_top1, curr_top5 = get_output_accuracy(output, labels, (1, 5))
            test_tot_top1_accuracy += curr_top1
            test_tot_top5_accuracy += curr_top5
    avg_test_top5_epoch_acc = test_tot_top5_accuracy / len(test_loader)
    avg_test_top1_epoch_acc = test_tot_top1_accuracy / len(test_loader)
    avg_top5_epoch_acc = tot_top5_accuracy / len(data)
    avg_top1_epoch_acc = tot_top1_accuracy / len(data)

    model.train() #Set to train when returning to function

    return avg_top1_epoch_acc, avg_top5_epoch_acc, avg_test_top1_epoch_acc, avg_test_top5_epoch_acc

def train(n_epochs, optimizer, model, loss_fn, train_loader, scheduler, device,
          decoder_save=None, plot_file=None, pickleLosses = None,
          starting_epoch=1, evaluate_epochs=False, folder='./', store_data=False, test_loader=None):

    model.train()
    total_losses = []
    total_top1_accuracy = []
    total_top5_accuracy = []
    test_total_top1_accuracy = []
    test_total_top5_accuracy = []

    final_loss = 0.0
    t_1 = t.time()
    print("\n=======================================")
    print("Training started at Epoch {}\n".format(starting_epoch))

    # loading_saved pickle data
    # if pickleLosses is not None:
    #     try:
    #         with open(pickleLosses, 'rb') as file:
    #             loaded_losses = pickle.load(file)
    #             (total_losses, content_losses, style_losses) = loaded_losses
    #             print("Loaded saved losses from file successfully: \n{} \n{} \n{}"
    #                   .format(total_losses, content_losses, style_losses))
    #     except Exception as e:
    #         print(f"An error occurred while loading arrays: {str(e)}")

    for epoch in range(starting_epoch, n_epochs + 1):
        print("Starting Epoch: ", epoch)
        # Losses for current epoch
        total_loss = 0.0
        t_2 = t.time()

        for idx, data in enumerate(train_loader):
            t_3 = t.time()
            imgs, labels = data[0].to(device), data[1].to(device)
            optimizer.zero_grad()

            outputs = model(imgs)
            loss = loss_fn(outputs, labels)
            loss.backward()
            optimizer.step()

            total_loss += loss.item()
            print('Batch #{}/{}         Time: {}'.format(idx + 1, len(train_loader), (t.time() - t_3)))

        if decoder_save is not None and (epoch % 50 == 0 or epoch == n_epochs):
            torch.save(model.decoder.state_dict(), decoder_save)
            print("Saved frontend model under name: {}".format(decoder_save))

        scheduler.step(total_loss)
        total_losses.append(total_loss / len(train_loader))
        final_loss = total_loss / len(train_loader)

        # Store loss data in pickled file
        # pickleSave = os.path.join(os.path.abspath(folder), "pickled_dump.pk1")
        # if pickleLosses is not None:
        #     pickleSave = os.path.join(os.path.abspath(folder), pickleLosses)
        # try:
        #     with open(pickleSave, 'wb') as file:
        #         pickle.dump((total_losses, epoch), file)
        #         print("Saved losses to '{}'".format(pickleSave))
        # except Exception as e:
        #     print(f"An error occurred while saving arrays: {str(e)}")

        # Plot loss data each epoch
        if plot_file is not None:
            plot_save = os.path.join(os.path.abspath(folder), plot_file)
            plt.figure(2, figsize=(12, 7))
            plt.clf()
            plt.plot(total_losses, label='Total')
            plt.xlabel('epoch')
            plt.ylabel('loss')
            plt.legend(loc=1)
            plt.savefig(plot_save)

        print('Epoch {}, Training loss {}, Time  {}'.format(epoch, final_loss,(t.time() - t_2)))
        if evaluate_epochs:
            top_1_acc, top_5_acc, test_top_1_acc, test_top_5_acc = evaluate_epoch_top1_5(model, train_loader, device, test_loader=test_loader)
            total_top1_accuracy.append(float(top_1_acc))
            total_top5_accuracy.append(float(top_5_acc))
            print("Accuracy (Training Data) = TOP-1:   {}  |  TOP-5:    {}".format(top_1_acc, top_5_acc))
            if test_loader is not None:
                test_total_top1_accuracy.append(float(test_top_1_acc))
                test_total_top5_accuracy.append(float(test_top_5_acc))
            print("Accuracy (Test Data) = TOP-1:   {}  |  TOP-5:    {}".format(test_top_1_acc, test_top_5_acc))

    print('Total Training loss {}, Time  {}'.format(final_loss, (t.time() - t_1)))

    # if store_data:
    #     filename = os.path.join(os.path.abspath(folder), 'output_results')
    #     print("Length: {}    {}   {}    {}   {}".format(total_top5_accuracy, total_top1_accuracy, test_total_top1_accuracy, test_total_top5_accuracy, total_losses))
    #     try:
    #         with open(filename, 'w') as file:
    #             file.write("Training results for run with the following hyper parameters:\n{}\n".format(arguments))
    #             if train_loader is not None: file.write("Accuracy & loss results per Epoch: (Train_Top 1,  Train_Top 5, "
    #                                                     "Test_Top_1, Test_Train_5,  Loss)\n")
    #             else: file.write("Accuracy & loss results per Epoch: (Top 1,    Top 5,   Loss)\n")
    #             for i in range(n_epochs):
    #                 top_1_item = total_top1_accuracy[i]
    #                 top_5_item = total_top5_accuracy[i]
    #                 test_top_5, test_top_1 = "", ""
    #                 if train_loader is not None:
    #                     test_top_1 = test_total_top1_accuracy[i]
    #                     test_top_5 = test_total_top5_accuracy[i]
    #
    #                 loss = total_losses[i]
    #                 file.write("Epoch {}:  ([{},  {}],   [{},   {}],   {})\n".format(i + 1, top_1_item, top_5_item, test_top_1, test_top_5, loss))
    #             file.close()
    #
    #     except Exception as e:
    #         print(f"An error occurred while loading arrays: {str(e)}")
    return final_loss

if __name__ == '__main__':

    # ssl._create_default_https_context = ssl._create_unverified_context

    device = 'cpu'
    # Setup parser
    parser = argparse.ArgumentParser()
    parser.add_argument('-dir', '--dir', type=str, required=False, default='./data',
                        help='Directory path to a batch of content images')
    # training options
    parser.add_argument('-e', '--e', type=int, default=50)
    parser.add_argument('-b', '--b', type=int, default=8, help="Batch size")
    parser.add_argument('-l', '--l', help="Encoder.pth", required=True)
    parser.add_argument('-s', '--s', help="Decoder.pth")
    parser.add_argument('-p', '--p', help="decoder.png")
    parser.add_argument('-cuda', '--cuda', default='Y')
    parser.add_argument('-type', '--type', type=int, default=0)

    parser.add_argument('-lr', '--lr', type=float, default=0.001)
    parser.add_argument('-wd', '--wd', type=float, default=0.00001)
    parser.add_argument('-minlr', '--minlr', type=float, default=0.001)
    parser.add_argument('-out', '--out', default=None, help="Output folder to put all files for training run")
    parser.add_argument('-gamma', '--gamma', type=float, default=0.9)
    parser.add_argument('-dataset', '--dataset', type=int, default=1)

    args = parser.parse_args()

    print("Beginning Training for model. Using the following parameters passed (Some default)\n")
    print("")

    if args.type == 0:
        model_type = vanilla
    else:
        model_type = moded

    # Import the dataset & get num of batches
    train_transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    if args.dataset == 0:
        cifar_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=train_transform)
        cifar_test_dataset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True,
                                                           transform=train_transform)
        cifar_data = DataLoader(cifar_dataset, batch_size=args.b, shuffle=True)
        cifar_test_data = DataLoader(cifar_test_dataset, batch_size=args.b, shuffle=True)
        frontend = model_type.encoder_decoder.frontend_10

    else:
        cifar_dataset = torchvision.datasets.CIFAR100(root='./data', train=True, download=True, transform=train_transform)
        cifar_test_dataset = torchvision.datasets.CIFAR100(root='./data', train=False, download=True, transform=train_transform)
        cifar_data = DataLoader(cifar_dataset, batch_size=args.b, shuffle=True)
        cifar_test_data = DataLoader(cifar_test_dataset, batch_size=args.b, shuffle=True)
        frontend = model_type.encoder_decoder.frontend

    encoder = model_type.encoder_decoder.encoder
    encoder.load_state_dict(torch.load(args.l))

    # Creating folder to store all files made during training
    if args.out != None:
        if os.path.exists(os.path.abspath(args.out)):
            print("Saving all files created to folder '{}'".format(args.out))
        else:
            os.mkdir(args.out)
            print("Created folder {} to save all files made during training".format(args.out))
    else:
        args.out = "./"

    # Set the device (GPU if available, otherwise CPU)
    if torch.cuda.is_available() and args.cuda == 'Y':
        if torch.cuda.device_count() > 1:
            device = torch.device("cuda:1")
        else:
            device = torch.device("cuda:0")
    else:
        device = torch.device("cpu")

    print("Using device: {}".format(device))

    loss_fn = nn.functional.cross_entropy
    model = model_type.model(encoder, frontend)
    model.to(device)

    # Define optimizer and learning rate scheduler
    optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.wd)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=2, verbose=True, factor=0.1,
                                                     min_lr=args.minlr)

    # Train the model
    train(args.e, optimizer, model, loss_fn, cifar_data, scheduler, device, args.s,
          args.p, evaluate_epochs=False, folder=args.out, store_data=True, test_loader=cifar_test_data)


"""dm_actinf.py

Developmental model: *actinf*, smp_graphs configured for base
experiment

Variations:
- actinf variant: M1: goal and goal error, M2: goal error only, M3: ?
- data dimensionality: 1, 2, 3, ..., N?
- proprioceptive space upwards: include extero
- Temporal embedding configured with lag* parameters to system and
  models

TODO:
- FIXME: target function properties like: frequency matched to body
  (attainable / not attainable)
- FIXME: learning rate: eta
- FIXME: priors: limits, learn the sensorimotor limits
- FIXME: priors: tapping, learn tapping
- FIXME: x insert measure_mse, x insert measure_moments, x insert
  measure_accumulated_error

Log:
- renamed from dm_actinf_m1_goal_error_ND_embedding.py
- ported from actinf/active_inference_basic.py --mode m1_goal_error_nd
"""

import copy
from functools import partial

from smp_base.plot import histogramnd
from smp_graphs.block_plot import SnsMatrixPlotBlock2, ImgPlotBlock2
from smp_graphs.block import dBlock2, IBlock2, SliceBlock2, DelayBlock2, StackBlock2
from smp_graphs.block_meas import XCorrBlock2, MSEBlock2, MomentBlock2
from smp_graphs.block_meas_infth import JHBlock2, MIBlock2, InfoDistBlock2, TEBlock2, CTEBlock2, MIMVBlock2, TEMVBlock2
from smp_graphs.block_models import ModelBlock2
from smp_graphs.block_cls import PointmassBlock2, SimplearmBlock2, BhasimulatedBlock2
from smp_graphs.block_cls_ros import STDRCircularBlock2, LPZBarrelBlock2, SpheroBlock2

from smp_graphs.common import escape_backslash
from smp_graphs.funcs import f_meshgrid, f_meshgrid_mdl, f_random_uniform, f_sin_noise

from smp_graphs.utils_conf import get_systemblock

# execution
saveplot = False
recurrent = True
debug = False
showplot = True
ros = False # True

# experiment
commandline_args = ['numsteps']
randseed = 12348

################################################################################
# experiment variations
# - algo
# - system
# - system order
# - dimensions
# - number of modalities

lconf = {
    # execution and global
    'numsteps': int(10000/5),
    # system
    'sys': {
        'name': 'pm',
        'lag': 0,
        'dim_s0': 2,
        'dim_s1': 2,
    },
    # motivation
    'motivation_i': 0,
    # models
    # learning modulation
    'lm' : {
        'washout': 200,
        'thr_predict': 200,
        'fit_offset': 200,
        # start stop in ratios of episode length numsteps
        'r_washout': (0.0, 0.1),
        'r_train': (0.1, 0.8),
        'r_test': (0.8, 1.0),
    }
}
lconf['systemblock'] = get_systemblock[lconf['sys']['name']](**lconf['sys'])

numsteps = lconf['numsteps']
loopblocksize = numsteps

# learning modulation
for k in list(lconf['lm']):
    if not k.startswith('r_'): continue
    k_ = k.replace('r_', 'n_')
    r0 = lconf['lm'][k][0] * lconf['numsteps']
    r1 = lconf['lm'][k][1] * lconf['numsteps']
    # lconf['lm'][k_] = range(int(r0), int(r1))
    lconf['lm'][k_] = (int(r0), int(r1))

sysname = lconf['sys']['name']
# sysname = 'pm'
# sysname = 'sa'
# sysname = 'bha'
# sysname = 'stdr'
# sysname = 'lpzbarrel'
# sysname = 'sphero'
# dim = 3 # 2, 1
# dim = 9 # bha

outputs = {
    'latex': {'type': 'latex',},
}

# get systemblock
systemblock = lconf['systemblock']

# print "systemblock", systemblock
dim_s0 = systemblock['params']['dims']['s0']['dim']
print(("dim_s0", dim_s0))
dim_s1 = systemblock['params']['dims']['s1']['dim']
print(("dim_s1", dim_s1))
m_mins = np.array([systemblock['params']['m_mins']]).T
print(("m_mins", m_mins))
m_maxs = np.array([systemblock['params']['m_maxs']]).T
print(("m_maxs", m_maxs))

lag = systemblock['params']['lag']
lag_past = systemblock['params']['lag_past']
lag_future = systemblock['params']['lag_future']

dt = systemblock['params']['dt']

algo_conf = {
    'knn': {
        'name': 'k-nearest neighbors',
    },
}

algo = 'knn' #
# algo = 'gmm' #
# algo = 'igmm' #
# algo = 'hebbsom'
# algo = 'soesgp'
# algo = 'storkgp'
# algo = 'resrls'
# algo = 'homeokinesis'

# lag_past = (-21, -20)
# lag_past = (-11, -10)
# lag_past = (-7, -5)
# lag_past = (-6, -5)
# lag_past = (-5, -4)
# lag_past = (-4, -3)
# lag_past = (-3, -2)
# lag_past = (-2, -1)
# lpzbarrel non-overlapping seems important
# lag_past = (-6, -2)
# lag_future = (-1, 0)

# # pm experiment with lag = 5
# lag_past = (-7, -5)
# lag_future = (-2, 0)

minlag = 1 # -lag_future[1]
# maxlag = -lag_past[0] # + lag_future[1]
maxlag = max(20, -lag_past[0])
laglen = maxlag # - minlag

# eta = 0.99
# eta = 0.95
eta = 0.7
# eta = 0.3
# eta = 0.25
# eta = 0.15
# eta = 0.1
# eta = 0.01

# print "dm_actinf_m1_goal_error_ND_embedding.py: dim_s0 = %s" % (dim_s0, )

# motivations
from smp_graphs.utils_conf import dm_motivations
motivations = dm_motivations(m_mins, m_maxs, dim_s0, dt)
motivation_i = lconf['motivation_i']
goal = motivations[motivation_i][1]['params']['models']['goal']['type']

desc = """An exemplary experiment involving a single developmental
episode of {0} steps of an actinf agent. The episode starts with an
initial washout period of 1/10th of the episode length during which
the top-down prediction is applied and the model is allowed to predict
but its is only predicting from its bootstrapping state. Starting with
time step {1}, the model is being updated with the incoming
measurements, which almost immediately has the effect of bringing the
system into the state predicted by the goal. The goal or target
function is a {2} sequence of proprioceptive states. In this case,
every consecutive change of goal leads to a new mini-episode of
learning for new combinations of goal prediction and local prediction
error. This is a simple type of \\emph{{goal babbling}}, and it is
only partially completed when learning is stopped at time step
{3}. The agent able to get close to most of the goal predictions based
on the existing knowledge but the remaining residual after the goal
change transient is not being corrected anymore.""".format(
    numsteps, lconf['lm']['n_train'][0], escape_backslash(goal),
    lconf['lm']['n_test'][0])

def plot_timeseries_block(l0 = 'pre_l0', l1 = 'pre_l1', blocksize = 1):
    global partial
    global PlotBlock2, numsteps, timeseries, algo, dim_s1, dim_s0, sysname, lag, lag_past, lag_future, saveplot
    global motivations, motivation_i, goal
    # goal = motivations[motivation_i][1]['params']['models']['goal']['type']
    return {
    'block': PlotBlock2,
    'params': {
        'blocksize': blocksize,
        'saveplot': saveplot,
        # 'debug': True,
        'title': '%s\nalgo %s, sys %s(dim_p=%d), goal = %s, lag %d, tap- %s, tap+ %s' % (
            'dm actinf', algo, sysname, dim_s0, goal, lag, lag_past, lag_future),
        'desc': """An {1} agent learning to control a {0}-dimensional {2}
            system using the {3} low-level algorithm.""".format(dim_s0, 'actinf', sysname, algo),
        'inputs': {
            'goals': {'bus': '%s/pre' % (l1,), 'shape': (dim_s0, blocksize)},
            'pre':   {'bus': '%s/pre' % (l0,), 'shape': (dim_s0, blocksize)},
            'err':   {'bus': '%s/err' % (l0,), 'shape': (dim_s0, blocksize)},
            'tgt':   {'bus': '%s/tgt' % (l0,), 'shape': (dim_s0, blocksize)},
            'X_fit': {'bus': '%s/X_fit' % (l0,), 'shape': (dim_s0 * (lag_past[1] - lag_past[0]) * 2, blocksize)},
            's0': {'bus': 'robot1/s0', 'shape': (dim_s0, blocksize)},
            's1': {'bus': 'robot1/s1',  'shape': (dim_s1, blocksize)},
            'mse_s_p_accum': {'bus': 'mse_s_p_accum/Ix', 'shape': (dim_s0, blocksize)},
            },
        'hspace': 0.2,
        'subplots': [
            [
                {
                    'input': ['err',], 'plot': partial(timeseries, marker='.'),
                    'title': 'Momentary prediction error',
                    'legend': {'e(s_p)': 0},
                    'xticks': False,
                },
            ],

            [
                {
                    'input': ['goals', 's0', 'pre'],
                    'plot': [partial(timeseries, marker='.')] * 3,
                    'title': 'Goal, state, and prediction',
                    'legend': {'Goal': 0, 'State': dim_s0, 'State_p': 2*dim_s0},
                    'xticks': False,
                }
            ],

            # [
            #     {
            #         'input': ['goals', 'pre', 's0'],
            #         'plot': partial(timeseries, marker='.')
            #     },
            # ],

            [
                {
                    'input': ['X_fit'], 'plot': partial(timeseries, marker='.'),
                    'title': 'Model %s input $\mathbf{X}$' % (algo),
                    'legend': {'X_pre_l1': 0, 'X_prerr_l0': dim_s0},
                    'xticks': False,
                },
            ],

            [
                {
                    'input': ['tgt'], 'plot': partial(timeseries, marker='.'),
                    'title': 'Model %s input $\mathbf{Y}$' % (algo),
                    'legend': {'Y_tgt': 0},
                    'xticks': False,
                },
            ],

            [
                {
                    'input': ['mse_s_p_accum',],
                    'plot': partial(timeseries, marker='.'),
                    'title': 'Accumulated error',
                    'legend': {'$\sum$ err': 0},
                }
            ],

            # [
            #     {'input': ['s0', 's1'], 'plot': timeseries},
            # ],

            # [
            #     {
            #         'input': ['hidden'], 'plot': partial(timeseries, marker = '.'),
            #         'title': 'Model %s hidden activation $\mathbf{Z}$' % (algo),
            #         # FIXME: knn particulars
            #         'legend': {'Z_dist': 0, 'E(Z_dist)': 5, 'Z_idx': 6},
            #         'xticks': False,
            #     },
            # ],

            # [
            #     {
            #         'input': ['wo_norm', 'wo_norm_fwd'], 'plot': partial(timeseries, marker = '.'),
            #         'title': 'Model %s parameter norm (accumulated adaptation)' % (algo),
            #         'legend': {'|W|': 0}
            #     },
            # ],

        ],
    }}

"""
sweep system subgraph
 - sweep block is an open-loop system (data source) itself
 - system block is the system we want to sweep
"""
sweepsys_steps = 40 # 6
sweepsys_input_flat = np.power(sweepsys_steps, dim_s0)
sweepsys = ('robot0', copy.deepcopy(systemblock))
sweepsys[1]['params']['blocksize'] = sweepsys_input_flat
sweepsys[1]['params']['debug'] = False
sweepsys[1]['params']['inputs'] = {'u': {'bus': 'sweepsys_grid/meshgrid'}}
sweepsys[1]['params']['outputs']['s0']['shape'] = (dim_s0, sweepsys_input_flat)
sweepsys[1]['params']['outputs']['s1']['shape']  = (dim_s1, sweepsys_input_flat)

sweepmdl_steps = 1000
sweepmdl_input_flat = sweepmdl_steps # np.power(sweepmdl_steps, dim_s0 * 2)
sweepmdl_func = f_random_uniform

# sweepmdl_steps = 3
# sweepmdl_input_flat = np.power(sweepmdl_steps, dim_s0 * 2)
# sweepmdl_func = f_meshgrid

loopblock = {
    'block': Block2,
    'params': {
        'id': 'bhier',
        'debug': False,
        'topblock': False,
        'logging': False,
        'numsteps': sweepsys_input_flat,  # inner numsteps when used as loopblock (sideways time)
        'blocksize': 1,           # compute single steps, has to be 1 so inner cnt is correct etc
        'blockphase': [0],        # phase = 0
        'outputs': {'meshgrid': {'shape': (dim_s0, sweepsys_input_flat), 'buscopy': 'sweepsys_grid/meshgrid'}},
        # subgraph
        'graph': OrderedDict([
            ('sweepsys_grid', {
                'block': FuncBlock2,
                'params': {
                    'debug': False,
                    'blocksize': sweepsys_input_flat,
                    'inputs': {
                        'ranges': {'val': np.array([[-1, 1]] * dim_s0)},
                        'steps':  {'val': sweepsys_steps},
                        },
                    'outputs': {'meshgrid': {'shape': (dim_s0, sweepsys_input_flat)}},
                    # 'func': f_meshgrid
                    'func': f_random_uniform,
                    },
                }),

                # sys to sweep
                sweepsys,

            # ('sweepsys_grid', {
            #     'block': FuncBlock2,
            #     'params': {
            #         'debug': False,
            #         'blocksize': sweepsys_input_flat,
            #         'inputs': {
            #             'ranges': {'val': np.array([[-1, 1]] * dim_s0)},
            #             'steps':  {'val': sweepsys_steps},
            #             },
            #         'outputs': {'meshgrid': {'shape': (dim_s0, sweepsys_input_flat)}},
            #         'func': f_meshgrid
            #         },
            #     }),

            #     # sys to sweep
            #     sweepsys,

            ]),
        }
    }

loopblock_model = {
    'block': Block2,
    'params': {
        'id': 'bhier',
        'debug': False,
        'topblock': False,
        'logging': False,
        'numsteps': sweepmdl_input_flat,  # inner numsteps when used as loopblock (sideways time)
        'blocksize': 1,           # compute single steps, has to be 1 so inner cnt is correct etc
        'blockphase': [0],        # phase = 0
        'outputs': {
            'pre': {
                'shape': (dim_s0 * 2, sweepmdl_input_flat),
                'buscopy': 'sweepmdl_grid/meshgrid'}},
        # subgraph
        'graph': OrderedDict([
            # model sweep input
            ('sweepmdl_grid', {
                'block': FuncBlock2,
                'params': {
                    'debug': False,
                    'blocksize': sweepmdl_input_flat,
                    'inputs': {
                        # 'ranges': {'val': np.array([[-1, 1], [-1e-0, 1e-0]])},
                        'ranges': {
                            'val': np.array([[m_mins[0], m_maxs[0]]] * dim_s0 * 2)},
                            # 'val': np.vstack((
                            #     np.array([[m_mins[0], m_maxs[0]]] * dim_s0),
                            #     np.array([[-2.0,      1.0]]       * dim_s0),
                            #     ))},
                        'steps':  {'val': sweepmdl_steps},
                        },
                    'outputs': {
                        'meshgrid': {
                            'shape': (
                                dim_s0 * 2,
                                sweepmdl_input_flat,)}},
                    'func': sweepmdl_func,
                    },
                }),

            # # model sweep input
            # ('sweepmdl_grid_goal', {
            #     'block': FuncBlock2,
            #     'params': {
            #         'debug': False,
            #         'blocksize': sweepmdl_input_flat,
            #         'inputs': {
            #             'ranges': {'val': np.array([[-1, 1]])},
            #             'steps':  {'val': sweepmdl_steps},
            #             },
            #         'outputs': {'pre': {'shape': (dim_s0, sweepmdl_input_flat)}},
            #         'func': f_meshgrid,
            #         },
            #     }),

            # # model sweep input
            # ('sweepmdl_grid_err', {
            #     'block': FuncBlock2,
            #     'params': {
            #         'debug': False,
            #         'blocksize': sweepmdl_input_flat,
            #         'inputs': {
            #             'ranges': {'val': np.array([[-1e-0, 1e-0]])},
            #             'steps':  {'val': sweepmdl_steps},
            #             },
            #         'outputs': {'pre': {'shape': (dim_s0, sweepmdl_input_flat)}},
            #         'func': f_meshgrid,
            #         },
            #     }),

            ('sweep_slice', {
                'block': SliceBlock2,
                'params': {
                    'id': 'sweep_slice',
                    'blocksize': sweepmdl_input_flat,
                    # puppy sensors
                    'inputs': {
                        'x': {
                            'bus': 'sweepmdl_grid/meshgrid',
                            'shape': (dim_s0 * 2, sweepmdl_input_flat)}},
                    'slices': {
                        'x': {
                            'goals': slice(0, dim_s0),
                            'errs':  slice(dim_s0, dim_s0 * 2)},
                        },
                    # 'slices': ,
                    }
                }),

            # model sweep
            # learner: basic actinf predictor proprio space learn_proprio_base_0
            ('pre_l0_test', {
                'block': ModelBlock2,
                'params': {
                    'blocksize': sweepmdl_input_flat,
                    'blockphase': [0],
                    # 'debug': True,
                    'lag': minlag,
                    'eta': eta, # 0.3,
                    'inputs': {
                        # descending prediction
                        'pre_l1': {
                            'bus': 'sweep_slice/x_goals',
                            'shape': (dim_s0,sweepmdl_input_flat)},
                        # ascending prediction error
                        'pre_l0': {
                            'bus': 'sweep_slice/x_errs',
                            'shape': (dim_s0, sweepmdl_input_flat)},
                        # ascending prediction error
                        'prerr_l0': {
                            'bus': 'pre_l0_test/err',
                            'shape': (dim_s0, minlag+1), 'lag': minlag},
                        # measurement
                        'meas_l0': {
                            'val': np.array([[-np.inf for i in range(sweepmdl_input_flat)]] * dim_s0),
                            # 'bus': 'robot1/s0',
                            'shape': (dim_s0, sweepmdl_input_flat)}},
                    'outputs': {
                        'pre': {'shape': (dim_s0, sweepmdl_input_flat)},
                        'err': {'shape': (dim_s0, sweepmdl_input_flat)},
                        'tgt': {'shape': (dim_s0, sweepmdl_input_flat)},
                        },
                    'models': {
                        # 'fwd': {'type': 'actinf_m1', 'algo': algo, 'idim': dim_s0 * 2, 'odim': dim},
                        'fwd': {
                            'type': 'actinf_m1', 'algo': 'copy',
                            'copyid': 'pre_l0', 'idim': dim_s0 * 2, 'odim': dim_s0},
                        },
                    'rate': 1,
                    },
                }),

            ('pre_l0_combined', {
                'block': StackBlock2,
                'params': {
                    'inputs': {
                        # 'goals': {
                        #     'bus': 'sweepmdl_grid_goal/pre',
                        #     'shape': (dim_s0, sweepmdl_input_flat),
                        #     },
                        # 'errs':  {
                        #     'bus': 'sweepmdl_grid_err/pre',
                        #     'shape': (dim_s0, sweepmdl_input_flat),
                        #     },
                        'goalerrs': {
                            'bus': 'sweepmdl_grid/meshgrid',
                            'shape': (dim_s0 * 2, sweepmdl_input_flat),
                            },
                        'pres':  {
                            'bus': 'pre_l0_test/pre',
                            'shape': (dim_s0, sweepmdl_input_flat),
                            },
                        },
                    'outputs': {
                        'y': {'shape': (dim_s0 * 3, sweepmdl_input_flat)}},
                    }
                }),

            # # plot timeseries
            # ('plot_ts',
            #      plot_timeseries_block(
            #          l0 = 'pre_l0_test',
            #          l1 = 'sweepmdl_grid_goal',
            #          blocksize = sweepmdl_input_flat)),

            # # plot model sweep 1d
            # ('plot_model_sweep', {
            #     'block': ImgPlotBlock2,
            #     'params': {
            #         'id': 'plot_model_sweep',
            #         'logging': False,
            #         'saveplot': saveplot,
            #         'debug': False,
            #         'blocksize': sweepmdl_input_flat,
            #         'inputs': {
            #             # 'sweepin_goal': {
            #             #     'bus': 'sweepmdl_grid_goal/pre',
            #             #     'shape': (dim_s0, sweepmdl_input_flat)},
            #             # 'sweepin_err':  {
            #             #     'bus': 'sweepmdl_grid_err/pre',
            #             #     'shape': (dim_s0, sweepmdl_input_flat)},
            #             # 'sweepout_mdl':  {
            #             #     'bus': 'pre_l0_test/pre',
            #             #     'shape': (dim_s0, sweepmdl_input_flat)},
            #             'all': {
            #                 'bus': 'pre_l0_combined/y',
            #                 'shape': (dim_s0 * 3, sweepmdl_input_flat),
            #                 },
            #             },
            #         'outputs': {},
            #         'wspace': 0.5,
            #         'hspace': 0.5,
            #         # with one subplot and reshape
            #         'subplots': [
            #             [
            #                 {
            #                     'input': ['all'],
            #                     'shape': (dim_s0 * 3, sweepmdl_input_flat),
            #                     'ndslice': [(slice(None), slice(None))],
            #                     # 'vmin': -1.0, 'vmax': 1.0,
            #                     # 'vmin': 0.1, 'vmax': 0.3,
            #                     'cmap': 'RdGy',
            #                     'dimstack': {
            #                         'x': range(2*dim-1, dim - 1, -1),
            #                         'y': range(dim-1,   -1     , -1)},
            #                     'digitize': {'argdims': range(0, dim_s0 * 2), 'valdim': 2*dim+i, 'numbins': 2},
            #                 } for i in range(dim_s0)],

            #             ],
            #         },
            #     }),
            ]),
        }
    }

# main graph
graph = OrderedDict([

    # # sweep system
    # ("sweepsys", {
    #     'debug': False,
    #     'block': SeqLoopBlock2,
    #     'params': {
    #         'id': 'sweepsys',
    #         # loop specification, check hierarchical block to completely
    #         # pass on the contained in/out space?
    #         'blocksize': numsteps, # execution cycle, same as global numsteps
    #         'blockphase': [1],     # execute on first time step only
    #         'numsteps':  numsteps,          # numsteps      / loopblocksize = looplength
    #         'loopblocksize': loopblocksize, # loopblocksize * looplength    = numsteps
    #         # can't do this dynamically yet without changing init passes
    #         'outputs': {'meshgrid': {'shape': (dim_s0, sweepsys_input_flat)}},
    #         'loop': [('none', {})], # lambda ref, i, obj: ('none', {}),
    #         'loopmode': 'sequential',
    #         'loopblock': loopblock,
    #     },
    # }),

    # # plot the system sweep result
    # ('plot_sweep', {
    #     'block': PlotBlock2,
    #     'params': {
    #         'debug': False,
    #         'blocksize': numsteps, # sweepsys_input_flat,
    #         'title': 'system sweep',
    #         'inputs': {
    #             'meshgrid': {
    #                 'bus': 'sweepsys_grid/meshgrid',
    #                 'shape': (dim_s0, sweepsys_input_flat)},
    #             's0': {
    #                 'bus': 'robot0/s0',
    #                 'shape': (dim_s0, sweepsys_input_flat)},
    #             's1': {
    #                 'bus': 'robot0/s1',
    #                 'shape': (dim_s1, sweepsys_input_flat)},
    #             },
    #             'hspace': 0.2,
    #             'subplots': [
    #                 [
    #                     {'input': ['meshgrid'], 'plot': timeseries},
    #                 ],
    #                 [
    #                     {'input': ['s0'], 'plot': timeseries},
    #                 ],
    #                 [
    #                     {'input': ['s1'], 'plot': timeseries},
    #                 ],
    #                 ],
    #         }
    #     }),

    # # sns matrix plot
    # ('plot2', {
    #     'block': SnsMatrixPlotBlock2,
    #     'params': {
    #         'id': 'plot2',
    #         'logging': False,
    #         'debug': False,
    #         'saveplot': saveplot,
    #         'blocksize': numsteps,
    #         'inputs': {
    #             'meshgrid': {
    #                 'bus': 'sweepsys_grid/meshgrid',
    #                 'shape': (dim_s0, sweepsys_input_flat)},
    #             's0': {
    #                 'bus': 'robot0/s0',
    #                 'shape': (dim_s0, sweepsys_input_flat)},
    #             's1': {
    #                 'bus': 'robot0/s1',
    #                 'shape': (dim_s1, sweepsys_input_flat)},
    #             },
    #         'outputs': {},#'x': {'shape': 3, 1)}},
    #         'subplots': [
    #             [
    #                 # stack inputs into one vector (stack, combine, concat
    #                 {
    #                     'input': ['meshgrid', 's0', 's1'],
    #                     'mode': 'stack',
    #                     'plot': histogramnd
    #                 },
    #             ],
    #         ],
    #     },
    # }),

    # system block from definition elsewhere
    ('robot1', systemblock),

    # learning experiment
    ('brain_learn_proprio', {
        'block': Block2,
        'params': {
            'numsteps': numsteps,
            'blocksize': 1,
            'graph': OrderedDict([

                ('cnt', {
                    'block': CountBlock2,
                    'params': {
                        'blocksize': 1,
                        # 'debug': True,
                        # 'inputs': {},
                        'outputs': {'x': {'shape': (dim_s0, 1)}},
                        },
                }),

                # generic block configured with subgraph on demand
                ('pre_l1', {
                    'block': Block2,
                    'params': {
                        'numsteps': numsteps,
                        'blocksize': 1,
                        # 'subgraph': OrderedDict([motivations[0]]),
                        'subgraph': OrderedDict([motivations[motivation_i]]),
                        'subgraph_rewrite_id': False,
                    }
                }),

                # dev model actinf_m1: learner is basic actinf predictor proprio space learn_proprio_base_0
                ('pre_l0', {
                    'block': ModelBlock2,
                    'params': {
                        'blocksize': 1,
                        'blockphase': [0],
                        'debug': True,
                        'lag': minlag,
                        'eta': eta, # 3.7,
                        'ros': ros,
                        # FIXME: relative shift = minlag, block length the maxlag
                        'inputs': {
                            # descending prediction
                            'pre_l1': {
                                'bus': 'pre_l1/pre',
                                # 'shape': (dim_s0, maxlag), 'lag': range(-maxlag, -minlag)},
                                # FIXME: check correctness here
                                # training on past goal, prediction with current goal
                                # should be possible inside step_model to use current
                                # goal
                                'shape': (dim_s0, maxlag), 'lag': list(range(lag_past[0], lag_past[1]))},
                            # ascending prediction error
                            'pre_l0': {
                                'bus': 'pre_l0/pre',
                                # 'shape': (dim_s0, maxlag), 'lag': range(-maxlag + 1, -minlag + 1)},
                                'shape': (dim_s0, maxlag), 'lag': list(range(lag_past[0] + 1, lag_past[1] + 1))},
                            # ascending prediction error
                            'prerr_l0': {
                                'bus': 'pre_l0/err',
                                # 'shape': (dim_s0, maxlag), 'lag': range(-maxlag + 1, -minlag + 1)},
                                'shape': (dim_s0, maxlag), 'lag': list(range(lag_past[0] + 1, lag_past[1] + 1))},
                            # measurement
                            'meas_l0': {
                                'bus': 'robot1/s0',
                                # 'shape': (dim_s0, maxlag), 'lag': range(-laglen, 0)}
                                'shape': (dim_s0, maxlag), 'lag': list(range(lag_future[0], lag_future[1]))}
                            },
                        'outputs': {
                            'pre': {'shape': (dim_s0, 1)},
                            'err': {'shape': (dim_s0, 1)},
                            'tgt': {'shape': (dim_s0, 1)},
                            'X_fit': {'shape': (dim_s0 * (lag_past[1] - lag_past[0]) * 2, 1)},
                            },
                        'models': {

                            'm1': {
                                'type': 'actinf_m1',
                                'algo': algo,
                                'lag_past': lag_past,
                                'lag_future': lag_future,
                                'idim': dim_s0 * (lag_past[1] - lag_past[0]) * 2, # laglen
                                'odim': dim_s0 * (lag_future[1] - lag_future[0]), # laglen,
                                'laglen': laglen,
                                'eta': eta,
                                # 'laglen_past': lag_past[1] - lag_past[0],
                                # 'laglen_future': lag_future[1] - lag_future[0],
                                # learning modulation
                                'n_washout': lconf['lm']['n_washout'],
                                'n_train': lconf['lm']['n_train'],
                                'n_test': lconf['lm']['n_test'],
                            },

                            # 'm2': {
                            #     'type': 'actinf_m2',
                            #     'algo': algo,
                            #     'idim': dim_s0 * laglen,
                            #     'odim': dim_s0 * laglen,
                            #     'laglen': laglen,
                            #     'eta': eta
                            #     },

                            # 'm3': {
                            #     'type': 'actinf_m3',
                            #     'algo': algo,
                            #     'idim': dim_s0 * laglen * 2,
                            #     'odim': dim_s0 * laglen,
                            #     'laglen': laglen,
                            #     'eta': eta
                            #     },

                            },
                        'rate': 1,
                        },
                    }),

                # # dev model uniform random sampler
                # ('pre_l0', {
                #     'block': ModelBlock2,
                #     'params': {
                #         'blocksize': 1,
                #         'blockphase': [0],
                #         'inputs': {
                #             'lo': {'val': m_mins, 'shape': (dim_s0, 1)},
                #             'hi': {'val': m_maxs, 'shape': (dim_s0, 1)},
                #             },
                #         'outputs': {
                #             'pre': {'shape': (dim_s0, 1)},
                #             'err': {'val': np.zeros((dim_s0, 1)), 'shape': (dim_s0, 1)},
                #             'tgt': {'val': np.zeros((dim_s0, 1)), 'shape': (dim_s0, 1)},
                #             },
                #         'models': {
                #             'goal': {'type': 'random_uniform'}
                #             },
                #         'rate': 50,
                #         },
                #     }),

                # # dev model homeokinesis
                # ('pre_l0', {
                #     'block': ModelBlock2,
                #     'params': {
                #         'blocksize': 1,
                #         'blockphase': [0],
                #         'debug': False,
                #         'lag': minlag,
                #         'eta': eta, # 3.7,
                #         'ros': ros,
                #         'inputs': {
                #             # descending prediction
                #             'pre_l1': {
                #                 'bus': 'pre_l1/pre',
                #                 'shape': (dim_s0, maxlag), 'lag': minlag},
                #             # ascending prediction error
                #             'pre_l0': {
                #                 'bus': 'pre_l0/pre',
                #                 'shape': (dim_s0, maxlag), 'lag': minlag},
                #             # ascending prediction error
                #             'prerr_l0': {
                #                 'bus': 'pre_l0/err',
                #                 'shape': (dim_s0, maxlag), 'lag': minlag},
                #             # measurement
                #             'meas_l0': {
                #                 'bus': 'robot1/s0', 'shape': (dim_s0, maxlag)}},
                #         'outputs': {
                #             'pre': {'shape': (dim_s0, 1)},
                #             'err': {'shape': (dim_s0, 1)},
                #             'tgt': {'shape': (dim_s0, 1)},
                #             },
                #         'models': {
                #             'hk': {
                #                 'type': 'homeokinesis', 'algo': algo, 'mode': 'hk', 'idim': dim_s0, 'odim': dim_s0,
                #                 'minlag': minlag, 'maxlag': maxlag, 'laglen': laglen, 'm_mins': m_mins, 'm_maxs': m_maxs,
                #                 'epsA': 0.01, 'epsC': 0.03, 'creativity': 0.8},
                #             },
                #         'rate': 1,
                #         },
                #     }),

                # learn_proprio_e2p2e

                # measure and introspect
                # total MSE goal - state
                ('mse_s_p', {
                    'block': MSEBlock2,
                    'params': {
                        'blocksize': 1, # numsteps,
                        'inputs': {
                            'x': {'bus': 'pre_l1/pre', 'shape': (dim_s0, 1)},
                            'x_': {'bus': 'robot1/s0', 'shape': (dim_s0, 1)},
                            },
                        'outputs': {
                            'y': {'shape': (dim_s0, 1)},
                            },
                        },
                    }),

                # accumulated MSE goal - state over time
                ('mse_s_p_accum', {
                    'block': IBlock2,
                    'params': {
                        'blocksize': numsteps,
                        'inputs': {
                            'x': {'bus': 'mse_s_p/y', 'shape': (dim_s0, numsteps)},
                            },
                        'outputs': {
                            'Ix': {'shape': (dim_s0, numsteps)},
                            },
                        },
                    }),

                # end brain
                ]),
            }
        }),

    ################################################################################
    # use a sequential loop block to insert probes running orthogonally in time
    # blocksize:  seqloops blocksize from the outside, same as overall experiment
    # blockphase: points in the cnt % numsteps space when to execute
    # numsteps:
    # loopblocksize: number of loop iterations = numsteps/loopblocksize
    # ("sweepmodel", {
    #     'debug': False,
    #     'block': SeqLoopBlock2,
    #     'params': {
    #         'id': 'sweepmodel',
    #         'blocksize': numsteps, # execution cycle, same as global numsteps
    #         #                        execution phase, on first time step only
    #         # 'blockphase': [numsteps/2, numsteps-10],
    #         # 'blockphase': [int(i * numsteps)-1 for i in np.linspace(1.0/2, 1, 2)],
    #         'blockphase': [int(i * numsteps)-1 for i in np.linspace(1.0/1, 1, 1)],
    #         # 'blockphase': [0],
    #         'numsteps':  1, # numsteps,          # numsteps      / loopblocksize = looplength
    #         'loopblocksize': 1, #loopblocksize, # loopblocksize * looplength    = numsteps
    #         # can't do this dynamically yet without changing init passes
    #         'outputs': {'pre': {'shape': (dim_s0 * 2, sweepmdl_input_flat)}},
    #         'loop': [('none', {}) for i in range(2)], # lambda ref, i, obj: ('none', {}),
    #         'loopmode': 'sequential',
    #         'loopblock': loopblock_model,
    #     },
    # }),

    # plot timeseries
    ('plot_ts', plot_timeseries_block(l0 = 'pre_l0', blocksize = numsteps)),

    # # plot / print measures
    # ('plot_meas', {
    #     'block': PlotBlock2,
    #     'params': {

    #         },
    #     }),

    ])


# Take a gzipped pileup file and a .tsv files from oncotator reformat (/data/ipadioleau/code/oncotator_extract_information_from_MAF_tumor_Vs_normal.py)
# Add strand coverage of variant and reference allele to the tsv file
import sys
from collections import Counter
import gzip
import io

# MPILEUP = gzip.open(sys.argv[1], 'rb')
ONCOTATOR_FILE = open(sys.argv[2], 'r')
COMPLETED_ONCOTATOR_FILES = open(sys.argv[3], 'w')

# A function that return a dictionnary with all alternatives count on forward and reverse strand
def return_alternativ_count(seq):
	nucleotides_list = {}
	acceptable = ['A', 'a', 'C', 'c', 'G', 'g', 'T', 't', '+', '-', '1', '2', '3', '4', '5', '6', '7', '8', '9', '0']
	getsize = 0
	indel = []
	tmp_size = []
	size = 0
	for nuc in seq:
		# skip all non informative value
		if nuc in acceptable:
			# If there is an indel set getsize to 1 so we get it's size next turn
			if (nuc=='+' or nuc=='-'):
				getsize = 1
				# if previous variant was a indel, save information
				if indel != []:
					indel_seq = ''.join(indel)
					indel = []
					if not indel_seq.upper() in nucleotides_list:
						nucleotides_list[indel_seq.upper()] = [0,0]
					if indel_seq.isupper():
						nucleotides_list[indel_seq.upper()][0] += 1
					else:
						nucleotides_list[indel_seq.upper()][1] += 1
				signe = nuc
			# Get indel size
			elif getsize == 1:
				if nuc.isdigit() :
					tmp_size.append(nuc)
				else:
					if tmp_size == []:
						indel.append(signe + nuc)
					else:
						size = int(''.join(tmp_size)) - 1
						tmp_size = []
						getsize = 0
						indel.append(signe + nuc)
			# If we are processing an indel, get next nucleotide
			elif size > 0 :
				indel.append(nuc)
				size -= 1
			# Save count for each alternative in a dictionary
			else:
				if indel != []:
					indel_seq = ''.join(indel)
					indel = []
					signe=''
					if not indel_seq.upper() in nucleotides_list:
						nucleotides_list[indel_seq.upper()] = [0,0]
					if indel_seq.isupper():
						nucleotides_list[indel_seq.upper()][0] += 1
					else:
						nucleotides_list[indel_seq.upper()][1] += 1
				if  not nuc.isdigit():
					if not nuc.upper() in nucleotides_list:
						nucleotides_list[nuc.upper()] = [0,0]
					if nuc.isupper() :
						nucleotides_list[nuc.upper()][0] += 1
					else:
						nucleotides_list[nuc.upper()][1] += 1
	if indel != []:
		indel_seq = ''.join(indel)
		indel = []
		if not indel_seq.upper() in nucleotides_list:
			nucleotides_list[indel_seq.upper()] = [0,0]
		if indel_seq.isupper():
			nucleotides_list[indel_seq.upper()][0] += 1
		else:
			nucleotides_list[indel_seq.upper()][1] += 1
	return(nucleotides_list)

variant_dict = {}
with io.TextIOWrapper(io.BufferedReader(gzip.open(sys.argv[1]))) as MPILEUP:
	for line in MPILEUP:
		info = line.split('\t')
		nucleotides = Counter(info[4])
		ref_forward = '0'
		ref_reverse = '0'
		if '.' in nucleotides:
			ref_forward = str(nucleotides['.'])
		if ',' in nucleotides:
			ref_reverse = str(nucleotides[','])
		alternativ = return_alternativ_count(info[4])
		for alt in alternativ:
			if alt[0] == '-' :
				info[1] = str(int(info[1]) + 1)
			chr_pos = info[0] + '_' + info[1]
			if chr_pos in variant_dict:
				variant_dict[chr_pos][alt] = [alternativ[alt][0],alternativ[alt][1], ref_forward, ref_reverse]
			else:
				variant_dict[chr_pos] = {}
				variant_dict[chr_pos][alt] = [alternativ[alt][0],alternativ[alt][1], ref_forward, ref_reverse]
			if alt[0] == '-' :
				info[1] = str(int(info[1]) - 1)
			#OUTFILE.write(info[0] + '\t' + info[1] + '\t' + info[2] + '\t' + alt + '\t' + ref_forward + '\t' + ref_reverse + '\t' + str(alternativ[alt][0]) + '\t' + str(alternativ[alt][1]) + '\n')
			#print(info[0] + '\t' + info[1] + '\t' + info[2] + '\t' + alt + '\t' + ref_forward + '\t' + ref_reverse + '\t' + str(alternativ[alt][0]) + '\t' + str(alternativ[alt][1]) + '\n')

#####
header = ONCOTATOR_FILE.readline().strip('\n')
COMPLETED_ONCOTATOR_FILES.write(header + '\tt_alt_forward_strand\tt_alt_reverse_strand\tt_ref_forward_strand\tt_ref_reverse_strand\n')
header_info = header.split('\t')
Variant_Type_index = header_info.index('Variant_Type')
Tumor_Seq_Allele2_index = header_info.index('Tumor_Seq_Allele2')

Start_position_index = header_info.index('Start_position')
Chromosome_index = header_info.index('Chromosome')
Tumor_Seq_Allele1_index = header_info.index('Tumor_Seq_Allele1')

for line in ONCOTATOR_FILE:
	info = line.strip().split('\t')
	chr_pos = info[Chromosome_index] + '_' + info[Start_position_index]
	variant_counts = [9999999999,9999999999, 9999999999, 9999999999]
	# My variant checker script do not detect DNP or TNP
	# To report DNP/TNP counts we report the minimum counts per strand for all alternatives nucleotides composing the variant
	if info[Variant_Type_index] == 'DNP' or info[Variant_Type_index] == 'TNP' :
		nucleotides = info[Tumor_Seq_Allele2_index]
		#print(nucleotides)
		tmp_pos = int(info[Start_position_index])
		for tmp_nuc in nucleotides:
			tmp_chr_pos = info[Chromosome_index] + '_' + str(tmp_pos)
			if (tmp_chr_pos in variant_dict) and (tmp_nuc in variant_dict[tmp_chr_pos]):
				if (int(variant_dict[tmp_chr_pos][tmp_nuc][0]) < variant_counts[0]):
					variant_counts[0] = int(variant_dict[tmp_chr_pos][tmp_nuc][0])
				if (int(variant_dict[tmp_chr_pos][tmp_nuc][1]) < variant_counts[1]):
					variant_counts[1] = int(variant_dict[tmp_chr_pos][tmp_nuc][1])
				if (int(variant_dict[tmp_chr_pos][tmp_nuc][2]) < variant_counts[2]):
					variant_counts[2] = int(variant_dict[tmp_chr_pos][tmp_nuc][2])
				if (int(variant_dict[tmp_chr_pos][tmp_nuc][3]) < variant_counts[3]):
					variant_counts[3] = int(variant_dict[tmp_chr_pos][tmp_nuc][3])
			tmp_pos += 1
	elif info[Variant_Type_index] == 'DEL' or info[Variant_Type_index] == 'ONP' :
		seq = '-' + info[Tumor_Seq_Allele1_index]
		if (chr_pos in variant_dict) and (seq in variant_dict[chr_pos]):
			if (int(variant_dict[chr_pos][seq][0]) < variant_counts[0]):
				variant_counts[0] = variant_dict[chr_pos][seq][0]
			if (int(variant_dict[chr_pos][seq][1]) < variant_counts[1]):
				variant_counts[1] = variant_dict[chr_pos][seq][1]
			if (int(variant_dict[chr_pos][seq][2]) < variant_counts[2]):
				variant_counts[2] = variant_dict[chr_pos][seq][2]
			if (int(variant_dict[chr_pos][seq][3]) < variant_counts[3]):
				variant_counts[3] = variant_dict[chr_pos][seq][3]
	elif info[Variant_Type_index] == 'INS' :
		seq = '+' + info[Tumor_Seq_Allele2_index]
		if (chr_pos in variant_dict) and (seq in variant_dict[chr_pos]):
			if (int(variant_dict[chr_pos][seq][0]) < variant_counts[0]):
				variant_counts[0] = variant_dict[chr_pos][seq][0]
			if (int(variant_dict[chr_pos][seq][1]) < variant_counts[1]):
				variant_counts[1] = variant_dict[chr_pos][seq][1]
			if (int(variant_dict[chr_pos][seq][2]) < variant_counts[2]):
				variant_counts[2] = variant_dict[chr_pos][seq][2]
			if (int(variant_dict[chr_pos][seq][3]) < variant_counts[3]):
				variant_counts[3] = variant_dict[chr_pos][seq][3]
	else:
		seq = info[Tumor_Seq_Allele2_index]
		if (chr_pos in variant_dict) and (seq in variant_dict[chr_pos]):
			if (int(variant_dict[chr_pos][seq][0]) < variant_counts[0]):
				variant_counts[0] = variant_dict[chr_pos][seq][0]
			if (int(variant_dict[chr_pos][seq][1]) < variant_counts[1]):
				variant_counts[1] = variant_dict[chr_pos][seq][1]
			if (int(variant_dict[chr_pos][seq][2]) < variant_counts[2]):
				variant_counts[2] = variant_dict[chr_pos][seq][2]
			if (int(variant_dict[chr_pos][seq][3]) < variant_counts[3]):
				variant_counts[3] = variant_dict[chr_pos][seq][3]
	if variant_counts[0] == 9999999999 and variant_counts[1] == 9999999999:
		if variant_counts[2] == 9999999999 :
			variant_counts[2] = 0
		if variant_counts[3] == 9999999999 :
			variant_counts[3] = 0
		variant_counts = ['0','0', str(variant_counts[2]), str(variant_counts[3])]

	COMPLETED_ONCOTATOR_FILES.write(line.strip('\n') + '\t' + str(variant_counts[0]) + '\t' + str(variant_counts[1]) + '\t' + str(variant_counts[2]) + '\t' + str(variant_counts[3]) + '\n')

COMPLETED_ONCOTATOR_FILES.close()
ONCOTATOR_FILE.close()
MPILEUP.close()


from urllib.parse import unquote

import dash  # type: ignore
from dash import dcc, html  # type: ignore
from dash_extensions.EventListener import EventListener  # type: ignore

import lecture_search.app.pages.asset_readers.reader_slides as reader_slides
import lecture_search.app.pages.asset_readers.reader_video as reader_video

dash.register_page(__name__, path_template="/<course_name>/<lecture_name>/<asset_name>")
event = {"event": "scroll", "props": ["target.scrollTop", "target.scrollLeft"]}

def layout(course_name, lecture_name, asset_name, page=None, startTime=None):
    if page is None:
        page = 0
    if startTime is None:
        startTime = 0

    app = dash.get_app()
    # decode url and special characters
    course_name = unquote(course_name)
    lecture_name = unquote(lecture_name)
    asset_name = unquote(asset_name)
    path = f"{course_name}/{lecture_name}/{asset_name}"
    print(path)
    if asset_name.endswith(".pdf"):
        content = reader_slides.bp.embed(app)
    elif asset_name.endswith(".mp4"):
        content = reader_video.bp.embed(app)
    else:
        content = html.Div("Not implemented yet")
    return html.Div(
        [
            dcc.Store("file_path", data=path),
            dcc.Store("page", data=page),
            dcc.Store("startTime", data=startTime),
            html.H2(lecture_name),
            EventListener(
                content,
                events=[event],
                logging=True,
                id="el",
            ),
        ]
    )


#chp
from PyQt5.QtWidgets import *
from PyQt5.uic import *
from PyQt5.QtGui import *
#from matplotlib import pyplot
#from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg
#from PyQt5.QtCore import *

class MyApp(QMainWindow):
    def __init__(self):
        super().__init__()
        loadUi("test2.ui",self)
        self.main()

    def main(self):
        self.img = cv2.imread('')
        self.img = self.processingImage(self.img)
        self.printImage(self.img, self.pic)

    def processingImage(self, img):

        #img = cv2.blur(img,(55,55))

        #img = cv2.

        img = cv2.cvtColor(img,cv2.COLOR_RGB2RRAY)
        kernal = numpy.ones((3,3), numpy.uint8)
        img = cv2.morphologyEx(img,cv2.MORPH_GRADIENT, kernal)

        return img

    def printImage(self, img):
        imgRGB = cv2.cvtColor(imgRGB, cv2.COLOR_RGB2RGB)
        h, w, byte = imgRGB.shape
        img = QImage(imgRGB,w,h,byte *w, QImage.Format_RGB888)
        pic.setPixmap(QPixmap(img))

app = QApplication([])
win = MyApp()
win.show()
app.exec()


# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, x):
#         self.val = x
#         self.next = None

class Solution:
    def reverseKGroup(self, head: ListNode, k: int) -> ListNode:
        if k == 1: return head
        if not head: return head
        node,pre = head,head
        for _ in range(k-1):
            if not node.next: return head
            node = node.next
        temp = node.next
        cur = pre.next
        for _ in range(k-1):
            nxt = cur.next
            cur.next = pre
            pre,cur = cur,nxt
        head.next = self.reverseKGroup(temp,k)
        return node


import os
import sqlite3

from pathlib import Path

def default_db(user) -> Path:
    currentdir = os.path.dirname(os.path.abspath(__file__))
    parentdir = os.path.dirname(currentdir)
    db_path = Path(parentdir + f'/dbs/{user}.babel.db')
    return db_path.resolve()

class BabelDatabase:
    def __init__(self, db_name: Path):
        self.__db_name = db_name

    def add_change_request(self, user_account, provider_account,
                           consumer_account, schema_id, perms, p_nonce,
                           p_sig, public_key):
        self.__open_con()
        self.__c.execute(f"INSERT INTO perm_change_requests "
                         f"VALUES (NULL,"
                         f"'{user_account}', "
                         f"'{provider_account}', "
                         f"'{consumer_account}',"
                         f"'{schema_id}',"
                         f"'{perms}',"
                         f"'{p_nonce}',"
                         f"'{p_sig}',"
                         f"'{public_key}',"
                         f"0,"
                         f"NULL,"
                         f"NULL)")

        self.__conn.commit()

        request_id = self.__c.lastrowid

        self.__close_con()

        return request_id

    def update_provider_process_id(self, request_id, provider_process_id):
        self.__open_con()

        self.__c.execute(f"UPDATE perm_change_requests "
                         f"SET provider_process_id='{provider_process_id}' "
                         f"WHERE req_id='{request_id}'")

        self.__conn.commit()
        self.__close_con()

    def update_processed(self, request_id, proof_tx=None):
        self.__open_con()

        self.__c.execute(f"UPDATE perm_change_requests "
                         f"SET processed='1',"
                         f"proof_tx='{proof_tx}' "
                         f"WHERE req_id='{request_id}'")

        self.__conn.commit()
        self.__close_con()

    def get_unprocessed(self, num=10):
        self.__open_con()
        self.__c.execute(f'SELECT * FROM perm_change_requests '
                         f'WHERE processed=0 '
                         f'LIMIT {num}')
        columns = [d[0] for d in self.__c.description]
        res = [dict(zip(columns, row)) for row in self.__c.fetchall()]

        self.__close_con()

        return res

    def get_request_by_id(self, request_id, user_account):
        op_data = None
        self.__open_con()
        self.__c.execute(f'SELECT * FROM perm_change_requests '
                         f"WHERE req_id='{request_id}' "
                         f"AND end_user_account='{user_account}' "
                         f'LIMIT 1')
        columns = [d[0] for d in self.__c.description]
        res = [dict(zip(columns, row)) for row in self.__c.fetchall()]

        self.__close_con()

        if res:
            op_data = res[0]

        return op_data

    def __open_con(self):
        self.__conn = sqlite3.connect(str(self.__db_name))
        self.__c = self.__conn.cursor()

    def __close_con(self):
        self.__conn.close()


import math
import os
import modules.shared as shared
import modules.sd_models
import gradio as gr
from scripts import postprocessing_upscale
from .prompt_util import readJsonPrompt
import asyncio

def fix_env_Path_ffprobe():
    envpath = os.environ["PATH"]
    ffppath = shared.opts.data.get("infzoom_ffprobepath", "")

    if ffppath and not ffppath in envpath:
        path_sep = ";" if os.name == "nt" else ":"
        os.environ["PATH"] = envpath + path_sep + ffppath

def closest_upper_divisible_by_eight(num):
    if num % 8 == 0:
        return num
    else:
        return math.ceil(num / 8) * 8

def load_model_from_setting(model_field_name, progress, progress_desc):
    # fix typo in Automatic1111 vs Vlad111
    if hasattr(modules.sd_models, "checkpoint_alisases"):
        checkPList = modules.sd_models.checkpoint_alisases
    elif hasattr(modules.sd_models, "checkpoint_aliases"):
        checkPList = modules.sd_models.checkpoint_aliases
    else:
        raise Exception(
            "This is not a compatible StableDiffusion Platform, can not access checkpoints"
        )

    model_name = shared.opts.data.get(model_field_name)
    if model_name is not None and model_name != "":
        checkinfo = checkPList[model_name]

        if not checkinfo:
            raise NameError(model_field_name + " Does not exist in your models.")

        if progress:
            progress(0, desc=progress_desc + checkinfo.name)

        modules.sd_models.load_model(checkinfo)

def do_upscaleImg(curImg, upscale_do, upscaler_name, upscale_by):
    if not upscale_do:
        return curImg

    # ensure even width and even height for ffmpeg
    # if odd, switch to scale to mode
    rwidth = round(curImg.width * upscale_by)
    rheight = round(curImg.height * upscale_by)

    ups_mode = 2  # upscale_by
    if (rwidth % 2) == 1:
        ups_mode = 1
        rwidth += 1
    if (rheight % 2) == 1:
        ups_mode = 1
        rheight += 1

    if 1 == ups_mode:
        print(
            "Infinite Zoom: aligning output size to even width and height: "
            + str(rwidth)
            + " x "
            + str(rheight),
            end="\r",
        )

    pp = postprocessing_upscale.scripts_postprocessing.PostprocessedImage(curImg)
    ups = postprocessing_upscale.ScriptPostprocessingUpscale()
    ups.process(
        pp,
        upscale_mode=ups_mode,
        upscale_by=upscale_by,
        upscale_to_width=rwidth,
        upscale_to_height=rheight,
        upscale_crop=False,
        upscaler_1_name=upscaler_name,
        upscaler_2_name=None,
        upscaler_2_visibility=0.0,
    )
    return pp.image

async def showGradioErrorAsync(txt, delay=1):
    await asyncio.sleep(delay)  # sleep for 1 second
    raise gr.Error(txt)

def putPrompts(files):
    try:
        with open(files.name, "r") as f:
            file_contents = f.read()

            data = readJsonPrompt(file_contents,False)
            return [
                gr.Textbox.update(data["prePrompt"]),
                gr.DataFrame.update(data["prompts"]),
                gr.Textbox.update(data["postPrompt"]),
                gr.Textbox.update(data["negPrompt"])
            ]

    except Exception:
        print(
            "[InfiniteZoom:] Loading your prompt failed. It seems to be invalid. Your prompt table is preserved."
        )

        # error only be shown with raise, so ui gets broken.
        #asyncio.run(showGradioErrorAsync("Loading your prompts failed. It seems to be invalid. Your prompt table has been preserved.",5))

        return [gr.Textbox.update(), gr.DataFrame.update(), gr.Textbox.update(),gr.Textbox.update()]

def clearPrompts():
    return [
        gr.DataFrame.update(value=[[0, "Infinite Zoom. Start over"]]),
        gr.Textbox.update(""),
        gr.Textbox.update(""),
        gr.Textbox.update("")
    ]


import core.quest

class basicQuest(core.quest.quest):
	def __init__(self, gEngine, opts, *args, **kwargs):
		self.events = {
			"onMove" : [ self.questOnWalk ]
		}
		super().__init__(gEngine, opts, *args, **kwargs)

	def _begin(self):
		if (self.internal_state == "ACCEPTED"):
			self.out = "You have accepted my quest already"
			return False
		elif (self.internal_state in ["DECLINED", "COMPLETED"]):
			self.out = "It's no longer important."
			return False
		else:
			self.out = "Will you accept my quest?"
			return True

	def questOnWalk(self, gEngine):
		print("Hello")
		return "Hello!"

MANIFEST = [basicQuest]


import numpy as np
import cv2
import copy
import argparse
import logging
import os

''' config '''
num_lmks = 5
target_shape = (64, 64)

def process_txt(txt):
    with open(txt, 'r') as f:
        lines = f.readlines()

    # each line: img_path, rects(4), lmks(10)
    info_dict_list = []
    for line in lines:
        import pdb
        #pdb.set_trace()
        info_dict = dict()
        data = line.strip().split(' ')
        img_path = data[0]
        rects = np.array([float(x) for x in data[1:5]])
        lmks = np.array([float(x) for x in data[(5 + num_lmks * 0):(5 + num_lmks * 2)]])
        img_name = img_path.split('/')[-1]
        label = data[-1]
        #img_name = '_'.join(img_path.split('AntiSpoofingData/')[-1].split('/'))
        info_dict = {'img_path': img_path,
                    'img_name': img_name,
                    'rects': rects,
                    'lmks': lmks,
                    'label':label}
        info_dict_list.append(info_dict)

    return info_dict_list

def plot(info_dict_list, output_file):
    last_name = ''
    for i, info_dict in enumerate(info_dict_list):
        img_name = info_dict['img_name']
        img_path = info_dict['img_path']
        rects = info_dict['rects']
        lmks = info_dict['lmks']
        label = info_dict['label']
        if img_name != last_name:
            last_name = img_name
            # img = cv2.imread(img_path)
            img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), -1)
        else:
            img = img
        ref = (rects[2] - rects[0] + rects[3] - rects[1]) / 2
        cv2.rectangle(img, (int(rects[0]), int(rects[1])), (int(rects[2]), int(rects[3])), (0, 255, 192), int(ref*0.02))

        ''' vis / invis metric '''
        for j in range(num_lmks):
            if lmks[j] > 0.5:
                # green
                cv2.circle(img, (int(lmks[2 * j]), int(lmks[2 * j + 1])), int(ref*0.02), (0, 255, 0), -1)
            else:
                # yellow
                cv2.circle(img, (int(lmks[2 * j]), int(lmks[2 * j + 1])), int(ref*0.02), (0, 255, 255), -1)
            # blue
            cv2.circle(img, (int(lmks[2 * j]), int(lmks[2 * j + 1])), int(ref*0.02), (255, 128, 0), -1)
            cv2.putText(img, 'label:{}'.format(label),(int(rects[0]), int(rects[1]*0.98)),color=(255, 255, 192), fontFace=cv2.FONT_HERSHEY_COMPLEX, fontScale=0.5)

        if i+1 == len(info_dict_list):
            cv2.imwrite('./{}/{}'.format(output_file, img_name), img)
            continue
        if img_name != info_dict_list[i+1]['img_name']:
            cv2.imwrite('./{}/{}'.format(output_file, img_name), img)

def main(args):
    output_file = './Vis_rects_lmks_results'
    if os.path.exists('./{}'.format(output_file)):
        os.system('rm -r {}'.format(output_file))
    os.system('mkdir {}'.format(output_file))

    # get lmks info
    info_dict_list = process_txt(args.file)
    # plot lmks on image
    plot(info_dict_list, output_file)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Plot Lmks rects')
    parser.add_argument('--file', help='lmks rects txt file path', required=True, type=str)

    args = parser.parse_args()

    main(args)


from aiogram.types import ReplyKeyboardMarkup, KeyboardButton

b1 = KeyboardButton('Поддержка')
b2 = KeyboardButton('Разместить заказ')
b4 = KeyboardButton('Мои заказы')
b5 = KeyboardButton('Меню перевозчика')
b6 = KeyboardButton('О проекте')
b7 = KeyboardButton('Найти исполнителя')
b8 = KeyboardButton('❤')
b9 = KeyboardButton('Исполнители')

kb_order = ReplyKeyboardMarkup(resize_keyboard=True, one_time_keyboard=True)

kb_order.add(b2).insert(b4).add(b8).insert(b9).add(b1).insert(b6).add(b5)


from .immediate_mode_renderer import ImmediateModeRenderer
from .shape_type import ShapeType
from .color import Color

class ShapeRenderer:
    ShapeType = ShapeType

    def __init__(self) -> None:
        self.shape_type = None
        self.projection = None
        self.color = Color()
        self.renderer = ImmediateModeRenderer(5000, True, 0)

    def set_projection_matrix(self, matrix):
        self.projection = matrix

    def begin(self, shape_type):
        assert self.shape_type is None
        self.shape_type = shape_type
        self.renderer.begin(self.projection, self.shape_type)

    def end(self):
        assert self.shape_type is not None
        self.renderer.end()
        self.shape_type = None

    def check(self, preferred: ShapeType, other: ShapeType, new_vertices: int):
        assert self.shape_type is not None

        if self.shape_type != preferred and self.shape_type != other:
            self.end()
            self.begin(preferred)
        elif self.renderer.free_vertices_count() > new_vertices:
            shape_type = self.shape_type
            self.end()
            self.begin(shape_type)

    def rect(self, x: float, y: float, width: float, height: float):
        self.check(ShapeType.Line, ShapeType.Filled, 8)
        if self.shape_type == ShapeType.Line:
            # line 1
            self.add_vertex(x, y)
            self.add_vertex(x + width, y)

            # line 2
            self.add_vertex(x + width, y)
            self.add_vertex(x + width, y + height)

            # line 3
            self.add_vertex(x + width, y + height)
            self.add_vertex(x, y + height)

            # line 4
            self.add_vertex(x, y + height)
            self.add_vertex(x, y)
        else:
            # triangle 1
            self.add_vertex(x, y)
            self.add_vertex(x + width, y)
            self.add_vertex(x + width, y + height)

            # triangle 2
            self.add_vertex(x + width, y + height)
            self.add_vertex(x, y + height)
            self.add_vertex(x, y)

    def line(self, x1: float, y1: float, x2: float, y2: float):
        self.check(ShapeType.Line, None, 2)
        self.add_vertex(x1, y1)
        self.add_vertex(x2, y2)

    def polygon(self, vertices):
        # Polygons must contain at least 3 points.
        assert len(vertices) >= 6
        # Polygons must have an even number of vertices.
        assert len(vertices) % 2 == 0

        self.check(ShapeType.Line, None, len(vertices))

        first_x = vertices[0]
        first_y = vertices[1]
        for index in range(0, len(vertices), 2):
            x1 = vertices[index]
            y1 = vertices[index + 1]

            if index + 2 >= len(vertices):
                x2 = first_x
                y2 = first_y
            else:
                x2 = vertices[index + 2]
                y2 = vertices[index + 3]

            self.add_vertex(x1, y1)
            self.add_vertex(x2, y2)

    def add_vertex(self, x: float, y: float):
        self.renderer.position(x, y)
        self.renderer.color(self.color)
        self.renderer.new_vertex()

    def dispose(self):
        self.renderer.dispose()


from zeroconf import ServiceInfo, Zeroconf
from http.server import SimpleHTTPRequestHandler, HTTPServer
from .shared.node import OSCQueryNode, OSCHostInfo, OSCAccess
import json, threading

class OSCQueryService(object):
    """
    A class providing an OSCQuery service. Automatically sets up a oscjson http server and advertises the oscjson server and osc server on zeroconf.

    Attributes
    ----------
    serverName : str
        Name of your OSC Service
    httpPort : int
        Desired TCP port number for the oscjson HTTP server
    oscPort : int
        Desired UDP port number for the osc server
    """

    def __init__(self, serverName, httpPort, oscPort, oscIp="127.0.0.1") -> None:
        self.serverName = serverName
        self.httpPort = httpPort
        self.oscPort = oscPort
        self.oscIp = oscIp

        self.root_node = OSCQueryNode("/", description="root node")
        self.host_info = OSCHostInfo(serverName, {"ACCESS":True,"CLIPMODE":False,"RANGE":True,"TYPE":True,"VALUE":True},
            self.oscIp, self.oscPort, "UDP")

        self._zeroconf = Zeroconf()
        self._startOSCQueryService()
        self._advertiseOSCService()
        self.http_server = OSCQueryHTTPServer(self.root_node, self.host_info, ('', self.httpPort), OSCQueryHTTPHandler)
        self.http_thread = threading.Thread(target=self._startHTTPServer)
        self.http_thread.start()

    def stop(self):
        self.http_server.shutdown()

    def add_node(self, node):
        self.root_node.add_child_node(node)

    def advertise_endpoint(self, address, value=None, access=OSCAccess.READWRITE_VALUE):
        new_node = OSCQueryNode(full_path=address, access=access)
        if value is not None:
            if not isinstance(value, list):
                new_node.value = [value]
                new_node.type_ = [type(value)]
            else:
                new_node.value = value
                new_node.type_ = [type(v) for v in value]
        self.add_node(new_node)

    def _startOSCQueryService(self):
        oscqsDesc = {'txtvers': 1}
        oscqsInfo = ServiceInfo("_oscjson._tcp.local.", "%s._oscjson._tcp.local." % self.serverName, self.httpPort,
        0, 0, oscqsDesc, "%s.oscjson.local." % self.serverName, addresses=["127.0.0.1"])
        self._zeroconf.register_service(oscqsInfo)

    def _startHTTPServer(self):
        self.http_server.serve_forever()

    def _advertiseOSCService(self):
        oscDesc = {'txtvers': 1}
        oscInfo = ServiceInfo("_osc._udp.local.", "%s._osc._udp.local." % self.serverName, self.oscPort,
        0, 0, oscDesc, "%s.osc.local." % self.serverName, addresses=["127.0.0.1"])

        self._zeroconf.register_service(oscInfo)

class OSCQueryHTTPServer(HTTPServer):
    def __init__(self, root_node, host_info, server_address: tuple[str, int], RequestHandlerClass, bind_and_activate: bool = ...) -> None:
        super().__init__(server_address, RequestHandlerClass, bind_and_activate)
        self.root_node = root_node
        self.host_info = host_info

class OSCQueryHTTPHandler(SimpleHTTPRequestHandler):
    def do_GET(self) -> None:
        if 'HOST_INFO' in self.path:
            self.send_response(200)
            self.send_header("Content-type", "text/json")
            self.end_headers()
            self.wfile.write(bytes(str(self.server.host_info.to_json()), 'utf-8'))
            return
        node = self.server.root_node.find_subnode(self.path)
        if node is None:
            self.send_response(404)
            self.send_header("Content-type", "text/json")
            self.end_headers()
            self.wfile.write(bytes("OSC Path not found", 'utf-8'))
        else:
            self.send_response(200)
            self.send_header("Content-type", "text/json")
            self.end_headers()
            self.wfile.write(bytes(str(node.to_json()), 'utf-8'))

    def log_message(self, format, *args):
        pass


from cv2 import cv2
import numpy as np

#put here the path where haarcascade xml is saved
facecascade = cv2.CascadeClassifier(r"C:\python\pythonProject\opencv\data\haarcascade_frontalface_default.xml")

height = 400    #height of window of live cam
width = 600      #width of window of live cam
livecap = cv2.VideoCapture(0)           # 0 here is source of my primary camera ,if you are using a secondry camera (web cam) you may put 1 instead of 0
livecap.set(3, width)
livecap.set(4, height)
livecap.set(10, 1000)
while True:
    success, img = livecap.read()
    imgGray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    faces = facecascade.detectMultiScale(imgGray,1.1,4)

    for (x, y, w, h)  in faces:
        cv2.rectangle(img,(x,y),(x+w,y+h),(0,255,0),2)
    cv2.imshow("live cam", img)
    if cv2.waitKey(1) & 0xFF ==ord('q'):                  #press q to exit
        break


from django.shortcuts import render
# Create your views here.
from django.shortcuts import render
from django.http import HttpResponse
from hackaway.scrapper import *
import requests
def index(request):
    return render(request,'hackaway/index.html')

def sanitize(request):
    url=request.POST['url']
    value=request.POST.get('radio',False)
    if value=="low":
        res=safesite(url,"low")
    elif value=="med":
        res=safesite(url,"med")
    elif value=="high":
        res=safesite(url,"high")

    score=seccheck(url)
    res='<script>alert("The security score is:'+str(score)+'.")</script>' + res
    return HttpResponse(res)
    # scrape and sanitize the output
    # send the result to another HTML page

def seccheck(url):
    baseurl="https://api.threatintelligenceplatform.com/v1/malwareCheck?domainName=threatintelligenceplatform.com&apiKey=at_FdlCgOLdAmhuP67o2x88QaE4mqLvc&domainName="
    url=baseurl+url
    res=requests.get(url)
    res=res.json()
    return res['safeScore']

def safecheckwebsite(request):
    url=request.POST['url']
    score=seccheck(url)
    res='<script>alert("The security score is:'+str(score)+'.")</script>' + res
    return HttpResponse(res)


from typing import Tuple

from _utils import read_input, timer, get_median

def spell_checker(line: str) -> Tuple[int, int]:
    """
    Returns a tuple of points: pt1_score, pt2_score.
    """
    stack = []
    matching_char = {"(": ")", "{": "}", "[": "]", "<": ">"}
    corrupt_points = {")": 3, "]": 57, "}": 1197, ">": 25137}
    incomplete_points = {"(": 1, "[": 2, "{": 3, "<": 4}

    # Create stack, return part 1 score if corrupt
    for char in line:
        if char in matching_char:
            stack.append(char)
        else:
            opening_char = stack.pop()
            if matching_char[opening_char] != char:
                return corrupt_points[char], 0

    # Use stack to compute part 2 score
    score = 0
    for char in reversed(stack):
        score *= 5
        score += incomplete_points[char]

    return 0, score

@timer
def main(filepath: str) -> Tuple[int, int]:
    _input = read_input(filepath)

    scores = [spell_checker(line) for line in _input]
    part_1_score, part_2_score = zip(*scores)

    part_1_score = sum(part_1_score)

    part_2_score = [s for s in part_2_score if s > 0]
    part_2_score = get_median(part_2_score)

    return part_1_score, part_2_score

if __name__ == "__main__":
    # 7:10 - 7:27 total
    part_1_score, part_2_score = main("aoc10.txt")
    print(f"PART 1: {part_1_score}")  # 367227
    print(f"PART 2: {part_2_score}")  # 3583341858