import re, os
from Tester import Tester
from TestHarness import util

class RunApp(Tester):

    @staticmethod
    def validParams():
        params = Tester.validParams()
        params.addParam('input',              "The input file to use for this test.")
        params.addParam('test_name',          "The name of the test - populated automatically")
        params.addParam('input_switch', '-i', "The default switch used for indicating an input to the executable")
        params.addParam('errors',             ['ERROR', 'command not found', 'erminate called after throwing an instance of'], "The error messages to detect a failed run")
        params.addParam('expect_out',         "A regular expression that must occur in the input in order for the test to be considered passing.")
        params.addParam('match_literal', False, "Treat expect_out as a string not a regular expression.")
        params.addParam('absent_out',         "A regular expression that must be *absent* from the output for the test to pass.")
        params.addParam('should_crash', False, "Inidicates that the test is expected to crash or otherwise terminate early")
        params.addParam('executable_pattern', "A test that only runs if the exectuable name matches the given pattern")
        params.addParam('delete_output_before_running',  True, "Delete pre-existing output files before running test. Only set to False if you know what you're doing!")
        params.addParam('delete_output_folders', True, "Delete output folders before running")

        # RunApp can also run arbitrary commands. If the "command" parameter is supplied
        # it'll be used in lieu of building up the command automatically
        params.addParam('command',            "The command line to execute for this test.")

        # Parallel/Thread testing
        params.addParam('max_parallel', 1000, "Maximum number of MPI processes this test can be run with      (Default: 1000)")
        params.addParam('min_parallel',    1, "Minimum number of MPI processes that this test can be run with (Default: 1)")
        params.addParam('max_threads',    16, "Max number of threads (Default: 16)")
        params.addParam('min_threads',     1, "Min number of threads (Default: 1)")
        params.addParam('redirect_output',  False, "Redirect stdout to files. Neccessary when expecting an error when using parallel options")

        params.addParam('allow_warnings',   True, "Whether or not warnings are allowed.  If this is False then a warning will be treated as an error.  Can be globally overridden by setting 'allow_warnings = False' in the testroot file.");
        params.addParam('allow_unused',   True, "Whether or not unused parameters are allowed in the input file.  Can be globally overridden by setting 'allow_unused = False' in the testroot file.");
        params.addParam('allow_override', True, "Whether or not overriding a parameter/block in the input file generates an error.  Can be globally overridden by setting 'allow_override = False' in the testroot file.");
        params.addParam('allow_deprecated', True, "Whether or not deprecated warnings are allowed.  Setting to False will cause deprecation warnings to be treated as test failures.  We do NOT recommend you globally set this permanently to False!  Deprecations are a part of the normal development flow and _SHOULD_ be allowed!")

        # Valgrind
        params.addParam('valgrind', 'NORMAL', "Set to (NONE, NORMAL, HEAVY) to determine which configurations where valgrind will run.")

        return params

    def __init__(self, name, params):
        Tester.__init__(self, name, params)
        if os.environ.has_key("MOOSE_MPI_COMMAND"):
            self.mpi_command = os.environ['MOOSE_MPI_COMMAND']
            self.force_mpi = True
        else:
            self.mpi_command = 'mpiexec'
            self.force_mpi = False

        # Make sure that either input or command is supplied
        if not (params.isValid('input') or params.isValid('command')):
            raise Exception('Either "input" or "command" must be supplied for a RunApp test')

    def getInputFile(self):
        if self.specs.isValid('input'):
            return self.specs['input'].strip()
        else:
            return None # Not all testers that inherit from RunApp have an input file

    def checkRunnable(self, options):
        if options.enable_recover:
            if self.specs.isValid('expect_out') or self.specs.isValid('absent_out') or self.specs['should_crash'] == True:
                self.addCaveats('expect_out RECOVER')
                self.setStatus(self.bucket_skip.status, self.bucket_skip)
                return False

        if self.specs.isValid('executable_pattern') and re.search(self.specs['executable_pattern'], self.specs['executable']) == None:
            self.addCaveats('EXECUTABLE PATTERN')
            self.setStatus(self.bucket_skip.status, self.bucket_skip)
            return False

        return True

    def getThreads(self, options):
        #Set number of threads to be used lower bound
        nthreads = max(options.nthreads, int(self.specs['min_threads']))
        #Set number of threads to be used upper bound
        nthreads = min(nthreads, int(self.specs['max_threads']))
        return nthreads

    def getProcs(self, options):
        if options.parallel == None:
            default_ncpus = 1
        else:
            default_ncpus = options.parallel

        # Raise the floor
        ncpus = max(default_ncpus, int(self.specs['min_parallel']))
        # Lower the ceiling
        ncpus = min(ncpus, int(self.specs['max_parallel']))
        return ncpus

    def getCommand(self, options):
        specs = self.specs

        # Just return an arbitrary command if one is supplied
        if specs.isValid('command'):
            return os.path.join(specs['test_dir'], specs['command']) + ' ' + ' '.join(specs['cli_args'])

        # Create the additional command line arguments list
        cli_args = list(specs['cli_args'])

        # Check for built application
        if not options.dry_run and not os.path.exists(specs['executable']):
            self.setStatus('Application not found', self.bucket_fail)
            return ''

        if (options.parallel_mesh or options.distributed_mesh) and ('--parallel-mesh' not in cli_args or '--distributed-mesh' not in cli_args):
            # The user has passed the parallel-mesh option to the test harness
            # and it is NOT supplied already in the cli-args option
            cli_args.append('--distributed-mesh')

        if '--error' not in cli_args and (not specs["allow_warnings"] or options.error):
            cli_args.append('--error')

        if '--error-unused' not in cli_args and (not specs["allow_unused"] or options.error_unused):
            cli_args.append('--error-unused')

        if '--error-override' not in cli_args and not specs["allow_override"]:
            cli_args.append('--error-override')

        if '--error-deprecated' not in cli_args and (not specs["allow_deprecated"] or options.error_deprecated):
            cli_args.append('--error-deprecated')

        if self.getCheckInput():
            cli_args.append('--check-input')

        timing_string = ' '
        if options.timing and specs["timing"]:
            cli_args.append('--timing')
            cli_args.append('Outputs/print_perf_log=true')

        if options.colored == False:
            cli_args.append('--color off')

        if options.cli_args:
            cli_args.insert(0, options.cli_args)

        if options.scaling and specs['scale_refine'] > 0:
            cli_args.insert(0, ' -r ' + str(specs['scale_refine']))

        # The test harness should never use GDB backtraces: they don't
        # work well when dozens of expect_err jobs run at the same time.
        cli_args.append('--no-gdb-backtrace')

        # Get the number of processors and threads the Tester requires
        ncpus = self.getProcs(options)
        nthreads = self.getThreads(options)

        if options.parallel == None:
            default_ncpus = 1
        else:
            default_ncpus = options.parallel

        if specs['redirect_output'] and ncpus > 1:
            cli_args.append('--keep-cout --redirect-output ' + self.name())

        caveats = []
        if nthreads > options.nthreads:
            caveats.append('min_threads=' + str(nthreads))
        elif nthreads < options.nthreads:
            caveats.append('max_threads=' + str(nthreads))
        # TODO: Refactor this caveats business
        if ncpus > default_ncpus:
            caveats.append('min_cpus=' + str(ncpus))
        elif ncpus < default_ncpus:
            caveats.append('max_cpus=' + str(ncpus))

        self.addCaveats(*caveats)

        if self.force_mpi or options.parallel or ncpus > 1 or nthreads > 1:
            command = self.mpi_command + ' -n ' + str(ncpus) + ' ' + specs['executable'] + ' --n-threads=' + str(nthreads) + ' ' + specs['input_switch'] + ' ' + specs['input'] + ' ' +  ' '.join(cli_args)
        elif options.valgrind_mode.upper() == specs['valgrind'].upper() or options.valgrind_mode.upper() == 'HEAVY' and specs['valgrind'].upper() == 'NORMAL':
            command = 'valgrind --suppressions=' + os.path.join(specs['moose_dir'], 'python', 'TestHarness', 'suppressions', 'errors.supp') + ' --leak-check=full --tool=memcheck --dsymutil=yes --track-origins=yes --demangle=yes -v ' + specs['executable'] + ' ' + specs['input_switch'] + ' ' + specs['input'] + ' ' + ' '.join(cli_args)
        else:
            command = specs['executable'] + timing_string + specs['input_switch'] + ' ' + specs['input'] + ' ' + ' '.join(cli_args)

        return command

    def testFileOutput(self, moose_dir, options, output):
        """ Set a failure status for expressions found in output """
        reason = ''
        specs = self.specs
        if specs.isValid('expect_out'):
            mode = ""
            if specs['match_literal']:
                have_expected_out = util.checkOutputForLiteral(output, specs['expect_out'])
                mode = 'literal'
            else:
                have_expected_out = util.checkOutputForPattern(output, specs['expect_out'])
                mode = 'pattern'

            if (not have_expected_out):
                reason = 'EXPECTED OUTPUT MISSING'
                output += "#"*80 + "\n\nUnable to match the following " + mode + " against the program's output:\n\n" + specs['expect_out'] + "\n"

        if reason == '' and specs.isValid('absent_out'):
            have_absent_out = util.checkOutputForPattern(output, specs['absent_out'])
            if (have_absent_out):
                reason = 'OUTPUT NOT ABSENT'
                output += "#"*80 + "\n\nMatched the following pattern, which we did NOT expect:\n\n" + specs['absent_out'] + "\n"

        if reason == '':
            # We won't pay attention to the ERROR strings if EXPECT_ERR is set (from the derived class)
            # since a message to standard error might actually be a real error.  This case should be handled
            # in the derived class.
            if options.valgrind_mode == '' and not specs.isValid('expect_err') and len( filter( lambda x: x in output, specs['errors'] ) ) > 0:
                reason = 'ERRMSG'
            elif self.exit_code == 0 and specs['should_crash'] == True:
                reason = 'NO CRASH'
            elif self.exit_code != 0 and specs['should_crash'] == False:
                reason = 'CRASH'
            # Valgrind runs
            elif self.exit_code == 0 and self.shouldExecute() and options.valgrind_mode != '' and 'ERROR SUMMARY: 0 errors' not in output:
                reason = 'MEMORY ERROR'

        if reason != '':
            self.setStatus(reason, self.bucket_fail)

        return reason

    def processResults(self, moose_dir, options, output):
        """
        Wrapper method for testFileOutput.

        testFileOutput does not set a success status, while processResults does.
        For testers that are RunApp types, they would call this method. While
        all other tester types (like exodiff) will call testFileOutput. This is
        to prevent derived testers from having a successful status set, before
        actually running their derived processResults method.
        """
        reason = self.testFileOutput(moose_dir, options, output)

        # Populate the bucket
        if reason != '':
            self.setStatus(reason, self.bucket_fail)
        else:
            self.setStatus(self.success_message, self.bucket_success)

        return output


import pygame, os, pickle, io, time, traceback
import Core
from Core import CntMng
from Core import MAIN
from Core import AppData
import Core as tge
from OneTrack.MAIN.Screens import Editor
from OneTrack.MAIN import LagIndicator
from OneTrack.MAIN import UI
from OneTrack.MAIN.Screens.Editor import InstanceVar as var

