import nestpy
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.axes as axes
import os

#Detector identification
detector = nestpy.DetectorExample_XENON10()
# detector = nestpy.VDetector()
# Performing NEST calculations according to the given detector example
nc = nestpy.NESTcalc(detector) #can also be left empty

#GetInteractionObject grabs the number for the interaction you want so you don't have to always reference the dictionary. Just type e.g., 'ion'
#It just changes the name to a number for nestpy to do its work.
def GetInteractionObject(name):
    name = name.lower()

    if name == 'er':
        raise ValueError("For 'er', specify either 'gammaray' or 'beta'")

    nest_interaction_number = dict(
        nr=0,
        wimp=1,
        b8=2,
        dd=3,
        ambe=4,
        cf=5,
        ion=6,
        gammaray=7,
        beta=8,
        ch3t=9,
        c14=10,
        kr83m=11,
        nonetype=12,
    )

    interaction_object = nestpy.INTERACTION_TYPE(nest_interaction_number[name])
    return interaction_object

#Once you have interaction, you can get yields
#Begin with np.vectorize so that we can get yields over a range of interaction types/energies/drift fields.
@np.vectorize
def GetYieldsVectorized(interaction, yield_type, **kwargs):
    # This function does nc.GetYields for the various interactions and arguments we pass into it
    # TODO: Look at docstrings

    interaction_object = GetInteractionObject(interaction)

    if 'energy' in kwargs.keys():
        if interaction_object == GetInteractionObject('nr') and kwargs['energy'] > 3e2:
            return np.nan
    if interaction_object == GetInteractionObject('gammaray') and kwargs['energy'] > 3e3:
            return np.nan
    if interaction_object == GetInteractionObject('beta') and kwargs['energy'] > 3e3:
            return np.nan
    yield_object = nc.GetYields(interaction = interaction_object, **kwargs)
    #this returns the yields for the type of yield we are considering be it ElectronYield or PhotonYield (an attribute of yield)
    return getattr(yield_object, yield_type)

#Gives us photon yield values
def PhotonYield(**kwargs):
    return GetYieldsVectorized(yield_type = 'PhotonYield', **kwargs)
#Gives electron yields
def ElectronYield(**kwargs):
    return GetYieldsVectorized(yield_type = 'ElectronYield', **kwargs)

def Yield(**kwargs):
    return {'photon' : PhotonYield(**kwargs),
            'electron' : ElectronYield(**kwargs),
           # What is missing?  Aren't there other parts of YieldObject?
           }

#we are able to do nestpy with 13 different interaction types and that's all we're going to use here.
# interaction_types = np.array(['nr','wimp','b8','dd','ambe','cf','ion', 'gammaray', 'beta', 'ch3t', 'c14', 'kr83m', 'nonetype'])

fields = [1, 2, 5, 10, 20, 50, 100, 200, 500, 1000, 2000, 5000, 10000, 20000]
energies = np.logspace(-1, 4, 2000)
energies = np.broadcast_to(energies, (len(fields), len(energies)))

file_path = './Images/'
directory = os.path.dirname(file_path)
if not os.path.exists(directory):
    os.makedirs(directory)

# ## Nuclear Recoils
def nr_subplot(x, y_photons, y_electrons, driftFields):
    plt.figure(1, figsize=(9,6))
    subplot1 = plt.subplot(111)
    plt.figure(2, figsize=(9,6))
    subplot2 = plt.subplot(111)
    for i in range(0, len(driftFields)-2):
        subplot1.plot(x[i,:], y_photons[i,:], label="{0} V/cm".format(driftFields[i]))
        subplot2.plot(x[i,:], y_electrons[i,:], label="{0} V/cm".format(driftFields[i]))

    #Just formatting to look presentable
    subplot1.set_xscale('log')
    subplot2.set_xscale('log')
    subplot1.set_ylim(bottom=0)
    subplot2.set_ylim(bottom=0)

    subplot1.legend(loc='best', fontsize= 9, ncol=3)
    subplot1.set_xlabel('Recoil Energy [keV]')
    subplot1.set_ylabel('Light Yields [n$_\gamma$/keV]')
    subplot1.set_title('Light Yields for Nuclear Recoils')
    subplot1.margins(0)
    subplot2.legend(loc='best', fontsize= 9, ncol=3)
    subplot2.set_xlabel('Recoil Energy [keV]')
    subplot2.set_title('Charge Yields for Nuclear Recoils')
    subplot2.set_ylabel('Charge Yield [n$_e$/keV]')
    subplot2.margins(0)

    #Here's where we are actually saving our plots!
    plt.figure(1)
    plt.tight_layout()
    plt.savefig(os.path.join(file_path, 'nr_LY.png'))
    plt.figure(2)
    plt.tight_layout()
    plt.savefig(os.path.join(file_path, 'nr_QY.png'))

def beta_subplot(x, y_photons, y_electrons, driftFields):
    plt.figure(3, figsize=(9,6))
    subplot1 = plt.subplot(111)
    plt.figure(4, figsize=(9,6))
    subplot2 = plt.subplot(111)
    for i in range(0, len(driftFields)):
        subplot1.plot(x[i,:], y_photons[i,:], label="{0} V/cm".format(driftFields[i]))
        subplot2.plot(x[i,:], y_electrons[i,:], label="{0} V/cm".format(driftFields[i]))

    #Just formatting to look presentable
    subplot1.set_xscale('log')
    subplot2.set_xscale('log')
    subplot1.set_ylim(bottom=0)
    subplot2.set_ylim(bottom=0)

    subplot1.legend(loc='best', fontsize= 9, ncol=1)
    subplot1.set_xlabel('Energy [keV]')
    subplot1.set_ylabel('Light Yields [n$_\gamma$/keV]')
    subplot1.set_title('Light Yields for $\\beta$ Electron Recoils')
    subplot1.margins(0)
    subplot2.legend(loc='best', fontsize= 9, ncol=1)
    subplot2.set_xlabel('Energy [keV]')
    subplot2.set_title('Charge Yields for $\\beta$ Electron Recoils')
    subplot2.set_ylabel('Charge Yield [n$_e$/keV]')
    subplot2.margins(0)

    #Here's where we are actually saving our plots!
    plt.figure(3)
    plt.tight_layout()
    plt.savefig(os.path.join(file_path, 'beta_LY.png'))
    plt.figure(4)
    plt.tight_layout()
    plt.savefig(os.path.join(file_path, 'beta_QY.png'))

# ## $\gamma$ electron recoils
def gamma_subplot(x, y_photons, y_electrons, driftFields):
    plt.figure(5, figsize=(9,6))
    subplot1 = plt.subplot(111)
    plt.figure(6, figsize=(9,6))
    subplot2 = plt.subplot(111)
    for i in range(0, len(driftFields)):
        subplot1.plot(x[i,:], y_photons[i,:], label="{0} V/cm".format(driftFields[i]))
        subplot2.plot(x[i,:], y_electrons[i,:], label="{0} V/cm".format(driftFields[i]))

    #Just formatting to look presentable
    subplot1.set_xscale('log')
    subplot2.set_xscale('log')
    subplot1.set_ylim(bottom=0)
    subplot2.set_ylim(bottom=0)

    subplot1.legend(loc='best', fontsize= 9, ncol=1)
    subplot1.set_xlabel('Energy [keV]')
    subplot1.set_ylabel('Light Yields [n$_\gamma$/keV]')
    subplot1.set_title('Light Yields for $\\gamma$ Electron Recoils')
    subplot1.margins(0)
    subplot2.legend(loc='best', fontsize= 9, ncol=1)
    subplot2.set_xlabel('Energy [keV]')
    subplot2.set_title('Charge Yields for $\\gamma$ Electron Recoils')
    subplot2.set_ylabel('Charge Yield [n$_e$/keV]')
    subplot2.margins(0)

    #Here's where we are actually saving our plots!
    plt.figure(5)
    plt.tight_layout()
    plt.savefig(os.path.join(file_path, 'gamma_LY.png'))
    plt.figure(6)
    plt.tight_layout()
    plt.savefig(os.path.join(file_path, 'gamma_QY.png'))

# ## $\alpha$-particle recoils
def alpha_subplot(x, y_photons, y_electrons, driftFields):
    plt.figure(7, figsize=(9,6))
    subplot1 = plt.subplot(111)
    plt.figure(8, figsize=(9,6))
    subplot2 = plt.subplot(111)
    for i in range(0, len(driftFields)):
        subplot1.plot(x[i,:], y_photons[i,:], label="{0} V/cm".format(driftFields[i]))
        subplot2.plot(x[i,:], y_electrons[i,:], label="{0} V/cm".format(driftFields[i]))

    #Just formatting to look presentable
    subplot1.set_xscale('log')
    subplot2.set_xscale('log')
    subplot1.set_ylim(bottom=0)
    subplot2.set_ylim(bottom=0)
    subplot1.set_xlim(1, 1e4)
    subplot2.set_xlim(1, 1e4)

    subplot1.legend(loc='best', fontsize= 9, ncol=2)
    subplot1.set_xlabel('Energy [keV]')
    subplot1.set_ylabel('Light Yields [n$_\gamma$/keV]')
    subplot1.set_title('Light Yields for $\\alpha$-Particle Nuclear Recoils')
    subplot1.margins(0)
    subplot2.legend(loc='best', fontsize= 9, ncol=2)
    subplot2.set_xlabel('Energy [keV]')
    subplot2.set_title('Charge Yields for $\\alpha$-Particle Nuclear Recoils')
    subplot2.set_ylabel('Charge Yield [n$_e$/keV]')
    subplot2.margins(0)

    #Here's where we are actually saving our plots!
    plt.figure(7)
    plt.tight_layout()
    plt.savefig(os.path.join(file_path, 'alpha_LY.png'))
    plt.figure(8)
    plt.tight_layout()
    plt.savefig(os.path.join(file_path, 'alpha_QY.png'))

# ## $^{206}$Pb nuclear recoils
def Pb_subplot(x, y_photons, y_electrons, driftFields):
    plt.figure(9, figsize=(9,6))
    subplot1 = plt.subplot(111)
    plt.figure(10, figsize=(9,6))
    subplot2 = plt.subplot(111)
    for i in range(0, len(driftFields)):
        subplot1.plot(x[i,:], y_photons[i,:], label="{0} V/cm".format(driftFields[i]))
        subplot2.plot(x[i,:], y_electrons[i,:], label="{0} V/cm".format(driftFields[i]))

    subplot1.set_xscale('log')
    subplot2.set_xscale('log')
    subplot1.set_ylim(0, 10)
    subplot1.set_xlim(1, 1e2)
    subplot2.set_ylim(bottom=0)
    subplot2.set_xlim(1, 1e2)

    subplot1.legend(loc='best', fontsize= 10, ncol=3)
    subplot1.set_xlabel('Energy [keV]')
    subplot1.set_ylabel('Light Yields [n$_\gamma$/keV]')
    subplot1.set_title('Light Yields for Nuclear Recoils from $^{206}$Pb')
    subplot2.legend(loc='best', fontsize= 10, ncol=3)
    subplot2.set_xlabel('Energy [keV]')
    subplot2.set_title('Charge Yields for Nuclear Recoils from $^{206}$Pb')
    subplot2.set_ylabel('Charge Yield [n$_e$/keV]')

    #Here's where we are actually saving our plots!
    plt.figure(9)
    plt.tight_layout()
    plt.savefig(os.path.join(file_path, '208Pb_LY.png'))
    plt.figure(10)
    plt.tight_layout()
    plt.savefig(os.path.join(file_path, '208Pb_QY.png'))

