# Noysim -- Noise simulation tools for Aimsun.
# Copyright (c) 2010-2011 by Bert De Coensel, Ghent University & Griffith University.
#
# Noise emission model functions and classes

import warnings
import random

import numpy
import pylab

from geo import Point, Direction, asPoint, asDirection
from acoustics import LOWDB, fromdB, TertsBandSpectrum
import numeric

#---------------------------------------------------------------------------------------------------
# Basic classes
#---------------------------------------------------------------------------------------------------

class Vehicle(object):
  """ base class defining a vehicle, containing all properties that could have an influence on noise emissions """
  def __init__(self, vid, length, width, height, weight, cat, axles, doublemount, studs, fuel, cc, maxspeed, maxdecel, maxaccel,
                     position = Point(0.0,0.0,0.0), direction = Direction(0.0,0.0), speed = 0.0, acceleration = 0.0):
    object.__init__(self)
    # unique vehicle ID
    self._vid          = vid
    # size properties
    self._length       = length       # in m
    self._width        = width        # in m
    self._height       = height       # in m
    self._weight       = weight       # in kg
    # emission related properties
    self._cat          = cat          # emission category (usually an integer), to be interpreted by the emission models
    self._axles        = axles        # integer, number of axles of vehicle
    self._doublemount  = doublemount  # boolean, true for double-mounted tires (for trucks)
    self._studs        = studs        # boolean, true if the vehicle has winter tires
    self._fuel         = fuel         # the type of fuel the vehicle uses (string, e.g. 'petrol', 'diesel', 'electric')
    self._cc           = cc           # the cilinder capacity (cc)
    # kinematic properties
    self._maxspeed     = maxspeed     # in km/h
    self._maxdecel     = maxdecel     # in m/s^2 (negative value)
    self._maxaccel     = maxaccel     # in m/s^2
    # dynamic properties
    self._position     = position     # position Point(x,y,z) - center of bottom plane - in m
    self._direction    = direction    # Direction(bearing,gradient) of the vehicle, in degrees (0-360)
    self._speed        = speed        # velocity (km/h)
    self._acceleration = acceleration # acceleration (m/s^2)

  def vid(self):
    return self._vid

  def length(self):
    return self._length
  def width(self):
    return self._width
  def height(self):
    return self._height
  def weight(self):
    return self._weight

  def cat(self):
    return self._cat
  def axles(self):
    return self._axles
  def doublemount(self):
    return self._doublemount
  def studs(self):
    return self._studs
  def fuel(self):
    return self._fuel
  def cc(self):
    return self._cc

  def maxspeed(self):
    return self._maxspeed
  def maxdecel(self):
    return self._maxdecel
  def maxaccel(self):
    return self._maxaccel

  def position(self):
    return self._position
  def direction(self):
    return self._direction
  def speed(self):
    return self._speed
  def acceleration(self):
    return self._acceleration

  def __str__(self):
    """ return a string representation of the vehicle """
    s = '[ID %s: %.1fx%.1fx%.1fm, %dkg, ' % (str(self.vid()), self.length(), self.width(), self.height(), int(self.weight()))
    s += 'c=%s, ax=%d, dm=%s, ' % (str(self.cat()), self.axles(), str(self.doublemount()))
    s += 'stud=%s, %s, %dcc,' % (str(self.studs()), self.fuel(), int(self.cc()))
    s += '\n max=(%.1f,%.1f,%.1f), ' % (self.maxspeed(), self.maxdecel(), self.maxaccel())
    s += 'p=%s, d=%s, v=%.1f, a=%.1f]' % (str(self.position()), str(self.direction()), self.speed(), self.acceleration())
    return s

  def copy(self, cls = None):
    """ create a copy of the vehicle """
    if cls == None:
      cls = Vehicle
    return cls(vid = self._vid, length = self._length, width = self._width, height = self._height, weight = self._weight,
               cat = self._cat, axles = self._axles, doublemount = self._doublemount, studs = self._studs, fuel = self._fuel,
               cc = self._cc, maxspeed = self._maxspeed, maxdecel = self._maxdecel, maxaccel = self._maxaccel,
               position = self._position.copy(), direction = self._direction.copy(), speed = self._speed, acceleration = self._acceleration)

  def move(self, dx = 0.0, dy = 0.0, dz = 0.0):
    """ move the location of the vehicle """
    self._position.x += dx
    self._position.y += dy
    self._position.z += dz

class Roadsurface(object):
  """ base class defining a road surface, containing all properties that could have an influence on noise emissions """
  def __init__(self, cat, temperature, chipsize, age, wet, tc):
    object.__init__(self)
    self._cat         = cat         # surface category, to be interpreted by the various emission models
    self._temperature = temperature # temperature of air above road surface (degrees celsius)
    self._chipsize    = chipsize    # road surface chip size (in mm)
    self._age         = age         # road surface age (in years)
    self._wet         = wet         # boolean, true if the surface is wet
    self._tc          = tc          # road surface temperature coefficient (in dB/degree)

  def cat(self):
    return self._cat
  def temperature(self):
    return self._temperature
  def chipsize(self):
    return self._chipsize
  def age(self):
    return self._age
  def wet(self):
    return self._wet
  def tc(self):
    return self._tc

  def __str__(self):
    """ return a string representation of the road surface """
    s = '[c=%s, t=%.1fC, chip=%dmm, ' % (str(self.cat()), self.temperature(), int(self.chipsize()))
    s += 'age=%dy, wet=%s, tc=%.2f]' % (int(self.age()), str(self.wet()), self.tc())
    return s

class Source(object):
  """ base source class, which serves as the output of an emission model """
  def __init__(self, position, direction, emission, directivity = lambda theta, phi, f: 0.0):
    object.__init__(self)
    self.position = position # Point(x,y,z) in m
    self.direction = direction # Direction(bearing,gradient) of the source, in degrees (0-360)
    self.emission = emission # source emission spectrum
    self.directivity = directivity # function (or function object) for the directivity pattern of the source
                                   # (correction on source power level as a function of theta, phi and frequency)

  def __str__(self):
    """ return a string representation of the source """
    # only prints the A-weighted level
    return '[p=%s, d=%s, e=%.2f dBA]' % (str(self.position), str(self.direction), self.emission.laeq())

  def copy(self):
    """ return a copy of the source """
    return Source(position = self.position.copy(),
                  direction = self.direction.copy(),
                  emission = self.emission.copy(),
                  directivity = self.directivity)

class EmissionModel(object):
  """ base emission model interface - to be filled in by derived classes """
  def __init__(self):
    object.__init__(self)

  def __str__(self):
    """ return a string representation of the emission model """
    return '[EmissionModel]'

  def categoryNames(self):
    """ return a list with the names of the vehicle categories that the model supports """
    raise NotImplementedError

  def sources(self, vehicle, road):
    """ calculate the list of emission sources associated with the given vehicle and road surface """
    raise NotImplementedError

  def emission(self, vehicle, road):
    """ shorthand function, calculating the total emission associated with the given vehicle and road surface """
    raise NotImplementedError

class Viewport(object):
  """ viewport function object, acting as a vehicle filter
      should return True for each vehicle that has to be taken into account in subsequent processing
  """
  def __init__(self):
    object.__init__(self)

  def __str__(self):
    """ return a string representation of the viewport """
    return '[Viewport]'

  def __call__(self, vehicle):
    """ should return True only if the vehicle has to be taken into account """
    return True # the default is to consider all vehicles

#---------------------------------------------------------------------------------------------------
# Specialized vehicle and road surface classes
#---------------------------------------------------------------------------------------------------

# The constants in these predefined classes were taken from the Queensland vehicles defined in the Griffith Univ. Aimsun file,
# with missing values (e.g. weight or cc) filled in by best guess (these additional values do not have an influence on noise anyhow)

class QLDCar(Vehicle):
  """ Queensland car (CAR_QLD) with default values for all parameters (Imagine category 1) """
  def __init__(self, vid = 0, length = 4.0, width = 2.0, height = 1.5, weight = 800.0,
                     cat = 1, axles = 2, doublemount = False, studs = False, fuel = 'petrol', cc = 1600,
                     maxspeed = 110.0, maxdecel = -6.5, maxaccel = 2.8,
                     position = Point(0.0,0.0,0.0), direction = Direction(0.0,0.0), speed = 0.0, acceleration = 0.0):
    Vehicle.__init__(self, vid, length, width, height, weight, cat, axles, doublemount, studs, fuel, cc,
                     maxspeed, maxdecel, maxaccel, position, direction, speed, acceleration)
  def copy(self):
    return Vehicle.copy(self, cls = QLDCar)

class QLDVan(Vehicle):
  """ Queensland van (HV) with default values for all parameters (Imagine category 1) """
  def __init__(self, vid = 0, length = 7.0, width = 2.0, height = 2.0, weight = 2000.0,
                     cat = 1, axles = 2, doublemount = False, studs = False, fuel = 'petrol', cc = 2500,
                     maxspeed = 110.0, maxdecel = -5.0, maxaccel = 2.5,
                     position = Point(0.0,0.0,0.0), direction = Direction(0.0,0.0), speed = 0.0, acceleration = 0.0):
    Vehicle.__init__(self, vid, length, width, height, weight, cat, axles, doublemount, studs, fuel, cc,
                     maxspeed, maxdecel, maxaccel, position, direction, speed, acceleration)
  def copy(self):
    return Vehicle.copy(self, cls = QLDVan)

class QLDLightTruck(Vehicle):
  """ Queensland light truck (LT) with default values for all parameters (Imagine category 2) """
  def __init__(self, vid = 0, length = 12.0, width = 2.5, height = 2.5, weight = 5000.0,
                     cat = 2, axles = 2, doublemount = True, studs = False, fuel = 'petrol', cc = 4000,
                     maxspeed = 110.0, maxdecel = -3.0, maxaccel = 1.5,
                     position = Point(0.0,0.0,0.0), direction = Direction(0.0,0.0), speed = 0.0, acceleration = 0.0):
    Vehicle.__init__(self, vid, length, width, height, weight, cat, axles, doublemount, studs, fuel, cc,
                     maxspeed, maxdecel, maxaccel, position, direction, speed, acceleration)
  def copy(self):
    return Vehicle.copy(self, cls = QLDLightTruck)

class QLDSemiTrailer(Vehicle):
  """ Queensland semi-trailer (ST) with default values for all parameters (Imagine category 3) """
  def __init__(self, vid = 0, length = 19.0, width = 2.5, height = 3.0, weight = 10000.0,
                     cat = 3, axles = 4, doublemount = True, studs = False, fuel = 'petrol', cc = 5000,
                     maxspeed = 110.0, maxdecel = -2.9, maxaccel = 1.0,
                     position = Point(0.0,0.0,0.0), direction = Direction(0.0,0.0), speed = 0.0, acceleration = 0.0):
    Vehicle.__init__(self, vid, length, width, height, weight, cat, axles, doublemount, studs, fuel, cc,
                     maxspeed, maxdecel, maxaccel, position, direction, speed, acceleration)
  def copy(self):
    return Vehicle.copy(self, cls = QLDSemiTrailer)

class QLDBDouble(Vehicle):
  """ Queensland B-Double heavy vehicle (BD) with default values for all parameters (Imagine category 3) """
  def __init__(self, vid = 0, length = 25.0, width = 2.5, height = 3.5, weight = 30000.0,
                     cat = 3, axles = 4, doublemount = True, studs = False, fuel = 'petrol', cc = 8000,
                     maxspeed = 110.0, maxdecel = -2.75, maxaccel = 0.8,
                     position = Point(0.0,0.0,0.0), direction = Direction(0.0,0.0), speed = 0.0, acceleration = 0.0):
    Vehicle.__init__(self, vid, length, width, height, weight, cat, axles, doublemount, studs, fuel, cc,
                     maxspeed, maxdecel, maxaccel, position, direction, speed, acceleration)
  def copy(self):
    return Vehicle.copy(self, cls = QLDBDouble)

class QLDMotorcycle(Vehicle):
  """ Queensland motorcycle with default values for all parameters (Imagine category 4b - here 5) """
  def __init__(self, vid = 0, length = 1.5, width = 0.5, height = 1.5, weight = 200.0,
                     cat = 5, axles = 2, doublemount = False, studs = False, fuel = 'petrol', cc = 1000,
                     maxspeed = 110.0, maxdecel = -6.0, maxaccel = 4.0,
                     position = Point(0.0,0.0,0.0), direction = Direction(0.0,0.0), speed = 0.0, acceleration = 0.0):
    Vehicle.__init__(self, vid, length, width, height, weight, cat, axles, doublemount, studs, fuel, cc,
                     maxspeed, maxdecel, maxaccel, position, direction, speed, acceleration)
  def copy(self):
    return Vehicle.copy(self, cls = QLDMotorcycle)