class Process():
    def __init__(self, pPID, pProcessName, pROOT_MODULE, pInitArgs):
        self.PID = pPID
        self.NAME = pProcessName
        self.ROOT_MODULE = pROOT_MODULE
        self.IS_GRAPHICAL = True
        self.INIT_ARGS = pInitArgs
        self.DISPLAY = pygame.Surface((Core.MAIN.ScreenWidth, Core.MAIN.ScreenHeight))
        self.LAST_SURFACE = self.DISPLAY.copy()
        self.APPLICATION_HAS_FOCUS = True
        self.POSITION = (0, 0)
        self.FULLSCREEN = True
        self.TITLEBAR_RECTANGLE = pygame.Rect(self.POSITION[0], self.POSITION[1], self.DISPLAY.get_width(), 15)
        self.TITLEBAR_TEXT = "OneTrack"
        self.WindowDragEnable = False
        self.DeleteInstanceOnFirstCycle = False
        self.DialogPID = -1
        self.DefaultContents = CntMng.ContentManager()

    def CheckForAnotherInstances(self):
        # Check if there is not another instance of OneTrack
        for process in Core.MAIN.ProcessList:
            if process.PID != self.PID:
                if process.IS_GRAPHICAL and process.FULLSCREEN == self.FULLSCREEN:
                    if process.EXECUTABLE_PATH == self.EXECUTABLE_PATH:
                        # Another instance detected
                        self.DialogPID = self.GreyDialog("Multiple instance detection", "OneTrack does not support multiple instances.\n\nOnly 1 instance of OneTrack is allowed.", "warn")
                        self.DeleteInstanceOnFirstCycle = True
                        return True
        return False

    def Initialize(self):
        var.ProcessReference = self

        print("Initializing {0}...".format(self.TITLEBAR_TEXT))

        # Initialize Variables
        self.CurrentScreenToUpdate = Editor

        # Initialize Content Manager
        self.DefaultContents.SetSourceFolder("OneTrack/")
        self.DefaultContents.SetFontPath("Data/fonts")
        self.DefaultContents.SetImageFolder("Data/img")
        self.DefaultContents.SetRegKeysPath("Data/reg")

        self.DefaultContents.LoadRegKeysInFolder()
        self.DefaultContents.LoadImagesInFolder()

        self.DefaultContents.InitSoundSystem()

        if self.CheckForAnotherInstances():
            return

        # Set the default content manager for the UI
        UI.ContentManager = self.DefaultContents
        var.DefaultContent = self.DefaultContents

        # Load UI Theme
        try:
            UI.ThemesManager_LoadTheme(self.DefaultContents.Get_RegKey("/selected_theme"))
        except:
            print("Error while loading selected theme\nUsing fallback theme")
            self.DefaultContents.Write_RegKey("/selected_theme", "default")

            UI.ThemesManager_LoadTheme(self.DefaultContents.Get_RegKey("/selected_theme"))

        MAIN.ReceiveCommand(0, 60)

        self.TITLEBAR_TEXT = "OneTrack v{0}".format(self.DefaultContents.Get_RegKey("/version"))

        var.LoadDefaultValues()
        Editor.Initialize()
        LagIndicator.Initialize()

    def GreyDialog(self, Title, Text, Icon="none"):
        var.AwaysUpdate = False
        return Core.MAIN.CreateProcess("OneTrack/UnatachedDialog", "OneTrack Dialog", (var.ProcessReference, "DIALOG_OK", "{0};{1}".format(Title, Text), "icon:{0},".format(Icon)))

    def Draw(self):
        self.CurrentScreenToUpdate.GameDraw(self.DISPLAY)

        LagIndicator.Draw(self.DISPLAY)

        self.LAST_SURFACE = self.DISPLAY.copy()
        return self.DISPLAY

    def Update(self):
        if self.DeleteInstanceOnFirstCycle and not self.DialogPID in Core.MAIN.ProcessList_PID:
            Core.MAIN.KillProcessByPID(self.PID)

        if not self.APPLICATION_HAS_FOCUS and var.AwaysUpdate is False:
            return

        self.CurrentScreenToUpdate.Update()

        LagIndicator.Update()

    def EventUpdate(self, event):
        self.CurrentScreenToUpdate.EventUpdate(event)

        if event.type == pygame.KEYUP:
            if event.key == 1073741895:
                print("Theme & RegKeys reloaded")
                UI.ContentManager.ReloadRegistry()
                UI.ThemesManager_LoadTheme(self.DefaultContents.Get_RegKey("/selected_theme"))


from ecluster.node import Node

class NodeList():
	def __init__(self, name, provider):
		self.name = name
		self.nodes = []
		self.provider = provider

	# Give next index number for type
	def num_type(self, type, deleted=False):
		num = 1

		for node in self.nodes:
			if deleted:
				# Count only deleted nodes
				if node.type == type and node.deleted:
					num += 1
			else:
				if node.type == type and not node.deleted:
					num += 1

		return num

	def add_node(self, type, deleted=False, nodeid=None, instance=None):
		if not nodeid:
			nodeid = self.num_type(type)

		node = Node(type, nodeid, self.provider, self.name)

		# We have some nodes to be deleted
		if self.num_type(type, deleted=True) > 1:
			node.replace = True

		if instance:
			node.instance = instance
		else:
			node.locate_in_cloud()

		node.deleted = deleted

		self.nodes += [node]

	def get_node_names(self):
		ret = []
		for node in self.nodes:
			ret += [str(node.name)]
		return ret

	def get_instance_ids(self):
		ret = []
		for node in self.nodes:
			if node.instance:
				ret += [str(node.instance.id)]
		return ret

	def needsync(self):
		for node in self.nodes:
			if node.changed:
				return True
			if node.deleted:
				return True
			if not node.created:
				return True
		return False

	def instance_in_list(self, instance):
		local_ids = self.get_instance_ids()
		if not instance.id in local_ids:
			return False

		return True


import unittest
from nalaf.structures.data import Dataset, Document, Part
from nalaf.preprocessing.spliters import NLTKSplitter

