from octavia_lib.api.drivers import data_models as driver_dm
from octavia_lib.api.drivers import exceptions
from oslo_config import cfg
from oslo_log import log as logging

from octavia.api.drivers.amphora_driver.v1 import driver
from octavia.common import constants as consts
from octavia.common import exceptions as api_exceptions
from octavia.db import api as db_apis

from octavia_f5.api.drivers.f5_driver import arbiter
from octavia_f5.common import constants as f5_consts
from octavia_f5.utils import driver_utils

CONF = cfg.CONF
CONF.import_group('oslo_messaging', 'octavia.common.config')
LOG = logging.getLogger(__name__)

class F5ProviderDriver(driver.AmphoraProviderDriver,
                       arbiter.MigrationArbiter,
                       arbiter.RescheduleMixin):
    """Octavia plugin for the F5 driver.

    Callbacks need to be overwritten for setting the _server_ attribute
    in RPC calls to specify the target rpc receiver (e.g. the scheduled
    host)
    """

    def __init__(self):
        super(F5ProviderDriver, self).__init__()

    def _get_server(self, loadbalancer_id):
        """ Get scheduled host of the loadbalancer.
        :param loadbalancer_id: loadbalancer id
        :return: scheduled host
        """
        loadbalancer = self.repositories.load_balancer.get(
            db_apis.get_session(), id=loadbalancer_id)
        if loadbalancer.server_group_id:
            return loadbalancer.server_group_id

        # fetch scheduled server from VIP port
        network_driver = driver_utils.get_network_driver()
        return network_driver.get_scheduled_host(loadbalancer.vip.port_id)

    def loadbalancer_create(self, loadbalancer):
        if loadbalancer.flavor == driver_dm.Unset:
            loadbalancer.flavor = None

        network_driver = driver_utils.get_network_driver()
        host = network_driver.get_scheduled_host(loadbalancer.vip_port_id)

        LOG.info("Scheduling loadbalancer %s to %s",
                 loadbalancer.loadbalancer_id,
                 host)
        payload = {consts.LOAD_BALANCER_ID: loadbalancer.loadbalancer_id,
                   consts.FLAVOR: loadbalancer.flavor}
        client = self.client.prepare(server=host)
        client.cast({}, 'create_load_balancer', **payload)

    def loadbalancer_delete(self, loadbalancer, cascade=False):
        loadbalancer_id = loadbalancer.loadbalancer_id
        payload = {consts.LOAD_BALANCER_ID: loadbalancer_id,
                   'cascade': cascade}
        client = self.client.prepare(server=self._get_server(loadbalancer_id))
        client.cast({}, 'delete_load_balancer', **payload)

    def loadbalancer_failover(self, loadbalancer_id):
        payload = {consts.LOAD_BALANCER_ID: loadbalancer_id}
        client = self.client.prepare(server=self._get_server(loadbalancer_id))
        client.cast({}, 'failover_load_balancer', **payload)

    def loadbalancer_reschedule(self, loadbalancer_id, target_host):
        store = {consts.LOADBALANCER_ID: loadbalancer_id, "candidate": target_host}
        self.run_flow(self.get_reschedule_flow, store=store)

    def loadbalancer_add(self, loadbalancer_id, target_host):
        payload = {consts.LOAD_BALANCER_ID: loadbalancer_id}
        client = self.client.prepare(server=target_host)
        client.call({}, 'add_loadbalancer', **payload)

    def loadbalancer_remove(self, loadbalancer_id, target_host):
        payload = {consts.LOAD_BALANCER_ID: loadbalancer_id}
        client = self.client.prepare(server=target_host)
        client.call({}, 'remove_loadbalancer', **payload)

    def loadbalancer_update(self, old_loadbalancer, new_loadbalancer):
        lb_id = new_loadbalancer.loadbalancer_id
        payload = {consts.LOAD_BALANCER_ID: lb_id,
                   consts.LOAD_BALANCER_UPDATES: {}}
        client = self.client.prepare(server=self._get_server(lb_id))
        client.cast({}, 'update_load_balancer', **payload)

    # Listener
    def listener_create(self, listener):
        payload = {consts.LISTENER_ID: listener.listener_id}
        client = self.client.prepare(server=self._get_server(listener.loadbalancer_id))
        client.cast({}, 'create_listener', **payload)

    def listener_delete(self, listener):
        listener_id = listener.listener_id
        payload = {consts.LISTENER_ID: listener_id}
        client = self.client.prepare(server=self._get_server(listener.loadbalancer_id))
        client.cast({}, 'delete_listener', **payload)

    def listener_update(self, old_listener, new_listener):
        listener_id = old_listener.listener_id
        payload = {consts.LISTENER_ID: listener_id,
                   consts.LISTENER_UPDATES: {}}
        client = self.client.prepare(server=self._get_server(old_listener.loadbalancer_id))
        client.cast({}, 'update_listener', **payload)

    # Pool
    def pool_create(self, pool):
        payload = {consts.POOL_ID: pool.pool_id}
        client = self.client.prepare(server=self._get_server(pool.loadbalancer_id))
        client.cast({}, 'create_pool', **payload)

    def pool_delete(self, pool):
        pool_id = pool.pool_id
        payload = {consts.POOL_ID: pool_id}
        client = self.client.prepare(server=self._get_server(pool.loadbalancer_id))
        client.cast({}, 'delete_pool', **payload)

    def pool_update(self, old_pool, new_pool):
        pool_id = new_pool.pool_id
        payload = {consts.POOL_ID: pool_id,
                   consts.POOL_UPDATES: {}}
        client = self.client.prepare(server=self._get_server(old_pool.loadbalancer_id))
        client.cast({}, 'update_pool', **payload)

    # Member
    def member_create(self, member):
        db_pool = self.repositories.pool.get(db_apis.get_session(),
                                             id=member.pool_id)
        payload = {consts.MEMBER_ID: member.member_id}
        db_pool = self.repositories.pool.get(db_apis.get_session(), id=member.pool_id)
        client = self.client.prepare(server=self._get_server(db_pool.load_balancer_id))
        client.cast({}, 'create_member', **payload)

    def member_delete(self, member):
        db_pool = self.repositories.pool.get(db_apis.get_session(), id=member.pool_id)
        payload = {consts.MEMBER_ID: member.member_id}
        client = self.client.prepare(server=self._get_server(db_pool.load_balancer_id))
        client.cast({}, 'delete_member', **payload)

    def member_update(self, old_member, new_member):
        db_pool = self.repositories.pool.get(db_apis.get_session(), id=old_member.pool_id)
        payload = {consts.MEMBER_ID: new_member.member_id,
                   consts.MEMBER_UPDATES: {}}
        client = self.client.prepare(server=self._get_server(db_pool.load_balancer_id))
        client.cast({}, 'update_member', **payload)

    def member_batch_update(self, pool_id, members):
        db_pool = self.repositories.pool.get(db_apis.get_session(), id=pool_id)
        payload = {'old_member_ids': [],
                   'new_member_ids': [],
                   'updated_members': []}
        client = self.client.prepare(server=self._get_server(db_pool.load_balancer_id))
        client.cast({}, 'batch_update_members', **payload)

    # Health Monitor
    def _health_monitor_check(self, healthmonitor):
        delay = healthmonitor.delay
        if not delay:
            return
        if delay > f5_consts.HEALTH_MONITOR_DELAY_MAX:
            raise api_exceptions.ValidationException(
                detail='Delay value for health monitor too high. Must not be higher than {}.'.format(
                    f5_consts.HEALTH_MONITOR_DELAY_MAX))

    def health_monitor_create(self, healthmonitor):
        self._health_monitor_check(healthmonitor)
        db_pool = self.repositories.pool.get(db_apis.get_session(), id=healthmonitor.pool_id)
        payload = {consts.HEALTH_MONITOR_ID: healthmonitor.healthmonitor_id}
        client = self.client.prepare(server=self._get_server(db_pool.load_balancer_id))
        client.cast({}, 'create_health_monitor', **payload)

    def health_monitor_delete(self, healthmonitor):
        db_pool = self.repositories.pool.get(db_apis.get_session(), id=healthmonitor.pool_id)
        payload = {consts.HEALTH_MONITOR_ID: healthmonitor.healthmonitor_id}
        client = self.client.prepare(server=self._get_server(db_pool.load_balancer_id))
        client.cast({}, 'delete_health_monitor', **payload)

    def health_monitor_update(self, old_healthmonitor, new_healthmonitor):
        self._health_monitor_check(new_healthmonitor)
        db_pool = self.repositories.pool.get(db_apis.get_session(), id=old_healthmonitor.pool_id)
        payload = {consts.HEALTH_MONITOR_ID: new_healthmonitor.healthmonitor_id,
                   consts.HEALTH_MONITOR_UPDATES: {}}
        client = self.client.prepare(server=self._get_server(db_pool.load_balancer_id))
        client.cast({}, 'update_health_monitor', **payload)

    # L7 Policy
    def l7policy_create(self, l7policy):
        db_listener = self.repositories.listener.get(db_apis.get_session(), id=l7policy.listener_id)
        payload = {consts.L7POLICY_ID: l7policy.l7policy_id}
        client = self.client.prepare(server=self._get_server(db_listener.load_balancer_id))
        client.cast({}, 'create_l7policy', **payload)

    def l7policy_delete(self, l7policy):
        db_listener = self.repositories.listener.get(db_apis.get_session(), id=l7policy.listener_id)
        payload = {consts.L7POLICY_ID: l7policy.l7policy_id}
        client = self.client.prepare(server=self._get_server(db_listener.load_balancer_id))
        client.cast({}, 'delete_l7policy', **payload)

    def l7policy_update(self, old_l7policy, new_l7policy):
        db_listener = self.repositories.listener.get(db_apis.get_session(), id=old_l7policy.listener_id)
        payload = {consts.L7POLICY_ID: new_l7policy.l7policy_id,
                   consts.L7POLICY_UPDATES: {}}
        client = self.client.prepare(server=self._get_server(db_listener.load_balancer_id))
        self.client.cast({}, 'update_l7policy', **payload)

    # L7 Rule
    def l7rule_create(self, l7rule):
        db_l7 = self.repositories.l7policy.get(db_apis.get_session(), id=l7rule.l7policy_id)

        payload = {consts.L7RULE_ID: l7rule.l7rule_id}
        client = self.client.prepare(server=self._get_server(db_l7.listener.load_balancer_id))
        client.cast({}, 'create_l7rule', **payload)

    def l7rule_delete(self, l7rule):
        db_l7 = self.repositories.l7policy.get(db_apis.get_session(), id=l7rule.l7policy_id)

        payload = {consts.L7RULE_ID: l7rule.l7rule_id}
        client = self.client.prepare(server=self._get_server(db_l7.listener.load_balancer_id))
        client.cast({}, 'delete_l7rule', **payload)

    def l7rule_update(self, old_l7rule, new_l7rule):
        db_l7 = self.repositories.l7policy.get(db_apis.get_session(), id=old_l7rule.l7policy_id)

        payload = {consts.L7RULE_ID: new_l7rule.l7rule_id,
                   consts.L7RULE_UPDATES: {}}
        client = self.client.prepare(server=self._get_server(db_l7.listener.load_balancer_id))
        client.cast({}, 'update_l7rule', **payload)

    def create_vip_port(self, loadbalancer_id, project_id, vip_dictionary):
        raise exceptions.NotImplementedError(
            operator_fault_string="F5 provider not creating VIP port. It will be created by Octavia API.")

    def get_supported_flavor_metadata(self):
        """Returns the valid flavor metadata keys and descriptions.

        This extracts the valid flavor metadata keys and descriptions
        from the JSON validation schema and returns it as a dictionary.

        :return: Dictionary of flavor metadata keys and descriptions.
        :raises DriverError: An unexpected error occurred.
        """
        raise exceptions.NotImplementedError()

    def validate_flavor(self, flavor_metadata):
        """Validates if driver can support flavor as defined in flavor_metadata.

        :param flavor_metadata (dict): Dictionary with flavor metadata.
        :return: Nothing if the flavor is valid and supported.
        :raises DriverError: An unexpected error occurred in the driver.
        :raises NotImplementedError: The driver does not support flavors.
        :raises UnsupportedOptionError: if driver does not
              support one of the configuration options.
        """
        raise exceptions.NotImplementedError()

    def validate_availability_zone(self, availability_zone_dict):
        """Validates availability zone profile data.

        This will validate an availability zone profile dataset against the
        availability zone settings the F5 driver supports.

        :param availability_zone_dict: The availability zone dict to validate.
        :type availability_zone_dict: dict
        :return: None
        :raises DriverError: An unexpected error occurred.
        :raises UnsupportedOptionError: If the driver does not support
          one of the availability zone settings.
        """
        if 'hosts' in availability_zone_dict:
            return

        raise exceptions.UnsupportedOptionError(
                user_fault_string='Failed to get the supported availability '
                                  'zone metadata.',
                operator_fault_string='Failed to get hosts from availability '
                                  'zone metadata: {}'.format(availability_zone_dict))