#Doing the actual functions
# The following are the energy and field ranges for each interaction in the detector we care about showing.
def makeplots():
    num_interactions = 5 #number of types of interactions we're studying, so we can make enough plots.
    nr_electrons = ElectronYield(interaction='nr', energy=energies.T, drift_field = fields).T/energies
    nr_photons = PhotonYield(interaction='nr', energy=energies.T, drift_field = fields).T/energies
    nr_subplot(energies, nr_photons, nr_electrons, fields)

    beta_electrons = ElectronYield(interaction='beta', energy=energies.T, drift_field = fields).T/energies
    beta_photons = PhotonYield(interaction='beta', energy=energies.T, drift_field = fields).T/energies
    beta_subplot(energies, beta_photons, beta_electrons, fields)

    gamma_electrons = ElectronYield(interaction='gammaray', energy=energies.T, drift_field = fields).T/energies
    gamma_photons = PhotonYield(interaction='gammaray', energy=energies.T, drift_field = fields).T/energies
    gamma_subplot(energies, gamma_photons, gamma_electrons, fields)

    alpha_electrons = ElectronYield(interaction='ion', Z=2, A = 4, energy=energies.T, drift_field = fields).T/energies
    alpha_photons = PhotonYield(interaction='ion',Z=2, A=4, energy=energies.T, drift_field = fields).T/energies
    alpha_subplot(energies, alpha_photons, alpha_electrons, fields)

    Pb_electrons = ElectronYield(interaction='ion', Z=82, A = 206, energy=energies.T, drift_field = fields).T/energies
    Pb_photons = PhotonYield(interaction='ion',Z=82, A=206, energy=energies.T, drift_field = fields).T/energies
    Pb_subplot(energies, Pb_photons, Pb_electrons, fields)

# makeplots()


#!/usr/bin/python3

from statistics import median

values = [] #array to store names
counter=0 #counter to check number of entries
while counter<13:
    value = int(input("Enter values: ")) #taking input from the user
    counter=counter+1 #incrementing the counter variable by 1
    values.append(value) #appending the entered value to the values array

maximum = max(values)
minimum = min(values)
average = sum(values)/len(values)
med = median(values)
sum = sum(values)
print('Max: ', maximum)
print("Min: ", minimum)
print("Avg: ", average)
print("Med: ", med)
print("Sum: ", sum)


from selenium.webdriver.common.by import By
import allure
from selenium.webdriver.support import expected_conditions
from selenium.webdriver.support.wait import WebDriverWait
import selenium
from model.project import Project
from datetime import datetime
import re

class ManageProjectPageHelper:
    def __init__(self, app):
        self.app = app
    def create_new_project(self, project):
        self.init_create_new_project()
        self.fill_project_fields(project)
        self.submit_create_project()

    def init_create_new_project(self):
        wd = self.app.wd
        with allure.step('Нажать Create New Project'):
            wd.find_element(By.XPATH, "//input[@value='Create New Project']").click()
            WebDriverWait(wd, 5).until(expected_conditions.presence_of_element_located((By.XPATH, "//td[contains(text(), 'Add Project')]")))

    def fill_project_fields(self, project):
        wd = self.app.wd
        with allure.step('Ввести уникальный Project Name'):
            wd.find_element(By.NAME, "name").send_keys(project.project_name)

    def submit_create_project(self):
        wd = self.app.wd
        with allure.step('Нажать Create New Project'):
            wd.find_element(By.XPATH, "//input[@type='submit' and @value='Add Project']").click()
            WebDriverWait(wd, 5).until(expected_conditions.invisibility_of_element_located((By.XPATH, "//td[contains(text(), 'Add Project')]")))

    def open_project(self, project_name):
        wd = self.app.wd
        with allure.step('Нажать на ссылку с названием проекта'):
            WebDriverWait(wd, 10).until(expected_conditions.element_to_be_clickable(
                (By.LINK_TEXT, f"{project_name}")))
            wd.find_element(By.LINK_TEXT, f"{project_name}").click()
            WebDriverWait(wd, 10).until(expected_conditions.presence_of_element_located(
                (By.NAME, "project_id")))

    def submit_delete_project(self):
        wd = self.app.wd
        with allure.step('Подтвердить удаление'):
            wd.find_element(By.XPATH, "//form[@action='manage_proj_delete.php']//input[@value='Delete Project']").click()
            WebDriverWait(wd, 10).until(expected_conditions.presence_of_element_located(
                (By.XPATH, "//table[@class='width100']//td[@class='form-title']")))
    def init_delete_project(self):
        wd = self.app.wd
        with allure.step('Нажать кнопку удаления проекта'):
            wd.find_element(By.XPATH, "//input[@value='Delete Project']").click()
            WebDriverWait(wd, 10).until(expected_conditions.element_to_be_clickable(
                (By.XPATH, "//form[@action='manage_proj_delete.php']//input[@value='Delete Project']")))

    def delete_project(self, project):
        self.open_project(project.project_name)
        self.init_delete_project()
        self.submit_delete_project()


import json
from random import randint

import vk_api
from vk_api.longpoll import VkLongPoll, VkEventType

import dialogflow_v2 as dialogflow

import os
from dotenv import load_dotenv

load_dotenv()

TOKEN_VK = os.environ.get('TOKEN_VK')
GOOGLE_APPLICATION_CREDENTIALS = os.environ.get('GOOGLE_APPLICATION_CREDENTIALS')

def handle_message(event, vk_api):
    client = dialogflow.SessionsClient()

    with open(GOOGLE_APPLICATION_CREDENTIALS, 'r') as file:
        config = json.load(file)

    session = client.session_path(config['project_id'], event.user_id)
    text_input = dialogflow.types.TextInput(
        text=event.text,
        language_code='ru'
    )
    query_input = dialogflow.types.QueryInput(text=text_input)

    response = client.detect_intent(
        session=session,
        query_input=query_input
    )

    if not response.query_result.intent.is_fallback:
        vk_api.messages.send(
            user_id=event.user_id,
            message=response.query_result.fulfillment_text,
            random_id=randint(1, 1000)
        )

if __name__ == "__main__":
    vk_session = vk_api.VkApi(token=TOKEN_VK)
    vk_api = vk_session.get_api()
    longpoll = VkLongPoll(vk_session)
    for event in longpoll.listen():
        if event.type == VkEventType.MESSAGE_NEW and event.to_me:
            handle_message(event, vk_api)


import select
import signal
import serial
from time import sleep
#from binascii import *

# export filtering
__all__ = ['GPS_reader']

def LOG(msg=''):
    print('[GPS_reader]%s' % msg)

class GPS_reader(object):
    # debug level
    DEBUG = 1
    # for interrupt handler and looping control
    _THREADED = False
    _STOP_EVENT = None
    # serial port
    PORT = '/dev/ttyUSB0'
    BAUDRATE = 9600
    # number of last NMEA info to store
    NMEA_NUM = 3
    # type of NMEA info to collect
    NMEA_TYPE = ("GPBOD", "GPBWC", "GPGGA", "GPGLL", "GPGSA", "GPGSV", "GPHDT",
                 "GPR00", "GPRMA", "GPRMB", "GPRMC", "GPRTE", "GPTRF", "GPSTN",
                 "GPVBW", "GPVTG", "GPWPL", "GPXTE", "GPZDA")

    def __init__(self):
        try:
            self._ser = serial.Serial(port=self.PORT, baudrate=self.BAUDRATE)
        except:
            LOG('[ERR] cannot open %s' % self.PORT)
            self._ser = None
        #
        self.infos = {}
        for nmea_t in self.NMEA_TYPE:
            self.infos[nmea_t] = []
        #
        self._listening = False
        self._reading = False
        #
        if self._ser:
            LOG(' reading position information over %s' % self.PORT)
        #
        if not self._THREADED:
            def handle_int(signum, frame):
                self.stop()
                LOG(' SIGINT: quitting')
            signal.signal(signal.SIGINT, handle_int)

    def stop(self):
        self._listening = False
        if self._ser:
            while self._reading:
                sleep(0.001)
            self._ser.close()

    def looping(self):
        if not self._listening:
            return False
        else:
            if not self._THREADED:
                return True
            elif hasattr(self._STOP_EVENT, 'is_set') \
            and not self._STOP_EVENT.is_set():
                return True
            return False

    def listen(self):
        if not self._ser:
            return
        self._listening = True
        while self.looping():
            try:
                self._reading = True
                l = self._ser.readline()
                self._reading = False
            # catch various errors that happen when the serial port is not
            # available in some way
            except (OSError, ValueError, serial.SerialException):
                break
            else:
                self.process(l)

    def process(self, buf='\n'):
        if self.DEBUG > 1:
            LOG(' process buf: %s' % buf)
        # grep NMEA_TYPE we want to get
        nmea_t = buf[1:6]
        if nmea_t in self.NMEA_TYPE:
            self.infos[nmea_t].append(buf[7:-2])
            if len(self.infos[nmea_t]) > self.NMEA_NUM:
                self.infos[nmea_t].remove( self.infos[nmea_t][0] )
            if self.DEBUG:
                LOG(' got %s: %s' % (nmea_t, buf[7:-2]))

    def get_last_info(self, nmea_t='GPRMC'):
        if not self._ser:
            return ''
        if nmea_t in self.NMEA_TYPE:
            if len(self.infos[nmea_t]):
                return self.infos[nmea_t][-1]


"""OpenAPI Tool."""

from typing import List, Optional

import requests

from llama_index.readers.schema.base import Document
from llama_index.tools.tool_spec.base import BaseToolSpec

class OpenAPIToolSpec(BaseToolSpec):
    """OpenAPI Tool

    This tool can be used to parse an OpenAPI spec for endpoints and operations
    Use the RequestsToolSpec to automate requests to the openapi server
    """

    spec_functions = ["load_openapi_spec"]

    def __init__(self, spec: Optional[dict] = None, url: Optional[str] = None):
        import yaml

        if spec and url:
            raise ValueError("Only provide one of OpenAPI dict or url")
        elif spec:
            pass
        elif url:
            response = requests.get(url).text
            spec = yaml.load(response, Loader=yaml.Loader)
        else:
            raise ValueError("You must provide a url or OpenAPI spec as a dict")

        parsed_spec = self.process_api_spec(spec)
        self.spec = Document(text=str(parsed_spec))

    def load_openapi_spec(self) -> List[Document]:
        """
        You are an AI agent specifically designed to retrieve information by making web requests to an API based on an OpenAPI specification.

        Here's a step-by-step guide to assist you in answering questions:

        1. Determine the base URL required for making the request

        2. Identify the relevant paths necessary to address the question

        3. Find the required parameters for making the request

        4. Perform the necessary requests to obtain the answer

        Returns:
            Document: A List of Document objects.
        """
        return [self.spec]

    def process_api_spec(self, spec: dict) -> dict:
        """Perform simplification and reduction on an OpenAPI specification.

        The goal is to create a more concise and efficient representation
        for retrieval purposes.
        """

        def reduce_details(details: dict) -> dict:
            reduced = {}
            if details.get("description"):
                reduced["description"] = details.get("description")
            if details.get("parameters"):
                reduced["parameters"] = [
                    param
                    for param in details.get("parameters", [])
                    if param.get("required")
                ]
            if "200" in details["responses"]:
                reduced["responses"] = details["responses"]["200"]
            return reduced

        def dereference_openapi(openapi_doc):
            """Dereferences a Swagger/OpenAPI document by resolving all $ref pointers."""
            try:
                import jsonschema
            except ImportError:
                raise ImportError(
                    "The jsonschema library is required to parse OpenAPI documents. "
                    "Please install it with `pip install jsonschema`."
                )

            resolver = jsonschema.RefResolver.from_schema(openapi_doc)

            def _dereference(obj):
                if isinstance(obj, dict):
                    if "$ref" in obj:
                        with resolver.resolving(obj["$ref"]) as resolved:
                            return _dereference(resolved)
                    return {k: _dereference(v) for k, v in obj.items()}
                elif isinstance(obj, list):
                    return [_dereference(item) for item in obj]
                else:
                    return obj

            dereferenced_doc = _dereference(openapi_doc)
            return dereferenced_doc

        spec = dereference_openapi(spec)
        endpoints = []
        for route, operations in spec["paths"].items():
            for operation, details in operations.items():
                if operation in ["get", "post", "patch"]:
                    endpoint_name = f"{operation.upper()} {route}"
                    description = details.get("description")
                    endpoints.append(
                        (endpoint_name, description, reduce_details(details))
                    )

        return {
            "servers": spec["servers"],
            "description": spec["info"].get("description"),
            "endpoints": endpoints,
        }