class TestNLTKSplitter(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.dataset = Dataset()
        doc = Document()
        part = Part('This is one sentence. This is another one.\n This is the third one; here continues.')
        cls.dataset.documents['doc_1'] = doc
        doc.parts['part_1'] = part

    def test_split(self):
        NLTKSplitter().split(self.dataset)

        sentences_ = []
        for document in self.dataset.documents.values():
            for part in document.parts.values():
                sentences_ += part.sentences_

        expected = ['This is one sentence.', 'This is another one.', 'This is the third one; here continues.']

        self.assertEqual(sentences_, expected)

if __name__ == '__main__':
    unittest.main()


from django.test import TestCase
from django.urls import resolve

from lists.views import home_page
from lists.models import Item, List

class ListTestCase(TestCase):
    def test_root_url_resolves_to_home_page_view(self):
        found = resolve('/')
        self.assertEqual(found.func, home_page)

    def test_home_page_returns_correct_html(self):
        response = self.client.get('/')
        self.assertTemplateUsed(response, 'home.html')

    def test_only_saves_items_when_necessary(self):
        self.client.get('/')
        self.assertEqual(Item.objects.count(), 0)

class ItemModelTest(TestCase):
    def test_saving_and_retrieving_items(self):
        first_item = Item(text='Item 1 da lista', list=List.objects.create())
        first_item.save()
        item_id = first_item.id
        item = Item.objects.get(id=item_id)
        self.assertIn(first_item.text, item.text)

    def test_list_items(self):
        Item(text='item1', list=List.objects.create()).save()
        Item(text='item2', list=List.objects.create()).save()
        self.assertEqual(Item.objects.count(), 2)

class ListViewTest(TestCase):

    def test_display_only_items_for_that_list(self):
        correct_list = List.objects.create()
        Item.objects.create(text='itemey 1', list=correct_list)
        Item.objects.create(text='itemey 2', list=correct_list)

        other_list = List.objects.create()
        Item.objects.create(text='other itemey 1', list=other_list)
        Item.objects.create(text='other itemey 2', list=other_list)

        response = self.client.get('/lists/{}/'.format(correct_list.id))

        self.assertContains(response, 'itemey 1')
        self.assertContains(response, 'itemey 2')
        self.assertNotContains(response, 'other itemey 1')
        self.assertNotContains(response, 'other itemey 2')

    def test_uses_list_template(self):
        list_ = List.objects.create()
        response = self.client.get('/lists/{}/'.format(list_.id))
        self.assertTemplateUsed(response, 'list.html')

class NewListTest(TestCase):
    def test_can_save_a_POST_request(self):
        self.client.post('/lists/new', data={'item_text': 'A new list item'})
        self.assertEqual(Item.objects.count(), 1)
        new_item = Item.objects.first()
        self.assertEqual(new_item.text, 'A new list item')

    def test_redirects_after_POST(self):
        response = self.client.post('/lists/new', data={'item_text': 'A new list item'})
        new_list = List.objects.first()
        self.assertRedirects(response, '/lists/{}/'.format(new_list.id))

class ListAndItemModels(TestCase):
    def test_create_list(self):
        list_ = List()
        list_.save()
        saved_list = List.objects.first()
        self.assertEqual(saved_list, list_)

        first_item = Item(text='The first (ever) list item')
        first_item.list = list_
        first_item.save()

        second_item = Item(text='Item the second')
        second_item.list = list_
        second_item.save()

        saved_items = Item.objects.all()
        self.assertEqual(saved_items.count(), 2)

        self.assertEqual(Item.objects.filter(list=list_).count(), 2)

class NewItemTest(TestCase):

    def test_can_save_a_POST_request_to_an_existing_list(self):
        other_list = List.objects.create()
        correct_list = List.objects.create()

        request = self.client.post(
            '/lists/{}/add_item'.format(correct_list.id),
            data=dict(item_text='A new item for an existing list')
        )

        self.assertEqual(Item.objects.count(), 1)
        new_item = Item.objects.first()
        self.assertEqual(new_item.text, 'A new item for an existing list')
        self.assertEqual(new_item.list, correct_list)

    def test_redirect_to_list_view(self):
        other_list = List.objects.create()
        correct_list = List.objects.create()

        response = self.client.post(
            'lists/{}/add_items',
            data=dict(item_text='A new item for an existing list')
        )

        self.assertRedirects(response, 'lists/{}/'.format(correct_list.id))


import contextlib
import datetime as dt
import io
import json
import traceback
import urllib.parse

import flask.json
import wand.image
from flask import Blueprint, jsonify, current_app, request, stream_with_context
from werkzeug.exceptions import HTTPException

import acnh.dodo as dodo
import acnh.designs.api as designs_api
import acnh.designs.render as designs_render
import acnh.designs.db as designs_db
import utils
import tarfile_stream
from acnh.errors import (
	ACNHError,
	InvalidDesignCodeError,
	MissingLayerError,
	InvalidScaleFactorError,
	CannotScaleThumbnailError,
	InvalidImageError,
	InvalidImageIdError,
	InvalidImageArgument,
	InvalidProArgument,
	InvalidAuthorIdError,
	TiledImageTooBigError,
	InvalidPaginationError,
	InvalidPaginationLimitError,
)
from acnh.designs.db import PageSpecifier, PageDirection
from acnh.designs.encode import BasicDesign, Design
from utils import limiter

def init_app(app):
	app.register_blueprint(bp)

bp = Blueprint('api', __name__, url_prefix='/api/v0')

@bp.route('/host-session/<dodo_code>')
@limiter.limit('1 per 4 seconds')
def host_session(dodo_code):
	return dodo.search_dodo_code(dodo_code)

@bp.route('/design/<design_code>')
@limiter.limit('5 per second')
def design(design_code):
	InvalidDesignCodeError.validate(design_code)
	return designs_api.download_design(design_code)

def get_scale_factor():
	scale_factor = request.args.get('scale', '1')
	InvalidScaleFactorError.validate(scale_factor)
	return int(scale_factor)

def maybe_scale(image):
	scale_factor = get_scale_factor()
	if scale_factor == 1:
		return image

	return utils.xbrz_scale_wand_in_subprocess(image, scale_factor)

@bp.route('/design/<design_code>.tar')
@limiter.limit('2 per 10 seconds')
def design_archive(design_code):
	InvalidDesignCodeError.validate(design_code)
	render_internal = 'internal_layers' in request.args
	get_scale_factor()  # do the validation now since apparently it doesn't work in the generator
	data = designs_api.download_design(design_code)
	meta, body = data['mMeta'], data['mData']
	# pylint: disable=unused-variable
	type_code = meta['mMtUse']
	design_name = meta['mMtDNm']  # hungarian notation + camel case + abbreviations DO NOT mix well

	def gen():
		# pylint: disable=no-member  # pylint are you drunk?
		if type_code == BasicDesign.type_code or render_internal:
			layers = designs_render.render_layers(body)
		else:
			layers = Design.from_data(data).layer_images.items()

		yield from make_tar(design_name, data['updated_at'], layers)

	encoded_filename = urllib.parse.quote(design_name + '.tar')
	return current_app.response_class(
		stream_with_context(gen()),
		mimetype='application/x-tar',
		headers={'Content-Disposition': f"attachment; filename*=utf-8''{encoded_filename}"},
	)

def make_tar(design_name, updated_at, layers):
	tar = tarfile_stream.open(mode='w|')
	yield from tar.header()

	for name, image in layers:
		tarinfo = tarfile_stream.TarInfo(f'{design_name}/{name}.png')
		tarinfo.mtime = updated_at

		image = maybe_scale(image)
		out = io.BytesIO()
		with image.convert('png') as c:
			c.save(file=out)
		tarinfo.size = out.tell()
		out.seek(0)

		yield from tar.addfile(tarinfo, out)

	yield from tar.footer()

# no rate limit as we need to render the thumbnails for all of an author's designs quickly
@bp.route('/design/<design_code>/<layer>.png')
def design_layer(design_code, layer):
	InvalidDesignCodeError.validate(design_code)
	data = designs_api.download_design(design_code)
	meta, body = data['mMeta'], data['mData']
	design_name = meta['mMtDNm']

	if layer == 'thumbnail':
		if request.args.get('scale', '1') != '1':
			raise CannotScaleThumbnailError
		rendered = Design.from_data(data).net_image()
	else:
		try:
			int(layer)
		except ValueError:
			rendered = designs_render.render_layer_name(data, layer)
		else:
			rendered = designs_render.render_layer(body, layer)

	rendered = maybe_scale(rendered)
	out = rendered.make_blob('png')

	encoded_filename = urllib.parse.quote(f'{design_name}-{layer}.png')
	return current_app.response_class(out, mimetype='image/png', headers={
		'Content-Length': len(out),
		'Content-Disposition': f"inline; filename*=utf-8''{encoded_filename}"
	})

@bp.route('/designs/<author_id>')
@limiter.limit('5 per 1 seconds')
def list_designs(author_id):
	author_id = int(InvalidAuthorIdError.validate(author_id).replace('-', ''))
	pro = request.args.get('pro', 'false')
	InvalidProArgument.validate(pro)

	page = designs_api.list_designs(author_id, pro=pro)
	del page['offset'], page['count'], page['total']
	page['designs'] = page.pop('headers')
	page['creator_name'] = page['designs'][0]['design_player_name']
	page['author_id'] = page['designs'][0]['design_player_id']

	for d in page['designs']:
		d['design_code'] = designs_api.design_code(d['id'])
		del d['id']
		del d['design_player_name'], d['design_player_id'], d['digest']
		d['created_at'] = dt.datetime.utcfromtimestamp(d['created_at'])
		# designs cannot be updated, so why is this even here??
		del d['updated_at']

	return page

@bp.route('/images', methods=['POST'])
@limiter.limit('1 per 15s')
def create_image():
	gen = format_created_design_results(create_image_gen())
	# Note: this is currently the only method (other than the rendering methods) which does *not* return JSON.
	# This is due to its iterative nature. I considered using JSON anyway, but very few libraries
	# support iterative JSON decoding, and we don't need anything other than an array anyway.
	return current_app.response_class(stream_with_context(gen), mimetype='text/plain')

def create_image_gen():
	try:
		image_name = request.values['image_name']
	except KeyError:
		raise InvalidImageArgument('image_name')

	author_name = request.values.get('author_name') or 'Anonymous'  # we are legion

	design_type_name = request.values.get('design_type', 'basic-design')
	try:
		if design_type_name == 'basic-design':
			yield from create_basic_image(image_name, author_name)
		else:
			yield from create_pro_image(image_name, author_name, design_type_name)
	except ACNHError as ex:
		yield ex.to_dict()

def create_pro_image(image_name, author_name, design_type_name):
	try:
		layers = {
			filename: wand.image.Image(blob=file.read()).convert('PNG')
			for filename, file
			in request.files.items()
		}
	except wand.image.WandException:
		print('In:', request.path)
		traceback.print_exc()
		raise InvalidImageError

	design = Design(
		design_type_name,
		island_name=designs_db.island_name(),
		author_name=author_name,
		design_name=image_name,
		layers=layers,
	)

	with contextlib.ExitStack() as stack:
		for img in layers.values():
			stack.enter_context(img)
		yield from designs_db.create_image(design)

def create_basic_image(image_name, author_name):
	width = height = None

	def get_int_value(name):
		v = request.values.get(name)
		if not v:
			return None
		try:
			return int(v)
		except ValueError:
			raise InvalidImageArgument(name)

	try:
		width, height = map(int, request.values['resize'].split('x'))
	except KeyError:
		width, height = map(get_int_value, ('resize-width', 'resize-height'))
	except ValueError:
		raise InvalidImageArgument('resize')

	scale = 'scale' in request.values or request.values.get('mode') == 'scale'

	try:
		img = wand.image.Image(blob=request.files['0'].read()).convert('PNG')
	except wand.image.WandException as exc:
		print('In', request.path)
		traceback.print_exc()
		raise InvalidImageError
	except KeyError:
		# pylint: disable=no-member  # external_layers is defined dynamically
		raise MissingLayerError(BasicDesign.external_layers[0])

	if width is not None and not scale:
		img.transform(resize=f'{width}x{height}')
	# do this again now because custom exception handlers don't run for generators ¯\_(ツ)_/¯
	TiledImageTooBigError.validate(img)

	design = BasicDesign(
		island_name=designs_db.island_name(),
		author_name=author_name,
		design_name=image_name,
		layers={'0': img},
	)

	with img:
		yield from designs_db.create_image(design, scale=scale)

def format_created_design_results(gen, *, header=True):
	def maybe_error(row):
		if isinstance(row, dict):
			return 'error: ' + flask.json.dumps(row)
		return None

	if header:
		# pylint: disable=stop-iteration-return  # this will never raise StopIteration
		row = next(gen)
		yield maybe_error(row) or str(row) + '\n'

	for row in gen:
		err = maybe_error(row)
		if err:
			yield err
			return

		was_quantized, design_id = row
		yield f'{int(was_quantized)},{designs_api.design_code(design_id)}\n'

@bp.route('/images')
def images():
	page = parse_keyset_params()
	rv = designs_db.images_keyset(page)
	for i, image_info in enumerate(rv):
		rv[i] = image_info = dict(rv[i])
		# images are meant to be anonymous, with the author identified solely by their chosen name
		del image_info['author_id']
		image_info['design_type'] = Design(image_info.pop('type_code')).name
	return jsonify(rv)

def parse_keyset_params():
	# before='' means last
	# after='' means first
	before = request.args.get('before')
	after = request.args.get('after')
	limit = request.args.get('limit')
	if limit is not None:
		limit = int(InvalidPaginationLimitError.validate(limit))

	if before is None and after is None:
		# default to first page
		return PageSpecifier.first()

	if before is not None and after is not None:
		raise TwoPaginationReferencesPassedError

	if not before and not after:
		reference = None
	else:
		reference = int(InvalidImageIdError.validate(before or after))

	if before is not None:
		direction = PageDirection.before
	elif after is not None:
		direction = PageDirection.after

	return PageSpecifier(direction, reference, limit)

@bp.route('/image/<image_id>')
def image(image_id):
	rv = designs_db.image(int(InvalidImageIdError.validate(image_id)))
	# images are meant to be anonymous, with the author identified solely by their chosen name
	del rv['image']['author_id']
	rv['image']['design_type'] = Design(rv['image'].pop('type_code')).name
	return rv

@bp.route('/image/<image_id>.tar')
@limiter.limit('2 per 10 seconds')
def image_archive(image_id):
	image_id = int(InvalidImageIdError.validate(image_id))
	image_info = designs_db.image(image_id)['image']
	render_internal = 'internal_layers' in request.args
	layers = {}
	cls = Design(image_info['type_code'])
	for layer, image_blob in zip(cls.external_layers, image_info['layers']):
		if image_info['pro']:
			layers[layer.name] = img = layer.as_wand()
		else:
			layers[layer.name] = img = wand.image.Image(width=image_info['width'], height=image_info['height'])

		img.import_pixels(channel_map='RGBA', data=image_blob)

	# pylint: disable=not-callable
	design = cls(layers=layers)
	if render_internal:
		requested_layers = enumerate(design.internalize())
	else:
		requested_layers = layers.items()

	gen = make_tar(image_info['image_name'], image_info['created_at'].timestamp(), requested_layers)
	encoded_filename = urllib.parse.quote(image_info['image_name'] + '.tar')
	return current_app.response_class(
		stream_with_context(gen),
		mimetype='application/x-tar',
		headers={
			'Content-Disposition': f"attachment; filename*=utf-8''{encoded_filename}",
		},
	)

@bp.route('/image/<image_id>/refresh', methods=['POST'])
def refresh_image(image_id):
	gen = stream_with_context(format_created_design_results(_refresh_image(image_id), header=False))
	return current_app.response_class(gen, mimetype='text/plain')

def _refresh_image(image_id):
	return designs_db.refresh_image(int(InvalidImageIdError.validate(image_id)))

@bp.route('/image/<image_id>', methods=['DELETE'])
def delete_image(image_id):
	image_id = int(InvalidImageIdError.validate(image_id))
	designs_db.delete_image(image_id)
	return jsonify('OK')

@bp.errorhandler(HTTPException)
def handle_exception(ex):
	"""Return JSON instead of HTML for HTTP errors."""
	# start with the correct headers and status code from the error
	response = ex.get_response()
	# replace the body with JSON
	response.data = json.dumps({
		'http_status': ex.code,
		'http_status_name': ex.name,
		'http_status_description': ex.description,
	})
	response.content_type = 'application/json'
	return response


import os
import time
from pip._vendor import requests
from pip._vendor.requests import HTTPError, RequestException
import pandas as pd
import difflib

class Restaurant(object):
    '''
    a class that represent each restaurant.
    '''
    def __init__(self, id, yelp_link, pic_link, name, categories, price, location):
        self.id = id
        self.yelp_link = yelp_link
        self.pic_link = pic_link
        self.name = name
        self.categories = categories
        self.price = price
        self.location = location

class YelpWrapper:
    def __init__(self):
        '''
        initiate a wrapper class by retrieving API keys from .env file
        '''
        self.API_KEY = os.getenv("YELP_API")
        self.CLIENT_KEY = os.getenv("YELP_CLIENT")

    def search(self, location: str, categories: str, open_now: bool, price: str) -> [Restaurant]:
        '''
        :param location: location user input
        :param categories: category information we parsed from user input by our parse method below so we can use a
        category best matched to what the user desires.
        :param open_now: whether user want it to be open now or not
        :param price: price range user put in
        :return: A list of restaurants matching the query.
        Use YelpAPI to retrieve real time list of restaurants that user might be interested in.
        '''
        params = {
            "location": location,
            "categories": categories,
            "open_now": open_now,
            "sort_by": "rating",
            "price": price,
            "limit": 50
        }
        try:
            url = "https://api.yelp.com/v3/businesses/search"
            headers = {"Authorization": f"Bearer {self.API_KEY}"}
            for trial in range(3):
                response = requests.get(url, headers=headers, params=params)
                response.raise_for_status()
                result = response.json()
                res = []
                if result and result.get('businesses'):
                    for bizz in result.get('businesses'):
                        id = bizz['id']
                        name = bizz['name']
                        yelp_link = bizz['url']
                        pic_link = bizz['image_url']
                        cato = []
                        for category in bizz['categories']:
                            cato.append(category['title'])
                        loca = bizz['location']['display_address']
                        p = bizz['price']
                        res.append(Restaurant(id=id, yelp_link=yelp_link, pic_link=pic_link, name=name, categories=cato,
                                              price=p, location=loca))
                    return res
                time.sleep(1)
        except HTTPError as http_err:
            print(f"HTTP error occurred when retrieving businesses: {http_err}")
            raise http_err
        except RequestException as r_err:
            print(f"Other request error occurred when retrieving businesses: {r_err}")
            raise r_err
        except Exception as err:
            print(f"Other error occurred when retrieving businesses: {err}")
            raise err
        else:
            raise ConnectionError

    def get_all_categories(self) -> [str]:
        """
        :return: a list of all categories.
        Use API key to call all categoriess. This method is used in pre_process.py and used to generate cato.csv.
        """
        try:
            url = "https://api.yelp.com/v3/categories"
            headers = {"Authorization": f"Bearer {self.API_KEY}"}
            for trial in range(3):
                response = requests.get(url, headers=headers)
                response.raise_for_status()
                result = response.json()
                res = []
                if result and result.get('categories'):
                    for cato in result.get('categories'):
                        res.append(cato["alias"])
                    return res
                time.sleep(2)
        except HTTPError as http_err:
            print(f"HTTP error occurred when retrieving categories: {http_err}")
            raise http_err
        except RequestException as r_err:
            print(f"Other request error occurred when retrieving categories")
            raise r_err
        except Exception as err:
            print(f"Other error occurred when retrieving categories: {err}")
            raise err
        else:
            raise ConnectionError

    @staticmethod
    def parse(input_cate):
        '''
        :param input_cate: the category information user put in
        :return: a parsed category.
        In the event that the input category does not match any one in yelp database, yelp API itself does not do
        anything with it and return a bunch of random categories. Through doing this, we can parse the category and
        return restaurants users are most likely looking for.
        '''
        all_df = pd.read_csv("cato.csv")
        all_cato_lst = all_df["categories"].tolist()
        top10_similiar = difflib.get_close_matches(input_cate, all_cato_lst)
        return top10_similiar[0]


import requests

# Функция выполнения запроса текущей структуры фонда
def get_structure(ticker: str) -> dict:
    headers = {
        'Content-Type': 'application/x-www-form-urlencoded',
        'Accept': '*/*',
        'Sec-Fetch-Site': 'same-origin',
        'Accept-Language': 'ru',
        'Sec-Fetch-Mode': 'cors',
        'Host': 'www.tinkoff.ru',
        'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.4 Safari/605.1.15',
        'Connection': 'keep-alive',
        'Sec-Fetch-Dest': 'empty',
    }

    params = {
        'appName': 'invest',
        'appVersion': '1.312.0',
        'origin': 'web',
        'platform': 'web',
        'idKind': 'ticker',
    }

    response = requests.get(
        f'https://www.tinkoff.ru/api/invest-gw/fireg-advisory/api/web/v2/etf/assets/{ticker}/structure',
        params=params,

        headers=headers,
    )
    if response.status_code == 200:
        data = response.json()
    return data


from flask import (
    Blueprint, flash, g, redirect, render_template, request, url_for
)

from werkzeug.exceptions import abort

from eventosABC.auth import login_required
from eventosABC.db import get_db

bp = Blueprint('evento', __name__)

@bp.route('/')
@login_required
def index():
    db = get_db()
    eventos = db.execute(
        'SELECT e.id, creado, titulo, fecha, hora_inicio, author_id FROM '
        'evento e JOIN user u on e.author_id=u.id WHERE u.id =?'
        ' ORDER BY creado ASC', (g.user['id'],)
    ).fetchall()
    return render_template('eventos/index.html', eventos=eventos)

@bp.route('/create', methods=('GET', 'POST'))
@login_required
def create():
    if request.method == 'POST':
        titulo = request.form['titulo']
        fecha = request.form['fecha']
        hora_inicio = request.form['hora_inicio']
        hora_fin = request.form['hora_fin']
        descripcion = request.form['descripcion']

        error = None

        if not titulo:
            error = "Es necesario llenar todos los campos"
        if not fecha:
            error = "Es necesario llenar todos los campos"
        if not descripcion:
            error = "Es necesario llenar todos los campos"
        if not hora_inicio:
            error = "Es necesario llenar todos los campos"
        if not hora_fin:
            error = "Es necesario llenar todos los campos"

        if error is not None:
            flash(error)
        else:
            db = get_db()
            db.execute(
                'INSERT INTO evento (titulo, descripcion, author_id, '
                'fecha, hora_inicio, hora_fin) values(?, ?, ?, ?, ?, ?)',
                (titulo, fecha, g.user['id'], descripcion, hora_inicio, hora_fin, )
            )
            db.commit()
            return redirect(url_for('evento.index'))

    return render_template('eventos/create.html')

def get_evento(id, check_author=True):
    evento = get_db().execute(
        'SELECT e.id, titulo, descripcion, creado, fecha, hora_inicio, hora_fin, author_id'
        ' FROM evento e JOIN user u ON e.author_id=u.id WHERE e.id=?', (id,)
    ).fetchone()

    if evento is None:
        abort(404, "El evento que esta buscando no existe.")
    elif check_author and evento['author_id'] != g.user['id']:
        abort(403, "No tienes perimso para ver este evento.")

    return evento

@bp.route('/<int:id>/detail', methods=('GET', 'POST'))
@login_required
def detail_evento(id):
    evento = get_evento(id)
    return render_template('eventos/detail.html', evento=evento)

@bp.route('/<int:id>/update', methods=('GET', 'POST'))
@login_required
def update_evento(id):
    evento = get_evento(id)

    if request.method == 'POST':
        titulo = request.form['titulo']
        fecha = request.form['fecha']
        descripcion = request.form['descripcion']
        hora_inicio = request.form['hora_inicio']
        hora_fin = request.form['hora_fin']

        error = None

        if not titulo:
            error = "Es necesario llenar todos los campos"
        if not fecha:
            error = "Es necesario llenar todos los campos"
        if not descripcion:
            error = "Es necesario llenar todos los campos"
        if not hora_inicio:
            error = "Es necesario llenar todos los campos"
        if not hora_fin:
            error = "Es necesario llenar todos los campos"

        if error is not None:
            flash(error)
        else:
            db = get_db()
            db.execute(
                'UPDATE evento SET titulo=?, descripcion=?,'
                'fecha=?, hora_inicio=?, hora_fin=? WHERE id=?',
                (titulo, fecha, descripcion, hora_inicio, hora_fin, evento['id'])
            )
            db.commit()
            return redirect(url_for('evento.index'))

    return render_template('eventos/update.html', evento=evento)

@bp.route('/<int:id>/delete', methods=('POST',))
@login_required
def delete_evento(id):
    db = get_db()
    get_evento(id)
    db.execute('DELETE FROM evento WHERE id = ?', (id,))
    db.commit()
    return redirect(url_for('evento.index'))


import json
import os
import textwrap

import datalabs
from datalabs import get_task, TaskType

logger = datalabs.logging.get_logger(__name__)

_CITATION = """\
@inproceedings{faisal-etal-21-sdqa,
 title = {{SD-QA}: {S}poken {D}ialectal {Q}uestion {A}nswering for the {R}eal {W}orld},
  author = {Faisal, Fahim and Keshava, Sharlina and ibn Alam, Md Mahfuz and Anastasopoulos, Antonios},
  url={https://arxiv.org/abs/2109.12072},
  year = {2021},
  booktitle = {Findings of the 2021 Conference on Empirical Methods in Natural Language Processing (EMNLP Findings)},
  publisher = {Association for Computational Linguistics},
  month = {November},
}
"""

_DESCRIPTION = """\
SDQA: : Spoken Dialectal Question Answering for the Real World
is a multi-dialect, spoken QA benchmark on five languages (Arabic, Bengali, English, Korean, Kiswahili)
with more than 68k audio prompts in 22 dialects from 245 speakers. Google Speech API with regional units
(eg. en-US, sw-TZ) are used to perform speech to text conversion of the questions. These questions are
replaced in place of the original TyDiQA gold questions to prepare dialectal development and test set. As a benchmark
train set, a discarded subset of original TyDiQA training set is used.

"""

LANG_URLS = {
    "ara": "https://drive.google.com/uc?export=download&id=1AUpczDwqwjnw44Trfu7ku-DQ533cmfVG",
    "ben": "https://drive.google.com/uc?export=download&id=1iP0r3JbF4P3QDkrHCHdIXsVFv-pUac3N",
    "eng": "https://drive.google.com/uc?export=download&id=1-XdM9TcDFACj7OXLi0KlTXnpwE2EgjN9",
    "kor": "https://drive.google.com/uc?export=download&id=120YpOhUje06yuTM58kXdgs1aBmrmmHDG",
    "swa": "https://drive.google.com/uc?export=download&id=1KjCthk7MDVz2IlznIsgmcOzQ39Iqe33u",
}

class SdqaConfig(datalabs.BuilderConfig):
    """BuilderConfig for SQUAD."""

    def __init__(self, **kwargs):
        """BuilderConfig for SQUADV2.
        Args:
          **kwargs: keyword arguments forwarded to super.
        """
        super(SdqaConfig, self).__init__(**kwargs)

class Sdqa(datalabs.GeneratorBasedBuilder):
    """TODO(squad_v2): Short description of my dataset."""

    # TODO(squad_v2): Set up version.
    BUILDER_CONFIGS = [
        datalabs.BuilderConfig(
            name="{}".format(lang),
            version=datalabs.Version("2.0.0"),
            description="SDQA 1.0.0"
        )
        for lang in list(LANG_URLS.keys())
    ]

    def _info(self):
        # TODO(squad_v2): Specifies the datalab.DatasetInfo object
        return datalabs.DatasetInfo(
            # This is the description that will appear on the datalab page.
            description=_DESCRIPTION,
            # datalab.features.FeatureConnectors
            features=datalabs.Features(
                {
                    "id": datalabs.Value("string"),
                    "title": datalabs.Value("string"),
                    "context": datalabs.Value("string"),
                    "question": datalabs.Value("string"),
                    "answers":
                        {
                            "text": datalabs.features.Sequence(datalabs.Value("string")),
                            "answer_start": datalabs.features.Sequence(datalabs.Value("int32")),
                        }
                    # These are the features of your dataset like images, labels ...
                }
            ),
            # If there's a common (input, target) tuple from the features,
            # specify them here. They'll be used if as_supervised=True in
            # builder.as_dataset.
            supervised_keys=None,
            # Homepage of the dataset for documentation
            homepage="https://nlp.cs.gmu.edu/publication/faisal-etal-21-sdqa/",
            citation=_CITATION,
            languages=["eng","ara","ben","swa","kor"],
            task_templates=[
                get_task(TaskType.qa_extractive)(
                    question_column="question", context_column="context", answers_column="answers"
                )
            ],
        )

    def _split_generators(self, dl_manager):
        """Returns SplitGenerators."""
        # TODO(squad_v2): Downloads the data and defines the splits
        # dl_manager is a datalab.download.DownloadManager that can be used to
        # download and extract URLs
        lang = str(self.config.name)
        url =LANG_URLS[lang]
        data_dir = dl_manager.download_and_extract(url)
        return [
            datalabs.SplitGenerator(
                name=datalabs.Split.TRAIN,
                gen_kwargs={
                    "filepath": os.path.join(data_dir,lang,"sdqa-train-" + lang + ".json"),
                    # "filepath": os.path.join(data_dir,  "train-en.json" + lang + "_train.jsonl"),
                },
            ),
            datalabs.SplitGenerator(
                name=datalabs.Split.TEST,
                gen_kwargs={
                    "filepath": os.path.join(data_dir,lang, "sdqa-test-" + lang + ".json"),
                },
            ),
            datalabs.SplitGenerator(
                name=datalabs.Split.VALIDATION,
                gen_kwargs={
                    "filepath": os.path.join(data_dir,lang, "sdqa-dev-" + lang + ".json"),
                },
            ),
        ]

    def _generate_examples(self, filepath):
        """Yields examples."""
        # TODO(squad_v2): Yields (key, example) tuples from the dataset
        with open(filepath, encoding="utf-8") as f:
            squad = json.load(f)
            for example in squad["data"]:
                title = example.get("title", "")
                for paragraph in example["paragraphs"]:
                    context = paragraph["context"]  # do not strip leading blank spaces GH-2585
                    for qa in paragraph["qas"]:
                        question = qa["question"]
                        id_ = qa["id"]

                        answer_starts = [answer["answer_start"] for answer in qa["answers"]]
                        answers = [answer["text"] for answer in qa["answers"]]

                        # Features currently used are "context", "question", and "answers".
                        # Others are extracted here for the ease of future expansions.
                        yield id_, {
                            "title": title,
                            "context": context,
                            "question": question,
                            "id": id_,
                            "answers": {
                                "answer_start": answer_starts,
                                "text": answers,
                            },
                        }


from ament_index_python.packages import get_package_share_directory

from launch import LaunchDescription
from launch.actions.declare_launch_argument import DeclareLaunchArgument
from launch.substitutions import LaunchConfiguration, PathJoinSubstitution

from launch_ros.actions import Node

def generate_launch_description():
    pkg_originbot_bringup = get_package_share_directory('originbot_bringup')

    joy_config_cmd = DeclareLaunchArgument(
        'joy_config',
        default_value=PathJoinSubstitution(
            [pkg_originbot_bringup, 'config', 'joy.yaml']),
        description='Originbot Joy teleop param file'
    )

    joy_config = LaunchConfiguration('joy_config')

    joy_node = Node(
        package='joy_linux',
        executable='joy_linux_node',
        name='joy_linux_node'
    )

    teleop_twist_joy_node = Node(
        package='teleop_twist_joy',
        executable='teleop_node',
        name='teleop_twist_joy_node',
        parameters=[joy_config]
    )

    ld = LaunchDescription()
    ld.add_action(joy_config_cmd)
    ld.add_action(joy_node)
    ld.add_action(teleop_twist_joy_node)

    return ld


import os
from typing import Dict, Optional, Union

import numpy as np
import tensorflow as tf

import merlin.io
from merlin.models.io import save_merlin_metadata
from merlin.models.tf.core import combinators
from merlin.models.tf.core.base import NoOp
from merlin.models.tf.core.prediction import TopKPrediction
from merlin.models.tf.inputs.base import InputBlockV2
from merlin.models.tf.inputs.embedding import CombinerType, EmbeddingTable
from merlin.models.tf.loader import Loader
from merlin.models.tf.models.base import BaseModel, get_output_schema
from merlin.models.tf.outputs.topk import TopKOutput
from merlin.models.tf.transforms.features import PrepareFeatures
from merlin.models.tf.utils import tf_utils
from merlin.models.tf.utils.batch_utils import TFModelEncode
from merlin.schema import ColumnSchema, Schema, Tags

@tf.keras.utils.register_keras_serializable(package="merlin.models")
class Encoder(tf.keras.Model):
    """Block that can be used for prediction and evaluation but not for training

    Parameters
    ----------
    inputs: Union[Schema, tf.keras.layers.Layer]
        The input block or schema.
        When a schema is provided, a default input block will be created.
    *blocks: tf.keras.layers.Layer
        The blocks to use for encoding.
    pre: Optional[tf.keras.layers.Layer]
        A block to use before the main blocks
    post: Optional[tf.keras.layers.Layer]
        A block to use after the main blocks
    prep_features: Optional[bool]
        Whether this block should prepare list and scalar features
        from the dataloader format. By default True.
    """

    def __init__(
        self,
        inputs: Union[Schema, tf.keras.layers.Layer],
        *blocks: tf.keras.layers.Layer,
        pre: Optional[tf.keras.layers.Layer] = None,
        post: Optional[tf.keras.layers.Layer] = None,
        prep_features: Optional[bool] = True,
        **kwargs,
    ):
        super().__init__(**kwargs)
        if isinstance(inputs, Schema):
            input_block = InputBlockV2(inputs)
            self._schema = inputs
        else:
            input_block = inputs
            if not hasattr(inputs, "schema"):
                raise ValueError("inputs must have a schema")
            self._schema = inputs.schema

        self.blocks = [input_block] + list(blocks) if blocks else [input_block]
        self.pre = pre
        self.post = post

        self.prep_features = prep_features

        self._prepare_features = PrepareFeatures(self.schema) if self.prep_features else NoOp()

    def encode(
        self,
        dataset: Union[merlin.io.Dataset, Loader],
        index: Union[str, ColumnSchema, Schema, Tags],
        batch_size: int,
        **kwargs,
    ) -> merlin.io.Dataset:
        """Encodes the given dataset and index.

        Parameters
        ----------
        dataset: Union[merlin.io.Dataset, merlin.models.tf.loader.Loader]
            The dataset to encode.
        index: Union[str, ColumnSchema, Schema, Tags]
            The index to use for encoding.
        batch_size: int
            The batch size for encoding.

        Returns
        -------
        merlin.io.Dataset
            The encoded dataset.
        """
        if isinstance(index, Schema):
            output_schema = index
        elif isinstance(index, ColumnSchema):
            output_schema = Schema([index])
        elif isinstance(index, str):
            output_schema = Schema([self.schema[index]])
        elif isinstance(index, Tags):
            output_schema = self.schema.select_by_tag(index)
        else:
            raise ValueError(f"Invalid index: {index}")

        return self.batch_predict(
            dataset,
            batch_size=batch_size,
            output_schema=output_schema,
            index=index,
            output_concat_func=np.concatenate,
            **kwargs,
        )

    def batch_predict(
        self,
        dataset: Union[merlin.io.Dataset, Loader],
        batch_size: int,
        output_schema: Optional[Schema] = None,
        index: Optional[Union[str, ColumnSchema, Schema, Tags]] = None,
        **kwargs,
    ) -> merlin.io.Dataset:
        """Batched prediction using Dask.

        Parameters
        ----------
        dataset: Union[merlin.io.Dataset, merlin.models.tf.loader.Loader]
            Dataset or Loader to predict on.
        batch_size: int
            Batch size to use for prediction.

        Returns
        -------
        merlin.io.Dataset
        """

        if index:
            if isinstance(index, ColumnSchema):
                index = Schema([index])
            elif isinstance(index, str):
                index = Schema([self.schema[index]])
            elif isinstance(index, Tags):
                index = self.schema.select_by_tag(index)
            elif not isinstance(index, Schema):
                raise ValueError(f"Invalid index: {index}")

            if len(index) != 1:
                raise ValueError("Only one column can be used as index")
            index = index.first.name

        dataset_schema = None
        if hasattr(dataset, "schema"):
            dataset_schema = dataset.schema
            data_output_schema = dataset_schema
            if isinstance(dataset, Loader):
                data_output_schema = dataset.output_schema
            if not set(self.schema.column_names).issubset(set(data_output_schema.column_names)):
                raise ValueError(
                    f"Model schema {self.schema.column_names} does not match dataset schema"
                    + f" {data_output_schema.column_names}"
                )

        loader_transforms = None
        if isinstance(dataset, Loader):
            loader_transforms = dataset.transforms
            batch_size = dataset.batch_size
            dataset = dataset.dataset

        # Check if merlin-dataset is passed
        if hasattr(dataset, "to_ddf"):
            dataset = dataset.to_ddf()

        model_encode = TFModelEncode(
            self,
            batch_size=batch_size,
            loader_transforms=loader_transforms,
            schema=dataset_schema,
            **kwargs,
        )

        encode_kwargs = {}
        if output_schema:
            encode_kwargs["filter_input_columns"] = output_schema.column_names

        # Processing a sample of the dataset with the model encoder
        # to get the output dataframe dtypes
        sample_output = model_encode(dataset.head(), **encode_kwargs)
        output_dtypes = sample_output.dtypes.to_dict()

        predictions = dataset.map_partitions(model_encode, meta=output_dtypes, **encode_kwargs)
        if index:
            predictions = predictions.set_index(index)

        return merlin.io.Dataset(predictions)

    def call(self, inputs, *, targets=None, training=False, testing=False, **kwargs):
        """Calls the model on new inputs and returns the outputs as tensors.

        Parameters
        ----------
        inputs : tensor-like or dict/tuple of tensors.
            Tensors or dict/tuple of tensors representing the input batch.
        targets : tensor-like, optional
            Tensors representing the target data.
        training : bool, optional
            Whether the model is in training mode.
        testing : bool, optional
            Whether the model is in testing mode.

        Returns
        -------
        A tensor or dict of tensors corresponding to the result of calling the layer.
        """
        inputs = self._prepare_features(inputs, targets=targets)
        if isinstance(inputs, tuple):
            inputs, targets = inputs
        return combinators.call_sequentially(
            list(self.to_call),
            inputs=inputs,
            features=inputs,
            targets=targets,
            training=training,
            testing=testing,
            **kwargs,
        )

    def __call__(self, inputs, **kwargs):
        """Overrides the default __call__ method to remove "features" from inputs.

        Parameters
        ----------
        inputs : tensor-like or dict/tuple of tensors.
            Tensors or dict/tuple of tensors representing the input batch.

        Returns
        -------
        A tensor or dict of tensors corresponding to the result of calling the layer.
        """
        # We remove features here since we don't expect them at inference time
        # Inside the `call` method, we will add them back by assuming inputs=features
        if "features" in kwargs:
            kwargs.pop("features")

        return super().__call__(inputs, **kwargs)

    def build(self, input_shape):
        """Creates the variables of the layer.

        Parameters
        ----------
        input_shape: Tuple[int]
            The shape of the input data.
        """
        self._prepare_features.build(input_shape)
        input_shape = self._prepare_features.compute_output_shape(input_shape)

        combinators.build_sequentially(self, list(self.to_call), input_shape=input_shape)
        if not hasattr(self.build, "_is_default"):
            self._build_input_shape = input_shape

    def compute_output_shape(self, input_shape):
        """Computes the output shape of the layer.

        Parameters
        ----------
        input_shape: Tuple[int]
            The shape of the input data.

        Returns
        -------
        Tuple[int]
            The output shape of the layer.
        """
        input_shape = self._prepare_features.compute_output_shape(input_shape)
        return combinators.compute_output_shape_sequentially(list(self.to_call), input_shape)

    def train_step(self, data):
        """Performs a training step.

        Train step method is not implemented and Raises an error as the
        Encoder block is not meant to be trained by itself and can only be
        trained as part of a model.
        """
        raise NotImplementedError(
            "This block is not meant to be trained by itself. ",
            "It can only be trained as part of a model.",
        )

    def fit(self, *args, **kwargs):
        """Fits the model.

        Fit method is not implemented and Raises an error as the Encoder block
        is not meant to be trained by itself and can only be trained as part
        of a model.
        """
        raise NotImplementedError(
            "This block is not meant to be trained by itself. ",
            "It can only be trained as part of a model.",
        )

    def save(
        self,
        export_path: Union[str, os.PathLike],
        include_optimizer=True,
        save_traces=True,
    ) -> None:
        """Saves the model to export_path as a Tensorflow Saved Model.
        Along with merlin model metadata.

        Parameters
        ----------
        export_path : Union[str, os.PathLike]
            Path where model will be saved to
        include_optimizer : bool, optional
            If False, do not save the optimizer state, by default True
        save_traces : bool, optional
            When enabled, will store the function traces for each layer. This
            can be disabled, so that only the configs of each layer are
            stored, by default True
        """
        super().save(
            export_path,
            include_optimizer=include_optimizer,
            save_traces=save_traces,
            save_format="tf",
        )
        input_schema = self.schema
        output_schema = get_output_schema(export_path)
        save_merlin_metadata(export_path, input_schema, output_schema)

    @property
    def to_call(self):
        """Provides the list of blocks to be called during the execution of the model."""
        if self.pre:
            yield self.pre

        for block in self.blocks:
            yield block

        if self.post:
            yield self.post

    @property
    def has_schema(self) -> bool:
        """Returns True as this class does contain a schema."""
        return True

    @property
    def schema(self) -> Schema:
        """Returns the schema of the model."""
        return self._schema

    @property
    def first(self):
        """Returns the first block of the model."""
        return self.blocks[0]

    @property
    def last(self):
        """Returns the last block of the model."""
        return self.blocks[-1]

    @classmethod
    def from_config(cls, config, custom_objects=None):
        """Creates a new instance of the class by deserializing.

        Parameters
        ----------
        config: dict
            A dictionary, typically the output of get_config.
        custom_objects: dict, optional
            A dictionary mapping the names of layers to the corresponding
            functions and classes.

        Returns
        -------
        A new instance of Encoder.
        """
        pre = config.pop("pre", None)
        post = config.pop("post", None)
        layers = [
            tf.keras.layers.deserialize(conf, custom_objects=custom_objects)
            for conf in config.values()
        ]

        if pre is not None:
            pre = tf.keras.layers.deserialize(pre, custom_objects=custom_objects)

        if post is not None:
            post = tf.keras.layers.deserialize(post, custom_objects=custom_objects)

        output = Encoder(*layers, pre=pre, post=post)
        output.__class__ = cls

        return output

    def get_config(self):
        """Returns the configuration of the model as a dictionary.

        Returns
        -------
        dict
            The configuration of the model.
        """
        config = tf_utils.maybe_serialize_keras_objects(self, {}, ["pre", "post"])
        for i, layer in enumerate(self.blocks):
            config[i] = tf.keras.utils.serialize_keras_object(layer)

        return config

@tf.keras.utils.register_keras_serializable(package="merlin.models")
class TopKEncoder(Encoder, BaseModel):
    """Block that can be used for top-k prediction & evaluation, initialized
    from a trained retrieval model

    Parameters
    ----------
    query_encoder: Union[Encoder, tf.keras.layers.Layer],
        The layer to use for encoding the query features
    topk_layer: Union[str, tf.keras.layers.Layer, TopKOutput]
        The layer to use for computing the top-k predictions.
        You can also pass the `name` of registered top-k layer.
        The current supported strategies are [`brute-force-topk`]
        By default "brute-force-topk"
    candidates: Union[tf.Tensor, ~merlin.io.Dataset]
        The candidate embeddings to use for the Top-k index.
        You can pass a tensor of pre-trained embeddings
        or a merlin.io.Dataset of pre-trained embeddings, indexed by
        the candidates ids.
        This is required when `topk_layer` is a string
        By default None
    candidate_encoder:  Union[Encoder, tf.keras.layers.Layer],
        The layer to use for encoding the item features
    k: int, Optional
        Number of candidates to return, by default 10
    pre: Optional[tf.keras.layers.Layer]
        A block to use before encoding the input query
        By default None
    post: Optional[tf.keras.layers.Layer]
        A block to use after getting the top-k prediction scores
        By default None
    target: str, optional
        The name of the target. This is required when multiple targets are provided.
        By default None
    """

    def __init__(
        self,
        query_encoder: Union[Encoder, tf.keras.layers.Layer],
        topk_layer: Union[str, tf.keras.layers.Layer, TopKOutput] = "brute-force-topk",
        candidates: Union[tf.Tensor, merlin.io.Dataset] = None,
        candidate_encoder: Union[Encoder, tf.keras.layers.Layer] = None,
        k: int = 10,
        pre: Optional[tf.keras.layers.Layer] = None,
        post: Optional[tf.keras.layers.Layer] = None,
        target: str = None,
        **kwargs,
    ):
        if isinstance(topk_layer, TopKOutput):
            topk_output = topk_layer
        else:
            topk_output = TopKOutput(to_call=topk_layer, candidates=candidates, k=k, target=target)
        self.k = k

        Encoder.__init__(self, query_encoder, topk_output, pre=pre, post=post, **kwargs)
        # The base model is required for the evaluation step:
        BaseModel.__init__(self, **kwargs)

    @classmethod
    def from_candidate_dataset(
        cls,
        query_encoder: Union[Encoder, tf.keras.layers.Layer],
        candidate_encoder: Union[Encoder, tf.keras.layers.Layer],
        dataset: merlin.io.Dataset,
        top_k: int = 10,
        index_column: Optional[Union[str, ColumnSchema, Schema, Tags]] = None,
        **kwargs,
    ):
        """Class method to initialize a TopKEncoder from a dataset of
        raw candidates features.

        Parameters
        ----------
        query_encoder : Union[Encoder, tf.keras.layers.Layer]
            The encoder layer to use for computing the query embeddings.
        candidate_encoder : Union[Encoder, tf.keras.layers.Layer]
            The encoder layer to use for computing the candidates embeddings.
        dataset : merlin.io.Dataset
            Raw candidate features dataset
        index_column : Union[str, ColumnSchema, Schema, Tags], optional
            The column to use as candidates identifiers, this will be used
            for returning the topk ids of candidates with the highest scores.
            If not specified, the candidates indices will be used instead.
            by default None
        top_k : int, optional
            Number of candidates to return, by default 10

        Returns
        -------
        TopKEncoder
            a `TopKEncoder` indexed by the pre-trained embeddings of the candidates
            in the specified `dataset`
        """
        # TODO: Add related unit-test after RetrievalModelV2 is merged
        candidates = cls.encode_candidates(dataset, candidate_encoder)
        topk_output = TopKOutput(
            to_call="topk_layer", candidate_dataset=candidates, k=top_k, **kwargs
        )
        return cls(query_encoder, topk_output, **kwargs)

    def compile(
        self,
        optimizer="rmsprop",
        loss=None,
        metrics=None,
        loss_weights=None,
        weighted_metrics=None,
        run_eagerly=None,
        steps_per_execution=None,
        jit_compile=None,
        k: int = None,
        **kwargs,
    ):
        """Extend the compile method of `BaseModel` to set the threshold `k`
        of the top-k encoder.
        """
        if k is not None:
            self.topk_layer._k = k
            self.k = k
        BaseModel.compile(
            self,
            optimizer=optimizer,
            loss=loss,
            metrics=metrics,
            weighted_metrics=weighted_metrics,
            run_eagerly=run_eagerly,
            loss_weights=loss_weights,
            steps_per_execution=steps_per_execution,
            jit_compile=jit_compile,
            **kwargs,
        )

    @property
    def topk_layer(self):
        return self.blocks[-1].to_call

    def index_candidates(self, candidates, identifiers=None):
        self.topk_layer.index(candidates, identifiers=identifiers)
        return self

    def encode_candidates(
        self,
        dataset: merlin.io.Dataset,
        index_column: Optional[Union[str, ColumnSchema, Schema, Tags]] = None,
        candidate_encoder: Optional[Union[Encoder, tf.keras.layers.Layer]] = None,
        **kwargs,
    ) -> merlin.io.Dataset:
        """Method to generate candidates embeddings

        Parameters
        ----------
        dataset : merlin.io.Dataset
            Raw candidate features dataset
        index_column  : Union[str, ColumnSchema, Schema, Tags], optional
            The column to use as candidates identifiers, this will be used
            for returning the topk ids of candidates with the highest scores.
            If not specified, the candidates indices will be used instead.
            by default None
        candidate_encoder : Union[Encoder, tf.keras.layers.Layer], optional
            The encoder layer to use for computing the candidates embeddings.
            If not specified, the candidate_encoder set in the constructor
            will be used instead.
            by default None
        Returns
        -------
        merlin.io.Dataset
            A merlin dataset of candidates embeddings, indexed by index_column.
        """
        # TODO: Add related unit-test after RetrievalModelV2 is merged
        if not candidate_encoder:
            candidate_encoder = self.candidate_encoder
        assert candidate_encoder is not None, ValueError(
            "You should provide a `candidate_encoder` to compute candidates embeddings"
        )
        return candidate_encoder.encode(dataset=dataset, index=index_column, **kwargs)

    def batch_predict(
        self,
        dataset: Union[merlin.io.Dataset, Loader],
        batch_size: int,
        output_schema: Optional[Schema] = None,
        **kwargs,
    ) -> merlin.io.Dataset:
        """Batched top-k prediction using Dask.

        Parameters
        ----------
        dataset : Union[merlin.io.Dataset, merlin.models.tf.loader.Loader]
            Raw queries features dataset or Loader
        batch_size : int
            The number of queries to process at each prediction step
        output_schema: Schema, optional
            The columns to output from the input dataset
        Returns
        -------
        merlin.io.Dataset
            A merlin dataset with the top-k predictions, the
            candidates identifiers and related scores.
        """
        from merlin.models.tf.utils.batch_utils import TFModelEncode

        loader_transforms = None
        if isinstance(dataset, Loader):
            loader_transforms = dataset.transforms
            batch_size = dataset.batch_size
            dataset = dataset.dataset

        dataset_schema = dataset.schema
        dataset = dataset.to_ddf()

        model_encode = TFModelEncode(
            model=self,
            batch_size=batch_size,
            loader_transforms=loader_transforms,
            schema=dataset_schema,
            output_names=TopKPrediction.output_names(self.k),
            **kwargs,
        )

        encode_kwargs = {}
        if output_schema:
            encode_kwargs["filter_input_columns"] = output_schema.column_names

        # Processing a sample of the dataset with the model encoder
        # to get the output dataframe dtypes
        sample_output = model_encode(dataset.head(), **encode_kwargs)
        output_dtypes = sample_output.dtypes.to_dict()

        predictions = dataset.map_partitions(model_encode, meta=output_dtypes, **encode_kwargs)

        return merlin.io.Dataset(predictions)

    def fit(self, *args, **kwargs):
        """Fit model"""
        raise NotImplementedError(
            "This block is not meant to be trained by itself. ",
            "It can only be trained as part of a model.",
        )

@tf.keras.utils.register_keras_serializable(package="merlin.models")
class EmbeddingEncoder(Encoder):
    """Creates an Encoder from an EmbeddingTable.
    Typically used with RetrievalModelV2.

    Parameters
    ----------
    schema : Union[ColumnSchema, Schema]
        The ColumnSchema of the column for which the
        embedding table needs to be created.
        If a Schema is passed, only the first column
        is considered
    dim : int
        Dimension of the embeddings
    embeddings_initializer : Union[str, tf.keras.layers.Layer], optional
        Initializer for the `embeddings`
        matrix (see `keras.initializers`). By default "uniform"
    embeddings_regularizer : Union[str, tf.keras.layers.Layer], optional
        Regularizer function applied to
        the `embeddings` matrix (see `keras.regularizers`)., by default None
    activity_regularizer : Union[str, tf.keras.layers.Layer], optional
        Sets a layer that applies an update to the cost function based
        input activity, by default None
    embeddings_constraint : Union[str, tf.keras.layers.Layer], optional
        Constraint function applied to
        the `embeddings` matrix (see `keras.constraints`), by default None
    mask_zero : bool, optional
        Whether or not the input value 0 is a special "padding"
        value that should be masked out.
        This is useful when using recurrent layers
        which may take variable length input.
        If this is `True`, then all subsequent layers
        in the model need to support masking or an exception will be raised.
        If mask_zero is set to True, as a consequence, index 0 cannot be
        used in the vocabulary (input_dim should equal size of
        vocabulary + 1), by default False
    input_length : int, optional
        This argument is required if you are going to connect
        `Flatten` then `Dense` layers upstream
        (without it, the shape of the dense outputs cannot be computed),
        by default None
    sequence_combiner : Optional[CombinerType], optional
        A string specifying how to combine embedding results for each
        entry ("mean", "sqrtn" and "sum" are supported) or a layer.
        Default is None (no combiner used), by default None
    trainable : bool, optional
        Whether the layer's variables should be trainable, by default True
    name : str, optional
        String name of the layer, by default None
    dtype : optional
        The dtype of the layer's computations and weights. Can also be a
        `tf.keras.mixed_precision.Policy`, which allows the computation and weight
        dtype to differ. Default of `None` means to use
        `tf.keras.mixed_precision.global_policy()`, which is a float32 policy
        unless set to different value., by default None
    dynamic : bool, optional
        Set this to `True` if your layer should only be run eagerly, and
        should not be used to generate a static computation graph.
        This would be the case for a Tree-RNN or a recursive network,
        for example, or generally for any layer that manipulates tensors
        using Python control flow. If `False`, we assume that the layer can
        safely be used to generate a static computation graph., by default False
    embeddings_l2_batch_regularization : Optional[Union[float, Dict[str, float]]], optional
        Factor for L2 regularization of the embeddings vectors (from the current batch only)
        by default 0.0, by default 0.0
    post : Optional[tf.keras.layers.Layer], optional
        _description_, by default None
    **kwargs: Forwarded Encoder parameters
    """

    def __init__(
        self,
        schema: Union[ColumnSchema, Schema],
        dim: int,
        embeddings_initializer: Optional[Union[str, tf.keras.layers.Layer]] = "uniform",
        embeddings_regularizer: Optional[Union[str, tf.keras.layers.Layer]] = None,
        activity_regularizer: Optional[Union[str, tf.keras.layers.Layer]] = None,
        embeddings_constraint: Optional[Union[str, tf.keras.layers.Layer]] = None,
        mask_zero: bool = False,
        input_length: int = None,
        sequence_combiner: Optional[CombinerType] = None,
        trainable: bool = True,
        name: str = None,
        dtype=None,
        dynamic: bool = False,
        embeddings_l2_batch_regularization: Optional[Union[float, Dict[str, float]]] = 0.0,
        post: Optional[tf.keras.layers.Layer] = None,
        **kwargs,
    ):
        if isinstance(schema, ColumnSchema):
            col = schema
        else:
            col = schema.first
        col_name = col.name

        table = EmbeddingTable(
            dim,
            col,
            embeddings_initializer=embeddings_initializer,
            embeddings_regularizer=embeddings_regularizer,
            activity_regularizer=activity_regularizer,
            embeddings_constraint=embeddings_constraint,
            mask_zero=mask_zero,
            input_length=input_length,
            sequence_combiner=sequence_combiner,
            trainable=trainable,
            name=name,
            dtype=dtype,
            dynamic=dynamic,
            l2_batch_regularization_factor=embeddings_l2_batch_regularization,
        )

        super().__init__(table, tf.keras.layers.Lambda(lambda x: x[col_name]), post=post, **kwargs)

    def to_dataset(self, gpu=None) -> merlin.io.Dataset:
        return self.blocks[0].to_dataset(gpu=gpu)


import os
import time
import argparse
import numpy as np
import tensorflow as tf

from progress.bar import Bar
from ipywidgets import IntProgress
from IPython.display import display
import skimage.transform
import imageio

from model import VGG16
from utils import read_images, read_masks, read_list, label2rgb

TRAIN_DIR = "hw3-train-validation/train/"
VAL_DIR = "hw3-train-validation/validation/"

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--init_from', type=str, default='keras-vgg16.npy', help='pre-trained weights')
    parser.add_argument('--save_dir', type=str, default='save', help='directory to store checkpointed models')
    parser.add_argument('--mode', type=str, default='FCN32s', help='FCN mode: FCN32s, FCN16s, FCN8s')
    parser.add_argument('--lr', type=float, default=1e-4, help='starting learning rate')
    parser.add_argument('--decay', type=float, default=0.0, help='multiplier for weight decay')
    parser.add_argument('--keep_prob', type=float, default=1.0, help='dropout keep probability for fc layer')
    parser.add_argument('--note', type=str, default='', help='argument for taking notes')

    FLAG = parser.parse_args()

    print("===== create directory =====")
    if not os.path.exists(FLAG.save_dir):
        os.makedirs(FLAG.save_dir)

    train(FLAG)

def train(FLAG):
    print("Reading dataset...")
    # load data
    Xtrain, Ytrain = read_images(TRAIN_DIR), read_masks(TRAIN_DIR, onehot=True)
    Xtest, Ytest = read_images(VAL_DIR), read_masks(VAL_DIR, onehot=True)
    track = ["hw3-train-validation/validation/0008","hw3-train-validation/validation/0097","hw3-train-validation/validation/0107"]
    Xtrack, Ytrack = read_list(track)

    vgg16 = VGG16(classes=7, shape=(256,256,3))
    vgg16.build(vgg16_npy_path=FLAG.init_from, mode=FLAG.mode, keep_prob=FLAG.keep_prob)

    saver = tf.train.Saver(tf.global_variables(), max_to_keep=1)
    checkpoint_path = os.path.join(FLAG.save_dir, 'model.ckpt')

    def initialize_uninitialized(sess):
        global_vars = tf.global_variables()
        is_not_initialized = sess.run([tf.is_variable_initialized(var) for var in global_vars])
        not_initialized_vars = [v for (v,f) in zip(global_vars, is_not_initialized) if not f]
        if len(not_initialized_vars):
                sess.run(tf.variables_initializer(not_initialized_vars))

    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())

        # hyper parameters
        batch_size = 32
        epoch = 500
        early_stop_patience = 50
        min_delta = 0.0001
        opt_type = 'adam'

        # recorder
        epoch_counter = 0

        # optimizer
        global_step = tf.Variable(0, trainable=False)

        # Passing global_step to minimize() will increment it at each step.
        if opt_type is 'sgd':
            start_learning_rate = FLAG.lr
            half_cycle = 2000
            learning_rate = tf.train.exponential_decay(start_learning_rate, global_step, half_cycle, 0.5, staircase=True)
            opt = tf.train.MomentumOptimizer(learning_rate=learning_rate, momentum=0.9, use_nesterov=True)
        else:
            start_learning_rate = FLAG.lr
            half_cycle = 2000
            learning_rate = tf.train.exponential_decay(start_learning_rate, global_step, half_cycle, 0.5, staircase=True)
            opt = tf.train.AdamOptimizer(learning_rate=learning_rate)

        obj = vgg16.loss
        train_op = opt.minimize(obj, global_step=global_step)

        # progress bar
        ptrain = IntProgress()
        pval = IntProgress()
        display(ptrain)
        display(pval)
        ptrain.max = int(Xtrain.shape[0]/batch_size)
        pval.max = int(Xtest.shape[0]/batch_size)

        # re-initialize
        initialize_uninitialized(sess)

        # reset due to adding a new task
        patience_counter = 0
        current_best_val_loss = np.float('Inf')

        # optimize when the aggregated obj
        while(patience_counter < early_stop_patience and epoch_counter < epoch):

            # start training
            stime = time.time()
            bar_train = Bar('Training', max=int(Xtrain.shape[0]/batch_size), suffix='%(index)d/%(max)d - %(percent).1f%% - %(eta)ds')
            bar_val =  Bar('Validation', max=int(Xtest.shape[0]/batch_size), suffix='%(index)d/%(max)d - %(percent).1f%% - %(eta)ds')

            train_loss, train_accu = 0.0, 0.0
            for i in range(int(Xtrain.shape[0]/batch_size)):
                st = i*batch_size
                ed = (i+1)*batch_size
                loss, accu,_ = sess.run([obj, vgg16.accuracy, train_op],
                                    feed_dict={vgg16.x: Xtrain[st:ed,:],
                                                vgg16.y: Ytrain[st:ed,:],
                                                vgg16.is_train: True})
                train_loss += loss
                train_accu += accu
                ptrain.value +=1
                ptrain.description = "Training %s/%s" % (ptrain.value, ptrain.max)
            train_loss = train_loss/ptrain.value
            train_accu = train_accu/ptrain.value

            # validation
            val_loss = 0
            val_accu = 0
            for i in range(int(Xtest.shape[0]/batch_size)):
                st = i*batch_size
                ed = (i+1)*batch_size
                loss, accu = sess.run([obj, vgg16.accuracy],
                                    feed_dict={vgg16.x: Xtest[st:ed,:],
                                                vgg16.y: Ytest[st:ed,:],
                                                vgg16.is_train: False})
                val_loss += loss
                val_accu += accu
                pval.value += 1
                pval.description = "Testing %s/%s" % (pval.value, pval.value)
            val_loss = val_loss/pval.value
            val_accu = val_accu/pval.value

            # plot
            if epoch_counter%10 == 0:
                Xplot = sess.run(vgg16.pred,
                        feed_dict={vgg16.x: Xtrack[:,:],
                                    vgg16.y: Ytrack[:,:],
                                    vgg16.is_train: False})

                for i, fname in enumerate(track):
                    saveimg = skimage.transform.resize(Xplot[i],output_shape=(512,512),order=0,preserve_range=True,clip=False)
                    saveimg = label2rgb(saveimg)
                    imageio.imwrite(os.path.join(FLAG.save_dir,os.path.basename(fname)+"_pred_"+str(epoch_counter)+".png"), saveimg)
                    print(os.path.join(FLAG.save_dir,os.path.basename(fname)+"_pred_"+str(epoch_counter)+".png"))

            # early stopping check
            if (current_best_val_loss - val_loss) > min_delta:
                current_best_val_loss = val_loss
                patience_counter = 0
                saver.save(sess, checkpoint_path, global_step=epoch_counter)
                print("save in %s" % checkpoint_path)
            else:
                patience_counter += 1

            # shuffle Xtrain and Ytrain in the next epoch
            idx = np.random.permutation(Xtrain.shape[0])
            Xtrain, Ytrain = Xtrain[idx,:,:,:], Ytrain[idx,:]

            # epoch end
            epoch_counter += 1

            ptrain.value = 0
            pval.value = 0
            bar_train.finish()
            bar_val.finish()

            print("Epoch %s (%s), %s sec >> train loss: %.4f, train accu: %.4f, val loss: %.4f, val accu: %.4f" % (epoch_counter, patience_counter, round(time.time()-stime,2), train_loss, train_accu, val_loss, val_accu))

        # para_dict = sess.run(vgg16.para_dict)
        # np.save(os.path.join(FLAG.save_dir, "para_dict.npy"), para_dict)
        # print("save in %s" % os.path.join(FLAG.save_dir, "para_dict.npy"))

        # FLAG.optimizer = opt_type
        # FLAG.lr = start_learning_rate
        # FLAG.batch_size = batch_size
        # FLAG.epoch_end = epoch_counter
        # FLAG.val_loss = current_best_val_loss

        # header = ''
        # row = ''
        # for key in sorted(vars(FLAG)):
        #     if header is '':
        #         header = key
        #         row = str(getattr(FLAG, key))
        #     else:
        #         header += ","+key
        #         row += ","+str(getattr(FLAG,key))
        # row += "\n"
        # if os.path.exists("/home/cmchang/DLCV2018SPRING/hw3/model.csv"):
        #     with open("/home/cmchang/DLCV2018SPRING/hw3/model.csv", "a") as myfile:
        #         myfile.write(row)
        # else:
        #     with open("/home/cmchang/DLCV2018SPRING/hw3/model.csv", "w") as myfile:
        #         myfile.write(header)
        #         myfile.write(row)