class ReferenceRoadsurface(Roadsurface):
  """ Harmonoise/Imagine reference road surface """
  def __init__(self, cat = 'REF', temperature = 20.0, chipsize = 11.0, age = 2.0, wet = False, tc = 0.08):
    Roadsurface.__init__(self, cat, temperature, chipsize, age, wet, tc)

#---------------------------------------------------------------------------------------------------
# Specialized viewport classes
#---------------------------------------------------------------------------------------------------

class RectangularViewport(Viewport):
  """ viewport that takes into account the vehicles that are situated in a rectangular area """
  def __init__(self, minx, miny, maxx, maxy):
    Viewport.__init__(self)
    (self.minx, self.miny, self.maxx, self.maxy) = (minx, miny, maxx, maxy)

  def __str__(self):
    """ return a string representation of the viewport """
    return '[RectangularViewport (%s, %s, %s, %s)]' % (self.minx, self.miny, self.maxx, self.maxy)

  def __call__(self, vehicle):
    """ check if the vehicle is situated inside the rectangular area """
    p = vehicle.position()
    if (self.minx <= p.x) and (p.x <= self.maxx) and (self.miny <= p.y) and (p.y <= self.maxy):
      return True
    return False

# TODO:
# - define the class below
# - uncomment the line in Configuration.viewport so it is used
# class DynamicViewport(RectangularViewport):
#   """ special case of a rectangular viewport, for which the dimensions are automatically tailored
#       to give optimal tradeoff between speed and accuracy, based on various network and receiver properties
#   """
#   def __init__(self):
#     # perform calculations for minx, miny, maxx, maxy
#     # finally create the viewport
#     RectangularViewport.__init__(minx, miny, maxx, maxy)

#---------------------------------------------------------------------------------------------------
# Imagine road traffic noise emission model
#---------------------------------------------------------------------------------------------------

# model constants (LOWDB is used to fill in the values that are not supplied by Imagine, i.e. at 20 Hz and at 12kHz and above)
IMAGINE = {# Light motor vehicles
           1: {'AR': numpy.asarray([LOWDB, 69.9, 69.9, 69.9, 74.9, 74.9, 74.9, 79.3, 82.0, 81.2, 80.9, 78.9, 78.8, 80.5, 85.0, 87.9,
                                     90.9, 93.3, 92.8, 91.5, 88.5, 84.9, 81.8, 78.7, 74.9, 71.8, 69.1, 65.6,LOWDB,LOWDB,LOWDB]),
               'BR': numpy.asarray([  0.0, 33.0, 33.0, 33.0, 30.0, 30.0, 30.0, 41.0, 41.2, 42.3, 41.8, 38.6, 35.5, 32.9, 25.0, 25.0,
                                     27.0, 33.4, 36.7, 37.0, 37.5, 37.5, 38.6, 39.6, 40.0, 39.9, 40.2, 40.3,  0.0,  0.0,  0.0]),
               'AP': numpy.asarray([LOWDB, 87.0, 87.0, 87.0, 87.9, 90.8, 89.9, 86.9, 82.6, 81.9, 82.3, 83.9, 83.3, 82.4, 80.6, 80.2,
                                     77.8, 78.0, 81.4, 82.3, 82.6, 81.5, 80.2, 78.5, 75.6, 73.3, 71.0, 68.1,LOWDB,LOWDB,LOWDB]),
               'BP': numpy.asarray([  0.0,  0.0,  0.0,  0.0,  0.0, -3.0,  0.0,  8.0,  6.0,  6.0,  7.0,  8.0,  8.0,  8.0,  8.0,  8.0,
                                      8.0,  8.0,  8.0,  8.0,  8.0,  8.0,  8.0,  8.0,  8.0,  8.0,  8.0,  8.0,  0.0,  0.0,  0.0]),
               'CP': numpy.asarray([  0.0,  4.0,  4.0,  4.0,  7.0,  7.0,  7.0,  7.0,  7.0,  7.0,  7.0,  4.0,  4.0,  4.0,  4.0,  4.0,
                                      4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  0.0,  0.0,  0.0])},
           # Medium heavy vehicles:
           2: {'AR': numpy.asarray([LOWDB, 76.5, 76.5, 76.5, 78.5, 79.5, 79.5, 82.5, 84.3, 84.7, 84.3, 87.4, 87.8, 89.8, 91.6, 93.5,
                                     94.6, 92.4, 89.6, 88.1, 85.9, 82.7, 80.7, 78.8, 76.8, 76.7, 75.7, 74.5,LOWDB,LOWDB,LOWDB]),
               'BR': numpy.asarray([  0.0, 33.0, 33.0, 33.0, 30.0, 30.0, 30.0, 32.9, 35.9, 38.1, 36.5, 33.5, 30.6, 27.7, 21.9, 23.8,
                                     28.4, 31.1, 35.4, 35.9, 36.7, 36.3, 37.7, 38.5, 39.8, 39.9, 40.2, 40.3,  0.0,  0.0,  0.0]),
               'AP': numpy.asarray([LOWDB, 93.9, 93.9, 94.1, 95.0, 97.3, 96.1, 92.5, 91.9, 90.4, 93.4, 94.4, 94.2, 93.0, 90.8, 92.1,
                                     92.5, 94.1, 94.5, 92.4, 90.1, 87.6, 85.8, 83.8, 81.4, 80.0, 77.2, 75.4,LOWDB,LOWDB,LOWDB]),
               'BP': numpy.asarray([  0.0,  0.0,  0.0,  0.0,  0.0, -4.0,  0.0,  4.0,  5.0,  5.5,  6.0,  6.5,  6.5,  6.5,  6.5,  6.5,
                                      6.5,  6.5,  6.5,  6.5,  6.5,  6.5,  6.5,  6.5,  6.5,  6.5,  6.5,  6.5,  0.0,  0.0,  0.0]),
               'CP': numpy.asarray([  0.0,  5.0,  5.0,  5.0,  9.0,  9.0,  9.0,  9.0,  9.0,  9.0,  9.0,  5.0,  5.0,  5.0,  5.0,  5.0,
                                      5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  0.0,  0.0,  0.0])},
           # Heavy vehicles
           3: {'AR': numpy.asarray([LOWDB, 79.5, 79.5, 79.5, 81.5, 82.5, 82.5, 85.5, 87.3, 87.7, 87.3, 89.5, 90.5, 93.8, 95.9, 97.3,
                                     98.0, 95.6, 93.2, 91.9, 88.9, 85.5, 84.1, 82.2, 79.8, 78.6, 77.5, 76.8,LOWDB,LOWDB,LOWDB]),
               'BR': numpy.asarray([  0.0, 33.0, 33.0, 33.0, 30.0, 30.0, 30.0, 31.4, 32.8, 36.0, 34.6, 32.7, 29.3, 26.4, 24.2, 25.9,
                                     30.4, 32.3, 36.5, 36.8, 38.0, 36.8, 38.5, 38.9, 38.5, 40.2, 40.8, 41.0,  0.0,  0.0,  0.0]),
               'AP': numpy.asarray([LOWDB, 95.7, 94.9, 94.1, 96.8,101.8, 98.6, 95.5, 96.2, 95.7, 97.2, 96.3, 97.2, 95.8, 95.9, 96.8,
                                     95.1, 95.8, 95.0, 92.7, 91.2, 88.7, 87.6, 87.2, 84.2, 82.7, 79.7, 77.6,LOWDB,LOWDB,LOWDB]),
               'BP': numpy.asarray([  0.0,  0.0,  0.0,  0.0, -4.0,  0.0,  4.0,  3.0,  3.0,  3.0,  4.0,  5.0,  5.0,  5.0,  5.0,  5.0,
                                      5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  0.0,  0.0,  0.0]),
               'CP': numpy.asarray([  0.0,  5.0,  5.0,  5.0,  9.0,  9.0,  9.0,  9.0,  9.0,  9.0,  9.0,  5.0,  5.0,  5.0,  5.0,  5.0,
                                      5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  0.0,  0.0,  0.0])},
           # Powered two-wheelers 4a (mopeds < 50cc)
           4: {'AR': LOWDB*numpy.ones(31),
               'BR': numpy.zeros(31),
               'AP': numpy.asarray([LOWDB, 88.7, 87.6, 85.5, 85.8, 81.5, 80.7, 82.0, 85.6, 81.6, 81.4, 85.5, 86.3, 87.9, 88.7, 89.9,
                                     91.8, 91.2, 92.4, 95.0, 94.1, 92.9, 90.4, 89.1, 87.4, 84.9, 84.4, 82.2,LOWDB,LOWDB,LOWDB]),
               'BP': numpy.asarray([  0.0, -2.2, -0.1,  1.7,  5.9,  1.9,  3.3,  0.9, 17.3, 14.5,  5.0, 14.6,  9.9,  9.7, 12.7, 12.3,
                                     13.9, 16.6, 17.2, 17.9, 19.3, 20.6, 19.9, 20.8, 20.5, 21.0, 21.0, 19.3,  0.0,  0.0,  0.0]),
               'CP': numpy.asarray([  0.0,  4.0,  4.0,  4.0,  7.0,  7.0,  7.0,  7.0,  7.0,  7.0,  7.0,  4.0,  4.0,  4.0,  4.0,  4.0,
                                      4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  0.0,  0.0,  0.0])},
           # Powered two-wheelers 4b (motorcycles > 50cc)
           5: {'AR': LOWDB*numpy.ones(31),
               'BR': numpy.zeros(31),
               'AP': numpy.asarray([LOWDB, 90.8, 88.9, 89.2, 90.5, 89.2, 90.7, 93.2, 93.2, 90.0, 88.4, 87.6, 87.7, 87.0, 87.4, 89.4,
                                     89.9, 90.1, 89.7, 89.8, 88.2, 86.5, 85.8, 85.1, 85.1, 82.7, 81.7, 80.4,LOWDB,LOWDB,LOWDB]),
               'BP': numpy.asarray([  0.0,  2.1,  3.1,  1.2,  2.3,  2.8,  4.2,  6.2,  4.8,  7.3, 11.3, 10.6, 13.9, 13.5, 11.0, 10.8,
                                     11.4, 11.4, 11.7, 13.4, 11.6, 12.2, 10.9, 10.5, 12.0, 12.0, 12.0, 12.0,  0.0,  0.0,  0.0]),
               'CP': numpy.asarray([  0.0,  4.0,  4.0,  4.0,  7.0,  7.0,  7.0,  7.0,  7.0,  7.0,  7.0,  4.0,  4.0,  4.0,  4.0,  4.0,
                                      4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  4.0,  0.0,  0.0,  0.0])}}

# correction values for studded tires
STUDS = {'A':  numpy.asarray([ 0.0,  0.0,  0.0,  0.0,  0.0,  0.0,  0.0,  0.0,  0.3,  1.4,  1.5,  0.9,  1.2,  1.5,  1.9,  1.8,
                               0.8,  0.5,  0.2, -0.2, -0.4,  0.5,  0.8,  0.9,  2.1,  5.0,  7.3, 10.0,  0.0,  0.0,  0.0]),
         'B':  numpy.asarray([ 0.0,  0.0,  0.0,  0.0,  0.0,  0.0,  0.0,  0.0, -4.1, -6.0, -8.5, -4.1,  1.7,  0.6, -4.6, -3.9,
                              -2.7, -4.2,-11.7,-11.7,-14.9,-17.6,-21.8,-21.6,-19.2,-14.6, -9.9,-10.2,  0.0,  0.0,  0.0])}