from flask_sqlalchemy import SQLAlchemy
from sqlalchemy.ext.declarative import declarative_base

db = SQLAlchemy()
Base = declarative_base()

def init_app(app):
    app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = True
    app.config['SQLALCHEMY_DATABASE_URI'] = app.config.DATABASE_URL
    db.init_app(app)


import datetime

from django.shortcuts import render
from django.http import JsonResponse, HttpResponse, HttpResponseRedirect
from django.views.decorators.csrf import csrf_exempt, csrf_protect
from django.utils.translation import gettext as _

from clinic import clinic_bot, verification
from clinic.clinic_bot import is_num_in_contracts_info
from clinic.models import Doctor, DoctorCabinet, PhoneNumber, Person, Schedule, Recording, Review
from bot.models import ContractInfo
from datetime import date
from .forms import ContractForm

# Create your views here.
def index(request):
    return render(request, template_name="clinic/index.html", context={})

def staff(request):
    doctors = Doctor.objects.all()
    doc_cabs = DoctorCabinet.objects.all()
    dates = [(date.today() - doctor.work_start_date).days // 365 for doctor in doctors]
    phones = PhoneNumber.objects.all()
    return render(request, template_name="clinic/pages/staff.html", context={"doctors": doctors,
                                                                             "doc_cabs": doc_cabs,
                                                                             "phones": phones,
                                                                             "dates": dates})

def schedule(request):
    return render(request, template_name="clinic/pages/schedule-work.html", context={})

def contacts(request):
    return render(request, template_name="clinic/pages/contacts.html", context={})

def vacancy(request):
    return render(request, template_name="clinic/pages/vacancy.html", context={})

def reviews(request, doctor_id):
    dctr = Doctor.objects.get(id=doctor_id)
    revs = dctr.review_set.all()[::-1]
    return render(request, template_name="clinic/pages/reviews.html", context={"doctor": dctr,
                                                                               "revs": revs})

def make_appointment(request):
    return render(request, template_name="clinic/pages/making-an-appointment.html", context={})

@csrf_exempt
def doctor_info(request) -> JsonResponse:
    patient_data = request.POST.dict()
    try:
        person = Person.objects.get(lastname=patient_data["lastname"], name=patient_data["name"],
                                    contract_num=patient_data["contract_num"])
        doctor: Doctor = person.doctor
        schedule_for_doctor = Schedule.objects.filter(doctor=doctor)
        print(schedule_for_doctor)
        schedule_dct = {}
        for obj in schedule_for_doctor:
            if not obj.is_busy:
                dt_tm = str(obj.dt_tm).split()
                dt_tm[1] = dt_tm[1][:5]
                try:
                    schedule_dct[dt_tm[0]].append(dt_tm[1])
                except KeyError:
                    schedule_dct[dt_tm[0]] = [dt_tm[1]]
        return JsonResponse(
            {
                "status": "success",
                "personId": person.id,
                "doctorId": doctor.id,
                "lastname": doctor.lastname,
                "name": doctor.name,
                "patronymic": doctor.patronymic,
                "workSchedule": schedule_dct
            })
    except Person.DoesNotExist:
        return JsonResponse({"status": "undefined"})

@csrf_protect
def record(request):
    record_data = request.POST
    person = Person.objects.get(id=record_data["person_id"])
    doctor = Doctor.objects.get(id=record_data["doctor_id"])
    print(record_data)
    rec = Recording(person=person,
                    doctor=doctor,
                    datetime=datetime.datetime.strptime(record_data["dt_tm"], "%Y-%m-%d %H:%M"),
                    health_complaint=record_data["complaint"])
    rec.save()
    sch = Schedule.objects.get(doctor=rec.doctor, dt_tm=rec.datetime)
    sch.is_busy = True
    sch.save()
    clinic_bot.send_msg_about_recording_to_doctor(person, doctor, rec)
    return HttpResponse("successful")

@csrf_exempt
def check_person(request):
    person_data = request.POST.dict()
    try:
        p = Person.objects.get(lastname=person_data["lastname"],
                               name=person_data["name"],
                               contract_num=person_data["contract_num"])
        if p.doctor.id == int(person_data["doctor_id"]):
            return JsonResponse({"status": "canLeaveReview"})
    except Person.DoesNotExist:
        pass
    return JsonResponse({"status": "undefined"})

@csrf_protect
def add_review(request, doctor_id):
    review_data = request.POST
    dctr = Doctor.objects.get(id=review_data["doctor_id"])
    prsn = Person.objects.get(lastname=review_data["lastname"], name=review_data["name"])
    r = Review(doctor=dctr, person=prsn, text=review_data["review-text"])
    r.save()
    return HttpResponseRedirect(f"/reviews/{doctor_id}")

@csrf_protect
def contract_page(request):
    form = ContractForm()
    print(request.headers["Accept-Language"])
    if request.method == "POST":
        form.data = request.POST
        form.is_bound = True
        if form.is_valid():
            if form.data["chosen"] == "telegram" \
                    and not ContractInfo.objects.filter(person_number=form.data["person_phone_number"]) \
                    and not is_num_in_contracts_info(form.data["person_phone_number"]):
                clinic_bot.send_msg_about_contract_to_doctor(form.data.dict())
                return HttpResponseRedirect("https://t.me/StepFristBot")
            elif ContractInfo.objects.filter(person_number=form.data["person_phone_number"]):
                form.add_error("person_phone_number", "Користувач з таким номером телефону вже зареєстрований")
            elif is_num_in_contracts_info(form.data["person_phone_number"]):
                form.add_error("person_phone_number",
                               "Користувач з таким номером телефону вже відправив інформацію для контракту")
    return render(request, "clinic/pages/making-a-contract.html", {"form": form})

# @csrf_protect
# def code_submitting(request):
#     otp_code, phone_number = request.POST["verification_code"], request.POST["person_phone_number"]
#     if verification.check_verification_code(phone_number, otp_code):
#         form_data = json.loads(request.POST["form_data"])
#         clinic_bot.send_msg_about_contract_to_doctor(form_data)
#         return HttpResponseRedirect("https://t.me/StepFristBot")
#     else:
#         return render(request,
#                       template_name="clinic/pages/number_submitting.html",
#                       context={"phone": request.POST["person_phone_number"],
#                                "form_data": request.POST["form_data"]})

@csrf_exempt
def doctor_lfp(request):
    search_data = request.body.decode(encoding="utf-8")
    searched_doctors = Doctor.get_doctors_by_input(search_data)
    searched_doctors = list(map(lambda tup: " ".join(tup), searched_doctors))
    return JsonResponse({"searchedDoctors": searched_doctors})


#!/usr/bin/env python3
import collections
import os
import sys

import numpy as np
import tensorflow as tf
import tensorflow_addons as tfa

from mapping import Mapping
import mtool_evaluate
from sentence import Sentence

NODE_NO, NODE_PARENT, NODE_CHILD = range(3)

def get_layers(s):
    # Shortest paths from start, compute anchors
    paths = [0] * s.n_tokens() + [None] * s.n_nodes()
    anchors = [{i} if i < s.n_tokens() else set() for i in range(s.n_tokens() + s.n_nodes())]
    queue = collections.deque(range(s.n_tokens()))

    while queue:
        n = queue.popleft()
        for e in s.parents[n]:
            p = s.factors[s.EDGE_PARENTS][e]
            anchors[p] |= anchors[n]
            if not paths[p]:
                paths[p] = paths[n] + 1
                queue.append(p)

    # Compute a sequence of created nodes
    layers = []

    nodes = set(range(s.n_tokens()))
    while len(nodes) < s.n_nodes() + s.n_tokens():
        new_nodes, new_layer = set(), []
        for n in nodes:
            best, best_e = None, None
            for e in s.parents[n]:
                p = s.factors[s.EDGE_PARENTS][e]
                if p not in nodes and p not in new_nodes:
                    if not best or len(anchors[p]) < len(anchors[best]) or sorted(anchors[p]) < sorted(anchors[best]):
                        best, best_e = p, e
            for e in s.children[n]:
                p = s.factors[s.EDGE_CHILDREN][e]
                if paths[p]: continue
                if p not in nodes and p not in new_nodes:
                    if not best or len(anchors[p]) < len(anchors[best]) or sorted(anchors[p]) < sorted(anchors[best]):
                        best, best_e = p, e
            if best:
                new_nodes.add(best)
                new_layer.append((n, best, best_e))

        if not new_nodes: break
        nodes |= new_nodes
        layers.append(new_layer)

    if len(nodes) != s.n_nodes() + s.n_tokens():
        raise ValueError("Cannot get layers for sentence {}".format(s.id))

    return layers

class Network:
    def __init__(self, train, args):
        # Create required layers
        self.layers = tf.keras.Model()

        # Encoder
        self.layers.encoder_we = tf.keras.layers.Embedding(len(train.mappings[Sentence.TOKENS].strings), args.we_dim, mask_zero=True)
        self.layers.encoder_cle = tf.keras.layers.Embedding(len(train.mappings[Sentence.TOKENS].alphabet), args.cle_dim, mask_zero=True)
        self.layers.encoder_cle_rnn = tf.keras.layers.Bidirectional(tf.keras.layers.GRU(args.cle_dim), merge_mode="concat")
        self.layers.encoder_token_values = [
            tf.keras.layers.Embedding(len(train.mappings[Sentence.TOKEN_VALUES][prop].strings), 128, mask_zero=True)
            for prop in train.token_properties
        ]
        self.layers.encoder_rnns = [
            tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(args.encoder_dim, return_sequences=True), merge_mode="sum")
            for _ in range(args.encoder_rnn_layers)
        ]
        self.layers.encoder_node_values = [
            tf.keras.layers.Embedding(len(train.mappings[Sentence.NODE_VALUES][prop].strings), args.encoder_dim)
            for prop in train.node_properties
        ]
        self.edge_values = [len(train.mappings[Sentence.EDGE_VALUES][prop].strings) for prop in train.edge_properties]
        self.layers.encoder_edge_values = [
            tf.keras.layers.Embedding(2 * edge_values, args.encoder_dim)
            for edge_values in self.edge_values
        ]
        self.layers.encoder_token_transfer = [
            tf.keras.layers.Dense(args.encoder_dim, activation=tf.nn.tanh)
            for _ in range(args.decoder_iterations)]

        # Decoder
        self.layers.decoder_node_operation = [
            tf.keras.layers.Dense(3, activation=tf.nn.softmax) for _ in range(args.decoder_iterations)]
        self.layers.decoder_node_values = [
            [tf.keras.layers.Dense(len(train.mappings[Sentence.NODE_VALUES][prop].strings), activation=tf.nn.softmax) for prop in train.node_properties]
            for _ in range(args.decoder_iterations)]