import gym
from gym import Env
from gym.error import DependencyNotInstalled

from mazeGame import MazeGame
from gym import spaces
import numpy as np

metadata = {
    "render_modes": ["human", "rgb_array"],
}

class MazeEnv(Env):

    # observation: [distance of sensor] * 6 + [end direction] * 1
    # action: {0: turn left, 1: turn right, 2: walk forward, 3: walk backward}
    def __init__(self, maze_idx, render_mode=None):
        self.metadata = metadata

        self.mazeGame = MazeGame(maze_idx)
        self.action_space = spaces.Discrete(4)
        self.observation_space = spaces.Box(low=0, high=200, shape=(7,))

        self.end_distance = self.mazeGame.end.distance(self.mazeGame.robot_pos)
        self.TIME_PENALTY = -1

        self.render_mode = render_mode
        self.screen = None
        self.ZOOM = 4
        self.SCREEN_SIZE = 200 * self.ZOOM

    def step(self, action):
        self.mazeGame.action(action)
        terminated, truncated = self.mazeGame.isEnd()
        observation = self.mazeGame.observation()
        observation = self._go2eo(observation)
        reward = self._reward()

        info = {'robot_pos': self.mazeGame.robot_pos, 'robot_face': self.mazeGame.robot_face,
                'start': self.mazeGame.start, 'end': self.mazeGame.end, 'time': self.mazeGame.time,
                'time_limit': self.mazeGame.time_limit, 'TIME_PENALTY': self.TIME_PENALTY}
        self.end_distance = self.mazeGame.end.distance(self.mazeGame.robot_pos)
        return observation, reward, terminated, truncated, info

    def reset(self, *args, **kwargs):
        self.end_distance = self.mazeGame.end.distance(self.mazeGame.start)
        self.mazeGame.reset()
        observation = self.mazeGame.observation()
        observation = self._go2eo(observation)
        info = {'robot_pos': self.mazeGame.robot_pos, 'robot_face': self.mazeGame.robot_face,
                'start': self.mazeGame.start, 'end': self.mazeGame.end, 'time': self.mazeGame.time,
                'time_limit': self.mazeGame.time_limit, 'TIME_PENALTY': self.TIME_PENALTY}
        return observation, info

    def render(self, mode='human'):
        if self.render_mode is None:
            gym.logger.warn(
                "You are calling render method without specifying any render mode. "
                "You can specify the render_mode at initialization, "
                f'e.g. gym("{self.spec.id}", render_mode="rgb_array")'
            )
            return

        try:
            import pygame
            from pygame import gfxdraw
        except ImportError:
            raise DependencyNotInstalled(
                "pygame is not installed, run `pip install gym[classic_control]`"
            )
        if self.screen is None:
            pygame.init()
            if self.render_mode == "human":
                pygame.display.init()
                self.screen = pygame.display.set_mode(
                    (self.SCREEN_SIZE, self.SCREEN_SIZE)
                )
            else:  # mode == "rgb_array"
                self.screen = pygame.Surface((self.SCREEN_SIZE, self.SCREEN_SIZE))

        self.screen.fill((255, 255, 255))
        for line in self.mazeGame.maze_lines:
            pygame.draw.line(
                self.screen,
                (0, 0, 0),
                self.mc2pc(*line[:2]),
                self.mc2pc(*line[2:]),
                1
            )
        pygame.draw.circle(
            self.screen,
            (0, 0, 255),
            self.mc2pc(self.mazeGame.start.x, self.mazeGame.start.y),
            5 * self.ZOOM,
        )
        pygame.draw.circle(
            self.screen,
            (255, 0, 0),
            self.mc2pc(self.mazeGame.end.x, self.mazeGame.end.y),
            5 * self.ZOOM,
        )
        pygame.draw.circle(
            self.screen,
            (0, 255, 0),
            self.mc2pc(self.mazeGame.robot_pos.x, self.mazeGame.robot_pos.y),
            5 * self.ZOOM,
        )

        pygame.draw.line(
            self.screen,
            (255, 255, 255),
            self.mc2pc(self.mazeGame.robot_pos.x, self.mazeGame.robot_pos.y),
            self.mc2pc(self.mazeGame.robot_pos.x + self.mazeGame.robot_face[0] * 5, self.mazeGame.robot_pos.y + self.mazeGame.robot_face[1] * 5),
            1
        )

        if self.render_mode == "human":
            pygame.event.pump()
            pygame.display.flip()

        elif self.render_mode == "rgb_array":
            return np.transpose(
                np.array(pygame.surfarray.pixels3d(self.screen)), axes=(1, 0, 2)
            )

    # game observation to environment observation
    def _go2eo(self, observation):
        return np.array(observation["sensor"] + [observation["end_direction"]], dtype=np.float32)

    #
    def _reward(self):
        new_distance = self.mazeGame.end.distance(self.mazeGame.robot_pos)
        move_reward = self.end_distance - new_distance
        return move_reward + self.TIME_PENALTY

    # math coordinate to pygame coordinate
    def mc2pc(self, x, y):
        return x * self.ZOOM, self.SCREEN_SIZE - y * self.ZOOM


"""
This is just a raw importer for the sprint.
"""

import os
import psycopg2
import psycopg2.extras
import simplejson

from elasticsearch import Elasticsearch
from elasticsearch.helpers import bulk_index

es = Elasticsearch()

class NominatimExporter(object):

    ITER_BY = 100

    def __init__(self, credentials, itersize=1000, limit=None, **kwargs):
        self.credentials = credentials
        print('***************** Export init ***************')
        self.conn = self.create_connexion()
        print('Connected')
        self.cur = self.conn.cursor("nominatim", cursor_factory=psycopg2.extras.DictCursor)
        print('Cursor created')
        self.cur.itersize = itersize
        self.limit = limit
        self.kwargs = kwargs

    def create_connexion(self):
        return psycopg2.connect(**self.credentials)

    def __enter__(self):
        sql = """SELECT osm_type,osm_id,class as osm_key,type as osm_value,admin_level,rank_search,rank_address,
            place_id,parent_place_id,calculated_country_code as country_code, postcode, housenumber,
            (extratags->'ref') as ref, street,
            ST_X(ST_Centroid(geometry)) as lon,
            ST_Y(ST_Centroid(geometry)) as lat,
            name->'name' as name,
            name->'name:de' as name_de,
            name->'name:fr' as name_fr,
            name->'name:en' as name_en,
            name->'short_name' as short_name,
            name->'official_name' as official_name, name->'alt_name' as alt_name,
            (extratags->'place') as extra_place
            FROM placex
            ORDER BY geometry
            """
        self.cur.execute(sql)
        print('Query executed with itersize', self.cur.itersize)
        return self

    def get_name_clause(self):
        return "name->'name' as name"

    def add_parents(self, row):
        if not "context_name" in row:
            row['context_name'] = []
        self.add_parent(row, row)

    def add_parent(self, child, row):
        if child['parent_place_id']:
            sql = """SELECT parent_place_id, type as osm_value, class as osm_key, {name}, admin_level FROM placex WHERE place_id={parent_place_id}"""
            sql = sql.format(**{
                'parent_place_id': child['parent_place_id'],
                'name': self.get_name_clause()
            })
            cur = self.conn.cursor(str(child['parent_place_id']), cursor_factory=psycopg2.extras.DictCursor)
            cur.execute(sql)
            parent = cur.fetchone()
            cur.close()
            self.add_parent_data(parent, row)
            self.add_parent(parent, row)

    def add_parent_data(self, parent, row):
        if not parent:
            return
        parent = dict(parent)
        name = parent['name']
        if name and not name in row['context_name']:
            row["context_name"].append(name)
        if (parent.get('osm_key') == "boundary"
           and parent.get('osm_value') == "administrative"
           and parent.get('admin_level') == 8):
            row['city'] = {
                'default': name
            }

    def __exit__(self, *args):
        self.cur.close()
        self.conn.close()

    def __iter__(self):
        return self.cur.__iter__()

    def to_json(self, raw):
        row = dict(raw)
        self.add_parents(row)
        row['coordinate'] = {
            'lon': raw['lon'],
            'lat': raw['lat']
        }
        row['name'] = {
            'default': raw['name'],
            'de': raw['name_de'],
            'fr': raw['name_fr'],
            'en': raw['name_en'],
        }
        row['id'] = '_'.join([raw['osm_type'], str(raw['osm_id'])])
        row['importance'] = 0.75 - raw['rank_search'] / 40
        row['context'] = {
            'default': row['context_name']
        }
        return row

class BaseConsumer(object):

    def __iter__(self):
        credentials = {
            'dbname': os.environ.get('DBNAME', 'nominatim')
        }
        with NominatimExporter(credentials) as exporter:
            for raw in exporter:
                yield exporter.to_json(raw)

class ESImporter(BaseConsumer):

    INDEX_CHUNK_SIZE = 10000
    ES_INDEX = "photon"

    def __call__(self):
        count = 0
        data = []
        print('Starting with ES index', self.ES_INDEX)
        for row in self:
            data.append(row)
            count += 1
            if count >= self.INDEX_CHUNK_SIZE:
                self.index(data)
                count = 0
                data = []
        if data:
            self.index(data)
        print('Done!')

    def format(self, row):
        return row

    def suggest_payload(self, row):
        return {
            "latlon": row['latlon'],
            'type': row['type'],
            'source': row['source'],
            'class': row['class'],
        }

    def exclude_row(self, row):
        return False

    def index(self, data):
        print('Start indexing batch of', len(data))
        bulk_index(es, data, index=self.ES_INDEX, doc_type='place', refresh=True)
        print('End indexing of current batch')

    def join(self, l, sep=", "):
        return sep.join([str(s) for s in l if str(s).strip()])

    def set_id(self, row, *args):
        row['_id'] = self.join(args, "_")