class ImagineDirectivity(object):
  """ function object implementing the Harmonoise/Imagine directivity pattern
      once an object of this class is created, it behaves like a function (thanks to the '__call__' method)
  """
  def __init__(self, cat, h):
    """ create the directivity function, for the given vehicle category and source height """
    object.__init__(self)
    self.cat = cat # vehicle category
    self.h = h # source height

  def __call__(self, theta, phi, f):
    """ return the directivity correction for the given angles, at the given frequency """
    # theta: horizontal angle between travelling direction of vehicle and reciever (0 degrees means receiver in front of vehicle)
    # phi: vertical angle between travelling direction of vehicle and receiver (0 degrees means receiver in front of vehicle)
    # NOTE: for readability reasons, other symbols are used than in the Imagine document
    # restrict the angles to the correct range
    while theta < 0.0:
      theta += 360.0
    while theta > 360.0:
      theta -= 360.0
    while phi < -180.0:
      phi += 360.0
    while phi > 180.0:
      phi -= 360.0
    # enforce horizontal and vertical symmetry
    if theta > 180.0:
      theta = 360.0 - theta
    if phi < 0.0:
      phi = -phi
    if phi > 90.0:
      phi = 180.0 - phi
    # calculate shorthands
    thetaRad = (numpy.pi/180.0)*theta
    phiRad = (numpy.pi/180.0)*phi
    pi2theta = numpy.pi/2.0 - thetaRad
    sinpi2theta = numpy.sin(pi2theta)
    sqrtcosphi = numpy.sqrt(numpy.cos(phiRad))
    # calculate horizontal directivity (formulas from Harmonoise model, depending on source height, independent of vehicle category)
    if self.h == 0.01:
      # horn effect, only for certain frequency range
      if (1600.0 <= f) and (f <= 6300.0):
        horizontal = (-1.5 + 2.5*abs(sinpi2theta)) * sqrtcosphi
      else:
        horizontal = 0.0
    elif self.h == 0.30:
      horizontal = 0.0
    elif self.h == 0.75:
      # screening by the body of the heavy vehicle
      horizontal = (1.546*(pi2theta**3) - 1.425*(pi2theta**2) + 0.22*pi2theta + 0.6) * sqrtcosphi
    else:
      raise Exception('no directivity defined for sources at height %.2f' % self.h)
    # calculate vertical directivity (approximations from Imagine model, depending on vehicle category, independent of source height)
    if self.cat == 1:
      vertical = -abs(phi/20.0) # phi in degrees
    elif self.cat in (2, 3):
      vertical = -abs(phi/30.0)
    else:
      vertical = 0.0 # no corrections for categories 4 and 5
    return horizontal + vertical

class ImagineModel(EmissionModel):
  """ Imagine road traffic noise emission model """
  def __init__(self):
    EmissionModel.__init__(self)
    # parameters
    self.vref = 70.0
    self.vinterval = (0.1, 160.0)
    # flags for applying different types of corrections - to be adjusted after object creation if necessary
    self.correction = {'acceleration': True, 'gradient': False, 'temperature': False,
                       'surface': False, 'wetness': False, 'axles': True, 'fleet': False}
    self.fleet = {'diesel': 19.0, 'tirewidth': 187.0, 'vans': 10.5, 'iress': (1.0, 35.0), 'studs': False}

  def __str__(self):
    """ return a string representation of the emission model """
    return '[Imagine emission model]'

  def categoryNames(self):
    return ['Light', 'Medium', 'Heavy', 'Moped', 'Motorcycle']

  def rollingNoise(self, vehicle, road):
    # calculate base rolling noise
    v = numpy.clip(vehicle.speed(), *self.vinterval)
    z = IMAGINE[vehicle.cat()]['AR'] + IMAGINE[vehicle.cat()]['BR']*numpy.log10(v/self.vref)
    # apply corrections
    if self.correction['temperature'] == True:
      z += self.temperatureCorrection(vehicle, road)
    if self.correction['surface'] == True:
      z += self.surfaceCorrection(vehicle, road)
    if self.correction['wetness'] == True:
      z += self.wetnessCorrection(vehicle, road)
    if self.correction['axles'] == True:
      z += self.axlesCorrection(vehicle, road)
    if self.correction['fleet'] == True:
      z += self.fleetCorrectionRolling(vehicle, road)
    return z

  def propulsionNoise(self, vehicle, road):
    # calculate base propulsion noise
    v = numpy.clip(vehicle.speed(), *self.vinterval)
    z = IMAGINE[vehicle.cat()]['AP'] + IMAGINE[vehicle.cat()]['BP']*((v - self.vref)/self.vref)
    # apply corrections
    if self.correction['acceleration'] == True:
      z += self.accelerationCorrection(vehicle, road)
    if self.correction['gradient'] == True:
      z += self.gradientCorrection(vehicle, road)
    if self.correction['fleet'] == True:
      z += self.fleetCorrectionPropulsion(vehicle, road)
    return z

  def accelerationCorrection(self, vehicle, road):
    # restrict acceleration
    maxaccel = {1: 2.0, 2: 1.0, 3: 1.0, 4: 4.0, 5: 4.0}[vehicle.cat()]
    a = numpy.clip(vehicle.acceleration(), -1.0, maxaccel)
    return IMAGINE[vehicle.cat()]['CP'] * a

  def gradientCorrection(self, vehicle, road):
    # calculate gradient as a percentage
    alpha = 100.0*numpy.tan(asDirection(vehicle.direction()).gradientRadians())
    cp = IMAGINE[vehicle.cat()]['CP']
    g = 9.81
    if vehicle.cat() in [1, 4, 5]:
      if alpha >= -2.0:
        return cp * g * (alpha/100.0)
      if (-8.0 < alpha) and (alpha < -2.0):
        return cp * g * (-2.0/100.0)
      if alpha <= -8.0:
        return -cp * g * ((alpha + 10.0)/100.0)
    if vehicle.cat() in [2, 3]:
      if alpha >= -2.0:
        return cp * g * (alpha/100.0)
      if alpha < -2.0:
        return -cp * g * ((alpha + 4.0)/100.0)
    raise Exception('Imagine model: condition for gradient correction not found')

  def temperatureCorrection(self, vehicle, road):
    factor = {1: 1.0, 2: 0.5, 3: 0.5, 4: 1.0, 5: 1.0}[vehicle.cat()] # half temperature coefficient for cats 2 and 3
    return factor * road.tc() * (20.0 - road.temperature())

  def surfaceCorrection(self, vehicle, road):
    """ corrections for road surfaces belonging to the reference cluster (DAC and SMA) """
    # check for vehicle type and surface type
    if vehicle.cat() in [2, 3]:
      return 0.0 # no correction for heavy vehicles
    if not (road.cat() in ['REF', 'DAC', 'SMA']):
      return 0.0 # surface not in reference cluster
    chipsize = numpy.clip(road.chipsize(), 8.0, 16.0)
    # surface chipsize and type correction
    z = 0.25*(chipsize - 11.0)
    z += {'REF': 0.0, 'DAC': -0.3, 'SMA': 0.3}[road.cat()]
    return z

  def wetnessCorrection(self, vehicle, road):
    # check for vehicle type and wetness
    if not (road.wet() and (vehicle.cat() == 1)):
      return 0.0
    # retrieve applicable model constants
    v = numpy.clip(vehicle.speed(), *self.vinterval)
    return max(0.0, 15.0*numpy.log10(numpy.asarray(FTERTS)) - 12.0*numpy.log10(v/self.vref) - 48.0)

  def axlesCorrection(self, vehicle, road):
    if not (vehicle.cat() == 3):
      return 0.0 # correction only for heavy trucks
    if not vehicle.doublemount():
      return 6.8*numpy.log10(vehicle.axles()/4.0)
    else:
      return 9.1*numpy.log10(vehicle.axles()/4.0) + 0.8

  def fleetCorrectionRolling(self, vehicle, road):
    z = 0.0
    # correction for tire width
    if vehicle.cat() == 1:
      z += 0.04 * (self.fleet['tirewidth'] - 187.0)
    # correction for percentage of delivery vans
    if vehicle.cat() == 1:
      z += 1.0 * (self.fleet['vans'] - 10.5)/100.0
    # correction for winter/studded tires
    if (self.fleet['studs'] == True) and (vehicle.cat() == 1):
      v = numpy.clip(vehicle.speed(), 50.0, 90.0)
      z += STUDS['A'] + STUDS['B']*numpy.log10(v/self.vref)
    return z

  def fleetCorrectionPropulsion(self, vehicle, road):
    z = 0.0
    # correction for percentage of diesel vehicles
    if vehicle.cat() == 1:
      z += 3.0 * (self.fleet['diesel'] - 19.0)/100.0
    # correction for percentage of delivery vans
    if vehicle.cat() == 1:
      z += 5.0 * (self.fleet['vans'] - 10.5)/100.0
    # correction for illegal replacement exhaust silencer systems (IRESS)
    if vehicle.cat() in [1, 2, 3]:
      piress = (self.fleet['iress'][0] - 1.0)/100.0
    else:
      piress = (self.fleet['iress'][1] - 35.0)/100.0
    z += (29.0 * piress) - (24.0 * piress**2)
    return z

  def sources(self, vehicle, road = ReferenceRoadsurface()):
    # calculate position of sources
    pos = vehicle.position()
    h = {1: (0.01, 0.30), 2: (0.01, 0.75), 3: (0.01, 0.75), 4: (0.30, 0.30), 5: (0.30, 0.30)}[vehicle.cat()]
    lopos = Point(pos[0], pos[1], pos[2] + h[0])
    hipos = Point(pos[0], pos[1], pos[2] + h[1])
    # calculate rolling and propulsion noise
    r = fromdB(self.rollingNoise(vehicle, road))
    p = fromdB(self.propulsionNoise(vehicle, road))
    # calculate low and high noise
    lonoise = 10.0*numpy.log10(0.8*r + 0.2*p)
    hinoise = 10.0*numpy.log10(0.2*r + 0.8*p)
    # construct sources
    return [Source(position = lopos,
                   direction = asDirection(vehicle.direction()),
                   emission = TertsBandSpectrum(lonoise),
                   directivity = ImagineDirectivity(cat = vehicle.cat(), h = h[0])),
            Source(position = hipos,
                   direction = asDirection(vehicle.direction()),
                   emission = TertsBandSpectrum(hinoise),
                   directivity = ImagineDirectivity(cat = vehicle.cat(), h = h[1]))]

  def emission(self, vehicle, road = ReferenceRoadsurface()):
    """ return the total emission as a single spectrum, assuming that all is emitted by the same source """
    return TertsBandSpectrum(10.0*numpy.log10(fromdB(self.rollingNoise(vehicle,road)) + fromdB(self.propulsionNoise(vehicle,road))))

#---------------------------------------------------------------------------------------------------
# Random corrected road traffic noise emission models
#---------------------------------------------------------------------------------------------------

class ImagineCorrectionModel(ImagineModel):
  """ Road traffic noise emission model with random corrections on the emissions
      This model uses the Imagine model as a base, but adds a random (frequency-independent) sound power level correction
      to each vehicle, in order to model a realistic distribution of louder/quiter vehicles within the same emission class.
  """
  def __init__(self, seed = None):
    ImagineModel.__init__(self)
    numeric.seed(seed)
    self.corrections = {} # dict with a random correction for each vehicle, filled during simulation
    self.generators = {} # dict with correction generators for each vehicle category, to be filled in by subclasses

  def __str__(self):
    """ return a string representation of the emission model """
    return '[Random correction emission model]'

  def getCorrection(self, vehicle):
    """ return the correction for a given vehicle """
    (vid, cat) = (vehicle.vid(), vehicle.cat())
    if vid in self.corrections:
      # vehicle was encountered before, so use the same correction
      corr = self.corrections[vid]
    else:
      # new vehicle, so calculate a new correction and save it
      corr = self.generators[cat].generate()
      self.corrections[vid] = corr
    return corr

  def sources(self, vehicle, road = ReferenceRoadsurface()):
    """ overload the ImagineModel implementation """
    # calculate the sources using the Imagine model
    result = ImagineModel.sources(self, vehicle, road)
    # get the sound power level correction and apply it to the sources
    corr = self.getCorrection(vehicle)
    for source in result:
      source.emission.correct(corr)
    return result

  def emission(self, vehicle, road = ReferenceRoadsurface()):
    """ overload the ImagineModel implementation """
    # calculate the emission using the Imagine model
    result = ImagineModel.emission(self, vehicle, road)
    # get the sound power level correction and apply it to the emission
    corr = self.getCorrection(vehicle)
    result.correct(corr)
    return result

class SkewedNormalImagineCorrectionModel(ImagineCorrectionModel):
  """ Emission model that samples random corrections from a skewed normal distribution """
  def __init__(self, stdev = (0.0,0.0,0.0,0.0,0.0), skew = (0.0,0.0,0.0,0.0,0.0), seed = None):
    ImagineCorrectionModel.__init__(self, seed=seed)
    assert (len(stdev) == 5) and (len(skew) == 5)
    for i in range(5):
      self.generators[i+1] = numeric.SkewedNormalCorrectionGenerator(stdev = stdev[i], skew = skew[i])

  def __str__(self):
    """ return a string representation of the emission model """
    return '[Skewed normal correction emission model]'