if __name__ == '__main__':
    main()


'''**********************************************************************************
Given an integer array nums of length n and an integer target, find three integers in nums such that the sum is closest to target.

Return the sum of the three integers.

You may assume that each input would have exactly one solution.

Example 1:

Input: nums = [-1,2,1,-4], target = 1
Output: 2
Explanation: The sum that is closest to the target is 2. (-1 + 2 + 1 = 2).
**********************************************************************************'''

class Solution:
    def threeSumClosest(self, nums: List[int], target: int) -> int:
        nums.sort()
        ans = 0
        m = 111111111111
        for i in range(len(nums)-2):
            if i>0 and nums[i]==nums[i-1]:
                continue

            l = i+1
            r = len(nums)-1
            while(l<r):
                s = nums[i]+nums[l]+nums[r]
                if abs(target-s)<m:
                    m=abs(target-s)
                    ans = s
                if s<target:
                    l+=1
                elif s>target:
                    r-=1
                else:
                    return target
        return ans


import sys

if __name__ == '__main__':
    n, s = map(int, input().split())
    arr = list(map(int, input().split()))

    tmp = 0
    cnt = 0
    mn = sys.maxsize

    for i in range(n):
        tmp += arr[i]  # 현재 인덱스의 값을 더한다.
        cnt += 1
        if tmp >= s: # 만약 s를 넘는다면
            for j in range(i - cnt + 1, i, 1): # 더했던 인덱스들을 차례로 빼본다.
                if tmp - arr[j] >= s: # 만약 전의 인덱스를 빼고도 s를 넘는다면, 해당 인덱스 제외 (cnt, tmp 감소)
                    tmp -= arr[j]
                    cnt -= 1
                else: break
            mn = min(mn, cnt)

    if mn == sys.maxsize: print(0)
    else: print(mn)