class JSONBatchDump(BaseConsumer):
    """Format for ES batch import."""

    MAX_ROWS = 100000

    def __call__(self):
        data = []
        count = 0
        index = 0
        print('Starting export')

        def do_write(data, index):
            with open('dump{}.eson'.format(index), mode='w', encoding='utf-8') as f:
                f.write('\n'.join(data))

        for row in self:
            data.append('{"index": {}}')
            data.append(simplejson.dumps(row))
            count += 1
            if count % 10000 == 0:
                print('Processed:', count, index)
            if count >= self.MAX_ROWS:
                do_write(data, index)
                data = []
                count = 0
                index += 1
        do_write(data)
        print('Done!')

if __name__ == "__main__":
    importer = JSONBatchDump()
    importer()


# -*- coding: utf-8 -*-
#
# Copyright (c) 2020 Ian Ottoway <ian@ottoway.dev>
# Copyright (c) 2014 Agostino Ruscito <ruscito@gmail.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#

import logging
from typing import Union
from reprlib import repr as _r

from . import Packet, DataFormatType
from . import (ResponsePacket, SendUnitDataResponsePacket, ReadTagServiceResponsePacket, RegisterSessionResponsePacket,
               UnRegisterSessionResponsePacket, ListIdentityResponsePacket, SendRRDataResponsePacket,
               MultiServiceResponsePacket, ReadTagFragmentedServiceResponsePacket, WriteTagServiceResponsePacket,
               WriteTagFragmentedServiceResponsePacket, GenericUnconnectedResponsePacket,
               GenericConnectedResponsePacket)
from ..exceptions import CommError, RequestError
from ..bytes_ import Pack, print_bytes_msg
from ..const import (EncapsulationCommand, INSUFFICIENT_PACKETS, DataItem, AddressItem, EXTENDED_SYMBOL, ELEMENT_TYPE,
                     Services, CLASS_TYPE, INSTANCE_TYPE, DataType, DataTypeSize, ConnectionManagerService,
                     ClassCode, Services, STRUCTURE_READ_REPLY, PRIORITY, TIMEOUT_TICKS, ATTRIBUTE_TYPE)

class RequestPacket(Packet):
    __log = logging.getLogger(f'{__module__}.{__qualname__}')
    _message_type = None
    _address_type = None
    _timeout = b'\x0a\x00'  # 10
    _encap_command = None
    _response_class = ResponsePacket
    _response_args = ()
    _response_kwargs = {}
    type_ = None
    VERBOSE_DEBUG = False

    def __init__(self, plc):
        super().__init__()
        self._msg = []  # message data
        self._plc = plc
        self.error = None

    def add(self, *value: bytes):
        self._msg.extend(value)
        return self

    @property
    def message(self) -> bytes:
        return b''.join(self._msg)

    def _build_request(self):
        msg = self._build_common_packet_format(addr_data=self._plc._target_cid)
        header = self._build_header(self._encap_command, len(msg))
        return header + msg

    def _build_header(self, command, length) -> bytes:
        """ Build the encapsulate message header

        The header is 24 bytes fixed length, and includes the command and the length of the optional data portion.

         :return: the header
        """
        try:
            return b''.join([
                command,
                Pack.uint(length),  # Length UINT
                Pack.udint(self._plc._session),  # Session Handle UDINT
                b'\x00\x00\x00\x00',  # Status UDINT
                self._plc._cfg['context'],  # Sender Context 8 bytes
                Pack.udint(self._plc._cfg['option']),  # Option UDINT
            ])

        except Exception as err:
            raise CommError('Failed to build request header') from err

    def _build_common_packet_format(self, addr_data=None) -> bytes:
        addr_data = b'\x00\x00' if addr_data is None else Pack.uint(len(addr_data)) + addr_data
        msg = self.message
        return b''.join([
            b'\x00\x00\x00\x00',  # Interface Handle: shall be 0 for CIP
            self._timeout,
            b'\x02\x00',  # Item count: should be at list 2 (Address and Data)
            self._address_type,
            addr_data,
            self._message_type,
            Pack.uint(len(msg)),
            msg
        ])

    def _send(self, message):

        try:
            if self.VERBOSE_DEBUG:
                self.__log.debug(print_bytes_msg(message, '>>> SEND >>>'))
            self._plc._sock.send(message)
        except Exception as err:
            raise CommError('failed to send message') from err

    def _receive(self):
        """
        socket receive
        :return: reply data
        """
        try:
            reply = self._plc._sock.receive()
        except Exception as err:
            raise CommError('failed to receive reply') from err
        else:
            if self.VERBOSE_DEBUG:
                self.__log.debug(print_bytes_msg(reply, '<<< RECEIVE <<<'))
            return reply

    def send(self) -> ResponsePacket:
        if not self.error:
            self._send(self._build_request())
            self.__log.debug(f'Sent: {self!r}')
            reply = self._receive()
            response = self._response_class(reply, *self._response_args, **self._response_kwargs)
        else:
            response = self._response_class(*self._response_args, **self._response_kwargs)
            response._error = self.error
        self.__log.debug(f'Received: {response!r}')
        return response

    def __repr__(self):
        return f'{self.__class__.__name__}(message={_r(self._msg)})'

    __str__ = __repr__

class SendUnitDataRequestPacket(RequestPacket):
    __log = logging.getLogger(f'{__module__}.{__qualname__}')
    _message_type = DataItem.connected
    _address_type = AddressItem.connection
    _response_class = SendUnitDataResponsePacket
    _encap_command = EncapsulationCommand.send_unit_data

    def __init__(self, plc):
        super().__init__(plc)
        self._msg = [Pack.uint(plc._sequence), ]

class ReadTagServiceRequestPacket(SendUnitDataRequestPacket):
    __log = logging.getLogger(f'{__module__}.{__qualname__}')
    type_ = 'read'
    _response_class = ReadTagServiceResponsePacket

    def __init__(self, plc):
        super().__init__(plc)
        self.tag = None
        self.elements = None
        self.tag_info = None
        self.request_id = None

    def add(self, tag, request_path, elements, tag_info, request_id):
        self.tag = tag
        self.elements = elements
        self.tag_info = tag_info
        self.request_id = request_id
        if request_path is None:
            self.error = 'Invalid Tag Request Path'

        super().add(
            Services.read_tag,
            request_path,
            Pack.uint(self.elements),
        )

    def send(self):
        if not self.error:
            self._send(self._build_request())
            self.__log.debug(f'Sent: {self!r}')
            reply = self._receive()
            response = ReadTagServiceResponsePacket(reply, elements=self.elements, tag_info=self.tag_info, tag=self.tag)
        else:
            response = ReadTagServiceResponsePacket(tag=self.tag)
            response._error = self.error
        self.__log.debug(f'Received: {response!r}')
        return response

    def __repr__(self):
        return f'{self.__class__.__name__}(tag={self.tag!r}, elements={self.elements!r})'

class ReadTagFragmentedServiceRequestPacket(SendUnitDataRequestPacket):
    __log = logging.getLogger(f'{__module__}.{__qualname__}')
    type_ = 'read'
    _response_class = ReadTagFragmentedServiceResponsePacket

    def __init__(self, plc):
        super().__init__(plc)
        self.tag = None
        self.elements = None
        self.tag_info = None
        self.request_path = None
        self.request_id = None

    def add(self, tag, request_path, elements, tag_info, request_id):
        self.tag = tag
        self.elements = elements
        self.tag_info = tag_info
        self.request_path = request_path
        self.request_id = request_id
        if self.request_path is None:
            self.error = 'Invalid Tag Request Path'

    def send(self):
        if not self.error:
            offset = 0
            responses = []
            while offset is not None:
                self._msg.extend([Services.read_tag_fragmented,
                                  self.request_path,
                                  Pack.uint(self.elements),
                                  Pack.dint(offset)])
                self._send(self._build_request())
                self.__log.debug(f'Sent: {self!r} (offset={offset})')
                reply = self._receive()
                response = ReadTagFragmentedServiceResponsePacket(reply, self.tag_info, self.elements)
                self.__log.debug(f'Received: {response!r}')
                responses.append(response)
                if response.service_status == INSUFFICIENT_PACKETS:
                    offset += len(response.bytes_)
                    self._msg = [Pack.uint(self._plc._sequence)]
                else:
                    offset = None
            if all(responses):
                final_response = responses[-1]
                final_response.bytes_ = b''.join(resp.bytes_ for resp in responses)
                final_response.parse_bytes()
                self.__log.debug(f'Reassembled Response: {final_response!r}')
                return final_response

        failed_response = ReadTagServiceResponsePacket()
        failed_response._error = self.error or 'One or more fragment responses failed'
        self.__log.debug(f'Reassembled Response: {failed_response!r}')
        return failed_response

    def __repr__(self):
        return f'{self.__class__.__name__}(tag={self.tag!r}, elements={self.elements!r})'

class WriteTagServiceRequestPacket(SendUnitDataRequestPacket):
    __log = logging.getLogger(f'{__module__}.{__qualname__}')
    type_ = 'write'
    _response_class = WriteTagServiceResponsePacket

    def __init__(self, plc):
        super().__init__(plc)
        self.tag = None
        self.elements = None
        self.tag_info = None
        self.value = None
        self.data_type = None
        self.request_id = None

    def add(self, tag, request_path, value, elements, tag_info, request_id, bits_write=None):
        self.tag = tag
        self.elements = elements
        self.tag_info = tag_info
        self.value = value
        self.request_id = request_id
        request_path = request_path
        if request_path is None:
            self.error = 'Invalid Tag Request Path'

        else:
            if bits_write:
                request_path = _make_write_data_bit(tag_info, value, request_path)
                data_type = 'BOOL'
            else:
                request_path, data_type = _make_write_data_tag(tag_info, value, elements, request_path)

            super().add(
                request_path,
            )
            self.data_type = data_type

    def __repr__(self):
        return f'{self.__class__.__name__}(tag={self.tag!r}, value={_r(self.value)}, elements={self.elements!r})'

class WriteTagFragmentedServiceRequestPacket(SendUnitDataRequestPacket):
    __log = logging.getLogger(f'{__module__}.{__qualname__}')
    type_ = 'write'
    _response_class = WriteTagFragmentedServiceResponsePacket

    def __init__(self, plc):
        super().__init__(plc)
        self.tag = None
        self.value = None
        self.elements = None
        self.tag_info = None
        self.request_path = None
        self.data_type = None
        self.segment_size = None
        self.request_id = None
        self._packed_type = None

    def add(self, tag, request_path, value, elements, tag_info, request_id):
        try:
            self.data_type = tag_info['data_type_name']
            self.tag = tag
            self.value = value
            self.elements = elements
            self.tag_info = tag_info
            self.request_path = request_path
            self.request_id = request_id

            if tag_info['tag_type'] == 'struct':
                self._packed_type = STRUCTURE_READ_REPLY + Pack.uint(tag_info['data_type']['template']['structure_handle'])
            else:
                self._packed_type = Pack.uint(DataType[self.data_type])

            if self.request_path is None:
                self.error = 'Invalid Tag Request Path'
        except Exception as err:
            self.__log.exception('Failed adding request')
            self.error = err

    def send(self):
        if not self.error:
            responses = []
            segment_size = self._plc.connection_size - (len(self.request_path) + len(self._packed_type)
                                                        + 9)  # 9 = len of other stuff in the path

            pack_func = Pack[self.data_type] if self.tag_info['tag_type'] == 'atomic' else lambda x: x
            segments = (self.value[i:i+segment_size]
                        for i in range(0, len(self.value), segment_size))

            offset = 0
            elements_packed = Pack.uint(self.elements)

            for i, segment in enumerate(segments, start=1):
                segment_bytes = b''.join(pack_func(s) for s in segment) if not isinstance(segment, bytes) else segment
                self._msg.extend((
                    Services.write_tag_fragmented,
                    self.request_path,
                    self._packed_type,
                    elements_packed,
                    Pack.dint(offset),
                    segment_bytes
                ))

                self._send(self._build_request())
                self.__log.debug(f'Sent: {self!r} (part={i} offset={offset})')
                reply = self._receive()
                response = WriteTagFragmentedServiceResponsePacket(reply)
                self.__log.debug(f'Received: {response!r}')
                responses.append(response)
                offset += len(segment_bytes)
                self._msg = [Pack.uint(self._plc._sequence), ]

            if all(responses):
                final_response = responses[-1]
                self.__log.debug(f'Reassembled Response: {final_response!r}')
                return final_response

        failed_response = WriteTagFragmentedServiceResponsePacket()
        failed_response._error = self.error or 'One or more fragment responses failed'
        self.__log.debug(f'Reassembled Response: {failed_response!r}')
        return failed_response