# Distributions of sound power levels from Australian study
# (Brown & Tomerini, Road & Transport Research, 20(3):50-63, 2011)
AUSTDIST = {1: numpy.asarray([89, 188, 207, 386, 438, 437, 556, 596, 821, 797, 905, 1099, 1369, 1567, 1727, 2184, 2674, 3426, 3502, 3522, 3123, 2455, 1696, 1238, 832, 468, 362, 286, 195, 126, 74, 77, 32, 34, 14, 14, 11, 6, 6, 4, 5, 6, 3, 2, 5, 8, 0, 0, 0, 0, 0], dtype=float),
            2: numpy.asarray([4, 8, 12, 27, 33, 39, 36, 55, 55, 67, 82, 70, 103, 87, 91, 94, 152, 258, 301, 369, 459, 568, 687, 622, 574, 486, 344, 296, 241, 168, 134, 80, 61, 50, 34, 28, 21, 12, 4, 4, 2, 1, 5, 0, 2, 3, 0, 0, 0, 0, 0], dtype=float),
            3: numpy.asarray([14, 120, 81, 194, 307, 277, 306, 319, 328, 345, 404, 436, 488, 456, 507, 514, 473, 584, 611, 668, 679, 711, 919, 1083, 1430, 1662, 2007, 2018, 1872, 1696, 1268, 935, 547, 378, 261, 197, 102, 77, 60, 40, 24, 32, 22, 11, 13, 16, 1, 0, 0, 0, 2], dtype=float),
            4: numpy.asarray([0, 0, 2, 5, 8, 5, 9, 8, 9, 19, 22, 24, 20, 32, 34, 48, 54, 58, 78, 74, 72, 70, 50, 52, 34, 33, 43, 26, 21, 13, 14, 14, 4, 3, 5, 4, 5, 2, 0, 2, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0], dtype=float),
            5: numpy.asarray([0, 0, 2, 5, 8, 5, 9, 8, 9, 19, 22, 24, 20, 32, 34, 48, 54, 58, 78, 74, 72, 70, 50, 52, 34, 33, 43, 26, 21, 13, 14, 14, 4, 3, 5, 4, 5, 2, 0, 2, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0], dtype=float)}

class DistributionImagineCorrectionModel(ImagineCorrectionModel):
  """ Emission model that samples random corrections from a predefined distribution """
  def __init__(self, corrdist = AUSTDIST, seed = None, zeropoint='meanEn'):
    # corrdist = dict with correction distributions (in bins of 1dB) for each Imagine vehicle class
    ImagineCorrectionModel.__init__(self, seed=seed)
    for cat, values in corrdist.iteritems():
      self.generators[cat] = numeric.TabulatedCorrectionGenerator(values = values, dx = 1.0, zeropoint=zeropoint)

  def __str__(self):
    """ return a string representation of the emission model """
    return '[Distribution correction emission model]'

#---------------------------------------------------------------------------------------------------
# Lookup table road traffic noise emission model
#---------------------------------------------------------------------------------------------------

# TODO: Implement lookup table model, for which the following has to be done
# - create a class LookupModel, that is a subclass from ImagineModel, and overload the 'sources' and 'emission' methods
#   (same way as the DistributionModel class above); a subclass from the Imagine model is necessary because the Lookup model
#   may use the shape of the emission spectrum as from the Imagine model
# - add necessary values to the default configuration (CFGEMODEL, see above) and create a new default configuration file,
#   by running this script from the command line with the flag MAIN_CONFIG set to True
# - update the method 'Configuration.emodel' (see below) such that a LookupModel object is created when appropriate

#---------------------------------------------------------------------------------------------------
# Test code
#---------------------------------------------------------------------------------------------------

if __name__ == '__main__':

  # test emission correction generator
  if 1:
    vClass = 3
    pylab.figure()
    g = numeric.TabulatedCorrectionGenerator(values = AUSTDIST[vClass], zeropoint='meanEn')
    g.plotData() # measured distribution
    g.plot() # fitted distribution

  # plot of Imagine emission spectrum at different vehicle speeds
  if 0:
    road = ReferenceRoadsurface()
    model = ImagineModel()
    cat = 1
    pylab.figure()
    speeds = [30, 40, 50, 60, 70, 80, 90, 100, 110, 120]
    markers = ['circle grey', 'diamond black', 'square white', 'triangle grey', 'circle black',
               'diamond white', 'square grey', 'triangle black', 'circle white', 'diamond grey']
    for v, m in zip(speeds, markers):
      model.emission(QLDCar(cat = cat, speed = v, acceleration = 0.0)).plot(m = m, interval = (60.0, 110.0))
    pylab.legend([('%d km/h' % v) for v in speeds])
    pylab.title('Imagine emission spectrum at different vehicle speeds')

  # reproduction of Figure 25 of Imagine report (emission spectra for rolling and propulsion noise, and total emission spectrum)
  if 0:
    road = ReferenceRoadsurface()
    model = ImagineModel()
    vList = {1: range(10, 161), 2: range(10, 101), 3: range(10, 101)}
    yRange = {1: (70, 115), 2: (75, 115), 3: (75, 115)}
    pylab.figure()
    for cat in [1, 2, 3]:
      p = [TertsBandSpectrum(model.propulsionNoise(QLDCar(cat = cat, speed = v, acceleration = 0.0), road)).laeq() for v in vList[cat]]
      r = [TertsBandSpectrum(model.rollingNoise(QLDCar(cat = cat, speed = v, acceleration = 0.0), road)).laeq() for v in vList[cat]]
      e = [model.emission(QLDCar(cat = cat, speed = v, acceleration = 0.0)).laeq() for v in vList[cat]]
      pylab.subplot(2, 2, cat)
      pylab.plot(vList[cat], p)
      pylab.plot(vList[cat], r)
      pylab.plot(vList[cat], e)
      # adjust y-axis
      a = pylab.axis()
      pylab.axis([a[0], a[1], yRange[cat][0], yRange[cat][1]])

  # reproduction of Figure 5.32 of Harmonoise report (directivity functions for source heights 0.1m and 0.75m)
  if 0:
    cat = 1 # vehicle category of interest
    f = 2000.0 # frequency of interest
    thetas = numpy.arange(0.0, 180.0)
    pylab.figure()
    for h in [0.01, 0.75]:
      func = [ImagineDirectivity(cat = cat, h = h)(theta = theta, phi = 0.0, f = f) for theta in thetas]
      pylab.plot((numpy.pi/180.0)*thetas, func)
    # adjust y-axis
    a = pylab.axis()
    pylab.axis([0.0, 3.5, -10.0, 4.0])
    # draw grid
    pylab.gca().xaxis.grid(color='gray', linestyle='dashed')
    pylab.gca().yaxis.grid(color='gray', linestyle='dashed')

  # plot of Harmonoise/Imagine directivity patterns as polar plot
  if 0:
    warnings.simplefilter('ignore', UserWarning)
    cat = 1
    f = 2000.0
    pylab.figure()
    iplot = 0
    for h in [0.01, 0.30, 0.75]:
      for phi in [0.0, 30.0, 60.0, 90.0]:
        directivity = ImagineDirectivity(cat = cat, h = h)
        thetas = numpy.arange(0.0, 360.0)
        Lcorrs = [directivity(theta = theta, phi = phi, f = f) for theta in thetas] # level corrections
        iplot += 1
        pylab.subplot(3, 4, iplot, polar = True)
        pylab.polar((numpy.pi/180.0)*thetas, numpy.zeros(360), color = 'black', linestyle = '--') # draw the zero line
        pylab.polar((numpy.pi/180.0)*thetas, Lcorrs)
        pylab.ylim(-10.0, 4.0)

  try:
    pylab.show()
  except:
    pass


import numpy as np
from pathlib import Path
from mat4py import loadmat
from streamz import Stream
import pandas as pd
import os.path
import pygame
from PIL import Image,ImageFilter
import time
import sys

def open_im(filename):#opens the image
    try:
        OriImage = Image.open(filename)
    except FileNotFoundError:
        print('Imgae file not Found in directory')
    return OriImage

def NF_instance(parameter_blur,parameter_size):
    """
	this function will recieve the two parameters and control the blurring and size of image accorting to them.

   	Parameters
    	----------
    	parameter 1:int
        	First Neurofeedback parameter.
    	parameter 2: int
        	Second Neurofeedback parameter.
    	Returns
    	-------
        size: image size tuple.
	data: image.tobytes() array
	mode: image mode, integer.

    """
    boxImage = OriImage.filter(ImageFilter.BoxBlur(parameter_blur))
    resizedImage = boxImage.resize([int(parameter_size*s) for s in boxImage.size])
    mode = resizedImage.mode
    data = resizedImage.tobytes()
    size = resizedImage.size
    return data,size,mode
OriImage = open_im('im2.jpg')
pygame.init()
screen = pygame.display.set_mode((1000,600))
pygame.display.set_caption('Neurofeedback Training')
icon = pygame.image.load('brain.png')
pygame.display.set_icon(icon)
screen.fill((250,250,250))
pygame.display.update()
image = pygame.image.load(r'im2.jpg') #should be an image loaded in the same file of the script or an image address
screen.blit(image,(100,40))
pygame.display.update()

def main(parameter_blur, parameter_size):#updates the gui with the current image using NF_instance
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
            pygame.quit()
            sys.exit(0)
    data,size,mode = NF_instance(parameter_blur,parameter_size)
    resed = pygame.image.fromstring(data,size,mode)
    res_rec = resed.get_rect(center = screen.get_rect().center)
    screen.fill((250,250,250))
    pygame.display.update()
    screen.blit(resed,res_rec)#update gui with current image
    pygame.display.update()
    time.sleep(0.1)

def load_mat_data(mat_data:str): # This function loads MATLAB data to Python dictionary
    try:
        py_data = loadmat(mat_data)
    except FileNotFoundError:
        print(f'There is no {mat_data}.mat file in the directory ')
    return py_data

def dict_2_np(mat_dict:dict): # This function converts dictionary data to numpy array
    np_data=np.array(list(mat_dict.values()))
    return np_data

def data_giver(data_array:np): # This function uploads nuerofeedbak parametars to the visual interface
    np_data=np.array(list(mat_dict.values()))
    return np_data

def data_giver(data_array:np): # This function uploads nuerofeedbak parametars to the visual interface
    """Streams data of two neurofeedback parameters from np.array matrix,
    assuming they're both positive integers in a defined range.

    Parameters
    ----------
    data_array:np
        A matrix Positive integers

    Returns
    -------
    The function activates neurofeedback visual interface

    Raises
    ------
    ValueError
        If one of the parameters is negative or equels to zero or None.
        If one of the parameter is not in the defined range.

    TypeError
        If the data array is not np.ndarray
    """

    if not (isinstance(data_array, np.ndarray)):
        raise TypeError(f"Received {data_array}; expected numpy array")

    for column in data_array.T:

        if  (column[0]== None) or (column[1] == None):
            raise ValueError(f"Received {column[0], column[1]}; expected a numerical value")

        if (column[0] <= 0) or (column[1] <= 0):
            raise ValueError(f"Received {column[0], column[1]}; expected positive integers or floats")

        if  not (0<column[0]<=10) or not (0.1<=column[1]<=1):
            raise ValueError(f"Received {column[0], column[1]}; expected a value in defined range")

        if  (column[0] == bool) or (column[1] ==  bool):
            raise TypeError(f"Received {column[0], column[1]}; expected a numerical value")

        main(column[0],column[1]) # Call for the visual interface function
        time.sleep(1) # Data samplig rate

def streamer_func (datout_file:str, lenout_file:str, folder_path=Path().absolute()):
    """This function receives two parameters from the EEG Neurofeedback system (Curry7)
    in Matlab formt and stream it to python. The function pushes the Matlab data through a papline based
    on streamz library.

    Parameters
    ----------
    datout_file:str
        First Neurofeedback parameter vector file name.
    lenout_file:str
        Second Neurofeedback parameter vector file name.
    path: pathlib path
        Pathlib path to the folder with the data.

    Returns
    -------
        The function activates the visual interface
    """
    datout_source = Stream()
    a=datout_source.map(lambda x: x[0]).map(load_mat_data).map(dict_2_np)
    b=datout_source.map(lambda x: x[1]).map(load_mat_data).map(dict_2_np)

    a.zip(b).map(np.concatenate).sink(data_giver)

    datout_source.emit((datout_file, lenout_file))

streamer_func('datout.mat', 'lenout.mat')


#### Reward evaluation for agent in Reacher7Dof gym env using a single real-world env
## Wrtitten by : leopauly | cnlp@leeds.ac.uk
## Courtesy for DDPG implementation : Steven Spielberg Pon Kumar (github.com/stevenpjg)
####

##Imports
import gym
from gym.spaces import Box, Discrete
import numpy as np
np.set_printoptions(suppress=True)
import cv2
from ddpg import DDPG
from ou_noise import OUNoise
import matplotlib.pyplot as plt
import scipy.misc as misc
import skimage
import os
from threading import Thread, Lock
import sys
from six.moves import xrange
import PIL.Image as Image
import random
import numpy as np
import cv2
import time
import math
import matplotlib.pyplot as plt
import tensorflow as tf
import keras
from keras import backend as K
import lscript as lsp
import modelling as md