import socket

class FGFS_Interface:
  """
  Interface with FlightGear

  Connects to FGFS socket, and can read properties.

  To launch FGFS: --telnet=socket,out,60,localhost,5555,udp
  """

  def __init__(self):
    self.fgfs_sock = None
    self.fgfs_file = None

  def connect(self, host, port):
    if host == "":
      host = socket.gethostname()

    self.fgfs_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    self.fgfs_sock.connect((host, port))
    self.fgfs_file = self.fgfs_sock.makefile()
    self.fgfs_sock.sendall(b"data\r\n")

  #def setprop(self, prop, value):
  #  self.fgfs_sock.sendall(bytes("get {0}\r\n".format(prop),"utf8"))

  def getprop(self, prop):
    self.fgfs_sock.sendall(bytes("get {0}\r\n".format(prop),"utf8"))
    return self.fgfs_file.readline().strip()

def updateFlightParams(db, fgfs):
  """
  Function that updates the flight parameters from FGFS in the database
  """
  p = dict() # parameters
  p["airspeed_kt"] = "/velocities/airspeed-kt"
  p["latitude_deg"] = "/position/latitude-deg"
  p["longitude_deg"] = "/position/longitude-deg"
  p["altitude_ft"] = "/position/altitude-ft"
  p["vspeed_fps"] = "/velocities/vertical-speed-fps"
  p["heading_deg"] = "/orientation/heading-deg"

  for param,path in p.items():
    db.setFlightParam(param, fgfs.getprop(path))