#         self.layers.decoder_edge_values = [
#             [tf.keras.layers.Dense(len(train.mappings[Sentence.EDGE_VALUES][prop].strings), activation=tf.nn.softmax) for prop in train.edge_properties]
#             for _ in range(args.decoder_iterations)]
        self.layers.decoder_edge_parents = [
            tf.keras.layers.Dense(args.edge_dim)
            for _ in range(args.decoder_iterations)]
        self.layers.decoder_edge_children = [
            tf.keras.layers.Dense(args.edge_dim)
            for _ in range(args.decoder_iterations)]
        self.layers.decoder_edge_arc = [
            tf.keras.layers.Dense(2, activation=tf.nn.softmax)
            for _ in range(args.decoder_iterations)]
        if args.highway:
            self.layers.decoder_edge_highway = [
                tf.keras.layers.Dense(args.edge_dim, activation=tf.nn.tanh)
                for _ in range(args.decoder_iterations)]
        self.layers.decoder_deprel_parents = [
            tf.keras.layers.Dense(args.deprel_dim)
            for _ in range(args.decoder_iterations)]
        self.layers.decoder_deprel_children = [
            tf.keras.layers.Dense(args.deprel_dim)
            for _ in range(args.decoder_iterations)]
        self.layers.decoder_deprel_values = [
            [tf.keras.layers.Dense(len(train.mappings[Sentence.EDGE_VALUES][prop].strings), activation=tf.nn.softmax) for prop in train.edge_properties]
            for _ in range(args.decoder_iterations)]
        if args.highway:
            self.layers.decoder_deprel_highway = [
                tf.keras.layers.Dense(args.edge_dim, activation=tf.nn.tanh)
                for _ in range(args.decoder_iterations)]
        self.layers.decoder_tops = tf.keras.layers.Dense(2, activation=tf.nn.softmax)

        # Generic helpers
        self.layers.dropout = tf.keras.layers.Dropout(args.dropout)
        self.layers.concat = tf.keras.layers.Concatenate()
        self.layers.sum = tf.keras.layers.Add()
        self.word_dropout = args.word_dropout
        self.optimizer_step_sentences = len(train.sentence_lens) / ((len(train.sentence_lens) - 1) // (args.batch_size * args.batch_aggregation) + 1)

        if args.predict: return

        # Create optimizer, losses, summary writer and accuracies
        self.optimizer = tfa.optimizers.LazyAdam(beta_2=args.beta_2)
        self.loss_sce = tf.keras.losses.SparseCategoricalCrossentropy()
        self.writer = tf.summary.create_file_writer(args.logdir, flush_millis=10*1000)
        self.accuracies = dict([
            ("node/ops", tf.keras.metrics.SparseCategoricalAccuracy())] + [
            ("node/" + prop, tf.keras.metrics.SparseCategoricalAccuracy()) for prop in train.node_properties] + [
            ("edge/arc", tf.keras.metrics.SparseCategoricalAccuracy())] + [
            ("edge/" + prop, tf.keras.metrics.SparseCategoricalAccuracy()) for prop in train.edge_properties] + [
            ("edge/tops", tf.keras.metrics.SparseCategoricalAccuracy())])

    def load_weights(self, path):
        self.layers.load_weights(path)

    def save_weights(self, path):
        self.layers.save_weights(path)

    def _encoder(self, token_ids, token_charseqs, token_charseq_ids, token_values, token_additionals, training):
        if training and self.word_dropout:
            word_dropout_mask = tf.cast(tf.math.greater_equal(tf.random.uniform(token_ids.shape), self.word_dropout), tf.int32)
            unk_mask = tf.cast(tf.math.not_equal(token_ids, 0), tf.int32) * Mapping.UNK
            token_ids = word_dropout_mask * token_ids + (1 - word_dropout_mask) * unk_mask
        inputs = [self.layers.encoder_we(token_ids)]

        characters_embedded = self.layers.encoder_cle(token_charseqs)
        characters_embedded = self.layers.dropout(characters_embedded, training=training)
        charseqs = self.layers.encoder_cle_rnn(characters_embedded)
        charseqs = tf.gather(charseqs, token_charseq_ids)
        inputs.append(charseqs)

        for i, embedding in enumerate(self.layers.encoder_token_values):
            inputs.append(embedding(token_values[:, :, i]))

        inputs.extend(token_additionals)

        hidden = self.layers.concat(inputs)
        hidden = self.layers.dropout(hidden, training=training)
        for i, rnn in enumerate(self.layers.encoder_rnns):
            rnn_output = rnn(hidden)
            rnn_output = self.layers.dropout(rnn_output, training=training)
            if i == 0:
                hidden = rnn_output
            else:
                hidden = self.layers.sum([hidden, rnn_output])

        delattr(hidden, "_keras_mask")
        return hidden

    class State:
        def __init__(self, network, sentence, tokens, training):
            self.network = network
            self.sentence = sentence
            self.tokens = tokens
            self.training = training

            self.nodes = tokens
            self.n_nodes = sentence.n_tokens()
            self.nodes_token = list(range(sentence.n_tokens()))

            self.n_nodes_cached = self.n_nodes
            self.node_mapping = {i: i for i in range(self.n_nodes)}

        def add_node(self, values, predecessor, mapping=None):
            self.sentence.add_node(values)
            self.nodes_token.append(self.nodes_token[predecessor])
            if mapping is not None:
                self.node_mapping[mapping] = self.n_nodes
            self.n_nodes += 1

        def add_edge(self, parent, child, values):
            if values[0] == Mapping.ANCHOR and child >= self.sentence.n_tokens(): return
            if values[0] == Mapping.ANCHOR and parent < self.sentence.n_tokens(): return
            self.sentence.add_edge(parent, child, values)

        @staticmethod
        def recompute_nodes(nodes, states, iteration):
            nodes_token, nodes_values = [], []
            start = 0
            for state in states:
                nodes_token.extend([start + t for t in state.nodes_token[state.n_nodes_cached:]])
                nodes_values.extend(state.sentence.factors[Sentence.NODE_VALUES][state.n_nodes_cached - state.sentence.n_tokens():])
                start += state.n_nodes_cached

            new_nodes = states[0].network.layers.encoder_token_transfer[iteration](tf.gather(nodes, np.array(nodes_token, np.int32)))
            for i, encoder in enumerate(states[0].network.layers.encoder_node_values):
                new_nodes += encoder(np.array([node_values[i] for node_values in nodes_values], np.int32))

            new_nodes = states[0].network.layers.dropout(new_nodes, training=states[0].training)

            start = 0
            for state in states:
                num_nodes = state.n_nodes - state.n_nodes_cached
                if num_nodes:
                    state.nodes = tf.concat([state.nodes, new_nodes[start:start + num_nodes]], axis=0)
                start += num_nodes
                state.n_nodes_cached = state.n_nodes

        @staticmethod
        def recompute_edges(ori_nodes, states, iteration):
            node_ids, edge_values = [], []
            start = 0
            for i, state in enumerate(states):
                for n in range(ori_nodes[i], state.n_nodes):
                    for e in state.sentence.children[n]:
                        node_ids.append(start + n - ori_nodes[i])
                        edge_values.append(state.sentence.factors[Sentence.EDGE_VALUES][e])
                    for e in state.sentence.parents[n]:
                        node_ids.append(start + n - ori_nodes[i])
                        edge_values.append([a + b for a, b in zip(
                            state.network.edge_values, state.sentence.factors[Sentence.EDGE_VALUES][e])])
                start += state.n_nodes - ori_nodes[i]

            edge_embeddings = 0
            for i, encoder in enumerate(states[0].network.layers.encoder_edge_values):
                edge_embeddings += encoder(np.array([edge_values[i] for edge_values in edge_values], np.int32))

            node_values = tf.math.segment_mean(edge_embeddings, node_ids)
            max_node_id = max(node_ids)
            if max_node_id + 1 < start:
                node_values = tf.pad(node_values, [[0, start - max_node_id - 1], [0, 0]])

            start = 0
            for i, state in enumerate(states):
                if state.n_nodes == ori_nodes[i]: continue
                state.nodes += tf.pad(node_values[start:start + state.n_nodes - ori_nodes[i]],
                                      [[ori_nodes[i], 0], [0, 0]])
                start += state.n_nodes - ori_nodes[i]

    def _summary_step(self):
        return int(int(self.optimizer.iterations) * self.optimizer_step_sentences)

    def train_epoch(self, train, learning_rate, args):
        self.optimizer.learning_rate = learning_rate

        skipped = 0
        num_gradients = 0
        while not train.epoch_finished():
            batch = train.next_batch(args.batch_size)
            sentences = batch["sentences"]

            layers = []
            for sentence in sentences:
                layers.append(get_layers(sentence))

            for accuracy in self.accuracies.values():
                accuracy.reset_states()

            with tf.GradientTape() as tape:
                encoded_tokens = self._encoder(
                    token_ids=batch["batch_factors"][Sentence.TOKENS],
                    token_charseqs=batch["charseqs"],
                    token_charseq_ids=batch["charseq_ids"],
                    token_values=batch["batch_factors"][Sentence.TOKEN_VALUES],
                    token_additionals=[batch[x] for x in ["bert", "fasttext", "word2vec"] if x in batch],
                    training=True)

                states = []
                for i, sentence in enumerate(sentences):
                    states.append(self.State(self, sentence.empty_copy(), encoded_tokens[i][:sentence.n_tokens()], True))

                loss = 0
                for iteration in range(args.decoder_iterations):
                    # 1) New nodes
                    ori_nodes = [s.n_nodes for s in states]
                    sum_nodes = sum(ori_nodes)
                    nodes = tf.concat([s.nodes for s in states], axis=0)

                    target_ops = np.zeros([sum_nodes], np.int32)
                    target_node_values = np.zeros([sum_nodes, len(train.node_properties)], np.int32)
                    target_node_values_mask = np.zeros([sum_nodes, 1], np.int32)
#                     target_edge_values = np.zeros([sum_nodes, len(train.edge_properties)], np.int32)
#                     target_edge_values_mask = np.zeros([sum_nodes, 1], np.int32)

                    start = 0
                    for i in range(len(sentences)):
                        if iteration < len(layers[i]):
                            state, layer, sentence = states[i], layers[i][iteration], sentences[i]
                            for n, target, edge in layer:
                                n = state.node_mapping[n]
                                target_ops[start + n] = NODE_PARENT if sentence.factors[Sentence.EDGE_PARENTS][edge] == target else NODE_CHILD
                                target_node_values[start + n] = sentence.factors[Sentence.NODE_VALUES][target - sentence.n_tokens()]
                                target_node_values_mask[start + n] = 1
#                                 target_edge_values[start + n] = sentence.factors[Sentence.EDGE_VALUES][edge]
#                                 target_edge_values_mask[start + n] = 1 if n >= state.sentence.n_tokens() else 0
                                state.add_node(target_node_values[start + n], n, target)
                        start += states[i].n_nodes_cached
                    self.State.recompute_nodes(nodes, states, iteration)

                    predictions = self.layers.decoder_node_operation[iteration](nodes)
                    loss += self.loss_sce(target_ops, predictions)
                    self.accuracies["node/ops"](target_ops, predictions)
                    for i, prop in enumerate(train.node_properties):
                        predictions = self.layers.decoder_node_values[iteration][i](nodes)
                        loss += self.loss_sce(target_node_values[:, i], predictions, target_node_values_mask)
                        self.accuracies["node/" + prop](target_node_values[:, i], predictions, target_node_values_mask)
#                     for i, prop in enumerate(train.edge_properties):
#                         predictions = self.layers.decoder_edge_values[iteration][i](nodes)
#                         loss += self.loss_sce(target_edge_values[:, i], predictions, target_edge_values_mask)
#                         self.accuracies["edge/" + prop](target_edge_values[:, i], predictions, target_edge_values_mask)

                    # 2) Edges
                    sum_nodes = sum(s.n_nodes for s in states)
                    nodes = tf.concat([s.nodes for s in states], axis=0)
                    target_indices_a, target_indices_b = [], []
                    target_a_parent, target_a_child = [], []
                    target_deprel_parents, target_deprel_children = [], []
                    target_deprel_values = []

                    start = 0
                    for i in range(len(sentences)):
                        state, sentence = states[i], sentences[i]
                        offset = len(target_indices_a)
                        for j in range(ori_nodes[i], state.n_nodes):
                            target_indices_a.append(start + np.repeat(np.int32(j), state.n_nodes))
                            target_indices_b.append(start + np.arange(0, state.n_nodes, dtype=np.int32))
                            target_a_parent.append(np.zeros(state.n_nodes, np.float32))
                            target_a_child.append(np.zeros(state.n_nodes, np.float32))
                        for n_ori, n in state.node_mapping.items():
                            if n >= ori_nodes[i]:
                                for e in sentence.parents[n_ori]:
                                    if sentence.factors[Sentence.EDGE_VALUES][e][0] == Mapping.ROOT: continue
                                    p = sentence.factors[sentence.EDGE_PARENTS][e]
                                    if p in state.node_mapping:
                                        p = state.node_mapping[p]
                                        if not args.no_anchors or not sentence.factors[Sentence.EDGE_VALUES][e][0] == Mapping.ANCHOR:
                                            target_a_child[offset + n - ori_nodes[i]][p] = 1
                                        target_deprel_parents.append(start + p)
                                        target_deprel_children.append(start + n)
                                        target_deprel_values.append(sentence.factors[Sentence.EDGE_VALUES][e])
                                        state.add_edge(p, n, target_deprel_values[-1])
                                for e in sentence.children[n_ori]:
                                    if sentence.factors[Sentence.EDGE_VALUES][e][0] == Mapping.ROOT: continue
                                    c = sentence.factors[sentence.EDGE_CHILDREN][e]
                                    if c in state.node_mapping:
                                        c = state.node_mapping[c]
                                        if not args.no_anchors or not sentence.factors[Sentence.EDGE_VALUES][e][0] == Mapping.ANCHOR:
                                            target_a_parent[offset + n-ori_nodes[i]][c] = 1
                                        target_deprel_parents.append(start + n)
                                        target_deprel_children.append(start + c)
                                        target_deprel_values.append(sentence.factors[Sentence.EDGE_VALUES][e])
                                        state.add_edge(n, c, target_deprel_values[-1])
                        start += state.n_nodes

                    if not target_indices_a or not target_deprel_parents:
                        continue
                    target_indices_a = np.concatenate(target_indices_a, axis=0)
                    target_indices_b = np.concatenate(target_indices_b, axis=0)
                    target_a_parent = np.concatenate(target_a_parent, axis=0)
                    target_a_child = np.concatenate(target_a_child, axis=0)
                    target_deprel_parents = np.array(target_deprel_parents, np.int32)
                    target_deprel_children = np.array(target_deprel_children, np.int32)
                    target_deprel_values = np.array(target_deprel_values, np.int32)

                    # 2.1) Compute arcs
                    edge_parents = self.layers.decoder_edge_parents[iteration](nodes)
                    edge_children = self.layers.decoder_edge_children[iteration](nodes)

                    a_parent = tf.nn.tanh(self.layers.sum([tf.gather(edge_parents, target_indices_a),
                                                           tf.gather(edge_children, target_indices_b)]))
                    if args.highway:
                        a_parent += self.layers.decoder_edge_highway[iteration](a_parent)
                    a_parent = self.layers.decoder_edge_arc[iteration](a_parent)
                    loss += self.loss_sce(target_a_parent, a_parent)
                    self.accuracies["edge/arc"](target_a_parent, a_parent)

                    a_child = tf.nn.tanh(self.layers.sum([tf.gather(edge_parents, target_indices_b),
                                                          tf.gather(edge_children, target_indices_a)]))
                    if args.highway:
                        a_child += self.layers.decoder_edge_highway[iteration](a_child)
                    a_child = self.layers.decoder_edge_arc[iteration](a_child)
                    loss += self.loss_sce(target_a_child, a_child)
                    self.accuracies["edge/arc"](target_a_child, a_child)

                    # 2.2) Compute deprels
                    deprel_parents = self.layers.decoder_deprel_parents[iteration](nodes)
                    deprel_children = self.layers.decoder_deprel_children[iteration](nodes)
                    deprel_weights = tf.nn.tanh(self.layers.sum(
                        [tf.gather(deprel_parents, target_deprel_parents),
                         tf.gather(deprel_children, target_deprel_children)]))
                    if args.highway:
                        deprel_weights += self.layers.decoder_deprel_highway[iteration](deprel_weights)
                    for i, prop in enumerate(train.edge_properties):
                        predictions = self.layers.decoder_deprel_values[iteration][i](deprel_weights)
                        loss += self.loss_sce(target_deprel_values[:, i], predictions)
                        self.accuracies["edge/" + prop](target_deprel_values[:, i], predictions)

                    self.State.recompute_edges(ori_nodes, states, iteration)

                # Tops
                sum_nodes = sum(s.n_nodes for s in states)
                nodes = tf.concat([s.nodes for s in states], axis=0)
                target_tops = np.zeros([sum_nodes], np.int32)
                start = 0
                for i, sentence in enumerate(sentences):
                    for e in sentence.children[0]:
                        if sentence.factors[Sentence.EDGE_VALUES][e][0] == Mapping.ROOT:
                            c = sentence.factors[Sentence.EDGE_CHILDREN][e]
                            if c in states[i].node_mapping:
                                target_tops[start + states[i].node_mapping[c]] = 1
                    start += states[i].n_nodes
                predictions = self.layers.decoder_tops(nodes)
                loss += self.loss_sce(target_tops, predictions)
                self.accuracies["edge/tops"](target_tops, predictions)

            tg = tape.gradient(loss, self.layers.trainable_variables)
            tg_none = [variable.name for g, variable in zip(tg, self.layers.trainable_variables) if g is None]
            if tg_none:
                print("Skipping a batch with None gradient for variables {}".format(tg_none), file=sys.stderr, flush=True)
                continue

            if num_gradients == 0:
                gradients = [g.numpy() if not isinstance(g, tf.IndexedSlices) else [(g.values.numpy(), g.indices.numpy())] for g in tg]
            else:
                for g, ng in zip(gradients, tg):
                    if isinstance(g, list):
                        g.append((ng.values.numpy(), ng.indices.numpy()))
                    else:
                        g += ng.numpy()
            num_gradients += 1
            if num_gradients == args.batch_aggregation or len(train._permutation) == 0:
                gradients = [tf.IndexedSlices(*map(np.concatenate, zip(*g))) if isinstance(g, list) else g for g in gradients]
                self.optimizer.apply_gradients(zip(gradients, self.layers.trainable_variables))
                num_gradients = 0
                if int(self.optimizer.iterations) % 100 == 0:
                    tf.summary.experimental.set_step(self._summary_step())
                    with self.writer.as_default():
                        for name, accuracy in self.accuracies.items():
                            tf.summary.scalar("train/" + name, accuracy.result())

        tf.summary.experimental.set_step(self._summary_step())
        with self.writer.as_default():
            tf.summary.scalar("train/skipped", skipped)

    def predict(self, data, args):
        output = []

        while not data.epoch_finished():
            batch = data.next_batch(args.batch_size)
            sentences = batch["sentences"]

            encoded_tokens = self._encoder(
                token_ids=batch["batch_factors"][Sentence.TOKENS],
                token_charseqs=batch["charseqs"],
                token_charseq_ids=batch["charseq_ids"],
                token_values=batch["batch_factors"][Sentence.TOKEN_VALUES],
                token_additionals=[batch[x] for x in ["bert", "fasttext", "word2vec"] if x in batch],
                training=False)

            states = []
            for i, sentence in enumerate(sentences):
                states.append(self.State(self, sentence.empty_copy(), encoded_tokens[i][:sentence.n_tokens()], False))

            for iteration in range(args.decoder_iterations):
                # 1) New nodes
                ori_nodes = [s.n_nodes for s in states]
                sum_nodes = sum(ori_nodes)
                nodes = tf.concat([s.nodes for s in states], axis=0)

                prediction_ops = tf.argmax(self.layers.decoder_node_operation[iteration](nodes), axis=1).numpy()
                prediction_node_values = np.stack(
                    [tf.argmax(decoder(nodes), axis=1).numpy() for decoder in self.layers.decoder_node_values[iteration]], axis=1)
#                 prediction_edge_values = np.stack(
#                     [tf.argmax(decoder(nodes), axis=1).numpy() for decoder in self.layers.decoder_edge_values[iteration]], axis=1)
                start = 0
                for i, state in enumerate(states):
                    for n in range(1, ori_nodes[i]):
                        if prediction_ops[start + n] == NODE_NO: continue
                        state.add_node(prediction_node_values[start + n], n)
#                         edge_values = prediction_edge_values[start + n] if n >= state.sentence.n_tokens() else [Mapping.ANCHOR] + [Mapping.NONE] * (len(data.edge_properties) - 1)
#                         if prediction_ops[start + n] == NODE_PARENT:
#                             state.add_edge(state.n_nodes - 1, n, edge_values)
#                         else:
#                             state.add_edge(n, state.n_nodes - 1, edge_values)
                    start += ori_nodes[i]
                self.State.recompute_nodes(nodes, states, iteration)

                # 2) Edges
                sum_nodes = sum(s.n_nodes for s in states)
                nodes = tf.concat([s.nodes for s in states], axis=0)
                indices_a, indices_b = [], []
                indices_parents, indices_children = [], []

                start = 0
                for i, state in enumerate(states):
                    for j in range(ori_nodes[i], state.n_nodes):
                        indices_a.append(start + np.repeat(np.int32(j), state.n_nodes))
                        indices_b.append(start + np.arange(0, state.n_nodes, dtype=np.int32))
                    start += state.n_nodes
                if not indices_a: continue

                indices_a = np.concatenate(indices_a, axis=0)
                indices_b = np.concatenate(indices_b, axis=0)

                # 2.1) Compute arcs
                edge_parents = self.layers.decoder_edge_parents[iteration](nodes)
                edge_children = self.layers.decoder_edge_children[iteration](nodes)

                a_parent = tf.nn.tanh(self.layers.sum([tf.gather(edge_parents, indices_a),
                                                       tf.gather(edge_children, indices_b)]))
                if args.highway:
                    a_parent += self.layers.decoder_edge_highway[iteration](a_parent)
                a_parent = tf.argmax(self.layers.decoder_edge_arc[iteration](a_parent), axis=1).numpy()
                a_child = tf.nn.tanh(self.layers.sum([tf.gather(edge_parents, indices_b),
                                                      tf.gather(edge_children, indices_a)]))
                if args.highway:
                    a_child += self.layers.decoder_edge_highway[iteration](a_child)
                a_child = tf.argmax(self.layers.decoder_edge_arc[iteration](a_child), axis=1).numpy()

                start, offset = 0, 0
                for i, state in enumerate(states):
                    node = 0
                    for n in range(1, ori_nodes[i]):
                        if prediction_ops[start + n] == NODE_NO: continue
                        if prediction_ops[start + n] == NODE_PARENT:
                            a_parent[offset + node * state.n_nodes + n] = 1
                        else:
                            a_child[offset + node * state.n_nodes + n] = 1
                        node += 1
                    start += ori_nodes[i]
                    offset += (state.n_nodes - ori_nodes[i]) * state.n_nodes
                for i in np.where(a_parent > 0)[0]:
                    indices_parents.append(indices_a[i])
                    indices_children.append(indices_b[i])
                for i in np.where(a_child > 0)[0]:
                    indices_parents.append(indices_b[i])
                    indices_children.append(indices_a[i])
                indices_parents = np.array(indices_parents, np.int32)
                indices_children = np.array(indices_children, np.int32)

                # 2.2) Compute deprels
                deprel_parents = self.layers.decoder_deprel_parents[iteration](nodes)
                deprel_children = self.layers.decoder_deprel_children[iteration](nodes)
                deprel_weights = tf.nn.tanh(self.layers.sum([tf.gather(deprel_parents, indices_parents),
                                                             tf.gather(deprel_children, indices_children)]))
                if args.highway:
                    deprel_weights += self.layers.decoder_deprel_highway[iteration](deprel_weights)
                edge_predictions = []
                for i, prop in enumerate(data.edge_properties):
                    edge_predictions.append(tf.argmax(
                        self.layers.decoder_deprel_values[iteration][i](deprel_weights), axis=1).numpy())
                edge_predictions = np.stack(edge_predictions, axis=1)

                offset = 0
                for _ in range(2):
                    start = 0
                    for state in states:
                        while offset < len(indices_parents) and indices_parents[offset] >= start and indices_parents[offset] < start + state.n_nodes:
                            state.add_edge(
                                indices_parents[offset] - start,
                                indices_children[offset] - start,
                                edge_predictions[offset])
                            offset += 1
                        start += state.n_nodes
                assert offset == len(indices_parents)

                self.State.recompute_edges(ori_nodes, states, iteration)

            # Tops
            sum_nodes = sum(s.n_nodes for s in states)
            nodes = tf.concat([s.nodes for s in states], axis=0)
            predictions = tf.argmax(self.layers.decoder_tops(nodes), axis=1)
            start = 0
            for i, state in enumerate(states):
                for n in range(1, state.n_nodes):
                    if predictions[start + n]:
                        state.add_edge(0, n, [Mapping.ROOT] + [Mapping.NONE] * (len(data.edge_properties) - 1))
                start += state.n_nodes

            for state in states:
                output.append(state.sentence.write(data.mappings))
                output.append("\n")

        return "".join(output)

    @staticmethod
    def predict_ensemble(networks, data, args):
        output = []

        while not data.epoch_finished():
            batch = data.next_batch(args.batch_size)
            sentences = batch["sentences"]

            for network in networks:
                encoded_tokens = network._encoder(
                    token_ids=batch["batch_factors"][Sentence.TOKENS],
                    token_charseqs=batch["charseqs"],
                    token_charseq_ids=batch["charseq_ids"],
                    token_values=batch["batch_factors"][Sentence.TOKEN_VALUES],
                    token_additionals=[batch[x] for x in ["bert", "fasttext", "word2vec"] if x in batch],
                    training=False)

                network.states = []
                for i, sentence in enumerate(sentences):
                    network.states.append(network.State(network, sentence.empty_copy(), encoded_tokens[i][:sentence.n_tokens()], False))

            for iteration in range(args.decoder_iterations):
                # 1) New nodes
                ori_nodes = [s.n_nodes for s in networks[0].states]
                sum_nodes = sum(ori_nodes)
                for network in networks:
                    network.nodes = tf.concat([s.nodes for s in network.states], axis=0)

                prediction_ops = np.argmax(
                    np.mean([network.layers.decoder_node_operation[iteration](network.nodes).numpy() for network in networks], axis=0),
                    axis=1)
                prediction_node_values = np.stack([
                    np.argmax(
                        np.mean([network.layers.decoder_node_values[iteration][i](network.nodes).numpy() for network in networks], axis=0),
                        axis=1)
                    for i in range(len(networks[0].layers.decoder_node_values[iteration]))
                ], axis=1)

                for network in networks:
                    start = 0
                    for i, state in enumerate(network.states):
                        for n in range(1, ori_nodes[i]):
                            if prediction_ops[start + n] == NODE_NO: continue
                            state.add_node(prediction_node_values[start + n], n)
                        start += ori_nodes[i]
                    network.State.recompute_nodes(network.nodes, network.states, iteration)

                # 2) Edges
                sum_nodes = sum(s.n_nodes for s in networks[0].states)
                for network in networks:
                    network.nodes = tf.concat([s.nodes for s in network.states], axis=0)
                indices_a, indices_b = [], []
                indices_parents, indices_children = [], []

                start = 0
                for i, state in enumerate(networks[0].states):
                    for j in range(ori_nodes[i], state.n_nodes):
                        indices_a.append(start + np.repeat(np.int32(j), state.n_nodes))
                        indices_b.append(start + np.arange(0, state.n_nodes, dtype=np.int32))
                    start += state.n_nodes
                if not indices_a: continue

                indices_a = np.concatenate(indices_a, axis=0)
                indices_b = np.concatenate(indices_b, axis=0)

                # 2.1) Compute arcs
                a_parent, a_child = [], []
                for network in networks:
                    edge_parents = network.layers.decoder_edge_parents[iteration](network.nodes)
                    edge_children = network.layers.decoder_edge_children[iteration](network.nodes)

                    a_parent.append(tf.nn.tanh(network.layers.sum([tf.gather(edge_parents, indices_a),
                                                                   tf.gather(edge_children, indices_b)])))
                    if args.highway:
                        a_parent[-1] += network.layers.decoder_edge_highway[iteration](a_parent[-1])
                    a_parent[-1] = network.layers.decoder_edge_arc[iteration](a_parent[-1]).numpy()

                    a_child.append(tf.nn.tanh(network.layers.sum([tf.gather(edge_parents, indices_b),
                                                                   tf.gather(edge_children, indices_a)])))
                    if args.highway:
                        a_child[-1] += network.layers.decoder_edge_highway[iteration](a_child[-1])
                    a_child[-1] = network.layers.decoder_edge_arc[iteration](a_child[-1]).numpy()
                a_parent = np.argmax(np.mean(a_parent, axis=0), axis=1)
                a_child = np.argmax(np.mean(a_child, axis=0), axis=1)

                start, offset = 0, 0
                for i, state in enumerate(networks[0].states):
                    node = 0
                    for n in range(1, ori_nodes[i]):
                        if prediction_ops[start + n] == NODE_NO: continue
                        if prediction_ops[start + n] == NODE_PARENT:
                            a_parent[offset + node * state.n_nodes + n] = 1
                        else:
                            a_child[offset + node * state.n_nodes + n] = 1
                        node += 1
                    start += ori_nodes[i]
                    offset += (state.n_nodes - ori_nodes[i]) * state.n_nodes
                for i in np.where(a_parent > 0)[0]:
                    indices_parents.append(indices_a[i])
                    indices_children.append(indices_b[i])
                for i in np.where(a_child > 0)[0]:
                    indices_parents.append(indices_b[i])
                    indices_children.append(indices_a[i])
                indices_parents = np.array(indices_parents, np.int32)
                indices_children = np.array(indices_children, np.int32)

                # 2.2) Compute deprels
                edge_predictions = []
                for i, prop in enumerate(data.edge_properties):
                    predictions = []
                    for network in networks:
                        deprel_parents = network.layers.decoder_deprel_parents[iteration](network.nodes)
                        deprel_children = network.layers.decoder_deprel_children[iteration](network.nodes)
                        deprel_weights = tf.nn.tanh(network.layers.sum([tf.gather(deprel_parents, indices_parents),
                                                                        tf.gather(deprel_children, indices_children)]))
                        if args.highway:
                            deprel_weights += network.layers.decoder_deprel_highway[iteration](deprel_weights)
                        predictions.append(network.layers.decoder_deprel_values[iteration][i](deprel_weights).numpy())
                    edge_predictions.append(np.argmax(np.mean(predictions, axis=0), axis=1))
                edge_predictions = np.stack(edge_predictions, axis=1)

                for network in networks:
                    offset = 0
                    for _ in range(2):
                        start = 0
                        for state in network.states:
                            while offset < len(indices_parents) and indices_parents[offset] >= start and indices_parents[offset] < start + state.n_nodes:
                                state.add_edge(
                                    indices_parents[offset] - start,
                                    indices_children[offset] - start,
                                    edge_predictions[offset])
                                offset += 1
                            start += state.n_nodes
                    assert offset == len(indices_parents)

                    network.State.recompute_edges(ori_nodes, network.states, iteration)

            # Tops
            sum_nodes = sum(s.n_nodes for s in networks[0].states)
            for network in networks:
                network.nodes = tf.concat([s.nodes for s in network.states], axis=0)
            predictions = np.argmax(
                np.mean([network.layers.decoder_tops(network.nodes).numpy() for network in networks], axis=0),
                axis=1)
            for network in networks:
                start = 0
                for i, state in enumerate(network.states):
                    for n in range(1, state.n_nodes):
                        if predictions[start + n]:
                            state.add_edge(0, n, [Mapping.ROOT] + [Mapping.NONE] * (len(data.edge_properties) - 1))
                    start += state.n_nodes

            for state in networks[0].states:
                output.append(state.sentence.write(data.mappings))
                output.append("\n")

        return "".join(output)

    def evaluate(self, dataset, data, gold_path, args):
        step = self._summary_step()

        predicted_path = os.path.join(args.logdireval, "{}_{:07d}.mrp".format(dataset, step))
        assert not os.path.exists(predicted_path)

        with open(predicted_path, "w", encoding="utf-8") as predicted_file:
            predicted = self.predict(data, args)
            print(predicted, end="", file=predicted_file)

        if dataset in ["dev", "bdev", "lpss"]:
            mtool_evaluate.mrp_eval(
                gold_path, predicted_path, "{}_eval".format(dataset), rrhc=4, mces=100000, parallelize=8, logdir=args.logdireval, logstep=step, distributed=True)


rub_dict = {1: ['a', 'd', 'g', 'j', 'm', 'p', 's', 'v', 'y', '.', ' '],
            2: ['b', 'e', 'h', 'k', 'n', 'q', 't', 'w', 'z', ','],
            3: ['c', 'f', 'i', 'l', 'o', 'r', 'u', 'x', '!']}
counter = 0

message = input()

for letter in message:
    for key in list(rub_dict.keys()):
        if letter in rub_dict[key]:
            counter += key
            break

print(counter)


strings = input().split(" ")
searched_palindrome = input()
palindromes = []

for word in strings:
    if word == "".join(reversed(word)):
        palindromes.append(word)

print(f'{palindromes}')
print(f'Found palindrome {palindromes.count(searched_palindrome)} times')


from __future__ import print_function, division
import numpy as np
import random
import json
import sys
import os
import argparse
from shutil import copyfile

import networkx as nx
from networkx.readwrite import json_graph
import pdb

def read_dict(dict_file):
    all_nodes_source = set()
    all_nodes_target = set()
    all_instances = []
    if os.path.exists(dict_file):
        with open(dict_file) as fp:
            for line in fp:
                ins = line.split()
                all_instances.append([ins[0], ins[1]])
                all_nodes_source.add(ins[0])
                all_nodes_target.add(ins[1])

    return all_instances, all_nodes_source, all_nodes_target

def read_data(dataset_dir, dataset_prefix):

    prefix = dataset_dir + "/" + dataset_prefix
    id_map = json.load(open(prefix + "-id_map.json"))
    feats = np.load(prefix + "-feats.npy")
    return id_map, feats

def check_dict_features(all_instances, source_id_map, source_feats, target_id_map, target_feats):
    correct_case = 0
    for instance in all_instances:
        source = np.argmax(source_feats[source_id_map[instance[0]]])
        target = np.argmax(target_feats[target_id_map[instance[1]]])
        if source == target:
            correct_case += 1

    print("Accuracy = {0}".format(correct_case / len(all_instances)))

def parse_args():
    parser = argparse.ArgumentParser(description="Randomly remove edges and generate dict.")
    parser.add_argument('--source_dataset_dir', default=None, help='Path to source dataset')
    parser.add_argument('--target_dataset_dir', default=None, help='Path to target dataset')
    parser.add_argument('--source_dataset_prefix', default=None, help='Source dataset prefix')
    parser.add_argument('--target_dataset_prefix', default=None, help='Target dataset prefix')
    parser.add_argument('--dict_file', default=None, help='Path to groundtruth file')
    parser.add_argument('--seed', type=int, default=121, help='Random seed')
    return parser.parse_args()

if __name__ == "__main__":
    args = parse_args()
    print(args)
    random.seed(args.seed)
    np.random.seed(args.seed)
    all_instances, all_nodes_source, all_nodes_target = read_dict(args.dict_file)
    source_id_map, source_feats = read_data(args.source_dataset_dir, args.source_dataset_prefix)
    target_id_map, target_feats = read_data(args.target_dataset_dir, args.target_dataset_prefix)
    check_dict_features(all_instances, source_id_map, source_feats, target_id_map, target_feats)


import json
import requests as req
import csv
import math
import re
import nltk
from nltk.corpus import stopwords
import pymorphy2

def CleanStopWords(stopWords3, line):
    stopWords = stopwords.words("russian")
    stopWords2 = stopwords.words("english")
    return [word for word in line if word not in stopWords and word not in stopWords2 and word not in stopWords3]

def CleanDigits(line):
    return [word for word in line if not word.isdigit()]

def Normalize(line):
    morph = pymorphy2.MorphAnalyzer()
    return [morph.parse(word)[0].normal_form for word in line if len(word) > 2 and
            morph.parse(word)[0].tag.POS != 'PREP' and
            morph.parse(word)[0].tag.POS != 'CONJ' and
            morph.parse(word)[0].tag.POS != 'PRCL' and
            morph.parse(word)[0].tag.POS != 'INTJ']

def Clean_Normalize(str):
    #additional stop list
    with open("stopwords.txt") as file:
        stopWords3 = [row.strip() for row in file]

    new_str1 = re.sub(r'\<[^>]*\>', '', str)
    new_str1 = new_str1.lower()
    new_str1 = re.sub(r'[^\w\d\s]', ' ', new_str1)
    tokenline = nltk.word_tokenize(new_str1)
    tokenline = CleanStopWords(stopWords3, tokenline)
    tokenline = Normalize(tokenline)
    tokenline = CleanDigits(tokenline)
    return tokenline

if __name__ == '__main__':

    f_write = csv.writer(open("1page_normalized_more.csv", "w+", newline='', encoding='cp1251'), delimiter=';')
    f = csv.reader(open("1page_notags.csv", "r", newline='', encoding='cp1251'), delimiter=';')

    i = 0
    for row in f:
        i += 1

        f_write.writerow([row[0],
                    row[1],
                    Clean_Normalize(row[2]),
                    Clean_Normalize(row[3]),
                    row[4],
                    row[5],
                    row[6],
                    row[7],
                    row[8],
                    row[9]])

        if (i % 10 == 0):
            print(' ',i)
        #print(' ',i)
    print('end')


#!/usr/bin/env python3
"""The function decomposes a number into prime factors."""

from math import sqrt

def factorise(number: int) -> list[int]:
    """Returns the prime factors of the number.

    Parameters
    ----------
    number : int

    Raises
    ------
    TypeError
        if the number if not integer.
    ValueError
        if the number is a negative integer.

    Returns
    -------
    list[int]
    """
    if not isinstance(number, int):
        raise TypeError(f"{number} is not integer")
    if number < 0:
        raise ValueError(f"{number} is negative")
    factors: list[int] = []
    while number and number % 2 == 0:
        factors.append(2)
        number //= 2
    factor: int = 3
    while factor <= int(sqrt(number)):
        while number % factor == 0:
            factors.append(factor)
            number //= factor
        factor += 2
    if number > 1:
        factors.append(number)
    return factors


import math

class Node:

    def __init__(self, key, value):
        self.key = key
        self.value = value

class HashTable:

    def __init__(self, table_size):  # add appropriate attributes, NO default size
        ''' Initializes an empty hash table with a size that is the smallest
            prime number that is bigger than table_size '''
        if self.is_prime(table_size):
            table_size = table_size
        else:
            table_size = self.next_prime(table_size)
        self.table_size = table_size
        self.hash_table = [None] * table_size
        self.entries = 0

    def insert(self, key, value=None):
        ''' Inserts an entry into the hash table (using Horner hash function to determine index,
        and quadratic probing to resolve collisions).
        The key is a string (a word) to be entered, and value is any object (e.g. Python List).
        If the key is not already in the table, the key is inserted along with the associated value
        If the key is is in the table, the new value replaces the existing value.
        If load factor is greater than 0.5 after an insertion, hash table size should be increased
        to the next prime greater than 2*table_size.'''
        entry = Node(key, value)
        index = self.horner_hash(key) % self.table_size
        start = index

        i = 0
        while self.hash_table[index] is not None:
            if self.hash_table[index].key == entry.key:
                self.hash_table[index].value = entry.value
                return
            i += 1
            index = start + (i ** 2)
            while index > self.table_size - 1:
                index = index - self.table_size
        self.hash_table[index] = entry
        self.entries += 1

        if self.get_load_factor() > 0.5:
            self.rehash()

    def rehash(self):
        temp = self.hash_table
        self.table_size = self.next_prime(self.table_size * 2)
        self.hash_table = [None] * self.table_size
        self.entries = 0
        for entry in temp:
            if entry is not None:
                self.insert(entry.key, entry.value)

    def horner_hash(self, key):
        ''' Compute the hash value by using Horner’s rule, as described in project specification. '''
        n = min(8, len(key))
        hash_value = 0
        for i in range(n):
            hash_value += (ord(key[i]) * 31 ** (n - 1 - i))
        return hash_value

    def next_prime(self, n):
        ''' Find the next prime number that is > n.'''
        if n <= 1:
            return 2

        prime = n
        found = False

        while not found:
            prime = prime + 1

            if self.is_prime(prime):
                found = True

        return prime

    def is_prime(self, n):
        if n <= 1:
            return False
        if n <= 3:
            return True

        if n % 2 == 0 or n % 3 == 0:
            return False

        for i in range(5, int(math.sqrt(n) + 1), 6):
            if n % i == 0 or n % (i + 2) == 0:
                return False

        return True

    def in_table(self, key):
        ''' Returns True if key is in an entry of the hash table, False otherwise.'''
        index = self.horner_hash(key) % self.table_size
        start = index
        i = 0
        while self.hash_table[index] is not None:
            if self.hash_table[index].key == key:
                return True
            else:
                i += 1
                index = start + (i ** 2)
                while index > self.table_size - 1:
                    index = index - self.table_size
        return False

    def get_index(self, key):
        ''' Returns the index of the hash table entry containing the provided key.
        If there is not an entry with the provided key, returns None. '''
        index = self.horner_hash(key) % self.table_size
        start = index
        i = 0
        while self.hash_table[index] is not None:
            if self.hash_table[index].key == key:
                return index
            else:
                i += 1
                index = start + (i ** 2)
                while index > self.table_size - 1:
                    index = index - self.table_size
        return None

    def get_all_keys(self):
        ''' Returns a Python list of all keys in the hash table.'''
        keys = []
        for entry in self.hash_table:
            if entry is not None:
                keys.append(entry.key)
        return keys

    def get_value(self, key):
        ''' Returns the value associated with the key.
        If key is not in hash table, returns None.'''
        index = self.horner_hash(key) % self.table_size
        start = index
        i = 0
        while self.hash_table[index] is not None:
            if self.hash_table[index].key == key:
                return self.hash_table[index].value
            else:
                i += 1
                index = start + (i ** 2)
                while index > self.table_size - 1:
                    index = index - self.table_size
        return None

    def get_num_items(self):
        ''' Returns the number of entries in the table.'''
        return self.entries

    def get_table_size(self):
        ''' Returns the size of the hash table.'''
        return self.table_size

    def get_load_factor(self):
        ''' Returns the load factor of the hash table (entries / table_size).'''
        return self.entries / self.table_size


import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import os
import time

flower_str_list = ['daisy', 'dandelion', 'rose', 'sunflower', 'tulip']

CURRENT_PATH = os.getcwd()
DATASET_PATH = os.path.join(CURRENT_PATH, 'drive',
                            'MyDrive', 'flower_dataset')
TEST_DATASET_PATH = os.path.join(CURRENT_PATH, 'drive',
                            'MyDrive', 'flower_test_dataset')

train_ds = tf.keras.preprocessing.image_dataset_from_directory(
    DATASET_PATH,
    image_size=(256, 256),
    batch_size=64,
    subset='training',
    validation_split=0.2,
    shuffle=True,
    seed=8282
)
val_ds = tf.keras.preprocessing.image_dataset_from_directory(
    DATASET_PATH,
    image_size=(256, 256),
    batch_size=64,
    subset='validation',
    validation_split=0.2,
    shuffle=True,
    seed=8282
)
test_ds = tf.keras.preprocessing.image_dataset_from_directory(
    TEST_DATASET_PATH,
    image_size=(256, 256),
    batch_size=64,
    shuffle=True,
    seed=8282
)

def get_overfitting_model():
    model = tf.keras.Sequential([
        tf.keras.layers.Conv2D(64, (3, 3), padding="same",
                               activation="relu", input_shape=(256, 256, 3)),
        tf.keras.layers.Conv2D(128, (3, 3), padding="same", activation="relu"),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(128, activation="relu"),
        tf.keras.layers.Dense(64, activation="relu"),
        tf.keras.layers.Dense(5, activation="softmax")
    ])
    return model

def get_not_alexnet_model():
    model = tf.keras.Sequential([
        tf.keras.layers.Conv2D(32, (3, 3), padding="same",
                               activation="relu", input_shape=(256, 256, 3)),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Conv2D(64, (3, 3), padding="same", activation="relu"),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Dropout(0.3),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(128, activation="relu"),
        tf.keras.layers.Dropout(0.2),
        tf.keras.layers.Dense(64, activation="relu"),
        tf.keras.layers.Dense(5, activation="softmax")
    ])
    return model

def get_alexnet_model():
    model = tf.keras.Sequential([
        tf.keras.layers.Conv2D(64, (3, 3), padding="same",
                               activation="relu", input_shape=(256, 256, 3)),
        tf.keras.layers.Conv2D(128, (3, 3), padding="same", activation="relu"),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Dropout(0.3),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(128, activation="relu"),
        tf.keras.layers.Dropout(0.2),
        tf.keras.layers.Dense(64, activation="relu"),
        tf.keras.layers.Dense(5, activation="softmax")
    ])
    return model

def get_callback(mode):
    save_callback = tf.keras.callbacks.ModelCheckpoint(
        filepath='flower/checkpoint/' + mode,
        monitor='val_acc',
        mode='max',
        save_weights_only=True,
        save_freq='epoch'
    )
    return save_callback

def simulation(mode, epochs, model):
    model_path = os.path.join(os.getcwd(), 'drive',
                              'MyDrive', 'flower_model', mode)
    logs_path = os.path.join(os.getcwd(), 'drive', 'MyDrive',
                             'flower_logs', mode)
    model.summary()
    model.compile(loss="sparse_categorical_crossentropy",
                  optimizer="adam", metrics=['accuracy'])
    tensorboard = tf.keras.callbacks.TensorBoard(log_dir=logs_path)
    save_callback = get_callback(mode)
    model.fit(train_ds, validation_data=val_ds,
              epochs=epochs, verbose=2, callbacks=[save_callback, tensorboard])
    model.save(model_path)
    result = model.evaluate(test_ds)
    print(result)

not_alexnet_model = get_not_alexnet_model()
alexnet_model = get_alexnet_model()
overfitting_model = get_overfitting_model()

simulation('alexnet_model', 100, alexnet_model)
simulation('not_alexnet_model', 100, not_alexnet_model)
simulation('overfitting_model', 100, overfitting_model)


#!/usr/bin/env python

"""
This module provides a single construct() function which produces a Smach state
machine that implements wallfollowing behavior.

The constructed state machine has three attached methods:
    * set_ranges(ranges): this function should be called to update the range
                          readings
    * get_twist(): returns a twist message that could be directly passed to the
                   velocity publisher
    * set_config(config): updates the machine userdata with the new config

The constructed state machine is preemptable, i.e. each state checks whether
a preemption is requested and returns 'preempted' if that is the case.
"""

PACKAGE = 'amr_bugs'

import roslib
import numpy as np
import matplotlib.pyplot as plt
roslib.load_manifest(PACKAGE)
import smach
from preemptable_state import PreemptableState
from math import copysign
from types import MethodType
from geometry_msgs.msg import Twist

__all__ = ['construct']

#=============================== YOUR CODE HERE ===============================
# Instructions: write a function for each state of wallfollower state machine.
#               The function should have exactly one argument (userdata
#               dictionary), which you should use to access the input ranges
#               and to provide the output velocity.
#               The function should have at least one 'return' statement, which
#               returns one of the possible outcomes of the state.
#               The function should not block (i.e. have infinite loops), but
#               rather it should implement just one iteration (check
#               conditions, compute velocity), because it will be called
#               regularly from the state machine.
#
# Hint: below is an example of a state that moves the robot forward until the
#       front sonar readings are less than the desired clearance. It assumes
#       that the corresponding variables ('front_min' and 'clearance') are
#       available in the userdata dictionary.
#
#           def search(ud):
#               if ud.front_min < ud.clearance:
#                   return 'found_obstacle'
#               ud.velocity = (1, 0, 0)

def find_wall(userdata):
    """Moves the robot forward until it finds a wall."""
    if userdata.front_distance > userdata.clearance:
        userdata.velocity = (0.5, 0, 0)
        return 'front_far'
    else:
        userdata.velocity = (0, 0, 0)
        return 'front_near'

def adjust_to_wall(userdata):
    """Rotates the robot such that it aligns itself with a wall."""
    if userdata.side_distance > userdata.clearance:
        userdata.velocity = (0, 0, -0.5)
        return 'side_sensor_far'
    else:
        userdata.velocity = (0, 0, 0)
        if userdata.front_distance < userdata.clearance:
            return 'side_near_front_near'
        else:
            return 'side_near_front_far'

def go_straight(userdata):
    """Moves the robot straight or returns a message that the robot should turn."""
    if userdata.front_distance < userdata.clearance:
        userdata.velocity = (0, 0, 0)
        return 'front_near'
    elif userdata.side_distance > userdata.clearance:
        if userdata.diag_front > userdata.clearance:
            userdata.velocity = (0.1, 0, 0.5)
            return 'side_far'
        else:
            userdata.velocity = (0.5, 0, 0)
            return 'go_straight'

    userdata.velocity = (0.5, 0, 0)
    return 'go_straight'

def turn_concave(userdata):
    """Turns the robot when it encounters concave obstacles."""
    if userdata.front_distance > (userdata.clearance + 0.1):
        userdata.velocity = (0, 0, 0)
        return 'front_far'
    else:
        userdata.velocity = (0.1, 0, -0.5)
        return 'turn_concave'

def turn_convex(userdata):
    """Turns the robot when it encounters convex obstacles."""
    if userdata.diag_front > userdata.clearance:
        userdata.velocity = (0.1, 0, 0.5)
        return 'turn_convex'
    elif userdata.front_distance > userdata.clearance and userdata.side_distance > userdata.clearance and userdata.back_distance > userdata.clearance and userdata.diag_front > (1.5 *userdata.clearance) and userdata.diag_back > (1.5 * userdata.clearance) and userdata.left_back_distance > userdata.clearance and userdata.right_back_distance > userdata.clearance and userdata.left_front_distance > userdata.clearance and userdata.right_front_distance > userdata.clearance:
        userdata.velocity = (0, 0, 0)
        return 'all_far'
    else:
        userdata.velocity = (0, 0, 0)
        return 'side_near'

#==============================================================================

def set_ranges(self, ranges):
    """
    This function will be attached to the constructed wallfollower machine.
    Its argument is a list of Range messages as received by a sonar callback.
    """
    #============================= YOUR CODE HERE =============================
    # Instructions: store the ranges from a ROS message into the userdata
    #               dictionary of the state machine.
    #               'ranges' is a list or Range messages (that should be
    #               familiar to you by now). It implies that to access the
    #               actual range reading of, say, sonar number 3, you need to
    #               write:
    #
    #                   ranges[3].range
    #
    #               For example, to create an item called 'front_min', which
    #               contains the minimum between the ranges reported by the two
    #               front sonars, you would write the following:
    #
    #                   self.userdata.front_min = min(ranges[3].range, ranges[4].range)
    #
    # Hint: you can just store the whole array of the range readings, but to
    #       simplify the code in your state functions, you may compute
    #       additional values, e.g. the difference between the reading of the
    #       side sonars, or the minimum of all sonar readings, etc.
    #
    # Hint: you can access all the variables stored in userdata. This includes
    #       the current settings of the wallfollower (that is clearance and the
    #       mode of wallfollowing). Think about how you could make your state
    #       functions independent of wallfollowing mode by smart preprocessing
    #       of the sonar readings.

    self.userdata.ranges = ranges
    self.userdata.left_back_distance = ranges[12].range
    self.userdata.right_back_distance = ranges[11].range
    self.userdata.left_front_distance = ranges[3].range
    self.userdata.right_front_distance = ranges[4].range
    if self.userdata.mode == 1:
        self.userdata.back_distance = min(ranges[12].range, ranges[13].range)
        self.userdata.front_distance = min(ranges[4].range, ranges[5].range)
        self.userdata.side_distance = min(ranges[7].range, ranges[8].range)
        self.userdata.diag_front = min(ranges[5].range, ranges[6].range)
        self.userdata.diag_back = min(ranges[9].range, ranges[10].range)
    elif self.userdata.mode == 0:
        self.userdata.back_distance = min(ranges[12].range, ranges[13].range)
        self.userdata.front_distance = min(ranges[2].range, ranges[3].range)
        self.userdata.side_distance = min(ranges[0].range, ranges[15].range)
        self.userdata.diag_front = min(ranges[1].range, ranges[2].range)
        self.userdata.diag_back = min(ranges[13].range, ranges[14].range)

    #==========================================================================

def get_twist(self):
    """
    This function will be attached to the constructed wallfollower machine.
    It creates a Twist message that could be directly published by a velocity
    publisher. The values for the velocity components are fetched from the
    machine userdata.
    """
    twist = Twist()
    twist.linear.x = self.userdata.velocity[0]
    twist.linear.y = self.userdata.velocity[1]
    twist.linear.z = 0
    twist.angular.x = 0
    twist.angular.y = 0
    twist.angular.z = self.userdata.velocity[2]
    #============================= YOUR CODE HERE =============================
    # Instructions: although this function is implemented, you may need to
    #               slightly tweak it if you decided to handle wallfolllowing
    #               mode in "the smart way".
    # Hint: state machine userdata is accessible in this function as well, for
    #       example you can read the current wallfollowing mode with
    #
    #           self.userdata.mode
    #

    if self.userdata.mode == 1:
        twist.angular.z = -self.userdata.velocity[2]

    #==========================================================================
    return twist

def set_config(self, config):
    """
    This function will be attached to the constructed wallfollower machine.
    It updates the relevant fields in the machine userdata.
    Its argument is the config object that comes from ROS dynamic reconfigure
    client.
    """
    self.userdata.mode = config['mode']
    self.userdata.clearance = config['clearance']
    return config

def construct():
    sm = smach.StateMachine(outcomes=['preempted'])
    # Attach helper functions
    sm.set_ranges = MethodType(set_ranges, sm, sm.__class__)
    sm.get_twist = MethodType(get_twist, sm, sm.__class__)
    sm.set_config = MethodType(set_config, sm, sm.__class__)
    # Set initial values in userdata
    sm.userdata.velocity = (0, 0, 0)
    sm.userdata.mode = 1
    sm.userdata.clearance = 0.6
    sm.userdata.ranges = None
    # Add states
    with sm:
        #=========================== YOUR CODE HERE ===========================
        # Instructions: construct the state machine by adding the states that
        #               you have implemented.
        #               Below is an example how to add a state:
        #
        #                   smach.StateMachine.add('SEARCH',
        #                                          PreemptableState(search,
        #                                                           input_keys=['front_min', 'clearance'],
        #                                                           output_keys=['velocity'],
        #                                                           outcomes=['found_obstacle']),
        #                                          transitions={'found_obstacle': 'ANOTHER_STATE'})
        #
        #               First argument is the state label, an arbitrary string
        #               (by convention should be uppercase). Second argument is
        #               an object that implements the state. In our case an
        #               instance of the helper class PreemptableState is
        #               created, and the state function in passed. Moreover,
        #               we have to specify which keys in the userdata the
        #               function will need to access for reading (input_keys)
        #               and for writing (output_keys), and the list of possible
        #               outcomes of the state. Finally, the transitions are
        #               specified. Normally you would have one transition per
        #               state outcome.
        #
        # Note: The first state that you add will become the initial state of
        #       the state machine.

        smach.StateMachine.add('FIND_WALL', PreemptableState(find_wall, input_keys=['clearance', 'front_distance'], output_keys=['velocity'], outcomes=['front_far', 'front_near']), transitions= {'front_far': 'FIND_WALL', 'front_near': 'ADJUST_TO_WALL'})

        smach.StateMachine.add('ADJUST_TO_WALL', PreemptableState(adjust_to_wall, input_keys=['clearance', 'front_distance', 'side_distance', 'back_distance', 'diag_front', 'diag_back'], output_keys=['velocity'], outcomes=['side_sensor_far', 'side_near_front_far', 'side_near_front_near', 'back_near_side_far']), transitions= {'side_sensor_far': 'ADJUST_TO_WALL', 'side_near_front_far': 'GO_STRAIGHT', 'side_near_front_near': 'TURN_CONCAVE', 'back_near_side_far': 'TURN_CONVEX'})

        smach.StateMachine.add('GO_STRAIGHT', PreemptableState(go_straight, input_keys=['clearance', 'front_distance', 'side_distance', 'diag_front'], output_keys=['velocity'], outcomes=['front_near', 'side_far', 'all_far', 'go_straight']), transitions= {'go_straight': 'GO_STRAIGHT', 'front_near': 'TURN_CONCAVE', 'side_far': 'TURN_CONVEX', 'all_far': 'FIND_WALL'})

        smach.StateMachine.add('TURN_CONCAVE', PreemptableState(turn_concave, input_keys=['clearance', 'front_distance', 'left_back_distance', 'right_back_distance', 'left_front_distance', 'right_front_distance', 'diag_back'], output_keys=['velocity'], outcomes=['turn_concave', 'front_far', 'far_from_wall']), transitions= {'turn_concave': 'TURN_CONCAVE', 'front_far': 'GO_STRAIGHT', 'far_from_wall': 'ADJUST_TO_WALL'})

        smach.StateMachine.add('TURN_CONVEX', PreemptableState(turn_convex, input_keys=['clearance', 'front_distance', 'back_distance', 'side_distance', 'diag_front', 'diag_back', 'left_back_distance', 'right_back_distance', 'left_front_distance', 'right_front_distance'], output_keys=['velocity'], outcomes=['turn_convex', 'side_near', 'all_far']), transitions= {'turn_convex': 'TURN_CONVEX', 'side_near': 'GO_STRAIGHT', 'all_far': 'FIND_WALL'})

        #======================================================================
    return sm


from pathlib import Path

from coxbuild import get_working_directory
from coxbuild.schema import depend, group, precond, run, task

from .. import ProjectSettings, withProject
from . import grouped
from .package import hasPackages, upgradePackages

subgrouped = group("test")

@grouped
@subgrouped
@precond(lambda: not hasPackages({"pytest": "*",  "pytest-asyncio": "*", "pytest-cov": "*", "coverage": "*"}))
@task
def restore():
    """Restore Python test tools."""
    upgradePackages("pytest", "pytest-asyncio", "coverage", "pytest-cov")

@grouped
@subgrouped
@withProject
@depend(restore)
@task
def pytest(src: Path | None = None, test: Path | None = None, *, project: ProjectSettings):
    """Use pytest to test Python code."""
    run(["pytest", "--cov-report=term-missing",
        "--cov-report=html", f"--cov={str(src or project.src)}"], cwd=str(test or project.test))

@grouped
@subgrouped
@depend(pytest)
@task
def test():
    """Test Python code."""
    pass