## Defining vars
num_episodes=80
steps=60 # No of actions taken in a roll out
is_batch_norm = False #batch normalization switch
xrange=range # For python3
start_training=64 # Buffer size, before starting to train the RL algorithm
height=112
width=112
channel=3
crop_size=112
cluster_length=16
nb_classes=2
feature_size=4608
#-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------#

switch=int(sys.argv[3])
if(switch==0):
    base_dir='/home/ironman2/Observation-Learning-Simulations/S2l/Thesis_Ch3/Exp1_reach3dof/Results/Reward_Eval_Conv5_80eps/rewardmap_misc/'
    demo_folder='./Demos/demo_reach_0deg_new/'
elif (switch==1):
    base_dir='/home/ironman2/Observation-Learning-Simulations/S2l/Thesis_Ch3/Exp1_reach3dof/Results/Reward_Eval_Conv5_80eps_rand/rewardmap_3Dview1/'
    demo_folder='./Demos/demo_reach_0deg_new/'
elif (switch==2):
    base_dir='/home/ironman2/Observation-Learning-Simulations/S2l/Thesis_Ch3/Exp1_reach3dof/Results/Reward_Eval_Conv5_80eps_rand/rewardmap_3Dview2/'
    demo_folder='./Demos/demo_reach_0deg_new/'
elif (switch==3):
    base_dir='/home/ironman2/Observation-Learning-Simulations/S2l/Thesis_Ch3/Exp1_reach3dof/Results/Rewmaps_Conv5_norm2pow5_80eps/random_weights_lr0/'
    demo_folder='./Demos/demo_reach_0deg_new/'
elif (switch==4):
    base_dir='/home/ironman2/Observation-Learning-Simulations/S2l/Thesis_Ch3/Exp1_reach3dof/Results/Rewmaps_Conv5_norm2pow5_80eps/random_weights_ll0/'
    demo_folder='./Demos/demo_reach_0deg_new/'
elif (switch==5):
    base_dir='/home/ironman2/Observation-Learning-Simulations/S2l/Thesis_Ch3/Exp1_reach3dof/Results/Reward_Eval_Conv5_80eps_rand/multi_target_close/'
    demo_folder='./Demos/demo_reach_0deg_new/'
elif (switch==6):
    base_dir='/home/ironman2/Observation-Learning-Simulations/S2l/Thesis_Ch3/Exp1_reach3dof/Results/Reward_Eval_Conv5_80eps_rand/multi_target_far/'
    demo_folder='./Demos/demo_reach_0deg_new/'
elif (switch==-2):
    base_dir='/home/ironman2/Observation-Learning-Simulations/S2l/Thesis_Ch3/Exp1_reach3dof/Results/Rewmaps_Conv5_norm2pow5_80eps/random_weights_V2/'
    demo_folder='./Demos/demo_reach_180deg_new/'
elif (switch==-4):
    base_dir='/home/ironman2/Observation-Learning-Simulations/S2l/Thesis_Ch3/Exp1_reach3dof/Results/Reward_Eval_Conv5_80eps_rand/rewardmap_Obj2_new/'
    demo_folder='./Demos/demo_reach_180deg_new/'
elif (switch==-3):
    base_dir='/home/ironman2/Observation-Learning-Simulations/S2l/Thesis_Ch3/Exp1_reach3dof/Results/Reward_Eval_Conv5_80eps_rand/rewardmap_Obj1/'
    demo_folder='./Demos/demo_reach_0deg_new/'
elif (switch==-5):
    base_dir='/home/ironman2/Observation-Learning-Simulations/S2l/Thesis_Ch3/Exp1_reach3dof/Results/Reward_Eval_Conv5_80eps_rand/rewardmap_BG/'
    demo_folder='./Demos/demo_reach_0deg_new/'
elif (switch==-6):
    base_dir='/home/ironman2/Observation-Learning-Simulations/S2l/Thesis_Ch3/Exp1_reach3dof/Results/Rewmaps_Conv5_norm2pow5_80eps/random_weights_M/'
    demo_folder='./Demos/demo_reach_0deg_h.s/'
else:
    base_dir='/home/ironman2/Observation-Learning-Simulations/S2l/Thesis_Ch3/Exp1_reach3dof/Results/Rewmaps_Conv5_norm2pow5_80eps/random_weights_baseline/'
    demo_folder='./Demos/demo_reach_0deg_new/'

os.system('mkdir %s' % base_dir)

#-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------#

## Defining env
env = gym.make('Pusher3DOFReal-v1')
env.switch=switch
env.initialize_env()
assert isinstance(env.observation_space, Box), "observation space must be continuous"
assert isinstance(env.action_space, Box), "action space must be continuous"

#-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------#

class Frame_Feature:
    def __init__(self):
         self.g=tf.Graph()
         with self.g.as_default():
              self.sess=tf.InteractiveSession(config=tf.ConfigProto(allow_soft_placement=True,log_device_placement=False))
              self.base_model=tf.keras.applications.vgg16.VGG16(include_top=False, weights='imagenet', input_tensor=None, input_shape=(height,width,channel), pooling=None, classes=1000)
              #print(tf.contrib.graph_editor.get_tensors(self.g))   #(tf.get_default_graph()))
              self.base_model._make_predict_function()
              print('VggNet loaded with Imagenet values')

    def frame_feature_extractor(self,frame_):
        frame= self.im_preprocess(frame_)
        frame=frame.reshape(-1,height,width,channel)
        frame_features=self.base_model.predict(frame)
        return frame_features

    def im_preprocess(self,im):
        im = np.float32(im)
        im[:,:,2] -= 103.939
        im[:,:,1] -= 116.779
        im[:,:,0] -= 123.68
        im = im[:, :, ::-1]  # change to BGR
        return im
#-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------#

### DEMO FEATURE EXTRACTION
def get_compress_frames_data(filename, num_frames_per_clip=cluster_length):
  ret_arr = []
  for parent, dirnames, filenames in os.walk(filename):

    filenames = sorted(filenames)
    jump=math.floor((len(filenames)/num_frames_per_clip))
    loop=0

    for i in range(0,len(filenames),jump):
      if (loop>15):
        break
      if (filenames[i].endswith('.png')):
        image_name = str(filename) + '/' + str(filenames[i])
        img = Image.open(image_name)
        img_data = np.array(img)
        ret_arr.append(img_data)
        loop=loop+1
  ret_arr=np.array(ret_arr)
  #ret_arr=ret_arr/255
  return ret_arr

def demo_array_extractor(demo_vid_path):
    demo_vid_array=get_compress_frames_data(demo_vid_path)
    return demo_vid_array

#-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------#

def sampling_obs(vid_robo_all,num_frames_per_clip=cluster_length):
    total_obs=len(vid_robo_all)
    jump=math.floor(total_obs/num_frames_per_clip)
    loop=0
    ret_arr=[]
    for i in range(0,total_obs,jump):
        if (loop>15):
            break
        img_data = vid_robo_all[i]
        ret_arr.append(img_data)
        loop=loop+1

    ret_arr=np.array(ret_arr)
    #ret_arr=ret_arr/255
    return ret_arr

#-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------#

### VIDEO FEATURE EXTRACTION
class Vid_Feature:

    def __init__(self):
        self.saved_path='/home/ironman2/S2l_storage/trained_activity_nets_thesis/saved/'
        self.network_name='activity_model.ckpt-67.meta'
        self.network_weigths_name='activity_model.ckpt-67'
        ### Activity_net
        self.g=tf.Graph()
        with self.g.as_default():

            self.sess = tf.InteractiveSession(config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=False))
            ## Restore model weights from previously saved model
            self.saver = tf.train.import_meta_graph(os.path.join(self.saved_path,self.network_name))
            self.saver.restore(self.sess, os.path.join(self.saved_path,self.network_weigths_name))
            #self.sess.run(tf.global_variables_initializer())
            print("Model restored from file: %s" % self.saved_path,flush=True)

    ## For extracting activity features
    def feature_extractor(self,vid_np):
        self.vid_=vid_np.reshape(-1,cluster_length,height,width,channel)
        f_v = self.sess.graph.get_tensor_by_name('flatten_1/Reshape:0')
        self.f_v_val=np.array(self.sess.run([f_v], feed_dict={'conv1_input:0':self.vid_,'Placeholder:0':self.vid_,'dropout_1/keras_learning_phase:0':0 }))
        self.features=np.reshape(self.f_v_val,(-1))
        return self.features

#-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------#

def distance(f_demo,f_robo):
    norm_val=2
    norm_pow=1
    distance_=np.linalg.norm(f_demo-f_robo,ord=norm_val)
    return pow(distance_,norm_pow)

#-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------#

def s2l(i_run):
    print('This is the ith run',i_run)

    #Randomly initialize critic,actor,target critic, target actor network  and replay buffer
    num_states = feature_size   #num_states = env.observation_space.shape[0]
    num_actions = env.action_space.shape[0]
    print ("Number of States:", num_states)
    print ("Number of Actions:", num_actions)

    agent = DDPG(env, is_batch_norm,num_states,num_actions)
    exploration_noise = OUNoise(env.action_space.shape[0])
    counter=0
    total_reward=0
    best_reward=-10000

    print ("Number of Steps per episode:", steps)
    reward_st_per_episode = np.array([0])  #saving reward
    man_pos_x_st_per_episode = np.array([0])
    man_pos_y_st_per_episode = np.array([0])
    eval_metric_st= np.array([0])
    eval_metric_st_per_episode= np.array([0])
    reward_st_per_step = np.array([0])  #saving reward after every step

    activity_obj=Vid_Feature()
    demo_vid_array=demo_array_extractor(demo_folder)
    if(i_run==0):
        plt.imshow(demo_vid_array[0])
        plt.savefig('demo_img'+str(switch) +'.png')
    demo_features=activity_obj.feature_extractor(demo_vid_array)
    frame_obj=Frame_Feature()

    for episode in range(num_episodes):
        print ("==== Starting episode no:",episode,"====","\n")
        env.reset()   # Reset env in the begining of each episode
        env.render()
        obs_img=env.render(mode='rgb_array')   # Get the observation
        if(i_run==0 and episode==0):
            plt.imshow(obs_img)
            plt.savefig('env_img'+str(switch) +'.png')
        obs_img=np.array(misc.imresize(obs_img,[112,112,3]))
        observation =np.array(frame_obj.frame_feature_extractor(obs_img))
        observation=observation.reshape(-1)

        reward_per_episode = 0
        vid_robo_=[]

        for i in range(steps):

            x = observation

            action = agent.evaluate_actor(np.reshape(x,[1,num_states]))
            action = action+1
            noise = exploration_noise.noise()#/(episode+1)
            action = action[0] + noise
            print ('Action at',i_run ,'episode-',episode, 'step-', i ," :",action)

            _,_,done,info=env.step(action)
            env.render()
            obs_robo_=env.render(mode='rgb_array')
            obs_robo=misc.imresize(obs_robo_,[112,112,3])
            vid_robo_.append(obs_robo)

            observation=np.array(frame_obj.frame_feature_extractor(np.array(obs_robo)))
            observation=observation.reshape(-1)

            if(i==(steps-1)):
                vid_robo_all=np.array(vid_robo_)
                vid_robo=sampling_obs(vid_robo_all)
                #for i in range(len(vid_robo)):
                #    plt.imshow(vid_robo[i])
                #    plt.savefig('obs_img'+str(i) +'.png')
                robo_features=activity_obj.feature_extractor(vid_robo)
                reward=-(distance(demo_features,robo_features))
                reward=np.array(reward)
                print('reward: ',reward)
            else:
                reward=0
                reward=np.array(reward)
                print('reward: ',reward)

            # Printing eval_metric after every step
            eval_metric=np.array(env.get_eval())
            eval_metric=eval_metric.reshape(-1)
            print('Distance to goal:',eval_metric)
            eval_metric_st = np.append(eval_metric_st,eval_metric)
            np.savetxt(base_dir+'eval_metric_per_step_run_'+str(i_run)+'.txt',eval_metric_st, newline="\n")

            # Storing reward after every step
            reward_st_per_step = np.append(reward_st_per_step,reward)
            np.savetxt(base_dir+'reward_per_step_run_'+str(i_run)+'.txt',reward_st_per_step, newline="\n")

            #add s_t,s_t+1,action,reward to experience memory
            agent.add_experience(x,observation,action,reward,False)
            counter+=1

            #train critic and actor network
            if counter > start_training:
                agent.train() #
            print ('\n\n')
            print('Episode: ',episode,' Manipulator position: ',env.get_man_pos())

            reward_per_step=reward
            reward_per_episode+=reward_per_step

        ## Printing and saving episode rewards
        print ('Episode: ',episode,' Episode Reward: ',reward_per_episode)
        exploration_noise.reset() #reinitializing random noise for action exploration
        reward_st_per_episode = np.append(reward_st_per_episode,reward_per_episode)
        np.savetxt(base_dir+'episode_reward_run_'+str(i_run)+'.txt',reward_st_per_episode, fmt='%f', newline="\n")

        ## Printing and saving final mnaipulator position
        manipulator_final_pos=env.get_man_pos()
        print('Episode: ',episode,' Manipulator position: ',manipulator_final_pos)

        man_pos_x_st_per_episode = np.append(man_pos_x_st_per_episode,manipulator_final_pos[0])
        np.savetxt(base_dir+'episode_man_x_pos_run_'+str(i_run)+'.txt',man_pos_x_st_per_episode, fmt='%f', newline="\n")

        man_pos_y_st_per_episode = np.append(man_pos_y_st_per_episode,manipulator_final_pos[1])
        np.savetxt(base_dir+'episode_man_y_pos_run_'+str(i_run)+'.txt',man_pos_y_st_per_episode, fmt='%f', newline="\n")

        ## Saving
        if (best_reward<reward_per_episode):
            best_reward=reward_per_episode
            print('best reward:',best_reward)
            print('current reward:',reward_per_episode)
            print('saving policy for episode..................:',episode)
            #agent.save_actor(episode,i_run)

        ## Printing eval_metric after every step
        eval_metric=np.array(env.get_eval())
        eval_metric=eval_metric.reshape(-1)
        print('Distance to goal at the end  of episode:',eval_metric)
        eval_metric_st_per_episode = np.append(eval_metric_st_per_episode,eval_metric)
        np.savetxt(base_dir+'eval_metric_per_epispde_run_'+str(i_run)+'.txt', eval_metric_st_per_episode, newline="\n")

        total_reward+=reward_per_episode

    print ("Average reward per episode {}".format(total_reward / num_episodes))
    print('Best episode reward',best_reward)
    print ('\n\n')

    del agent
    del activity_obj
    del frame_obj