"""Advent of Code 2022, Day 14."""

import doctest
from collections import namedtuple
from collections.abc import Callable, Iterable

Point = namedtuple("Point", ["x", "y"])

def part1(problem_input: Iterable[str]) -> int:
    """Solution to part 1.

    >>> part1(sample_input)
    24
    """
    rock_formations = [
        [Point(*[int(n) for n in p.split(",")]) for p in f.split("->")]
        for f in problem_input
    ]
    occupied_space: dict[Point, str] = {}
    for formation in rock_formations:
        while len(formation) >= 2:
            start, end = formation[:2]
            if start.x == end.x:
                for y in range(min(start.y, end.y), max(start.y, end.y) + 1):
                    occupied_space[Point(start.x, y)] = "#"
            else:
                for x in range(min(start.x, end.x), max(start.x, end.x) + 1):
                    occupied_space[Point(x, start.y)] = "#"
            formation = formation[1:]

    origin = Point(500, 0)
    floor_y = max(p.y for p in occupied_space)
    is_void_breached = False
    while not is_void_breached:
        sand = origin
        while True:
            if sand.y > floor_y:
                is_void_breached = True
                break
            if Point(sand.x, sand.y + 1) not in occupied_space:
                sand = Point(sand.x, sand.y + 1)
            elif Point(sand.x - 1, sand.y + 1) not in occupied_space:
                sand = Point(sand.x - 1, sand.y + 1)
            elif Point(sand.x + 1, sand.y + 1) not in occupied_space:
                sand = Point(sand.x + 1, sand.y + 1)
            else:
                occupied_space[sand] = "o"
                break
        if origin in occupied_space:
            break
    return len([x for x in occupied_space.values() if x == "o"])