class MultiServiceRequestPacket(SendUnitDataRequestPacket):
    __log = logging.getLogger(f'{__module__}.{__qualname__}')
    type_ = 'multi'
    _response_class = MultiServiceResponsePacket

    def __init__(self, plc):
        super().__init__(plc)
        self.tags = []
        self._msg.extend((
            Services.multiple_service_request,  # the Request Service
            Pack.usint(2),  # the Request Path Size length in word
            CLASS_TYPE["8-bit"],
            ClassCode.message_router,
            INSTANCE_TYPE["8-bit"],
            b'\x01',  # Instance 1
        ))
        self._message = None
        self._msg_errors = None

    @property
    def message(self) -> bytes:
        return self._message

    def build_message(self, tags):
        rp_list, errors = [], []
        for tag in tags:
            if tag['rp'] is None:
                errors.append(f'Unable to create request path {tag["tag"]}')
            else:
                rp_list.append(tag['rp'])

        offset = len(rp_list) * 2 + 2
        offsets = []
        for rp in rp_list:
            offsets.append(Pack.uint(offset))
            offset += len(rp)

        msg = self._msg + [Pack.uint(len(rp_list))] + offsets + rp_list
        return b''.join(msg)

    def add_read(self, tag, request_path, elements, tag_info, request_id):
        if request_path is not None:
            rp = Services.read_tag + request_path + Pack.uint(elements)
            _tag = {
                'tag': tag,
                'elements': elements,
                'tag_info': tag_info,
                'rp': rp,
                'service': 'read',
                'request_id': request_id
            }
            message = self.build_message(self.tags + [_tag])
            if len(message) < self._plc.connection_size:
                self._message = message
                self.tags.append(_tag)
                return True
            else:
                return False
        else:
            self.__log.error(f'Failed to create request path for {tag}')
            raise RequestError('Failed to create request path')

    def add_write(self, tag, request_path, value, elements, tag_info, request_id, bits_write=None):
        if request_path is not None:
            if bits_write:
                data_type = tag_info['data_type']
                request_path = _make_write_data_bit(tag_info, value, request_path)
            else:
                request_path, data_type = _make_write_data_tag(tag_info, value, elements, request_path)

            _tag = {'tag': tag, 'elements': elements, 'tag_info': tag_info, 'rp': request_path, 'service': 'write',
                    'value': value, 'data_type': data_type, 'request_id': request_id}

            message = self.build_message(self.tags + [_tag])
            if len(message) < self._plc.connection_size:
                self._message = message
                self.tags.append(_tag)
                return True
            else:
                return False

        else:
            self.__log.error(f'Failed to create request path for {tag}')
            raise RequestError('Failed to create request path')

    def send(self):
        if not self._msg_errors:
            request = self._build_request()
            self._send(request)
            self.__log.debug(f'Sent: {self!r}')
            reply = self._receive()
            response = MultiServiceResponsePacket(reply, tags=self.tags)
        else:
            self.error = f'Failed to create request path for: {", ".join(self._msg_errors)}'
            response = MultiServiceResponsePacket()
            response._error = self.error

        self.__log.debug(f'Received: {response!r}')
        return response

def _make_write_data_tag(tag_info, value, elements, request_path, fragmented=False):
    data_type = tag_info['data_type']
    if tag_info['tag_type'] == 'struct':
        if not isinstance(value, bytes):
            raise RequestError('Writing UDTs only supports bytes for value')
        _dt_value = b'\xA0\x02' + Pack.uint(tag_info['data_type']['template']['structure_handle'])
        data_type = tag_info['data_type']['name']

    elif data_type not in DataType:
        raise RequestError("Unsupported data type")
    else:
        _dt_value = Pack.uint(DataType[data_type])

    service = Services.write_tag_fragmented if fragmented else Services.write_tag

    rp = b''.join((service,
                   request_path,
                   _dt_value,
                   Pack.uint(elements),
                   value))
    return rp, data_type

def _make_write_data_bit(tag_info, value, request_path):
    mask_size = DataTypeSize.get(tag_info['data_type'])
    if mask_size is None:
        raise RequestError(f'Invalid data type {tag_info["data_type"]} for writing bits')

    or_mask, and_mask = value
    return b''.join((
        Services.read_modify_write,
        request_path,
        Pack.uint(mask_size),
        Pack.ulint(or_mask)[:mask_size],
        Pack.ulint(and_mask)[:mask_size]
        ))

class SendRRDataRequestPacket(RequestPacket):
    __log = logging.getLogger(f'{__module__}.{__qualname__}')
    _message_type = DataItem.unconnected
    _address_type = AddressItem.uccm
    _encap_command = EncapsulationCommand.send_rr_data
    _response_class = SendRRDataResponsePacket

    def _build_common_packet_format(self, addr_data=None) -> bytes:
        return super()._build_common_packet_format(addr_data=None)

class RegisterSessionRequestPacket(RequestPacket):
    __log = logging.getLogger(f'{__module__}.{__qualname__}')
    _encap_command = EncapsulationCommand.register_session
    _response_class = RegisterSessionResponsePacket

    def _build_common_packet_format(self, addr_data=None) -> bytes:
        return self.message

class UnRegisterSessionRequestPacket(RequestPacket):
    __log = logging.getLogger(f'{__module__}.{__qualname__}')
    _encap_command = EncapsulationCommand.unregister_session
    _response_class = UnRegisterSessionResponsePacket

    def _build_common_packet_format(self, addr_data=None) -> bytes:
        return b''

    def _receive(self):
        return b''

class ListIdentityRequestPacket(RequestPacket):
    __log = logging.getLogger(f'{__module__}.{__qualname__}')
    _encap_command = EncapsulationCommand.list_identity
    _response_class = ListIdentityResponsePacket

    def _build_common_packet_format(self, addr_data=None) -> bytes:
        return b''

class GenericConnectedRequestPacket(SendUnitDataRequestPacket):
    __log = logging.getLogger(f'{__module__}.{__qualname__}')
    _response_class = GenericConnectedResponsePacket

    def __init__(self, plc):
        super().__init__(plc)
        self.service = None
        self.class_code = None
        self.instance = None
        self.attribute = None
        self.request_data = None

    def build(self,
              service: bytes,
              class_code: bytes,
              instance: bytes,
              attribute: bytes = b'',
              request_data: bytes = b'',
              data_format: DataFormatType = None):

        self._response_kwargs = {'data_format': data_format}
        self.class_code = class_code
        self.instance = instance
        self.attribute = attribute
        self.service = service
        self.request_data = request_data
        req_path = request_path(class_code, instance, attribute)

        self.add(service, req_path, request_data)

class GenericUnconnectedRequestPacket(SendRRDataRequestPacket):
    __log = logging.getLogger(f'{__module__}.{__qualname__}')
    _response_class = GenericUnconnectedResponsePacket

    def __init__(self, plc):
        super().__init__(plc)
        self.service = None
        self.class_code = None
        self.instance = None
        self.attribute = None
        self.request_data = None

    def build(self,
              service: Union[int, bytes],
              class_code: Union[int, bytes],
              instance: Union[int, bytes],
              attribute: Union[int, bytes] = b'',
              request_data: bytes = b'',
              route_path: bytes = b'',
              unconnected_send: bool = False,
              data_format: DataFormatType = None):
        self._response_kwargs = {'data_format': data_format}
        self.class_code = class_code
        self.instance = instance
        self.attribute = attribute
        self.service = service
        self.request_data = request_data
        req_path = request_path(class_code, instance, attribute)

        if unconnected_send:
            self.add(wrap_unconnected_send(b''.join((service, req_path, request_data)), route_path))
        else:
            self.add(service, req_path, request_data, route_path)

def wrap_unconnected_send(message, route_path):
    rp = request_path(class_code=ClassCode.connection_manager, instance=b'\x01')
    msg_len = len(message)
    return b''.join(
        [
            ConnectionManagerService.unconnected_send,
            rp,
            PRIORITY,
            TIMEOUT_TICKS,
            Pack.uint(msg_len),
            message,
            b'\x00' if msg_len % 2 else b'',
            route_path
        ]
    )

def request_path(class_code: Union[int, bytes], instance: Union[int, bytes],
                 attribute: Union[int, bytes] = b'', data: bytes = b''):

    path = [encode_segment(class_code, CLASS_TYPE), encode_segment(instance, INSTANCE_TYPE)]

    if attribute:
        path.append(encode_segment(attribute, ATTRIBUTE_TYPE))

    if data:
        path.append(data)

    return Pack.epath(b''.join(path))

def encode_segment(segment: Union[bytes, int], segment_types: dict):
    if isinstance(segment, int):
        if segment <= 0xff:
            segment = Pack.usint(segment)
        elif segment <= 0xffff:
            segment = Pack.uint(segment)
        elif segment <= 0xfffffffff:
            segment = Pack.dint(segment)
        else:
            raise RequestError('Invalid segment value')

    _type = segment_types.get(len(segment))

    if _type is None:
        raise RequestError('Segment value not valid for segment type')

    return _type + segment


import os
import math
import torch
import itertools
import torchvision
from torch.autograd import Variable
from TripletNet import TripletNet

def shuffle_data(data, seed = 0):
	image_ids, labels = data
	shuffled_image_ids = []
	shuffled_labels = []
	num_images = len(image_ids)
	torch.manual_seed(seed)
	perm = list(torch.randperm(num_images))
	for i in range(num_images):
		shuffled_image_ids.append(image_ids[perm[i]])
		shuffled_labels.append(labels[perm[i]])
	return shuffled_image_ids, shuffled_labels