if __name__=='__main__':
    from datetime import datetime
    start_time=str(datetime.now())
    run_start=int(sys.argv[1])
    run_end=int(sys.argv[2])
    print('Start trial:',run_start,'End trail:',run_end-1)
    for i_run in range(run_start,run_end):
        s2l(i_run)
    print('Start to end time:',start_time,str(datetime.now()))


import re
import numpy as np
import editdistance as ed

def cer(hypothesis, groundtruth, **kwargs):
    err = 0
    tot = 0
    for p, t in zip(hypothesis, groundtruth):
        err += float(ed.eval(p, t))
        tot += len(t)
    return err / tot

def wer(hypothesis, groundtruth, **kwargs):
    err = 0
    tot = 0
    for p, t in zip(hypothesis, groundtruth):
        p = p.split(' ')
        t = t.split(' ')
        err += float(ed.eval(p, t))
        tot += len(t)
    return err / tot

def clean(ref):
    ref = re.sub(r'B\-(\S+) ', '', ref)
    ref = re.sub(r' E\-(\S+)', '', ref)
    return ref

def parse(hyp, ref):
    gex = re.compile(r'B\-(\S+) (.+?) E\-\1')

    hyp = re.sub(r' +', ' ', hyp)
    ref = re.sub(r' +', ' ', ref)

    hyp_slots = gex.findall(hyp)
    ref_slots = gex.findall(ref)

    ref_slots = ';'.join([':'.join([x[1], x[0]]) for x in ref_slots])
    if len(hyp_slots)>0:
        hyp_slots = ';'.join([':'.join([clean(x[1]), x[0]]) for x in hyp_slots])
    else:
        hyp_slots = ''

    ref = clean(ref)
    hyp = clean(hyp)

    return ref, hyp, ref_slots, hyp_slots

def slot_type_f1(hypothesis, groundtruth, **kwargs):
    F1s = []
    for p, t in zip(hypothesis, groundtruth):
        ref_text, hyp_text, ref_slots, hyp_slots = parse(p, t)
        ref_slots = ref_slots.split(';')
        hyp_slots = hyp_slots.split(';')
        unique_slots = []
        ref_dict = {}
        hyp_dict = {}
        if ref_slots[0] != '':
            for ref_slot in ref_slots:
                v, k = ref_slot.split(':')
                ref_dict.setdefault(k, [])
                ref_dict[k].append(v)
        if hyp_slots[0] != '':
            for hyp_slot in hyp_slots:
                v, k = hyp_slot.split(':')
                hyp_dict.setdefault(k, [])
                hyp_dict[k].append(v)
        # Slot Type F1 evaluation
        if len(hyp_dict.keys()) == 0 and len(ref_dict.keys()) == 0:
            F1 = 1.0
        elif len(hyp_dict.keys()) == 0:
            F1 = 0.0
        elif len(ref_dict.keys()) == 0:
            F1 = 0.0
        else:
            P, R = 0.0, 0.0
            for slot in ref_dict:
                if slot in hyp_dict:
                    R += 1
            R = R / len(ref_dict.keys())
            for slot in hyp_dict:
                if slot in ref_dict:
                    P += 1
            P = P / len(hyp_dict.keys())
            F1 = 2*P*R/(P+R) if (P+R) > 0 else 0.0
        F1s.append(F1)
    return sum(F1s) / len(F1s)

def slot_value_cer(hypothesis, groundtruth, **kwargs):
    value_hyps = []
    value_refs = []
    for p, t in zip(hypothesis, groundtruth):
        ref_text, hyp_text, ref_slots, hyp_slots = parse(p, t)
        ref_slots = ref_slots.split(';')
        hyp_slots = hyp_slots.split(';')
        unique_slots = []
        ref_dict = {}
        hyp_dict = {}
        if ref_slots[0] != '':
            for ref_slot in ref_slots:
                v, k = ref_slot.split(':')
                ref_dict.setdefault(k, [])
                ref_dict[k].append(v)
        if hyp_slots[0] != '':
            for hyp_slot in hyp_slots:
                v, k = hyp_slot.split(':')
                hyp_dict.setdefault(k, [])
                hyp_dict[k].append(v)
        # Slot Value WER/CER evaluation
        unique_slots = list(ref_dict.keys())
        for slot in unique_slots:
            for ref_i, ref_v in enumerate(ref_dict[slot]):
                if slot not in hyp_dict:
                    hyp_v = ''
                    value_refs.append(ref_v)
                    value_hyps.append(hyp_v)
                else:
                    min_cer = 100
                    best_hyp_v = ""
                    for hyp_v in hyp_dict[slot]:
                        tmp_cer = cer([hyp_v], [ref_v])
                        if min_cer > tmp_cer:
                            min_cer = tmp_cer
                            best_hyp_v = hyp_v
                    value_refs.append(ref_v)
                    value_hyps.append(best_hyp_v)

    return cer(value_hyps, value_refs)

def slot_value_wer(hypothesis, groundtruth, **kwargs):
    value_hyps = []
    value_refs = []
    for p, t in zip(hypothesis, groundtruth):
        ref_text, hyp_text, ref_slots, hyp_slots = parse(p, t)
        ref_slots = ref_slots.split(';')
        hyp_slots = hyp_slots.split(';')
        unique_slots = []
        ref_dict = {}
        hyp_dict = {}
        if ref_slots[0] != '':
            for ref_slot in ref_slots:
                v, k = ref_slot.split(':')
                ref_dict.setdefault(k, [])
                ref_dict[k].append(v)
        if hyp_slots[0] != '':
            for hyp_slot in hyp_slots:
                v, k = hyp_slot.split(':')
                hyp_dict.setdefault(k, [])
                hyp_dict[k].append(v)
        # Slot Value WER/CER evaluation
        unique_slots = list(ref_dict.keys())
        for slot in unique_slots:
            for ref_i, ref_v in enumerate(ref_dict[slot]):
                if slot not in hyp_dict:
                    hyp_v = ''
                    value_refs.append(ref_v)
                    value_hyps.append(hyp_v)
                else:
                    min_wer = 100
                    best_hyp_v = ""
                    for hyp_v in hyp_dict[slot]:
                        tmp_wer = wer([hyp_v], [ref_v])
                        if min_wer > tmp_wer:
                            min_wer = tmp_wer
                            best_hyp_v = hyp_v
                    value_refs.append(ref_v)
                    value_hyps.append(best_hyp_v)

    return wer(value_hyps, value_refs)

def slot_edit_f1(hypothesis, groundtruth, loop_over_all_slot, **kwargs):
    test_case, TPs, FNs, FPs = [], 0, 0, 0

    slot2F1  = {} #defaultdict(lambda: [0,0,0]) # TPs, FNs, FPs
    for p, t in zip(hypothesis, groundtruth):
        ref_text, hyp_text, ref_slots, hyp_slots = parse(p, t)
        ref_slots = ref_slots.split(';')
        hyp_slots = hyp_slots.split(';')
        unique_slots = []
        ref_dict = {}
        hyp_dict = {}
        if ref_slots[0] != '':
            for ref_slot in ref_slots:
                v, k = ref_slot.split(':')
                ref_dict.setdefault(k, [])
                ref_dict[k].append(v)
        if hyp_slots[0] != '':
            for hyp_slot in hyp_slots:
                v, k = hyp_slot.split(':')
                hyp_dict.setdefault(k, [])
                hyp_dict[k].append(v)
        # Collecting unique slots
        unique_slots = list(ref_dict.keys())
        if loop_over_all_slot:
            unique_slots += [x for x in hyp_dict if x not in ref_dict]
        # Evaluating slot edit F1
        for slot in unique_slots:
            TP = 0
            FP = 0
            FN = 0
            if slot not in ref_dict: # this never happens in list(ref_dict.keys())
                for hyp_v in hyp_dict[slot]:
                    FP += 1
            else:
                for ref_i, ref_v in enumerate(ref_dict[slot]):
                    if slot not in hyp_dict:
                        FN += 1
                    else:
                        match = False
                        for hyp_v in hyp_dict[slot]:
                            #if ref_i < len(hyp_dict[slot]):
                            #    hyp_v = hyp_dict[slot][ref_i]
                            if hyp_v == ref_v:
                                match = True
                                break
                        if match:
                            TP += 1
                        else:
                            FN += 1
                            FP += 1
            slot2F1.setdefault(slot, [0,0,0])
            slot2F1[slot][0] += TP
            slot2F1[slot][1] += FN
            slot2F1[slot][2] += FP

    all_TPs, all_FNs, all_FPs = 0, 0, 0
    for slot in slot2F1.keys():
        all_TPs += slot2F1[slot][0]
        all_FNs += slot2F1[slot][1]
        all_FPs += slot2F1[slot][2]

    return (100.0 * 2*all_TPs/(2*all_TPs + all_FPs + all_FNs))

def slot_edit_f1_full(hypothesis, groundtruth, **kwargs):
    return slot_edit_f1(hypothesis, groundtruth, loop_over_all_slot=True, **kwargs)

def slot_edit_f1_part(hypothesis, groundtruth, **kwargs):
    return slot_edit_f1(hypothesis, groundtruth, loop_over_all_slot=False, **kwargs)


# Copyright (c) 2020-2021 impersonator.org authors (Wen Liu and Zhixin Piao). All rights reserved.

import cv2
import torch
import numpy as np
import math
from operator import itemgetter

from .dataset import normalize, pad_width

def infer(net, img, scales, base_height, stride, device, num_kpts=25,
          pad_value=(0, 0, 0), img_mean=(128, 128, 128), img_scale=1 / 256):
    normed_img = normalize(img, img_mean, img_scale)
    height, width, _ = normed_img.shape
    scales_ratios = [scale * base_height / float(height) for scale in scales]
    avg_heatmaps = np.zeros((height, width, num_kpts + 1), dtype=np.float32)
    avg_pafs = np.zeros((height, width, num_kpts * 2 + 2), dtype=np.float32)

    for ratio in scales_ratios:
        scaled_img = cv2.resize(normed_img, (0, 0), fx=ratio, fy=ratio, interpolation=cv2.INTER_CUBIC)
        min_dims = [base_height, max(scaled_img.shape[1], base_height)]
        padded_img, pad = pad_width(scaled_img, stride, pad_value, min_dims)

        tensor_img = torch.from_numpy(padded_img).permute(2, 0, 1).unsqueeze(0).float().to(device)
        stages_output = net(tensor_img)

        stage2_heatmaps = stages_output[-2]
        heatmaps = np.transpose(stage2_heatmaps.squeeze().cpu().data.numpy(), (1, 2, 0))
        heatmaps = cv2.resize(heatmaps, (0, 0), fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC)
        heatmaps = heatmaps[pad[0]:heatmaps.shape[0] - pad[2], pad[1]:heatmaps.shape[1] - pad[3]:, :]
        heatmaps = cv2.resize(heatmaps, (width, height), interpolation=cv2.INTER_CUBIC)
        avg_heatmaps = avg_heatmaps + heatmaps / len(scales_ratios)

        stage2_pafs = stages_output[-1]
        pafs = np.transpose(stage2_pafs.squeeze().cpu().data.numpy(), (1, 2, 0))
        pafs = cv2.resize(pafs, (0, 0), fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC)
        pafs = pafs[pad[0]:pafs.shape[0] - pad[2], pad[1]:pafs.shape[1] - pad[3], :]
        pafs = cv2.resize(pafs, (width, height), interpolation=cv2.INTER_CUBIC)
        avg_pafs = avg_pafs + pafs / len(scales_ratios)

    outputs = {
        "heatmaps": avg_heatmaps,
        "pafs": avg_pafs
    }

    return outputs