def part2(problem_input: Iterable[str]) -> int:
    """Solution to part 2.

    >>> part2(sample_input)
    93
    """
    rock_formations = [
        [Point(*[int(n) for n in p.split(",")]) for p in f.split("->")]
        for f in problem_input
    ]
    occupied_space: dict[Point, str] = {}
    for formation in rock_formations:
        while len(formation) >= 2:
            start, end = formation[:2]
            if start.x == end.x:
                for y in range(min(start.y, end.y), max(start.y, end.y) + 1):
                    occupied_space[Point(start.x, y)] = "#"
            else:
                for x in range(min(start.x, end.x), max(start.x, end.x) + 1):
                    occupied_space[Point(x, start.y)] = "#"
            formation = formation[1:]

    origin = Point(500, 0)
    floor_y = max(p.y for p in occupied_space) + 2
    while True:
        sand = origin
        while True:
            down = Point(sand.x, sand.y + 1)
            down_left = Point(sand.x - 1, sand.y + 1)
            down_right = Point(sand.x + 1, sand.y + 1)
            if down not in occupied_space and down.y != floor_y:
                sand = down
            elif down_left not in occupied_space and down_left.y != floor_y:
                sand = down_left
            elif down_right not in occupied_space and down_right.y != floor_y:
                sand = down_right
            else:
                occupied_space[sand] = "o"
                break
        if origin in occupied_space:
            break
    return len([x for x in occupied_space.values() if x == "o"])

def solve(func: Callable[[Iterable[str]], int]) -> None:
    with open("problem_input/day14.txt") as f:
        print(func(f))

if __name__ == "__main__":
    sample_input = """
498,4 -> 498,6 -> 496,6
503,4 -> 502,4 -> 502,9 -> 494,9
""".strip().splitlines()

    doctest.testmod(verbose=True)
    solve(part1)
    solve(part2)