def make_minibatches(data, minibatch_size = 16,  seed = 0, shuffle = 'random'):
	X, Y = data
	m = len(X)
	minibatches = []
	if shuffle == 'sequential':
		shuffled_X, shuffled_Y = X, Y

	elif shuffle == 'random':
		shuffled_X, shuffled_Y = shuffle_data(data, seed = seed)

	num_complete_minibatches = math.floor(m/minibatch_size)
	for k in range(0, num_complete_minibatches):
		minibatch_X = shuffled_X[k * minibatch_size : k * minibatch_size + minibatch_size]
		minibatch_Y = shuffled_Y[k * minibatch_size : k * minibatch_size + minibatch_size]
		minibatches.append((minibatch_X, minibatch_Y))

	rem_size = m - num_complete_minibatches * minibatch_size
	if m % minibatch_size != 0:
		minibatch_X = shuffled_X[num_complete_minibatches * minibatch_size : m]
		minibatch_Y = shuffled_Y[num_complete_minibatches * minibatch_size : m]
		minibatches.append((minibatch_X, minibatch_Y))

	return minibatches

def batch2images_tensor(minibatch_X, dataloader, gpu_device):
	minibatch_size = len(minibatch_X)
	images_tensor = torch.zeros(minibatch_size, 3, 224, 224)
	for i in range(minibatch_size):
		x = minibatch_X[i]
		x_image = dataloader.get_image(x)
		images_tensor[i, :, :, :] = x_image
	images_tensor = Variable(images_tensor)
	if torch.cuda.is_available():
		with torch.cuda.device(gpu_device):
			images_tensor = images_tensor.cuda()
	return images_tensor

def gen_triplets(minibatch, id2embeds, embedding_dim, gpu_device, mode = 'all'):
	X, Y = minibatch
	Y_prod = itertools.product(Y, repeat=3)
	X_prod = itertools.product(X, repeat=3)
	triplet = []
	for x, y  in zip(X_prod, Y_prod):
		xa, xp, xn = x
		ya, yp, yn = y
		if (ya == yp) and (ya!=yn) and (xa!=xp):
			triplet.append((xa, xp, xn))

	num_triplets = len(triplet)
	anchor = torch.zeros(num_triplets, embedding_dim)
	positive = torch.zeros(num_triplets, embedding_dim)
	negative = torch.zeros(num_triplets, embedding_dim)
	if torch.cuda.is_available():
		with torch.cuda.device(gpu_device):
			anchor = anchor.cuda()
			positive = positive.cuda()
			negative = negative.cuda()

	for i in range(num_triplets):
		xa, xp, xn = triplet[i]
		anchor[i, :] = id2embeds[xa]
		positive[i, :] = id2embeds[xp]
		negative[i, :] = id2embeds[xn]

	return anchor, positive, negative

def label2embeds_list2dict(labels_list, embeds):
	label2embeds = {}
	num_labels = len(labels_list)
	for i in range(num_labels):
		label = labels_list[i]
		label2embeds[label] = embeds[i, :]
	return label2embeds

def final_label2embeds(triplet_net, train_dataloader, gpu_device):
	labels_list = []
	image_ids = []
	for label, images_list in train_dataloader.images_dict.items():
		image_ids.append(images_list[0])
		labels_list.append(label)

	images_tensor = batch2images_tensor(image_ids, train_dataloader, gpu_device)
	with torch.no_grad():
		embeds = triplet_net.embedding(images_tensor)
	label2embeds = label2embeds_list2dict(labels_list, embeds)
	return label2embeds

def who_is_it(label2embeds, embed):
	labels = []
	num_labels = len(label2embeds)
	embedding_dim = embed.shape[0]
	embeds = torch.zeros(num_labels, embedding_dim)
	i = 0
	for label, cur_embed in label2embeds.items():
		labels.append(label)
		embeds[i, :] = cur_embed
		i += 1
	dist = torch.pow(embeds - embed, 2).sum(dim = 1)
	index = torch.argmin(dist).tolist()
	return labels[index]

def accuracy(data, dataloader, label2embeds, triplet_net, gpu_device):
	image_ids, Y = data
	num_data = len(Y)
	embedding_dim = triplet_net.embedding.embedding_dim
	embeds = torch.zeros(num_data, embedding_dim)
	minibatch_size = 32
	minibatches = make_minibatches(data, minibatch_size = minibatch_size,  seed = 0, shuffle = 'sequential')
	start = 0
	end = 0
	for cur_minibatch in minibatches:
		minibatch_X, _ = cur_minibatch
		cur_minibatch_size = len(minibatch_X)
		end += cur_minibatch_size
		images_tensor = batch2images_tensor(minibatch_X, dataloader, gpu_device)
		with torch.no_grad():
			embeds[start:end, :] = triplet_net(images_tensor)
		start = end

	acc = 0
	# pred = []

	for i in range(num_data):
		embed = embeds[i]
		target_label = Y[i]
		predicted_label = who_is_it(label2embeds, embed)
		if predicted_label == target_label:
			acc += 1
			# pred.append(predicted_label)
	return acc, num_data


from __future__ import absolute_import
import logging
import os
import shutil
import json
from collections import deque
import time
import re
import datetime
import traceback
import numpy as np
from typing import Dict, Optional, List, Tuple, Union, Iterable, Any, Set
import pdb

import torch
import torch.optim.lr_scheduler
from torch.nn.parallel import replicate, parallel_apply
from torch.nn.parallel.scatter_gather import scatter_kwargs, gather
from tensorboardX import SummaryWriter

from allennlp.common import Params
from allennlp.common.checks import ConfigurationError
from allennlp.common.util import peak_memory_mb, gpu_memory_mb
from allennlp.common.tqdm import Tqdm
from allennlp.data.instance import Instance
from allennlp.data.iterators.data_iterator import DataIterator
from allennlp.models.model import Model
from allennlp.nn import util
from allennlp.training.learning_rate_schedulers import LearningRateScheduler
from allennlp.training.optimizers import Optimizer

from AttentionSegmentation.commons.trainer_utils import is_sparse,\
    sparse_clip_norm, move_optimizer_to_cuda, TensorboardWriter
# from AttentionSegmentation.visualization.visualize_attns \
#     import html_visualizer
from AttentionSegmentation.model.attn2labels import BasePredictionClass
logger = logging.getLogger(__name__)

TQDM_COLUMNS = 200