def infer_fast(net, img, net_input_height_size, stride, upsample_ratio, device,
               pad_value=(0, 0, 0), img_mean=(128, 128, 128), img_scale=1 / 256):
    height, width, _ = img.shape
    scale = net_input_height_size / height

    scaled_img = cv2.resize(img, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
    scaled_img = normalize(scaled_img, img_mean, img_scale)
    min_dims = [net_input_height_size, max(scaled_img.shape[1], net_input_height_size)]
    padded_img, pad = pad_width(scaled_img, stride, pad_value, min_dims)

    tensor_img = torch.from_numpy(padded_img).permute(2, 0, 1).unsqueeze(0).float()
    tensor_img = tensor_img.to(device)

    stages_output = net(tensor_img)

    stage2_heatmaps = stages_output[-2]
    heatmaps = np.transpose(stage2_heatmaps.squeeze().cpu().data.numpy(), (1, 2, 0))
    heatmaps = cv2.resize(heatmaps, (0, 0), fx=upsample_ratio, fy=upsample_ratio, interpolation=cv2.INTER_CUBIC)

    stage2_pafs = stages_output[-1]
    pafs = np.transpose(stage2_pafs.squeeze().cpu().data.numpy(), (1, 2, 0))
    pafs = cv2.resize(pafs, (0, 0), fx=upsample_ratio, fy=upsample_ratio, interpolation=cv2.INTER_CUBIC)

    outputs = {
        "heatmaps": heatmaps,
        "pafs": pafs,
        "pad": pad,
        "scale": scale,
        "upsample_ratio": upsample_ratio,
        "stride": stride,
        "orig_shape": (height, width)
    }

    return outputs

def linspace2d(start, stop, n=10):
    points = 1 / (n - 1) * (stop - start)
    return points[:, None] * np.arange(n) + start[:, None]

def extract_keypoints(heatmap, all_keypoints, total_keypoint_num):
    heatmap[heatmap < 0.1] = 0
    heatmap_with_borders = np.pad(heatmap, [(2, 2), (2, 2)], mode='constant')
    heatmap_center = heatmap_with_borders[1:heatmap_with_borders.shape[0] - 1, 1:heatmap_with_borders.shape[1] - 1]
    heatmap_left = heatmap_with_borders[1:heatmap_with_borders.shape[0] - 1, 2:heatmap_with_borders.shape[1]]
    heatmap_right = heatmap_with_borders[1:heatmap_with_borders.shape[0] - 1, 0:heatmap_with_borders.shape[1] - 2]
    heatmap_up = heatmap_with_borders[2:heatmap_with_borders.shape[0], 1:heatmap_with_borders.shape[1] - 1]
    heatmap_down = heatmap_with_borders[0:heatmap_with_borders.shape[0] - 2, 1:heatmap_with_borders.shape[1] - 1]

    heatmap_peaks = (heatmap_center > heatmap_left) & \
                    (heatmap_center > heatmap_right) & \
                    (heatmap_center > heatmap_up) & \
                    (heatmap_center > heatmap_down)
    heatmap_peaks = heatmap_peaks[1:heatmap_center.shape[0] - 1, 1:heatmap_center.shape[1] - 1]
    keypoints = list(zip(np.nonzero(heatmap_peaks)[1], np.nonzero(heatmap_peaks)[0]))  # (w, h)
    keypoints = sorted(keypoints, key=itemgetter(0))

    suppressed = np.zeros(len(keypoints), np.uint8)
    keypoints_with_score_and_id = []
    keypoint_num = 0
    for i in range(len(keypoints)):
        if suppressed[i]:
            continue
        for j in range(i + 1, len(keypoints)):
            if math.sqrt((keypoints[i][0] - keypoints[j][0]) ** 2 +
                         (keypoints[i][1] - keypoints[j][1]) ** 2) < 6:

                suppressed[j] = 1
        keypoint_with_score_and_id = (keypoints[i][0], keypoints[i][1], heatmap[keypoints[i][1], keypoints[i][0]],
                                      total_keypoint_num + keypoint_num)
        keypoints_with_score_and_id.append(keypoint_with_score_and_id)
        keypoint_num += 1
    all_keypoints.append(keypoints_with_score_and_id)
    return keypoint_num

def group_keypoints(all_keypoints_by_type, pafs, PoseClass, min_paf_score=0.05, demo=False):

    BODY_PARTS_KPT_IDS = PoseClass.BODY_PARTS_KPT_IDS
    BODY_PARTS_PAF_IDS = PoseClass.BODY_PARTS_PAF_IDS
    pose_entry_size = PoseClass.pose_entry_size

    pose_entries = []
    all_keypoints = np.array([item for sublist in all_keypoints_by_type for item in sublist])
    for part_id in range(len(BODY_PARTS_PAF_IDS)):
        part_pafs = pafs[:, :, BODY_PARTS_PAF_IDS[part_id]]
        kpt_a_id = BODY_PARTS_KPT_IDS[part_id][0]
        kpt_b_id = BODY_PARTS_KPT_IDS[part_id][1]
        kpts_a = all_keypoints_by_type[BODY_PARTS_KPT_IDS[part_id][0]]  # [(x, y, s, part_id), ..., (x, y, s, part_id)]
        kpts_b = all_keypoints_by_type[BODY_PARTS_KPT_IDS[part_id][1]]  # [(x, y, s, part_id), ..., (x, y, s, part_id)]
        num_kpts_a = len(kpts_a)
        num_kpts_b = len(kpts_b)

        # ipdb.set_trace()

        if num_kpts_a == 0 and num_kpts_b == 0:  # no keypoints for such body part
            continue
        elif num_kpts_a == 0:  # body part has just 'b' keypoints
            for i in range(num_kpts_b):
                num = 0
                for j in range(len(pose_entries)):  # check if already in some pose, was added by another body part
                    if pose_entries[j][kpt_b_id] == kpts_b[i][3]:
                        num += 1
                        continue
                if num == 0:
                    pose_entry = np.ones(pose_entry_size) * -1
                    pose_entry[kpt_b_id] = kpts_b[i][3]  # keypoint idx
                    pose_entry[-1] = 1  # num keypoints in pose
                    pose_entry[-2] = kpts_b[i][2]  # pose score
                    pose_entries.append(pose_entry)
            continue
        elif num_kpts_b == 0:  # body part has just 'a' keypoints
            for i in range(num_kpts_a):
                num = 0
                for j in range(len(pose_entries)):
                    if pose_entries[j][kpt_a_id] == kpts_a[i][3]:
                        num += 1
                        continue
                if num == 0:
                    pose_entry = np.ones(pose_entry_size) * -1
                    pose_entry[kpt_a_id] = kpts_a[i][3]
                    pose_entry[-1] = 1
                    pose_entry[-2] = kpts_a[i][2]
                    pose_entries.append(pose_entry)
            continue

        # ipdb.set_trace()

        connections = []
        for i in range(num_kpts_a):
            kpt_a = np.array(kpts_a[i][0:2])
            for j in range(num_kpts_b):
                kpt_b = np.array(kpts_b[j][0:2])
                mid_point = [(), ()]
                mid_point[0] = (int(round((kpt_a[0] + kpt_b[0]) * 0.5)),
                                int(round((kpt_a[1] + kpt_b[1]) * 0.5)))
                mid_point[1] = mid_point[0]

                vec = [kpt_b[0] - kpt_a[0], kpt_b[1] - kpt_a[1]]
                vec_norm = math.sqrt(vec[0] ** 2 + vec[1] ** 2)
                if vec_norm == 0:
                    continue
                vec[0] /= vec_norm
                vec[1] /= vec_norm
                cur_point_score = (vec[0] * part_pafs[mid_point[0][1], mid_point[0][0], 0] +
                                   vec[1] * part_pafs[mid_point[1][1], mid_point[1][0], 1])

                height_n = pafs.shape[0] // 2
                success_ratio = 0
                point_num = 10  # number of points to integration over paf
                if cur_point_score > -100:
                    passed_point_score = 0
                    passed_point_num = 0
                    x, y = linspace2d(kpt_a, kpt_b)
                    for point_idx in range(point_num):
                        if not demo:
                            px = int(round(x[point_idx]))
                            py = int(round(y[point_idx]))
                        else:
                            px = int(x[point_idx])
                            py = int(y[point_idx])
                        paf = part_pafs[py, px, 0:2]
                        cur_point_score = vec[0] * paf[0] + vec[1] * paf[1]
                        if cur_point_score > min_paf_score:
                            passed_point_score += cur_point_score
                            passed_point_num += 1
                    success_ratio = passed_point_num / point_num
                    ratio = 0
                    if passed_point_num > 0:
                        ratio = passed_point_score / passed_point_num
                    ratio += min(height_n / vec_norm - 1, 0)
                if ratio > 0 and success_ratio > 0.8:
                    score_all = ratio + kpts_a[i][2] + kpts_b[j][2]
                    connections.append([i, j, ratio, score_all])

        # ipdb.set_trace()

        if len(connections) > 0:
            connections = sorted(connections, key=itemgetter(2), reverse=True)

        num_connections = min(num_kpts_a, num_kpts_b)
        has_kpt_a = np.zeros(num_kpts_a, dtype=np.int32)
        has_kpt_b = np.zeros(num_kpts_b, dtype=np.int32)
        filtered_connections = []
        for row in range(len(connections)):
            if len(filtered_connections) == num_connections:
                break
            i, j, cur_point_score = connections[row][0:3]
            if not has_kpt_a[i] and not has_kpt_b[j]:
                filtered_connections.append([kpts_a[i][3], kpts_b[j][3], cur_point_score])
                has_kpt_a[i] = 1
                has_kpt_b[j] = 1
        connections = filtered_connections
        if len(connections) == 0:
            continue

        # ipdb.set_trace()

        if part_id == 0:
            pose_entries = [np.ones(pose_entry_size) * -1 for _ in range(len(connections))]
            for i in range(len(connections)):
                pose_entries[i][BODY_PARTS_KPT_IDS[0][0]] = connections[i][0]
                pose_entries[i][BODY_PARTS_KPT_IDS[0][1]] = connections[i][1]
                pose_entries[i][-1] = 2
                pose_entries[i][-2] = np.sum(all_keypoints[connections[i][0:2], 2]) + connections[i][2]
        # elif part_id == 17 or part_id == 18:
        elif part_id == 18 or part_id == 19:
            kpt_a_id = BODY_PARTS_KPT_IDS[part_id][0]
            kpt_b_id = BODY_PARTS_KPT_IDS[part_id][1]
            for i in range(len(connections)):
                for j in range(len(pose_entries)):
                    if pose_entries[j][kpt_a_id] == connections[i][0] and pose_entries[j][kpt_b_id] == -1:
                        pose_entries[j][kpt_b_id] = connections[i][1]
                    elif pose_entries[j][kpt_b_id] == connections[i][1] and pose_entries[j][kpt_a_id] == -1:
                        pose_entries[j][kpt_a_id] = connections[i][0]
            continue
        else:
            kpt_a_id = BODY_PARTS_KPT_IDS[part_id][0]
            kpt_b_id = BODY_PARTS_KPT_IDS[part_id][1]
            for i in range(len(connections)):
                num = 0
                for j in range(len(pose_entries)):
                    if pose_entries[j][kpt_a_id] == connections[i][0]:
                        pose_entries[j][kpt_b_id] = connections[i][1]
                        num += 1
                        pose_entries[j][-1] += 1
                        pose_entries[j][-2] += all_keypoints[connections[i][1], 2] + connections[i][2]
                if num == 0:
                    pose_entry = np.ones(pose_entry_size) * -1
                    pose_entry[kpt_a_id] = connections[i][0]
                    pose_entry[kpt_b_id] = connections[i][1]
                    pose_entry[-1] = 2
                    pose_entry[-2] = np.sum(all_keypoints[connections[i][0:2], 2]) + connections[i][2]
                    pose_entries.append(pose_entry)

    filtered_entries = []
    for i in range(len(pose_entries)):
        if pose_entries[i][-1] < 3 or (pose_entries[i][-2] / pose_entries[i][-1] < 0.2):
            continue
        filtered_entries.append(pose_entries[i])
    pose_entries = np.asarray(filtered_entries)
    return pose_entries, all_keypoints

def infer_fast_post_process(net_outputs, PoseClass):
    """

    Args:
        net_outputs (dict): the output of the networks, and it contains,
            --heatmaps:
            --pafs:

        PoseClass (type of tools.human_pose2d_estimators.utils.pose_utils.OpenPoseBody25):

    Returns:
        outputs (dict): the output results, and it contains the followings keys,
            --pose_entries:
            --all_keypoints:
            --current_poses:
    """

    heatmaps = net_outputs["heatmaps"]
    pafs = net_outputs["pafs"]
    pad = net_outputs["pad"]
    scale = net_outputs["scale"]
    stride = net_outputs["stride"]
    upsample_ratio = net_outputs["upsample_ratio"]
    height, width = net_outputs["orig_shape"]

    num_keypoints = PoseClass.num_kpts

    total_keypoints_num = 0
    all_keypoints_by_type = []
    for kpt_idx in range(num_keypoints):  # 19th for bg
        total_keypoints_num += extract_keypoints(heatmaps[:, :, kpt_idx], all_keypoints_by_type, total_keypoints_num)

    pose_entries, all_keypoints = group_keypoints(all_keypoints_by_type, pafs, PoseClass, demo=True)

    for kpt_id in range(all_keypoints.shape[0]):
        all_keypoints[kpt_id, 0] = (all_keypoints[kpt_id, 0] * stride / upsample_ratio - pad[1]) / scale
        all_keypoints[kpt_id, 1] = (all_keypoints[kpt_id, 1] * stride / upsample_ratio - pad[0]) / scale

    if len(all_keypoints):
        all_keypoints[:, 0] = np.clip(all_keypoints[:, 0], 0, width)
        all_keypoints[:, 1] = np.clip(all_keypoints[:, 1], 0, height)

    current_poses = []
    for n in range(len(pose_entries)):
        if len(pose_entries[n]) == 0:
            continue
        pose_keypoints = np.zeros((num_keypoints, 3), dtype=all_keypoints.dtype)
        for kpt_id in range(num_keypoints):
            kpt_num_id = int(pose_entries[n][kpt_id])
            if kpt_num_id != -1:  # keypoint was found
                pose_keypoints[kpt_id] = all_keypoints[kpt_num_id, 0:3]
            else:
                pose_keypoints[kpt_id, 0:2] = -1.0

        # print(n, pose_keypoints)
        pose = PoseClass(pose_keypoints, pose_entries[n][-2])
        current_poses.append(pose)

    outputs = {
        "pose_entries": pose_entries,
        "all_keypoints": all_keypoints,
        "current_poses": current_poses
    }

    return outputs

def infer_post_process(net_outputs, PoseClass):

    avg_heatmaps = net_outputs["heatmaps"]
    avg_pafs = net_outputs["pafs"]

    num_keypoints = PoseClass.num_kpts

    total_keypoints_num = 0
    all_keypoints_by_type = []
    for kpt_idx in range(num_keypoints):  # 19th for bg
        total_keypoints_num += extract_keypoints(avg_heatmaps[:, :, kpt_idx], all_keypoints_by_type,
                                                 total_keypoints_num)

    pose_entries, all_keypoints = group_keypoints(all_keypoints_by_type, avg_pafs, PoseClass, demo=True)

    current_poses = []
    for n in range(len(pose_entries)):
        if len(pose_entries[n]) == 0:
            continue
        pose_keypoints = np.zeros((num_keypoints, 3), dtype=all_keypoints.dtype)
        for kpt_id in range(num_keypoints):
            kpt_num_id = int(pose_entries[n][kpt_id])
            if kpt_num_id != -1:  # keypoint was found
                pose_keypoints[kpt_id] = all_keypoints[kpt_num_id, 0:3]
            else:
                pose_keypoints[kpt_id, 0:2] = -1.0

        # print(n, pose_keypoints)
        pose = PoseClass(pose_keypoints, pose_entries[n][-2])
        current_poses.append(pose)

    outputs = {
        "pose_entries": pose_entries,
        "all_keypoints": all_keypoints,
        "current_poses": current_poses
    }

    return outputs


# -*- coding: utf-8 -*-
"""
Created on Tue Jun  4 00:14:31 2019

@author: Administrator
"""

# -*- coding: utf-8 -*-
"""
Created on Wed Feb 13 03:50:30 2019

@author: Administrator
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as  optim
import torch.backends.cudnn as cudnn
from torch.utils.data import Dataset, DataLoader
from torch.autograd import Variable

import os
import sys
import csv
import numpy as np

from src.model import LinearModel, weight_init
from src.datasets.cmu_panoptic import CMU

from src import Bar
import src.utils as utils
import time
from opt import Options
from src.procrustes import get_transformation
import src.data_process as data_process
from src.poseVisualizer import visualizePose

#from create_data_norm import calcu_mean_std
torch.manual_seed(1)

def main(opt):
    start_epoch = 0
    err_best = 1000
    glob_step = 0
    lr_now = opt.lr

    # create model
    print(">>> creating model")
    model = LinearModel(joint_num=opt.joint_num)

    model=model.cuda()
    model.apply(weight_init)
    print(">>> total params: {:.2f}M".format(sum(p.numel() for p in model.parameters()) / 1000000.0))
    criterion = nn.MSELoss(reduction='elementwise_mean').cuda()
    optimizer = torch.optim.Adam(model.parameters(), lr=opt.lr)

   # load ckpt
    if opt.load:
        print(">>> loading ckpt from '{}'".format(opt.load))
        ckpt = torch.load(opt.load)
        start_epoch = ckpt['epoch']
        lr_now = ckpt['lr']
        glob_step = ckpt['step']
        err_best = ckpt['err']
        mean_pose = ckpt['mean_pose']
        model.load_state_dict(ckpt['state_dict'])
        optimizer.load_state_dict(ckpt['optimizer'])
        print(">>> ckpt loaded (epoch: {} | err: {})".format(start_epoch, err_best))

    # data loading
    print(">>> loading data")
    train_data = torch.load(opt.data_dir+'_train.pth')
    mean_pose = np.mean(train_data['tgt'], axis=0)
    mean_pose = np.reshape(mean_pose, (opt.joint_num, 3))
    test_data=CMU(data_path=opt.data_dir+'_test.pth',  use_hg=opt.use_hg)
    train_data=CMU(data_path=opt.data_dir+'_train.pth',  use_hg=opt.use_hg)

    test_loader = DataLoader(dataset=test_data, batch_size=opt.test_batch, shuffle=False,num_workers=opt.job)
    train_loader = DataLoader(dataset=train_data, batch_size=opt.train_batch, shuffle=True,num_workers=opt.job)
    print(">>> data loaded !")

    cudnn.benchmark = True
    for epoch in range(start_epoch, opt.epochs):
        print('==========================')
        print('>>> epoch: {} | lr: {:.5f}'.format(epoch, lr_now))

        # train
        glob_step, lr_now, loss_train = train(
            train_loader, model, criterion, optimizer, opt.joint_num,
            lr_init=opt.lr, lr_now=lr_now, glob_step=glob_step, lr_decay=opt.lr_decay, gamma=opt.lr_gamma,
            max_norm=opt.max_norm)
        print("loss_train:", loss_train)

        # test
        outputs_use, loss_test, err_test = test(test_loader, model, criterion, opt.joint_num, procrustes=opt.procrustes)

        # save best checkpoint
        is_best = err_test < err_best
        err_best = min(err_test, err_best)
        save_path = ''
        if is_best:
            print("Saved Check Point (error : {})".format(err_test))
            checkpoint = {'epoch': epoch,
                            'lr': lr_now,
                            'step': glob_step,
                            'err': err_best,
                            'state_dict': model.state_dict(),
                            'optimizer': optimizer.state_dict(),
                            'mean_pose': mean_pose}
            if save_path != '' and os.path.exists(save_path):
                os.remove(save_path)
            save_path = opt.ckpt+'_best.chkpt'
            torch.save(checkpoint, save_path)

        # write loss to log file
        log_train_file = opt.log + 'train.log'
        with open(log_train_file, 'a') as log_tr:
            log_tr.write('{},{},{},{},{}\n'.format(epoch, lr_now, loss_train, loss_test, err_test))

def train(train_loader, model, criterion, optimizer, joint_num,
          lr_init=None, lr_now=None, glob_step=None, lr_decay=None, gamma=None,
          max_norm=True):
    losses = utils.AverageMeter()

    model.train()

    start = time.time()
    batch_time = 0
    bar = Bar('>>>', fill='>', max=len(train_loader))

    for i, data in enumerate(train_loader):
        # Turn down Learning Rate
        glob_step += 1
        if glob_step % lr_decay == 0 or glob_step == 1:
            lr_now = utils.lr_decay(optimizer, glob_step, lr_init, lr_decay, gamma)

        joint2d, truth = data['joint2d'], data['truth']
        inputs=Variable(joint2d.cuda().type(torch.cuda.FloatTensor))
        targets=Variable(truth.cuda().type(torch.cuda.FloatTensor))

        outputs = model(inputs)
        outputs=torch.reshape(outputs,(-1,(joint_num)*3))
        targets=torch.reshape(targets,(-1,(joint_num)*3))

        # calculate loss
        optimizer.zero_grad()
        loss = criterion(outputs, targets)

        losses.update(loss.item(), inputs.size(0))
        loss.backward()

        if max_norm:
            nn.utils.clip_grad_norm_(model.parameters(), max_norm=1)
        optimizer.step()

        # update summary
        if (i + 1) % 100 == 0:
            batch_time = time.time() - start
            start = time.time()

        bar.suffix = '({batch}/{size}) | batch: {batchtime:.4}ms | Total: {ttl} | ETA: {eta:} | loss: {loss:.4f}' \
            .format(batch=i + 1,
                    size=len(train_loader),
                    batchtime=batch_time * 10.0,
                    ttl=bar.elapsed_td,
                    eta=bar.eta_td,
                    loss=losses.avg)
        bar.next()

    bar.finish()
    return glob_step, lr_now, losses.avg

def test(test_loader, model, criterion, joint_num, procrustes=False):
    losses = utils.AverageMeter()

    model.eval()

    all_dist = []
    start = time.time()
    batch_time = 0
    bar = Bar('>>>', fill='>', max=len(test_loader))

    for i, data in enumerate(test_loader):
        joint2d,truth=data['joint2d'],data['truth']

        inputs=Variable(joint2d.cuda().type(torch.cuda.FloatTensor))
        targets=Variable(truth.cuda().type(torch.cuda.FloatTensor))

        outputs = model(inputs)

        outputs=torch.reshape(outputs,(-1,(joint_num)*3))
        targets=torch.reshape(targets,(-1,(joint_num)*3))

        # calculate loss
        loss = criterion(outputs, targets)
        losses.update(loss.item(), inputs.size(0))

        sqerr = (outputs - targets) ** 2
        distance = np.zeros((sqerr.shape[0],joint_num+1))
        dist_idx = 0
        for k in np.arange(0, (joint_num+1) * 3, 3):
            distance[:, dist_idx] = torch.sqrt(torch.sum(sqerr[:, k:k + 3], axis=1)).to('cpu').detach().numpy()
            dist_idx += 1
        all_dist.append(distance)

        # update summary
        if (i + 1) % 100 == 0:
            batch_time = time.time() - start
            start = time.time()

        bar.suffix = '({batch}/{size}) | batch: {batchtime:.4}ms | Total: {ttl} | ETA: {eta:} | loss: {loss:.6f}' \
            .format(batch=i + 1,
                    size=len(test_loader),
                    batchtime=batch_time * 10.0,
                    ttl=bar.elapsed_td,
                    eta=bar.eta_td,
                    loss=losses.avg)
        bar.next()

    all_dist = np.vstack(all_dist)
#    joint_err = np.mean(all_dist, axis=0)
    ttl_err = np.mean(all_dist)
    bar.finish()
    print (">>> error: {} <<<".format(ttl_err))

    return targets, losses.avg, ttl_err

if __name__ == "__main__":

    option = Options().parse()
    main(option)