class Trainer(object):
    def __init__(self,
                 model: Model,
                 optimizer: torch.optim.Optimizer,
                 iterator: DataIterator,
                 train_dataset: Iterable[Instance],
                 validation_dataset: Optional[Iterable[Instance]] = None,
                 segmenter: Optional[BasePredictionClass] = None,
                 patience: Optional[int] = None,
                 validation_metric: str = "-loss",
                 num_epochs: int = 20,
                 base_dir: Optional[str] = None,
                 num_serialized_models_to_keep: Optional[int] = None,
                 cuda_device: Union[int, List] = -1,
                 grad_norm: Optional[float] = None,
                 grad_clipping: Optional[float] = None,
                 learning_rate_scheduler: Optional[LearningRateScheduler] = None
                 ) -> None:
        """
        Parameters
        ----------
        model : ``Model``, required.
            An AllenNLP model to be optimized. Pytorch Modules can also be optimized if
            their ``forward`` method returns a dictionary with a "loss" key, containing a
            scalar tensor representing the loss function to be optimized.
        optimizer : ``torch.nn.Optimizer``, required.
            An instance of a Pytorch Optimizer, instantiated with the parameters of the
            model to be optimized.
        iterator : ``DataIterator``, required.
            A method for iterating over a ``Dataset``, yielding padded indexed batches.
        train_dataset : ``Dataset``, required.
            A ``Dataset`` to train on. The dataset should have already been indexed.
        validation_dataset : ``Dataset``, optional, (default = None).
            A ``Dataset`` to evaluate on. The dataset should have already been indexed.
        segmenter : ``BasePredictionClass`` for converting attention segmentation to labels.
            None if we don't want to visualize
        patience : Optional[int] > 0, optional (default=None)
            Number of epochs to be patient before early stopping: the training is stopped
            after ``patience`` epochs with no improvement. If given, it must be ``> 0``.
            If None, early stopping is disabled.
        validation_metric : str, optional (default="loss")
            Validation metric to measure for whether to stop training using patience
            and whether to serialize an ``is_best`` model each epoch. The metric name
            must be prepended with either "+" or "-", which specifies whether the metric
            is an increasing or decreasing function.
            Allowed values are -loss, +precision-overall, +recall-overall, +f1-measure-overall
        num_epochs : int, optional (default = 20)
            Number of training epochs.
        base_dir : str, optional (default=None)
            Path to directory for the experiment. Models will be saved at
            {base_dir}/models, and won't be saved if this parameter is not passed.
        num_serialized_models_to_keep: int, optional (default=None)
            The number of models to keep during training
        cuda_device : ``int``, optional (default = -1)
            An integer specifying the CUDA device to use. If -1, the CPU is used.
        grad_norm : ``float``, optional, (default = None).
            If provided, gradient norms will be rescaled to have a maximum of this value.
        grad_clipping : ``float``, optional (default = ``None``).
            If provided, gradients will be clipped `during the backward pass` to have an (absolute)
            maximum of this value.  If you are getting ``NaNs`` in your gradients during training
            that are not solved by using ``grad_norm``, you may need this.
        learning_rate_scheduler : ``PytorchLRScheduler``, optional, (default = None)
            A Pytorch learning rate scheduler. The learning rate will be decayed with respect to
            this schedule at the end of each epoch. If you use
            :class:`torch.optim.lr_scheduler.ReduceLROnPlateau`, this will use the ``validation_metric``
            provided to determine if learning has plateaued.  To support updating the learning
            rate on every batch, this can optionally implement ``step_batch(batch_num_total)`` which
            updates the learning rate given the batch number.
        visualization_dirname : The directory to save the visualization in

        """
        self._model = model
        self._iterator = iterator
        self._optimizer = optimizer
        self._train_data = train_dataset
        self._validation_data = validation_dataset
        self._segmenter = segmenter
        self._base_dir = base_dir

        if patience is None:  # no early stopping
            if validation_dataset:
                logger.warning('You provided a validation dataset but patience was set to None, '
                               'meaning that early stopping is disabled')
        elif (not isinstance(patience, int)) or patience <= 0:
            raise ConfigurationError('{} is an invalid value for "patience": it must be a positive integer '
                                     'or None (if you want to disable early stopping)'.format(patience))
        self._patience = patience
        self._num_epochs = num_epochs

        serialization_dir = None
        if base_dir is not None:
            serialization_dir = os.path.join(base_dir, "models")
        self._serialization_dir = serialization_dir
        assert os.path.exists(self._serialization_dir), "Directory {0} does not exist".format(self._serialization_dir)

        self._serialized_paths = deque()
        # self._last_permanent_saved_checkpoint_time = time.time()
        # self._model_save_interval = model_save_interval

        self._grad_norm = grad_norm
        self._grad_clipping = grad_clipping
        self._learning_rate_scheduler = learning_rate_scheduler

        increase_or_decrease = validation_metric[0]
        if increase_or_decrease not in ["+", "-"]:
            raise ConfigurationError("Validation metrics must specify whether they should increase "
                                     "or decrease by pre-pending the metric name with a +/-.")
        self._validation_metric = validation_metric[1:]
        self._validation_metric_decreases = increase_or_decrease == "-"

        if not isinstance(cuda_device, int) and not isinstance(cuda_device, list):
            raise ConfigurationError("Expected an int or list for cuda_device, got {}".format(cuda_device))

        if isinstance(cuda_device, list):
            logger.info(f"WARNING: Multiple GPU support is experimental not recommended for use. "
                        "In some cases it may lead to incorrect results or undefined behavior.")
            self._multiple_gpu = True
            self._cuda_devices = cuda_device
            # data_parallel will take care of transfering to cuda devices,
            # so the iterator keeps data on CPU.
            self._iterator_device = -1
        else:
            self._multiple_gpu = False
            self._cuda_devices = [cuda_device]
            self._iterator_device = cuda_device

        if self._cuda_devices[0] != -1:
            self._model = self._model.cuda(self._cuda_devices[0])

        self._log_interval = 10  # seconds

        self._batch_num_total = 0

        self._last_log = 0.0  # time of last logging
        self._num_serialized_models_to_keep = num_serialized_models_to_keep

    def _enable_gradient_clipping(self) -> None:
        if self._grad_clipping is not None:
            # Pylint is unable to tell that we're in the case that _grad_clipping is not None...
            # pylint: disable=invalid-unary-operand-type
            clip_function = lambda grad: grad.clamp(-self._grad_clipping, self._grad_clipping)
            for parameter in self._model.parameters():
                if parameter.requires_grad:
                    parameter.register_hook(clip_function)

    def _rescale_gradients(self) -> Optional[float]:
        """
        Performs gradient rescaling.
        Is a no-op if gradient rescaling is not enabled.
        """
        if self._grad_norm:
            parameters_to_clip = [p for p in self._model.parameters()
                                  if p.grad is not None]
            return sparse_clip_norm(parameters_to_clip, self._grad_norm)
        return None

    def _data_parallel(self, batch):
        """
        Do the forward pass using multiple GPUs.  This is a simplification
        of torch.nn.parallel.data_parallel to support the allennlp model
        interface.
        """
        inputs, module_kwargs = scatter_kwargs((), batch,
                                               self._cuda_devices, 0)
        used_device_ids = self._cuda_devices[:len(inputs)]
        replicas = replicate(self._model, used_device_ids)
        outputs = parallel_apply(replicas, inputs,
                                 module_kwargs, used_device_ids)

        # Only the 'loss' is needed.
        # a (num_gpu, ) tensor with loss on each GPU
        losses = gather([output['loss']
                         for output in outputs], used_device_ids[0], 0)
        return {'loss': losses.mean()}

    def _batch_loss(self, batch: torch.Tensor,
                    for_training: bool) -> torch.Tensor:
        """
        Does a forward pass on the given batch
        and returns the ``loss`` value in the result.
        If ``for_training`` is `True` also applies regularization penalty.
        """
        if self._multiple_gpu:
            output_dict = self._data_parallel(batch)
        else:
            output_dict = self._model(**batch)

        try:
            loss = output_dict["loss"]
            if for_training:
                loss += self._model.get_regularization_penalty()
        except KeyError:
            if for_training:
                raise RuntimeError("The model you are trying to optimize does not contain a"
                                   " 'loss' key in the output of model.forward(inputs).")
            loss = None

        return loss

    def _get_metrics(self, total_loss: float,
                     num_batches: int, reset: bool = False) -> Dict[str, float]:
        """
        Gets the metrics but sets ``"loss"`` to
        the total loss divided by the ``num_batches`` so that
        the ``"loss"`` metric is "average loss per batch".
        """
        metrics = self._model.get_metrics(reset=reset)
        metrics["loss"] = float(total_loss / num_batches) if num_batches > 0 else 0.0
        return metrics

    def _train_epoch(self, epoch: int) -> Dict[str, float]:
        """
        Trains one epoch and returns metrics.
        """
        logger.info(f"Peak CPU memory usage MB: {peak_memory_mb()}")
        if torch.cuda.is_available():
            for gpu, memory in gpu_memory_mb().items():
                logger.info(f"GPU {gpu} memory usage MB: {memory}")

        train_loss = 0.0
        # Set the model to "train" mode.
        self._model.train()

        # Get tqdm for the training batches
        train_generator = self._iterator(self._train_data,
                                         num_epochs=1,
                                         cuda_device=self._iterator_device)
        num_training_batches = self._iterator.get_num_batches(self._train_data)
        train_generator_tqdm = Tqdm.tqdm(train_generator,
                                         total=num_training_batches
                                         )
        self._last_log = time.time()
        last_save_time = time.time()

        batches_this_epoch = 0
        if self._batch_num_total is None:
            self._batch_num_total = 0

        for batch in train_generator_tqdm:
            batches_this_epoch += 1
            self._batch_num_total += 1
            batch_num_total = self._batch_num_total

            self._optimizer.zero_grad()
            loss = self._batch_loss(batch, for_training=True)
            loss.backward()

            # Make sure Variable is on the cpu before converting to numpy.
            # .cpu() is a no-op if you aren't using GPUs.
            train_loss += loss.data.cpu().numpy()

            batch_grad_norm = self._rescale_gradients()

            # This does nothing if batch_num_total is None or you are using an
            # LRScheduler which doesn't update per batch.
            if self._learning_rate_scheduler:
                self._learning_rate_scheduler.step_batch(batch_num_total)
            self._optimizer.step()

            # Update the description with the latest metrics
            metrics = self._get_metrics(train_loss, batches_this_epoch)
            description = self._description_from_metrics(metrics)

            train_generator_tqdm.set_description(description, refresh=False)
            if hasattr(self, "_tf_params") and self._tf_params is not None:
                # We have TF logging
                if self._batch_num_total % self._tf_params["log_every"] == 0:
                    self._tf_log(metrics, self._batch_num_total)

        return self._get_metrics(train_loss, batches_this_epoch, reset=True)

    def _should_stop_early(self, metric_history: List[float]) -> bool:
        """
        uses patience and the validation metric to determine if training should stop early
        """
        if self._patience and self._patience < len(metric_history):
            # Pylint can't figure out that in this branch `self._patience` is an int.
            # pylint: disable=invalid-unary-operand-type

            # Is the best score in the past N epochs worse than the best score overall?
            if self._validation_metric_decreases:
                return min(metric_history[-self._patience:]) > min(metric_history)
            else:
                return max(metric_history[-self._patience:]) < max(metric_history)

        return False

    def _metrics_to_console(self,  # pylint: disable=no-self-use
                            train_metrics: dict,
                            val_metrics: dict = None) -> None:
        """
        Logs all of the train metrics (and validation metrics, if provided) to the console.
        """
        val_metrics = val_metrics or {}
        dual_message_template = "Training %s : %3f    Validation %s : %3f "
        message_template = "%s %s : %3f "

        metric_names = set(sorted(train_metrics.keys()))
        if val_metrics:
            metric_names.update(val_metrics.keys())

        write_buf = "Metrics:\n"

        for name in metric_names:
            train_metric = train_metrics.get(name)
            val_metric = val_metrics.get(name)
            if val_metric is not None and train_metric is not None:
                training = "Training"
                validation = "Validation"
                write_buf += f"{training:>32s} {name:18s}: {train_metric:4.3f}  "
                write_buf += f"{validation} {name:18s}: {val_metric:4.3f}\n"
            elif val_metric is not None:
                # logger.info(message_template, "Validation", name, val_metric)
                write_buf += "{0:>32s} {1:18s}: {2:4.3f}".format(
                    "Validation", name, val_metric)
            elif train_metric is not None:
                # logger.info(message_template, "Training", name, train_metric)
                write_buf += "{0:>32s} {1:18s}: {2:4.3f}".format(
                    "Training", name, train_metric)
        logger.info(write_buf)

    def test(self, data):
        """
            Computes the metrics for data
        """
        model_path = os.path.join(self._serialization_dir, "best.th")
        logger.info("Loading best model from {0}".format(model_path))
        model_state = torch.load(model_path,
                                 map_location=util.device_mapping(-1))
        self._model.load_state_dict(model_state)
        loss, num_batches = self._inference_loss(data, logger_string="Testing")
        metrics = self._get_metrics(loss, num_batches, reset=True)
        metric_names = metrics.keys()
        message_template = "%s %s : %3f"
        for name in metric_names:
            test_metric = metrics.get(name)
            logger.info(message_template, "Test", name, test_metric)
        return metrics

    def _inference_loss(self, data,
                        logger_string="validation") -> Tuple[float, int]:
        """
        Computes the loss on the data passed.
        Returns it and the number of batches.
        Is primarily used for validation and testing
        """
        logger.info("Starting with {0}".format(logger_string))

        self._model.eval()

        inference_generator = self._iterator(data,
                                             num_epochs=1,
                                             cuda_device=self._iterator_device,
                                             for_training=False)
        num_batches = self._iterator.get_num_batches(data)
        inference_generator_tqdm = Tqdm.tqdm(inference_generator,
                                             total=num_batches
                                             )
        batches_this_epoch = 0
        inference_loss = 0
        for batch in inference_generator_tqdm:

            loss = self._batch_loss(batch, for_training=False)
            if loss is not None:
                # You shouldn't necessarily have to compute a loss for validation, so we allow for
                # `loss` to be None.  We need to be careful, though - `batches_this_epoch` is
                # currently only used as the divisor for the loss function, so we can safely only
                # count those batches for which we actually have a loss.  If this variable ever
                # gets used for something else, we might need to change things around a bit.
                batches_this_epoch += 1
                inference_loss += loss.data.cpu().numpy()

            # Update the description with the latest metrics
            inference_metrics = self._get_metrics(inference_loss, batches_this_epoch)
            description = self._description_from_metrics(inference_metrics)
            inference_generator_tqdm.set_description(description, refresh=False)
        logger.info("{0} done".format(logger_string))

        return inference_loss, batches_this_epoch

    def _get_validation_metric(self, val_metrics):
        return val_metrics[self._validation_metric]

    def train(self, continue_training=False) -> Dict[str, Any]:
        """
        Trains the supplied model with the supplied parameters.

        Parameters
        ------------
        continue_training: bool (default=False)
            Load a checkpoint and start training
        """
        epoch_counter, validation_metric_per_epoch = 0, []
        if continue_training:
            try:
                epoch_counter, validation_metric_per_epoch = self._restore_checkpoint()
            except RuntimeError:
                traceback.print_exc()
                raise ConfigurationError("Could not recover training from the checkpoint.  Did you mean to output to "
                                         "a different serialization directory or delete the existing serialization "
                                         "directory?")

        self._enable_gradient_clipping()

        logger.info("Beginning training.")

        train_metrics: Dict[str, float] = {}
        val_metrics: Dict[str, float] = {}
        epochs_trained = 0
        training_start_time = time.time()
        for epoch in range(epoch_counter, self._num_epochs):
            logger.info("=" * 50)
            logger.info("Starting Training Epoch %d/%d",
                        epoch + 1, self._num_epochs)
            epoch_start_time = time.time()
            train_metrics = self._train_epoch(epoch)

            if self._validation_data is not None:
                # We have a validation set, so compute all the metrics on it.
                val_loss, num_batches = self._inference_loss(
                    self._validation_data, logger_string="Validation")
                val_metrics = self._get_metrics(
                    val_loss, num_batches, reset=True)

                # Check validation metric for early stopping
                this_epoch_val_metric = self._get_validation_metric(
                    val_metrics)
                validation_metric_per_epoch.append(this_epoch_val_metric)
                if self._should_stop_early(validation_metric_per_epoch):
                    logger.info("Ran out of patience. Stopping training.")
                    break

                # Check validation metric to see if it's the best so far
                if self._validation_metric_decreases:
                    is_best_so_far = this_epoch_val_metric == min(
                        validation_metric_per_epoch)
                else:
                    is_best_so_far = this_epoch_val_metric == max(
                        validation_metric_per_epoch)
            else:
                # No validation set, so just assume it's the best so far.
                is_best_so_far = True
                val_metrics = {}
                this_epoch_val_metric = None

            self._save_checkpoint(
                epoch, validation_metric_per_epoch, is_best=is_best_so_far)
            self._metrics_to_console(train_metrics, val_metrics)
            # Now the predictions and visualization
            if self._segmenter is not None and is_best_so_far:
                visualization_file = os.path.join(
                    self._base_dir, "visualization", "validation.html")
                prediction_file = os.path.join(
                    self._base_dir, "predictions.json")
                logger.info(
                    f"Writing validation visualization at {visualization_file}"
                )
                self._segmenter.get_predictions(
                    instances=self._validation_data,
                    model=self._model,
                    cuda_device=self._iterator_device,
                    prediction_file=prediction_file,
                    visualization_file=visualization_file,
                    verbose=True)
                logger.info(f"Writing predictions to {prediction_file}")

            if self._learning_rate_scheduler:
                # The LRScheduler API is agnostic to whether your schedule requires a validation metric -
                # if it doesn't, the validation metric passed here is ignored.
                old_lr = min(param_group['lr']
                             for param_group in self._optimizer.param_groups)
                self._learning_rate_scheduler.step(
                    this_epoch_val_metric, epoch)
                new_lr = min(param_group['lr']
                             for param_group in self._optimizer.param_groups)
                logger.info("Reducing LR: {0:.2e} -> {1:.2e}".format(
                    old_lr, new_lr))

            epoch_elapsed_time = time.time() - epoch_start_time
            logger.info("Epoch duration: %s",
                        time.strftime("%H:%M:%S",
                                      time.gmtime(epoch_elapsed_time)))

            if epoch < self._num_epochs - 1:
                training_elapsed_time = time.time() - training_start_time
                estimated_time_remaining = training_elapsed_time * \
                    ((self._num_epochs - epoch_counter) / float(
                        epoch - epoch_counter + 1) - 1)
                formatted_time = time.strftime(
                    "%H:%M:%S", time.gmtime(estimated_time_remaining))
                logger.info("Estimated training time remaining: %s",
                            formatted_time)

            epochs_trained += 1
            logger.info("=" * 50)

        training_elapsed_time = time.time() - training_start_time
        metrics = {
            "training_duration": time.strftime(
                "%H:%M:%S", time.gmtime(training_elapsed_time)),
            "training_start_epoch": epoch_counter,
            "training_epochs": epochs_trained
        }
        for key, value in train_metrics.items():
            metrics["training_" + key] = value
        for key, value in val_metrics.items():
            metrics["validation_" + key] = value

        if validation_metric_per_epoch:
            # We may not have had validation data, so we need to hide this behind an if.
            if self._validation_metric_decreases:
                best_validation_metric = min(validation_metric_per_epoch)
            else:
                best_validation_metric = max(validation_metric_per_epoch)
            best_val_key = f"best_validation_{self._validation_metric}"
            metrics[best_val_key] = best_validation_metric
            metrics['best_epoch'] = [
                i for i, value in enumerate(validation_metric_per_epoch)
                if value == best_validation_metric][-1]
        return metrics

    def _description_from_metrics(self, metrics: Dict[str, float]) -> str:
        # pylint: disable=no-self-use
        return ', '.join(["%s: %.4f" % (name, value)
                          for name, value in metrics.items()]) + " ||"

    def _save_checkpoint(self,
                         epoch: Union[int, str],
                         val_metric_per_epoch: List[float],
                         is_best: Optional[bool] = None) -> None:
        """
        Saves a checkpoint of the model to self._serialization_dir.
        Is a no-op if self._serialization_dir is None.

        Parameters
        ----------
        epoch : Union[int, str], required.
            The epoch of training.  If the checkpoint is saved in the middle
            of an epoch, the parameter is a string with the epoch and timestamp.
        is_best: bool, optional (default = None)
            A flag which causes the model weights at the given epoch to
            be copied to a "best.th" file. The value of this flag should
            be based on some validation metric computed by your model.
        """
        if self._serialization_dir is not None:
            model_path = os.path.join(self._serialization_dir,
                                      "model_state_epoch_{}.th".format(epoch))
            model_state = self._model.state_dict()
            torch.save(model_state, model_path)

            training_state = {'epoch': epoch,
                              'val_metric_per_epoch': val_metric_per_epoch,
                              'optimizer': self._optimizer.state_dict(),
                              'batch_num_total': self._batch_num_total}
            training_path = os.path.join(
                self._serialization_dir,
                "training_state_epoch_{}.th".format(epoch))
            torch.save(training_state, training_path)
            if is_best:
                logger.info("Best validation performance so far. "
                            "Copying weights to '%s/best.th'.", self._serialization_dir)
                shutil.copyfile(
                    model_path, os.path.join(
                        self._serialization_dir, "best.th"))

            if self._num_serialized_models_to_keep:
                self._serialized_paths.append(
                    [time.time(), model_path, training_path])
                if len(self._serialized_paths) > self._num_serialized_models_to_keep:
                    paths_to_remove = self._serialized_paths.popleft()
                    for fname in paths_to_remove[1:]:
                        os.remove(fname)

    def _restore_checkpoint(self) -> Tuple[int, List[float]]:
        """
        Restores a model from a serialization_dir to the last saved checkpoint.
        This includes an epoch count and optimizer state, which is serialized separately
        from  model parameters. This function should only be used to continue training -
        if you wish to load a model for inference/load parts of a model into a new
        computation graph, you should use the native Pytorch functions:
        `` model.load_state_dict(torch.load("/path/to/model/weights.th"))``

        If ``self._serialization_dir`` does not exist or does not contain any checkpointed weights,
        this function will do nothing and return 0.

        Returns
        -------
        epoch: int
            The epoch at which to resume training, which should be one after the epoch
            in the saved training state.
        """
        have_checkpoint = (
            self._serialization_dir is not None and
            any(
                "model_state_epoch_" in x
                for x in os.listdir(os.path.join(self._serialization_dir))
            )
        )
        if not have_checkpoint:
            # No checkpoint to restore, start at 0
            logger.info("Didn't find any models to restore."
                        "Starting training from scratch")
            return 0, []

        serialization_files = os.listdir(
            os.path.join(self._serialization_dir))
        model_checkpoints = [x for x in serialization_files
                             if "model_state_epoch" in x]
        # Get the last checkpoint file.  Epochs are specified as either an
        # int (for end of epoch files) or with epoch and timestamp for
        # within epoch checkpoints, e.g. 5.2018-02-02-15-33-42
        found_epochs = [
            # pylint: disable=anomalous-backslash-in-string
            re.search("model_state_epoch_([0-9\.\-]+)\.th", x).group(1)
            for x in model_checkpoints
        ]
        int_epochs: Any = []
        for epoch in found_epochs:
            pieces = epoch.split('.')
            if len(pieces) == 1:
                # Just a single epoch without timestamp
                int_epochs.append([int(pieces[0]), 0])
            else:
                # has a timestamp
                int_epochs.append([int(pieces[0]), pieces[1]])
        last_epoch = sorted(int_epochs, reverse=True)[0]
        if last_epoch[1] == 0:
            epoch_to_load = str(last_epoch[0])
        else:
            epoch_to_load = '{0}.{1}'.format(last_epoch[0], last_epoch[1])

        model_path = os.path.join(
            self._serialization_dir,
            "model_state_epoch_{}.th".format(epoch_to_load))
        training_state_path = os.path.join(
            self._serialization_dir,
            "training_state_epoch_{}.th".format(epoch_to_load))

        # Load the parameters onto CPU, then transfer to GPU.
        # This avoids potential OOM on GPU for large models that
        # load parameters onto GPU then make a new GPU copy into the parameter
        # buffer. The GPU transfer happens implicitly in load_state_dict.
        model_state = torch.load(
            model_path, map_location=util.device_mapping(-1))
        training_state = torch.load(
            training_state_path, map_location=util.device_mapping(-1))
        self._model.load_state_dict(model_state)
        self._optimizer.load_state_dict(training_state["optimizer"])
        move_optimizer_to_cuda(self._optimizer)

        # We didn't used to save `validation_metric_per_epoch`, so we can't assume
        # that it's part of the trainer state. If it's not there, an empty list is all
        # we can do.
        if "val_metric_per_epoch" not in training_state:
            logger.warning("trainer state `val_metric_per_epoch`"
                           "not found, using empty list")
            val_metric_per_epoch: List[float] = []
        else:
            val_metric_per_epoch = training_state["val_metric_per_epoch"]

        if isinstance(training_state["epoch"], int):
            epoch_to_return = training_state["epoch"] + 1
        else:
            epoch_to_return = int(training_state["epoch"].split('.')[0]) + 1

        # For older checkpoints with batch_num_total missing, default to old behavior where
        # it is unchanged.
        batch_num_total = training_state.get('batch_num_total')
        if batch_num_total is not None:
            self._batch_num_total = batch_num_total

        return epoch_to_return, val_metric_per_epoch

    @classmethod
    def get_args(cls,
                 model: Model,
                 base_dir: str,
                 iterator: DataIterator,
                 train_data: Iterable[Instance],
                 validation_data: Optional[Iterable[Instance]],
                 segmenter: Optional[BasePredictionClass],
                 params: Params) -> Dict[str, Any]:
        patience = params.pop_int("patience", None)
        validation_metric = params.pop("validation_metric", "-loss")
        num_epochs = params.pop_int("num_epochs", 20)
        cuda_device = params.pop_int("cuda_device", -1)
        grad_norm = params.pop_float("grad_norm", None)
        grad_clipping = params.pop_float("grad_clipping", None)
        lr_scheduler_params = params.pop("learning_rate_scheduler", None)
        num_serialized_models_to_keep = params.pop(
            "num_serialized_models_to_keep", None)

        if cuda_device >= 0:
            model = model.cuda(cuda_device)
        parameters = [[n, p] for n, p in model.named_parameters()
                      if p.requires_grad]
        optimizer = Optimizer.from_params(parameters, params.pop("optimizer"))

        if lr_scheduler_params:
            scheduler = LearningRateScheduler.from_params(
                optimizer, lr_scheduler_params)
        else:
            scheduler = None
        params.assert_empty(cls.__name__)
        kwargs = {}
        kwargs['model'] = model
        kwargs['optimizer'] = optimizer
        kwargs['iterator'] = iterator
        kwargs['train_dataset'] = train_data
        kwargs['validation_dataset'] = validation_data
        kwargs['segmenter'] = segmenter
        kwargs['patience'] = patience
        kwargs['validation_metric'] = validation_metric
        kwargs['num_epochs'] = num_epochs
        kwargs['base_dir'] = base_dir
        kwargs['cuda_device'] = cuda_device
        kwargs['grad_norm'] = grad_norm
        kwargs['grad_clipping'] = grad_clipping
        kwargs['learning_rate_scheduler'] = scheduler
        kwargs['num_serialized_models_to_keep'] = num_serialized_models_to_keep
        return kwargs

    @classmethod
    def from_params(cls, *args, **kwargs) -> 'Trainer':
        new_args = cls.get_args(*args, **kwargs)
        return Trainer(**new_args)