import unittest
import re
import sys
sys.path.append('../drippy')
from drippy import SprinklerGPIO

class OptionsInputTests(unittest.TestCase):

    def testGoodInput(self):
        goodInput = "I am searching for #waterMe"
        waterOptions = re.compile(r"#waterMe \d+|#waterMe")
        result = waterOptions.search(goodInput)
        self.assertEqual(result.group(), '#waterMe')

    def testGoodInputWithTime(self):
        goodInputWithTime = "I am searching for #waterMe 2"
        waterOptions = re.compile(r"#waterMe \d+|#waterMe")
        result = waterOptions.search(goodInputWithTime)
        self.assertEqual(result.group(), '#waterMe 2')

    def testBadInput(self):
        badInput = "Nothing here"
        waterOptions = re.compile(r"#waterMe \d+|#waterMe")
        result = waterOptions.search(badInput)
        self.assertEqual(result, None)

    def testFaultyTime(self):
        faultyTime = "#waterMe -2"
        waterOptions = re.compile(r"#waterMe \d+|#waterMe")
        result = waterOptions.search(faultyTime)
        self.assertEqual(result.group(), '#waterMe')

class GPIOTests(unittest.TestCase):

    def testNumberOfStations(self):
        dripController = SprinklerGPIO.SprinklerGPIO(1)
        self.assertEqual(dripController.numberOfStations, 1)

    def testOpenDrip(self):
        dripController = SprinklerGPIO.SprinklerGPIO(1)
        self.assertEqual(dripController.getStationStatus(0), 0)
        dripController.setStationStatus(0, 1)
        self.assertEqual(dripController.getStationStatus(0), 1)
        dripController.setStationStatus(0, 0)

    def testCloseDrip(self):
        dripController = SprinklerGPIO.SprinklerGPIO(1)
        self.assertEqual(dripController.getStationStatus(0), 0)
        dripController.setStationStatus(0, 1)
        dripController.setStationStatus(0, 0)
        self.assertEqual(dripController.getStationStatus(0), 0)


from django.shortcuts import render,redirect
from django.http import HttpResponse
from .models import UsersModal,Mobiles,Laptops,HeadSet,Camera,Powerbank,Refrigerator,Kettle,Television,WashingMachine
from http.cookies import SimpleCookie
from math import floor,ceil
from GRS.randkey import make_ratings
from django.core.paginator import Paginator

def GRSappHome(request):
    user = request.COOKIES.get('username')
    mobiles = Mobiles.objects.all()
    paginator1 = Paginator(mobiles,6)
    page1 = request.GET.get('page1')
    mobiles = paginator1.get_page(page1)

    laptops = Laptops.objects.all()
    paginator2 = Paginator(laptops,4)
    page2 = request.GET.get('page2')
    laptops = paginator2.get_page(page2)

    headsets = HeadSet.objects.all()
    paginator3 = Paginator(headsets,6)
    page3 = request.GET.get('page3')
    headsets = paginator3.get_page(page3)

    cameras = Camera.objects.all()
    paginator4 = Paginator(cameras,4)
    page4 = request.GET.get('page4')
    cameras = paginator4.get_page(page4)

    powerbanks = Powerbank.objects.all()
    paginator5 = Paginator(powerbanks,6)
    page5 = request.GET.get('page5')
    powerbanks = paginator5.get_page(page5)

    fridges = Refrigerator.objects.all()
    paginator6 = Paginator(fridges,6)
    page6 = request.GET.get('page6')
    fridges = paginator6.get_page(page6)

    kettles = Kettle.objects.all()
    paginator7 = Paginator(kettles,4)
    page7 = request.GET.get('page7')
    kettles = paginator7.get_page(page7)

    tvs = Television.objects.all()
    paginator8 = Paginator(tvs,8)
    page8 = request.GET.get('page8')
    tvs = paginator8.get_page(page8)

    wms = WashingMachine.objects.all()
    paginator9 = Paginator(wms,8)
    page9 = request.GET.get('page9')
    wms = paginator9.get_page(page9)

    params = {'user':user,'mobiles':mobiles,'laptops':laptops,'headsets':headsets,'cameras':cameras,'powerbanks':powerbanks,'fridges':fridges,'kettles':kettles,'tvs':tvs,'wms':wms}

    return render(request,'GRSapp/GRSHome.html',params)

def loaditem(request,slug):
    if slug=="Mobiles":
        return redirect('loadmobiles')
    else:
        return HttpResponse("Page Not Found")

def loadmobiles(request):
    user = request.COOKIES.get('username')
    mobiles = Mobiles.objects.all()
    mobiles = make_ratings(mobiles)

    params = {"user":user,"mobiles":mobiles}

    return render(request,'GRSapp/loadMobiles.html',params)


import vcs
import sys
import os

baseline = sys.argv[1]

pth = os.path.join(os.path.dirname(__file__), "..")
sys.path.append(pth)
import checkimage

canvas = vcs.init()
canvas.setantialiasing(0)
canvas.setbgoutputdimensions(792, 611, units="pixels")
canvas.drawlogooff()

fillarea = vcs.createfillarea()
fillarea.x = [[0, .33, .33, 0], [.33, .67, .67, .33], [.67, 1, 1, .67]]
fillarea.y = [[0, 0, 1, 1]] * 3
fillarea.style = ["solid", "pattern", "hatch"]
fillarea.index = [1, 5, 5]
fillarea.color = [50, 50, 50]

canvas.plot(fillarea, bg=1)

testImage = os.path.abspath("test_vcs_large_pattern_hatch.png")
canvas.png(testImage)

ret = checkimage.check_result_image(testImage, baseline,
                                    checkimage.defaultThreshold)

sys.exit(ret)


"""
This is a demo of how you can optimize machine learning algorithm results with
only linear programming tools like scipy.optimize
"""

import numpy as np
import pandas as pd
from scipy.optimize import minimize
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_absolute_error
from sklearn.model_selection import train_test_split

def mape(y, y_pred, weight=1):
    """Custom metric with bias possibilities. Weight parameter will multiply
    every prediction by it's value thus changing the mape error.

    Arguments:
        y {array} -- ground truth
        y_pred {array} -- prediction made by estimator

    Keyword Arguments:
        weight {int} -- custom bias for y_pred (default: {1})

    Returns:
        [float] -- mape error with added bias (default = 1)
    """
    mape = round(mean_absolute_error(y, y_pred * weight) / np.mean(y), 4)
    return mape

def objective(x):
    """Simple optimization function for sklearn.optimize

    Arguments:
        x {float} -- initial value to start optimization from

    Returns:
        float -- error minimized by sklearn.optimize function
    """
    return mean_absolute_error(y_train, y_pred_train * x)

# Let's load dataset and select only numerical data for simplicity sake
df = pd.read_csv('datasets/weatherHistory.csv')
num = df.select_dtypes(include='float64').columns
df = df[num]

# cut data between target value and features
y = df['Pressure (millibars)']
X = df.loc[:, df.columns != 'Pressure (millibars)']

# this requires no explanation
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.20, random_state=42)

# define model and predict train and test data
model = LinearRegression()
model.fit(X_train, y_train)
y_pred_train = model.predict(X_train)
y_pred_test = model.predict(X_test)

# let's further minimize error by calling minimize on our predictions
x0 = 1.0
sol = minimize(objective, x0, method='Nelder-Mead')
print('\n' + str(sol) + '\n')

# weight is the optimized value that will further reduce error on our prediction
# you have to multiply weight value by y_train to see how good it optimizes
# result. After that test it on test data
weight = sol.x[0]

print('mape on train', mape(y_train, y_pred_train))
print('mape on train scipy optimized', mape(y_train, y_pred_train, weight))
print('mape on test', mape(y_test, y_pred_test))
print('mape on test scipy optimized', mape(y_test, y_pred_test, weight))

# We can see that after optimizing our predictions error went down. This is done
# on toy dataset so reduction is small. Test it on real business problem to see
# better results


import cv2
import os

def get_all_images(folder_path):
    """
    loads all images from a given directory
    :param folder_path: path to folder containing loaded images, relative to main project folder
                        eg. 'images/lifeLines'
    :return: python list of images in opencv (numpy) format
    """
    images = []
    for filename in os.listdir(folder_path):
        image = cv2.imread(f'{folder_path}/{filename}')
        if image is not None:
            images.append(image)
    return images


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# BioSTEAM: The Biorefinery Simulation and Techno-Economic Analysis Modules
# Copyright (C) 2020-2021, Yoel Cortes-Pena <yoelcortes@gmail.com>
# Bioindustrial-Park: BioSTEAM's Premier Biorefinery Models and Results
# Copyright (C) 2021, Yalin Li <yalinli2@illinois.edu>
#
# This module is under the UIUC open-source license. See
# github.com/BioSTEAMDevelopmentGroup/biosteam/blob/master/LICENSE.txt
# for license details.

import math
import flexsolve as flx
import biosteam as bst

__all__ = ('SludgeHandling', 'BeltThickener', 'SludgeCentrifuge')

class SludgeHandling(bst.Unit):
    '''
    A generic class for handling of wastewater treatmnet sludge.

    The 0th outs is the water-rich supernatant (effluent) and
    the 1st outs is the solid-rich sludge.

    Two pumps (one for the supernatant and one for sludge) are included.

    Separation split is determined by the moisture (i.e., water) content
    of the sludge, soluble chemicals will have the same split as water,
    insolubles chemicals will all go to the retentate.

    Parameters
    ----------
    sludge_moisture : float
        Moisture content of the sludge, [wt% water].
    solubles : tuple
        IDs of the soluble chemicals.
        Note that all chemicals that are not included in this tuple and not
        locked as gas phase (i.e., `chemical.locked_state!='g'`) will be
        treated as solids in simulation.

    References
    ----------
    .. [1] Shoener et al., Design of Anaerobic Membrane Bioreactors for the
    Valorization of Dilute Organic Carbon Waste Streams.
    Energy Environ. Sci. 2016, 9 (3), 1102–1112.
    https://doi.org/10.1039/C5EE03715H.
    '''

    _graphics = bst.Splitter._graphics
    _ins_size_is_fixed = False
    _N_outs = 2
    auxiliary_unit_names = ('effluent_pump', 'sludge_pump')

    def __init__(self, ID='', ins=None, outs=(), thermo=None,
                 sludge_moisture=0.96, solubles=()):
        bst.Unit.__init__(self, ID, ins, outs, thermo)
        self.sludge_moisture = sludge_moisture
        self.solubles = tuple(solubles)
        self.solids = tuple(i.ID for i in self.chemicals
                            if (i.ID not in solubles) and (i.locked_state!='g'))
        self.effluent_pump = bst.Pump(f'{self.ID}_eff')
        self.sludge_pump = bst.Pump(f'{self.ID}_sludge')

    @staticmethod
    def _mc_at_split(split, solubles, mixed, eff, sludge, target_mc):
        eff.imass[solubles] = mixed.imass[solubles] * split
        sludge.imass[solubles] = mixed.imass[solubles] - eff.imass[solubles]
        mc = sludge.imass['Water'] / sludge.F_mass
        return mc-target_mc

    def _run(self):
        eff, sludge = self.outs
        solubles, solids = self.solubles, self.solids

        mixed = bst.Stream()
        mixed.mix_from(self.ins)
        eff.T = sludge.T = mixed.T
        self._mixed = mixed.copy()

        sludge.copy_flow(mixed, solids, remove=True) # all solids go to sludge
        eff.copy_flow(mixed, solubles)

        flx.IQ_interpolation(
                f=self._mc_at_split, x0=1e-3, x1=1.-1e-3,
                args=(solubles, mixed, eff, sludge, self.sludge_moisture),
                checkbounds=False)

    def _cost(self):
        pumps = (self.effluent_pump, self.sludge_pump)
        for i in range(2):
            pumps[i].ins[0] = self.outs[i].copy() # use `.proxy()` will interfere `_run`
            pumps[i].simulate()
            self.power_utility.rate += pumps[i].power_utility.rate

class BeltThickener(SludgeHandling):
    '''
    Gravity belt thickener (GBT) designed based on the manufacture specification
    data sheet. [1]_

    The 0th outs is the water-rich supernatant (effluent) and
    the 1st outs is the solid-rich sludge.

    Key parameters include:
        - Capacity: 80-100 m3/h.
        - Influent solids concentration: 0.2-1%.
        - Sludge cake moisture content: 90-96%.
        - Motor power: 3 (driving motor) and 1.1 agitator motor.
        - Belt width: 2.5 m.
        - Weight: 2350 kg.
        - Quote price: $3680 ea for three or more sets.

    The bare module (installation) factor is from Table 25 in Humbird et al. [2]_
    (solids handling equipment).

    Parameters
    ----------
    sludge_moisture : float
        Moisture content of the thickened sludge, [wt% water].
    solubles : tuple
        IDs of the soluble chemicals.
        Note that all chemicals that are not included in this tuple and not
        locked as gas phase (i.e., `chemical.locked_state!='g'`) will be
        treated as solids in simulation.
    max_capacity : float
        Maximum hydraulic loading per belt thickener, [m3/h].
    power_demand : float
        Total power demand of each belt thickener, [kW].

    References
    ----------
    .. [1] `Industrial filtering equipment gravity thickener rotary thickening belt filter press \
    <https://www.alibaba.com/product-detail/Industrial-filtering-equipment-gravity-thickener-rotary_60757627922.html?spm=a2700.galleryofferlist.normal_offer.d_title.78556be9t8szku>`_
    Data obtained on 7/21/2021.

    .. [2] Humbird et al., Process Design and Economics for Biochemical Conversion of
    Lignocellulosic Biomass to Ethanol: Dilute-Acid Pretreatment and Enzymatic
    Hydrolysis of Corn Stover; Technical Report NREL/TP-5100-47764;
    National Renewable Energy Lab (NREL), 2011.
    https://www.nrel.gov/docs/fy11osti/47764.pdf
    '''

    def __init__(self, ID='', ins=None, outs=(), thermo=None,
                 sludge_moisture=0.96, solubles=(),
                 max_capacity=100, power_demand=4.1):
        SludgeHandling.__init__(self, ID, ins, outs, thermo,
                                sludge_moisture=sludge_moisture,
                                solubles=solubles)
        self.max_capacity = max_capacity
        self.power_demand = power_demand

    def _design(self):
        self._N_thickener = N = math.ceil(self._mixed.F_vol/self.max_capacity)
        self.design_results['Number of thickners'] = N
        self.F_BM['Thickeners'] = 1.7 # ref [2]
        self.baseline_purchase_costs['Thickeners'] = 4000 * N
        self.power_utility.rate = self.power_demand * N

    @property
    def N_thickener(self):
        '''[int] Number of required belt thickeners.'''
        return self._N

class SludgeCentrifuge(SludgeHandling, bst.SolidsCentrifuge):
    '''
    Solid centrifuge for sludge dewatering.

    `_run` and `_cost` are based on `SludgeHandling` and `_design`
    is based on `SolidsCentrifuge`.

    The 0th outs is the water-rich supernatant (effluent) and
    the 1st outs is the solid-rich sludge.

    Parameters
    ----------
    sludge_moisture : float
        Moisture content of the thickened sludge, [wt% water].
    solubles : tuple
        IDs of the soluble chemicals.
        Note that all chemicals that are not included in this tuple and not
        locked as gas phase (i.e., `chemical.locked_state!='g'`) will be
        treated as solids in simulation.
    '''

    def __init__(self, ID='', ins=None, outs=(), thermo=None,
                 sludge_moisture=0.8, solubles=(),
                 centrifuge_type='scroll_solid_bowl'):
        SludgeHandling.__init__(self, ID, ins, outs, thermo,
                                sludge_moisture=sludge_moisture,
                                solubles=solubles)
        self.centrifuge_type = centrifuge_type

    _run = SludgeHandling._run

    _design = bst.SolidsCentrifuge._design

    _cost = SludgeHandling._cost


import ifish
import general_lib

# getting gps start
while True:
    print("hi")
    # log_result = ifish.send_log()
    loc = ifish.get_location()
    if loc == False:
        command = {"father_process":"gps","data":False}
        general_lib.epaper_pipe_write(command)
    # if log_result == True:
    while loc == False:
        loc = ifish.get_location()
    command={"father_process":"gps","data":loc}
    general_lib.pipe_write(command)


# Copyright 2020 Lorna Authors. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import torch
import torch.nn as nn

from .utils import BasicBlock
from .utils import Bottleneck
from .utils import conv1x1
from .utils import get_model_params
from .utils import load_pretrained_weights
from .utils import resnet_params

class ResNet(nn.Module):

    def __init__(self, layers=None, global_params=None):
        super(ResNet, self).__init__()
        assert isinstance(layers, tuple), "blocks_args should be a tuple"
        assert len(layers) > 0, "layers must be greater than 0"

        if global_params.norm_layer is None:
            norm_layer = nn.BatchNorm2d
        self._norm_layer = norm_layer

        self.inplanes = 64
        self.dilation = 1
        if global_params.replace_stride_with_dilation is None:
            # each element in the tuple indicates if we should replace
            # the 2x2 stride with a dilated convolution instead
            replace_stride_with_dilation = [False, False, False]
        if len(replace_stride_with_dilation) != 3:
            raise ValueError("replace_stride_with_dilation should be None "
                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
        self.groups = global_params.groups
        self.base_width = global_params.width_per_group
        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.bn1 = norm_layer(self.inplanes)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(global_params.block, 64, layers[0])
        self.layer2 = self._make_layer(global_params.block, 128, layers[1], stride=2,
                                       dilate=replace_stride_with_dilation[0])
        self.layer3 = self._make_layer(global_params.block, 256, layers[2], stride=2,
                                       dilate=replace_stride_with_dilation[1])
        self.layer4 = self._make_layer(global_params.block, 512, layers[3], stride=2,
                                       dilate=replace_stride_with_dilation[2])
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512 * global_params.block.expansion, global_params.num_classes)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

        # Zero-initialize the last BN in each residual branch,
        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
        if global_params.zero_init_residual:
            for m in self.modules():
                if isinstance(m, Bottleneck):
                    nn.init.constant_(m.bn3.weight, 0)
                elif isinstance(m, BasicBlock):
                    nn.init.constant_(m.bn2.weight, 0)

    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
        norm_layer = self._norm_layer
        downsample = None
        previous_dilation = self.dilation
        if dilate:
            self.dilation *= stride
            stride = 1
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                conv1x1(self.inplanes, planes * block.expansion, stride),
                norm_layer(planes * block.expansion),
            )

        layers = [block(self.inplanes, planes, stride, downsample, self.groups,
                        self.base_width, previous_dilation, norm_layer)]
        self.inplanes = planes * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.inplanes, planes, groups=self.groups,
                                base_width=self.base_width, dilation=self.dilation,
                                norm_layer=norm_layer))

        return nn.Sequential(*layers)

    def extract_features(self, inputs):
        """ Returns output of the final convolution layer """
        x = self.conv1(inputs)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        return x

    def forward(self, inputs):
        x = self.conv1(inputs)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)

        return x

    @classmethod
    def from_name(cls, model_name, override_params=None):
        cls._check_model_name_is_valid(model_name)
        layers, global_params = get_model_params(model_name, override_params)
        return cls(layers, global_params)

    @classmethod
    def from_pretrained(cls, model_name, num_classes=1000):
        model = cls.from_name(model_name, override_params={"num_classes": num_classes})
        load_pretrained_weights(model, model_name, load_fc=(num_classes == 1000))
        return model

    @classmethod
    def get_image_size(cls, model_name):
        cls._check_model_name_is_valid(model_name)
        _, res = resnet_params(model_name)
        return res

    @classmethod
    def _check_model_name_is_valid(cls, model_name):
        """ Validates model name. None that pretrained weights are only available for
        the first four models (resnet_pytorch{i} for i in 18,34,50,101,152) at the moment. """
        num_models = [18, 34, 50, 101, 152]
        valid_models = ["resnet" + str(i) for i in num_models]
        if model_name not in valid_models:
            raise ValueError("model_name should be one of: " + ", ".join(valid_models))


from datetime import datetime, timedelta
from server import db

class Sensor(db.Model):
    """Sensor model, DB table."""
    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String(16), unique=True, nullable=False)
    description = db.Column(db.Text, nullable=True)
    created_at = db.Column(db.DateTime, nullable=False)

    def __init__(self, name, description=None, created_at=None):
        self.name = name
        self.description = description
        if not created_at:
            created_at = datetime.utcnow()
        self.created_at = created_at

    def __repr__(self):
        return '<Sensor id={} name={} created_at={}>'.format(
            self.id, self.name, self.created_at
        )

    @property
    def serialize(self):
        """Simplify object -> JSON conversion."""
        return {
            'id': self.id,
            'name': self.name,
            'description': self.description,
            'created_at': self.created_at.isoformat()
        }

class Reading(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    value = db.Column(db.Float, nullable=False)
    timestamp = db.Column(db.DateTime, nullable=False)

    sensor_id = db.Column(db.Integer, db.ForeignKey('sensor.id'), nullable=False)
    sensor = db.relationship('Sensor',
        backref=db.backref('readings', lazy='dynamic', cascade='all, delete-orphan')
    )

    def __init__(self, sensor_name, value, timestamp=None):
        sensor = Sensor.query.filter_by(name=sensor_name).first()
        if sensor is None:
            sensor = Sensor(sensor_name)
            db.session.add(sensor)
            db.session.commit()

        self.sensor = sensor
        self.value = value
        if not timestamp:
            timestamp = datetime.utcnow()
        self.timestamp = timestamp

    def __repr__(self):
        return '<Data sensor={} value={} timestamp={}>'.format(
            self.sensor_id, self.value, self.timestamp
        )

    @property
    def serialize(self):
        """Simplify object -> JSON conversion."""
        return {
            'id': self.id,
            'value': self.value,
            'timestamp': self.timestamp.isoformat(),
            'sensor': self.sensor_id
        }


#
#    Copyright (C) 2017 Kenneth A. Giusti
#
#    Licensed to the Apache Software Foundation (ASF) under one
#    or more contributor license agreements.  See the NOTICE file
#    distributed with this work for additional information
#    regarding copyright ownership.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.
#

import abc
import json
import logging
import math
import os
import socket
import threading
import time
from time import time as now
try:
    import Queue as queue
except ImportError:
    import queue
import oslo_messaging as om
import uuid

__all__ = [
    "RPCTestClient",
    "RPCTestServer",
    "TestNotifier",
    "TestListener",
    "Controller"
]

# the types of oslo.messaging clients
RPC_CLIENT = 'RPCClient'
RPC_SERVER = 'RPCServer'
LISTENER = 'Listener'
NOTIFIER = 'Notifier'
MESSAGING_CLIENT_TYPES = [RPC_CLIENT, RPC_SERVER, LISTENER, NOTIFIER]

# addressing for control messages
CONTROL_EXCHANGE = 'ombt-control'
CONTROLLER_TOPIC = 'controller-%s'
CLIENT_TOPIC = "client-%s-%s"    # client-$type-$topic

# addressing for RPC tests
RPC_EXCHANGE = 'ombt-rpc-test'
RPC_TOPIC = "rpc-%s"

# addressing for Notification tests
NOTIFY_EXCHANGE = 'ombt-notify-test'
NOTIFY_TOPIC = "notify-%s"

def _wait_stabilize(delay, count_fn):
    # helper to wait until a counter has stabilized for delay seconds
    count = count_fn()
    timeout = delay
    while timeout > 0:
        time.sleep(1.0)
        new_count = count_fn()
        if count != new_count:
            # reset
            count = new_count
            timeout = delay
        else:
            timeout -= 1

class Stats(object):
    """Manage a single statistic"""
    def __init__(self, min=None, max=None, total=0, count=0,
                 sum_of_squares=0, distribution=None):
        self.min = min
        self.max = max
        self.total = total
        self.count = count
        self.sum_of_squares = sum_of_squares
        # distribution of values grouped by powers of 10
        self.distribution = distribution or dict()

    @classmethod
    def from_dict(cls, values):
        if 'distribution' in values:
            # hack alert!
            # when a Stats is passed via an RPC call it appears as if the
            # distribution map's keys are converted from int to str.
            # Fix that by re-indexing the distribution map:
            new_dict = dict()
            old_dict = values['distribution']
            for k in old_dict.keys():
                new_dict[int(k)] = old_dict[k]
            values['distribution'] = new_dict
        return Stats(**values)

    def to_dict(self):
        new_dict = dict()
        for a in ["min", "max", "total", "count", "sum_of_squares"]:
            new_dict[a] = getattr(self, a)
        new_dict["distribution"] = self.distribution.copy()
        return new_dict

    def update(self, value):
        self.total += value
        self.count += 1
        self.sum_of_squares += value**2
        self.min = min(self.min, value) if self.min else value
        self.max = max(self.max, value) if self.max else value
        log = int(math.log10(value)) if value >= 1.0 else 0
        base = 10**log
        index = int(value / base)  # 0..9
        if log not in self.distribution:
            self.distribution[log] = [0 for i in range(10)]
        self.distribution[log][index] += 1

    def reset(self):
        self.__init__()

    def average(self):
        return (self.total / float(self.count)) if self.count else 0

    def std_deviation(self):
        return math.sqrt((self.sum_of_squares / float(self.count)) -
                         (self.average() ** 2)) if self.count else -1

    def merge(self, stats):
        if stats.min is not None and self.min is not None:
            self.min = min(self.min, stats.min)
        else:
            self.min = self.min or stats.min
        if stats.max is not None and self.max is not None:
            self.max = max(self.max, stats.max)
        else:
            self.max = self.max or stats.max

        self.total += stats.total
        self.count += stats.count
        self.sum_of_squares += stats.sum_of_squares
        for k in stats.distribution.keys():
            if k in self.distribution:
                self.distribution[k] = [z for z in map(lambda a, b: a + b,
                                                       stats.distribution[k],
                                                       self.distribution[k])]
            else:
                self.distribution[k] = stats.distribution[k]

    def __str__(self):
        return "min=%i, max=%i, avg=%f, std-dev=%f" % (self.min, self.max,
                                                       self.average(),
                                                       self.std_deviation())

    def print_distribution(self):
        keys = list(self.distribution.keys())
        keys.sort()
        for order in keys:
            row = self.distribution[order]
            # order=0, index=0 is special case as it is < 1.0, for all orders >
            # 0, index 0 is ignored since everthing < 10^order is accounted for
            # in index 9 of the (order - 1) row
            index = 0 if order == 0 else 1
            while index < len(row):
                print("[%d..<%d):  %d" %
                      ((10 ** int(order)) * index,
                       (10 ** int(order)) * (index + 1),
                       row[index]))
                index += 1

class NullOutputter(object):
    """ Output handler used if no output is desired
    """

    def write(self, msg):
        pass

class FileOutputter(object):
    """Output handler used for sending output to a file
    """

    def __init__(self, filepath):
        self._fobj = open(filepath, 'w', -1)

    def write(self, msg):
        self._fobj.write(msg)

class TestResults(object):
    """Client results of a test run.
    """

    def __init__(self, start_time=None, stop_time=None, latency=None,
                 msgs_ok=0, msgs_fail=0, errors=None):
        super(TestResults, self).__init__()
        self.start_time = start_time
        self.stop_time = stop_time
        self.latency = latency or Stats()
        self.msgs_ok = msgs_ok  # count of successful msg transfers
        self.msgs_fail = msgs_fail  # count of failed msg transfers
        self.errors = errors or dict()  # error msgs and counts

    @classmethod
    def from_dict(cls, values):
        if 'latency' in values:
            values['latency'] = Stats.from_dict(values['latency'])
        if 'errors' in values:
            values['errors'] = values['errors'].copy()
        return TestResults(**values)

    def to_dict(self):
        new_dict = dict()
        for a in ['start_time', 'stop_time', 'msgs_ok', 'msgs_fail']:
            new_dict[a] = getattr(self, a)
        new_dict['latency'] = self.latency.to_dict()
        new_dict['errors'] = self.errors.copy()
        return new_dict

    def error(self, reason):
        key = str(reason)
        self.errors[key] = self.errors.get(key, 0) + 1

    def reset(self):
        self.__init__()

    def merge(self, results):
        self.start_time = (min(self.start_time, results.start_time)
                           if self.start_time and results.start_time
                           else (self.start_time or results.start_time))
        self.stop_time = (max(self.stop_time, results.stop_time)
                          if self.stop_time and results.stop_time
                          else (self.stop_time or results.stop_time))
        self.msgs_ok += results.msgs_ok
        self.msgs_fail += results.msgs_fail
        self.latency.merge(results.latency)
        for err in results.errors:
            self.errors[err] = self.errors.get(err, 0) + results.errors[err]

    def print_results(self):
        if self.msgs_fail:
            print("Error: %d message transfers failed"
                  % self.msgs_fail)
        if self.errors:
            print("Error: errors detected:")
            for err in self.errors:
                print("  '%s' (occurred %d times)" % (err, self.errors[err]))

        total = self.msgs_ok + self.msgs_fail
        print("Total Messages: %d" % total)

        delta_time = self.stop_time - self.start_time
        print("Test Interval: %f - %f (%f secs)" % (self.start_time,
                                                    self.stop_time,
                                                    delta_time))

        if delta_time > 0.0:
            print("Aggregate throughput: %f msgs/sec" %
                  (float(total) / delta_time))

        latency = self.latency
        if latency.count:
            print("Latency %d samples (msecs): Average %f StdDev %f"
                  " Min %f Max %f"
                  % (latency.count,
                     latency.average(), latency.std_deviation(),
                     latency.min, latency.max))
            print("Latency Distribution: ")
            latency.print_distribution()

class _Base(object):
    """Common base for all ombt2 processes.  Establishes a connection to the
    control message bus and a subscription for control messages
    """

    def __init__(self, cfg, ctl_url, topic, output, name, unique=False,
                 kind=None, timeout=None):
        super(_Base, self).__init__()
        self._finished = threading.Event()

        self._timeout = timeout
        if kind is None:
            ctl_topic = CONTROLLER_TOPIC % (topic
                                            if not unique else 'singleton')
            self.kind = "Controller"
        else:
            ctl_topic = CLIENT_TOPIC % ((kind, topic)
                                        if not unique else (kind, 'singleton'))
            self.kind = kind

        self.name = name or 'ombt-%s-%s-%s-%s-%s' % (topic,
                                                     kind,
                                                     socket.gethostname(),
                                                     os.getpid(),
                                                     uuid.uuid4().hex)
        self.ctl_url = ctl_url
        self.ctl_tport = om.get_rpc_transport(cfg.CONF,
                                              url=ctl_url)
        # My address and subscription for receiving control commands/responses
        self.ctl_target = om.Target(exchange=CONTROL_EXCHANGE,
                                    topic=ctl_topic,
                                    server=self.name)
        self._ctl_server = om.get_rpc_server(self.ctl_tport,
                                             target=self.ctl_target,
                                             endpoints=[self],
                                             executor="threading")
        self._ctl_server.start()

        if output is not None:
            try:
                self._output = FileOutputter(output)
            except Exception as exc:
                logging.error("Cannot open output file %s: %s!",
                              output, str(exc))
                self._output = NullOutputter()
        else:
            self._output = NullOutputter()

    def start(self):
        # blocks until connection to the control bus is active
        ready = False
        attempts = 0
        logging.debug("%s connecting to the control message bus...", self.name)
        # call my "self_ready" method until it returns successfully.
        # this indicates the connection to the control bus is active.
        client = om.RPCClient(self.ctl_tport,
                              target=self.ctl_target,
                              timeout=2)
        while not ready and attempts < 25:
            try:
                ready = client.call({}, 'self_ready')
            except om.MessagingTimeout:
                attempts += 1
        if not ready:
            raise Exception("Unable to contact message bus")
        logging.debug("%s is listening", self.name)

    def wait(self, timeout=None):
        # blocks until client completes shutdown
        return self._finished.wait(timeout)

    def _do_shutdown(self):
        self._ctl_server.stop()
        self._ctl_server.wait()
        self._finished.set()
        logging.debug("%s has shut down", self.name)

    #
    # RPC calls:
    #

    def shutdown(self, ctxt):
        # cannot synchronously shutdown server since this call is dispatched by
        # the server...
        threading.Thread(target=self._do_shutdown).start()

    def self_ready(self, ctxt):
        # simple ping to determine when message bus is connected
        return True

class _Client(_Base):
    """Common base for non-controller clients.  Defines RPC calls that are
    invoked by the Controller to control the tests.
    """
    def __init__(self, cfg, ctl_url, topic, kind, timeout, name,
                 unique=False, output=None):
        # listen on 'client-$topic' for controller commands:
        super(_Client, self).__init__(cfg, ctl_url, topic, output, name,
                                      unique, kind, timeout)
        self.topic = CLIENT_TOPIC % ((kind, topic)
                                     if not unique else (kind, 'singleton'))
        self.results = TestResults()
        self._results_lock = threading.Lock()
        self._output.write('{"client-name": "%(client)s",'
                           ' "kind": "%(kind)s"}\n'
                           % {'client': self.name,
                              'kind': kind})

    #
    # RPC Calls
    #

    def client_ping(self, ctxt, reply_addr):
        # invoked by controller via rpc-cast to roll-call available clients
        logging.debug("Client ping received (%s)", self.name)
        target = om.Target(**reply_addr)
        ctrl = om.RPCClient(self.ctl_tport, target=target,
                            timeout=self._timeout)
        try:
            ctrl.call({}, "client_pong", kind=self.kind, name=self.name)
        except Exception as exc:
            err = str(exc)
            logging.error("client pong call failed: %s", err)
            self.error(err)
        else:
            logging.debug("Client pong sent (%s) (%s)", self.name, target)

class _TestClient(_Client):
    """Base class for Notifier and RPC clients
    """
    def __init__(self, cfg, ctl_url, topic, kind, name, timeout, unique=False,
                 output=None):
        super(_TestClient, self).__init__(cfg, ctl_url, topic, kind, timeout,
                                          name, unique, output)

    # helper to execute func(timestamp) count times, pausing pause seconds
    # between invocation.  Provides extra logging if verbose
    def _execute(self, func, count=0, pause=0, verbose=False):
        stop = False
        msgid = uuid.uuid4().hex
        seq = 0
        self.results.start_time = now()
        while not stop:
            err = None
            ts = now()
            try:
                func(ts, "%s:%d" % (msgid, seq))
            except Exception as exc:
                self.results.msgs_fail += 1
                err = str(exc)
                self.results.error(err)
                logging.error("Test client failed to send message: %s", err)
            else:
                self.results.msgs_ok += 1
            done = now()
            self.results.latency.update((done - ts) * 1000)
            self._output.write('{"id": "%(msgid)s", "start": %(start)f,'
                               ' "stop": %(stop)f%(error)s}\n'
                               % {'msgid': "%s:%d" % (msgid, seq),
                                  'start': ts,
                                  'stop': done,
                                  'error':
                                  (', "error": "%s"' % err) if err else ""
                                  }
                               )
            seq += 1
            if pause:
                time.sleep(pause)
            if count and self.results.latency.count >= count:
                stop = True
        self.results.stop_time = now()

    #
    # RPC Calls
    #

    @abc.abstractmethod
    def run_test(self, ctxt, test, kwargs, reply_addr):
        """Called by the controller to have the client run test 'test' with
        arguments kwargs. When the test completes the client sends the results
        to the controller at 'reply_addr' by calling its 'client_result'
        method.  Note: this is an RPC method that is invoked by the test
        controller via a fanout 'cast' - not 'call' (the controller does not
        block for results)
        """

class _TestServer(_Client):
    """Base class for Listener and RPC servers
    """
    def __init__(self, cfg, ctl_url, topic, kind, name, timeout, unique=False,
                 output=None):
        super(_TestServer, self).__init__(cfg, ctl_url, topic, kind, timeout,
                                          name, unique, output)

    #
    # Controller RPC Calls
    #
    def get_server_results(self, ctxt, reply_addr):
        """Called by the controller to gather server side test data.  May be
        called repeatedly until the test completes.  Note this is invoked by
        the controller via a fanout 'cast' - not a 'call' (the controller does
        not block for results)
        """
        with self._results_lock:
            results = self.results.to_dict()
            self.results.reset()

        controller = om.RPCClient(self.ctl_tport,
                                  om.Target(**reply_addr),
                                  timeout=self._timeout)
        try:
            controller.call({}, 'client_result', name=self.name,
                            kind=self.kind, results=results)
        except Exception as exc:
            # I don't think recovery is possible as the call may be in-doubt.
            # For now simply let folks know the results may be invalid
            err = str(exc)
            logging.error("%s failed to report results!"
                          " Test results may be invalid!"
                          " Error: %s", err)
        else:
            logging.debug("Server %s test results sent", self.name)

class RPCTestClient(_TestClient):
    """Runs the RPC tests against the RPCTestServer
    """
    def __init__(self, cfg, ctl_url, test_url, topic, name, timeout,
                 unique=False, output=None):
        super(RPCTestClient, self).__init__(cfg, ctl_url, topic, RPC_CLIENT,
                                            name, timeout, unique, output)
        # for calling the test RPC server(s):
        target = om.Target(exchange=RPC_EXCHANGE,
                           topic=RPC_TOPIC % topic)
        fanout_target = om.Target(exchange=RPC_EXCHANGE,
                                  topic=RPC_TOPIC % topic,
                                  fanout=True)
        tport = (self.ctl_tport
                 if test_url == self.ctl_url
                 else om.get_rpc_transport(cfg.CONF, url=test_url))
        self._rpc_client = om.RPCClient(tport,
                                        target=target,
                                        timeout=timeout)
        self._fanout_client = om.RPCClient(tport,
                                           target=fanout_target,
                                           timeout=timeout)
    #
    # RPC Calls:
    #

    def run_test(self, ctxt, test, kwargs, reply_addr):
        func = None
        verbose = kwargs.get("verbose", False)
        pause = kwargs.get("pause", 0)
        data = kwargs.get("data", "")
        count = kwargs.get("count", 0)

        if test == "test_call":
            func = lambda ts, msgid: self._rpc_client.call({}, 'echo',
                                                           data=data,
                                                           timestamp=ts,
                                                           msgid=msgid)
        elif test == "test_cast":
            func = lambda ts, msgid: self._rpc_client.cast({}, 'noop',
                                                           data=data,
                                                           timestamp=ts,
                                                           msgid=msgid)
        elif test == "test_fanout":
            func = lambda ts, msgid: self._fanout_client.cast({}, 'noop',
                                                              data=data,
                                                              timestamp=ts,
                                                              msgid=msgid)
        else:
            logging.error("Client %s ignoring unknown test %s",
                          self.name, test)
            return

        self._output.write('{"test-name": "%(name)s",'
                           ' "test-start": %(start)f}\n'
                           % {'name': test, 'start': now()})

        controller = om.RPCClient(self.ctl_tport,
                                  om.Target(**reply_addr),
                                  timeout=self._timeout)
        logging.debug("Client %s starting test %s ...", self.name, test)

        # Before running the test, try to ping the server.  This will force a
        # link setup so the first latency-timed message will not be blocked
        # waiting for the setup to complete.
        try:
            self._rpc_client.call({}, "self_ready")
        except Exception as exc:
            logging.warning("Client %s is unable to reach RPC server: %s",
                            self.name, str(exc))
            self.results.error(str(exc))
            # keep going, perhaps the test will not fail...

        self._execute(func, count, pause, verbose)

        self._output.write('{"test-end": %f}\n' % now())

        logging.debug("Client %s test %s finished, sending results...",
                      self.name, test)
        try:
            controller.call({}, 'client_result', name=self.name,
                            kind=self.kind, results=self.results.to_dict())
        except Exception as exc:
            # I don't think recovery is possible as the call may be in-doubt.
            # For now simply let folks know the results may be invalid
            logging.error("%s failed to report results!"
                          " Test results may be invalid!"
                          " Error: %s", RPC_CLIENT, str(exc))
        else:
            logging.debug("Client %s test %s results sent", self.name, test)
        self.results.reset()

class RPCTestServer(_TestServer):
    """Response to RPC requests from RPCTestClient
    """
    def __init__(self, cfg, ctl_url, test_url, topic, executor, name, timeout,
                 unique=False, output=None):
        super(RPCTestServer, self).__init__(cfg, ctl_url, topic, RPC_SERVER,
                                            name, timeout, unique, output)
        target = om.Target(exchange=RPC_EXCHANGE,
                           topic=RPC_TOPIC % topic,
                           server=self.name)
        tport = (self.ctl_tport
                 if test_url == self.ctl_url
                 else om.get_rpc_transport(cfg.CONF, url=test_url))
        self._rpc_server = om.get_rpc_server(tport,
                                             target,
                                             [self],
                                             executor=executor)
        self._rpc_server.start()

    def _update_stats(self, timestamp, msgid):
        # given timestamp from arriving message
        ts = now()
        self._output.write('{"id": "%s", "start": %f, "recv": %f'
                           % (msgid, timestamp, ts))

        if timestamp > ts:
            logging.error("Clock error detected:"
                          " send time (%f) after arrival time (%f)"
                          " test results will be invalid!",
                          timestamp, ts)
            with self._results_lock:
                self.results.error("Clocks not synchronized")
                self.results.msgs_fail += 1
            self._output.write(', "error": "unsynchronized clocks"\n')
        else:
            with self._results_lock:
                self.results.start_time = (min(self.results.start_time, ts)
                                           if self.results.start_time else ts)
                self.results.stop_time = (max(self.results.stop_time, ts)
                                          if self.results.stop_time else ts)
                self.results.msgs_ok += 1
                self.results.latency.update((ts - timestamp) * 1000)
        self._output.write('}\n')

    #
    # Controller RPC Calls:
    #

    def shutdown(self, ctxt):
        self._rpc_server.stop()
        self._rpc_server.wait()
        super(RPCTestServer, self).shutdown(ctxt)

    #
    # Test RPC Calls:
    #

    def noop(self, ctxt, data, timestamp, msgid):
        # for cast testing - called by RPCTestClient, no return value
        self._update_stats(timestamp, msgid)
        logging.debug("RPCServer.noop(timestamp=%s)", timestamp)

    def echo(self, ctxt, data, timestamp, msgid):
        # for call testing - called by RPCTestClient
        self._update_stats(timestamp, msgid)
        logging.debug("RPCServer.echo(timestamp=%s)", timestamp)
        return data

class TestNotifier(_TestClient):
    """Client for issuing Notification calls to the TestListener
    """
    def __init__(self, cfg, ctl_url, test_url, topic, name, timeout,
                 output=None):
        super(TestNotifier, self).__init__(cfg,
                                           ctl_url,
                                           topic,
                                           NOTIFIER,
                                           name,
                                           timeout,
                                           output)
        # for notifying the test listener:
        om.set_transport_defaults(control_exchange=NOTIFY_EXCHANGE)
        tport = om.get_notification_transport(cfg.CONF, url=test_url)
        topic = NOTIFY_TOPIC % topic
        self._notifier = om.notify.notifier.Notifier(tport,
                                                     self.name,
                                                     driver='messaging',
                                                     topics=[topic])
    #
    # RPC Calls:
    #

    def run_test(self, ctxt, test, kwargs, reply_addr):
        if test != 'test_notify':
            # ignore other tests, like rpc-call, etc
            return
        verbose = kwargs.get("verbose", False)
        pause = kwargs.get("pause", 0)
        data = kwargs.get("data", "")
        count = kwargs.get("count", 0)
        severity = kwargs.get("severity", "debug")

        controller = om.RPCClient(self.ctl_tport,
                                  om.Target(**reply_addr),
                                  timeout=self._timeout)
        logging.debug("Client %s starting test %s ...", self.name, test)

        func = getattr(self._notifier, severity)
        payload = {'payload': data}

        def test_func(timestamp, msgid):
            payload['timestamp'] = timestamp
            payload['msgid'] = msgid
            func({}, "notification-test", payload)

        self._output.write('{"test-name": "test_notify",'
                           ' "test-start": %(start)f}\n'
                           % {'start': now()})

        self._execute(test_func, count, pause, verbose)

        self._output.write('{"test-end": %f}\n' % now())

        logging.debug("Client %s test %s finished, sending results...",
                      self.name, test)

        with self._results_lock:
            results = self.results.to_dict()
            self.results.reset()

        try:
            controller.call({}, 'client_result', name=self.name,
                            kind=self.kind, results=results)
        except Exception as exc:
            # I don't think recovery is possible as the call may be in-doubt.
            # For now simply let folks know the results may be invalid
            logging.error("%s failed to report results!"
                          " Test results may be invalid!"
                          " Error: %s", str(exc))
        else:
            logging.debug("Client %s test %s results sent", self.name, test)

class TestListener(_TestServer):
    def __init__(self, cfg, ctl_url, test_url, topic, executor, name, timeout,
                 pool=None, output=None):
        super(TestListener, self).__init__(cfg,
                                           ctl_url,
                                           topic,
                                           LISTENER,
                                           name,
                                           timeout, output)
        target = om.Target(exchange=NOTIFY_EXCHANGE,
                           topic=NOTIFY_TOPIC % topic,
                           server=self.name)
        om.set_transport_defaults(control_exchange=NOTIFY_EXCHANGE)
        tport = om.get_notification_transport(cfg.CONF, url=test_url)
        self._listener = om.get_notification_listener(tport,
                                                      [target],
                                                      [self],
                                                      executor=executor,
                                                      pool=pool)
        self._listener.start()

    #
    # Controller RPC Calls:
    #

    def shutdown(self, ctxt):
        self._listener.stop()
        self._listener.wait()
        super(TestListener, self).shutdown(ctxt)

    #
    # Notifications:
    #

    def _report(self, severity, ctx, publisher, event_type, payload, metadata):
        ts = now()
        logging.debug("%s Notification %s:%s:%s:%s:%s", self.name, severity,
                      publisher, event_type, payload, metadata)

        timestamp = payload['timestamp']
        msgid = payload['msgid']

        self._output.write('{"id": "%(msgid)s", "start": %(start)f,'
                           ' "recv": %(recv)f'
                           % {'msgid': msgid, 'start': timestamp, 'recv': ts})

        if timestamp > ts:
            logging.error("Clock error detected:"
                          " send time (%f) after arrival time (%f)"
                          " test results will be invalid!",
                          timestamp, ts)
            with self._results_lock:
                self.results.error("Clocks not synchronized")
                self.results.msgs_fail += 1
            self._output.write(', "error": "unsynchronized clocks"\n')
        else:
            with self._results_lock:
                self.results.start_time = (min(self.results.start_time, ts)
                                           if self.results.start_time else ts)
                self.results.stop_time = (max(self.results.stop_time, ts)
                                          if self.results.stop_time else ts)
                self.results.latency.update((ts - timestamp) * 1000)
                self.results.msgs_ok += 1
        self._output.write('}\n')

    def debug(self, ctx, publisher, event_type, payload, metadata):
        self._report("debug", ctx, publisher, event_type, payload, metadata)

    def audit(self, ctx, publisher, event_type, payload, metadata):
        self._report("audit", ctx, publisher, event_type, payload, metadata)

    def critical(self, ctx, publisher, event_type, payload, metadata):
        self._report("critical", ctx, publisher, event_type, payload, metadata)

    def error(self, ctx, publisher, event_type, payload, metadata):
        self._report("error", ctx, publisher, event_type, payload, metadata)

    def info(self, ctx, publisher, event_type, payload, metadata):
        self._report("info", ctx, publisher, event_type, payload, metadata)

    def warn(self, ctx, publisher, event_type, payload, metadata):
        self._report("warn", ctx, publisher, event_type, payload, metadata)

class Controller(_Base):
    """The test controller
    """
    def __init__(self, cfg, ctl_url, topic, timeout, unique=False, idle=2,
                 output=None):
        # each controller has a unique topic not to be confused
        # with future or past controller instances
        self.topic = topic
        self._idle = idle
        self.unique = unique
        super(Controller, self).__init__(cfg, ctl_url, topic, output,
                                         unique=False,
                                         name=None,
                                         kind=None,
                                         timeout=timeout)
        self._total_minions = 0
        self._queue = queue.Queue()

        # count of clients per type
        self._minions = dict([(k, 0) for k in MESSAGING_CLIENT_TYPES])
        # aggregated client results per type
        self._results = dict([(k, TestResults())
                              for k in MESSAGING_CLIENT_TYPES])
        # control rpc client for each type:
        self._clients = dict()

    def start(self):
        super(Controller, self).start()
        logging.debug("Polling for clients...")
        reply = {'exchange': self.ctl_target.exchange,
                 'topic': self.ctl_target.topic,
                 'server': self.ctl_target.server}
        for kind in MESSAGING_CLIENT_TYPES:
            target = om.Target(exchange=CONTROL_EXCHANGE,
                               topic=CLIENT_TOPIC %
                               ((kind, self.topic)
                                if not self.unique else (kind, 'singleton')),
                               fanout=True)
            self._clients[kind] = om.RPCClient(self.ctl_tport, target=target)
            self._clients[kind].cast({}, 'client_ping', reply_addr=reply)

        # wait until no more clients reply to the ping
        # (things are idle)
        _wait_stabilize(self._idle, lambda: self._total_minions)

    def shutdown(self):
        """Shutdown this Controller
        """
        super(Controller, self).shutdown({})
        self.wait()

    def shutdown_clients(self):
        """Shutdown all clients listening to $topic
        """
        for kind in MESSAGING_CLIENT_TYPES:
            self._clients[kind].cast({}, 'shutdown')
        time.sleep(1.0)

    def run_call_test(self, count, data, verbose, pause):
        clients = self._minions[RPC_CLIENT]
        servers = self._minions[RPC_SERVER]
        kwargs = {'verbose': verbose,
                  'pause': pause,
                  'data': data,
                  'count': count}
        self._run_test(RPC_CLIENT, 'test_call', kwargs)
        # note: set the poll time to 2x the client's pause between calls,
        # otherwise we're polling too frequently
        self._query_servers(RPC_SERVER, pause * 2, count * clients)

        print("RPC call test results")
        print("%d RPC clients, %d RPC Servers (%d total)"
              % (clients, servers, clients + servers))

        print("\n")
        print("Aggregated RPC Client results:")
        print("------------------------------")
        self._results[RPC_CLIENT].print_results()
        print("\n")
        print("Aggregated RPC Server results:")
        print("------------------------------")
        self._results[RPC_SERVER].print_results()

    def run_cast_test(self, count, data, verbose, pause, delay):
        clients = self._minions[RPC_CLIENT]
        servers = self._minions[RPC_SERVER]
        kwargs = {'verbose': verbose,
                  'pause': pause,
                  'data': data,
                  'count': count}
        self._run_test(RPC_CLIENT, 'test_cast', kwargs)
        # cast are async, wait a bit for msgs to propagate
        time.sleep(delay)
        # note: set the poll time to 2x the client's pause between calls,
        # otherwise we're polling too frequently
        self._query_servers(RPC_SERVER, pause * 2, count * clients)

        print("RPC cast test results")
        print("%d RPC clients, %d RPC Servers (%d total)"
              % (clients, servers, clients + servers))

        print("\n")
        print("Aggregated RPC Client results:")
        print("------------------------------")
        self._results[RPC_CLIENT].print_results()
        print("\n")
        print("Aggregated RPC Server results:")
        print("------------------------------")
        self._results[RPC_SERVER].print_results()

    def run_fanout_test(self, count, data, verbose, pause,
                        delay):
        clients = self._minions[RPC_CLIENT]
        servers = self._minions[RPC_SERVER]
        kwargs = {'verbose': verbose,
                  'pause': pause,
                  'data': data,
                  'count': count}
        self._run_test(RPC_CLIENT, 'test_fanout', kwargs)
        # fanouts are async, wait a bit for msgs to propagate
        time.sleep(delay)
        # note: set the poll time to 2x the client's pause between calls,
        # otherwise we're polling too frequently
        self._query_servers(RPC_SERVER, pause * 2,
                            count * clients * servers)

        print("RPC fanout test results")
        print("%d RPC clients, %d RPC Servers (%d total)"
              % (clients, servers, clients + servers))

        print("\n")
        print("Aggregated RPC Client results:")
        print("------------------------------")
        self._results[RPC_CLIENT].print_results()
        print("\n")
        print("Aggregated RPC Server results:")
        print("------------------------------")
        self._results[RPC_SERVER].print_results()

        start = self._results[RPC_CLIENT].start_time
        stop = self._results[RPC_SERVER].stop_time
        print("\n")
        print("Fanout propagation delay:")
        print("-------------------------")
        print("  First client transmit time: %f" % start)
        print("  Last server receive time: %f" % stop)
        print("  Duration (secs): %f" % (stop - start))

    def run_notification_test(self, count, data, severity, verbose, pause,
                              delay):
        clients = self._minions[NOTIFIER]
        servers = self._minions[LISTENER]
        kwargs = {'verbose': verbose,
                  'pause': pause,
                  'data': data,
                  'count': count,
                  'severity': severity}
        self._run_test(NOTIFIER, 'test_notify', kwargs)
        # notifications are async, wait a bit for msgs to propagate
        time.sleep(delay)
        # note: set the poll time to 2x the client's pause between calls,
        # otherwise we're polling too frequently
        self._query_servers(LISTENER, pause * 2, count * clients)

        print("Notification test results")
        print("%d Notifiers, %d Listeners (%d total)"
              % (clients, servers, clients + servers))

        print("\n")
        print("Aggregated Notifier (Client) results:")
        print("------------------------------------")
        self._results[NOTIFIER].print_results()
        print("\n")
        print("Aggregated Listener (Server) results:")
        print("-------------------------------------")
        self._results[LISTENER].print_results()

    def _run_test(self, kind, test, kwargs):
        """Tell the messaging clients to run a test.  When the client completes
        it will call the 'client_result' method below.
        """
        count = self._minions[kind]
        if count == 0:
            raise Exception("No %s clients visible" % kind)

        reply = {'exchange': self.ctl_target.exchange,
                 'topic': self.ctl_target.topic,
                 'server': self.ctl_target.server}

        # tell 'kind' clients to run the test
        self._clients[kind].cast({}, 'run_test',
                                 test=test,
                                 kwargs=kwargs,
                                 reply_addr=reply)

        results_per_client = dict()
        # wait for the clients to send results
        while count:
            try:
                name, ckind, results = self._queue.get(timeout=self._timeout)
            except queue.Empty:
                raise Exception("%s test timed out: no response from clients!"
                                % test)
            results_per_client[name] = results.to_dict()
            self._results[ckind].merge(results)
            if ckind == kind:
                count -= 1
            else:
                # TODO(kgiusti) uh, is this a problem?
                logging.warning("Huh? results from %s while expecting a %s",
                                ckind, kind)
        self._output.write(json.dumps(results_per_client) + '\n')

    def _query_servers(self, kind, pause, total):
        """Ask the servers for any data gathered during the test. The servers
        will respond by calling the 'client_result' method below.  Once the
        servers stop reporting new statistics the query is done.

        :param kind: the type of server to query - RPC or LISTENER
        :type kind: str
        :param pause: time in seconds to wait between each server poll.  Should
                      be at least twice the client's inter-message pause time.
        :type pause: float
        :param total: total number of messages expected to be received by all
                      servers
        """
        if self._minions[kind] == 0:
            raise Exception("No %s servers visible" % kind)

        # avoid hammering the servers if no pause given
        pause = max(pause, 0.250) if pause else 0.250

        server_results = self._results[kind]
        reply = {'exchange': self.ctl_target.exchange,
                 'topic': self.ctl_target.topic,
                 'server': self.ctl_target.server}

        logging.debug("Querying servers...")

        results_per_server = dict()
        done = False
        start = now()
        while not done and abs(now() - start) < self._timeout:

            # tell 'kind' servers to return results
            self._clients[kind].cast({}, 'get_server_results',
                                     reply_addr=reply)
            # wait for the servers to send results
            count = self._minions[kind]
            seen = 0
            while count:
                try:
                    _ = self._queue.get(timeout=self._timeout)
                    name, ckind, results = _
                except queue.Empty:
                    raise Exception("%s test timed out: no response from"
                                    " servers!" % kind)

                if ckind != kind:
                    # TODO(kgiusti): uh, is this a problem?
                    logging.warning("Huh? results from %s while expecting %s",
                                    ckind, kind)
                    continue
                if name not in results_per_server:
                    results_per_server[name] = results
                else:
                    results_per_server[name].merge(results)

                server_results.merge(results)
                seen += results.msgs_ok + results.msgs_fail
                count -= 1

            # exit loop once the expected number of messages have been received
            # by the servers, or the test times out
            done = server_results.msgs_ok + server_results.msgs_fail >= total
            if not done:
                time.sleep(pause)
                if seen == 0:
                    # no replies: try pausing longer next time
                    pause = min(pause * 2.0, self._timeout)

        if not done:
            logging.error("Test timed out - not all messages accounted for")
        results_per_server = dict([[k, v.to_dict()]
                                   for k, v in results_per_server.items()])
        self._output.write(json.dumps(results_per_server) + '\n')

        logging.debug("... servers queried")

    #
    # RPC calls:
    #
    def client_pong(self, ctxt, kind, name):
        # A client 'name' is checking in
        if kind not in self._minions:
            self._minions[kind] = 0
        self._minions[kind] += 1
        self._total_minions += 1
        logging.debug("New %s detected (%s) - %d total clients found",
                      kind, name, self._total_minions)
        return True

    def client_result(self, ctxt, name, kind, results):
        # A test client is reporting a test result in response to the above
        # _run_test method.
        logging.debug("%s results received from %s", (kind, name))
        try:
            self._queue.put((name, kind, TestResults.from_dict(results)))
        except Exception as exc:
            logging.error("Invalid TestResult from %s:%s (%s)",
                          kind, str(exc), str(results))


# -*- coding: utf-8 -*-
'''
Created on 2013-4-18

@author: seven
'''
from src import argparser, define
import re
import time
from src.modules.mod_io import S_VALUE
from src.output import Output

def S_VALUE(m, n, p):
    d = round(p, 4)
    return (int(n) - int(m)) / d if d != 0 else 0.0

class tcp_stat(object):
    def __init__(self, line):
        self.time = time.time()
        # Tcp:
        # RtoAlgorithm
        # RtoMin
        # RtoMax
        # MaxConn
        # ActiveOpens
        # PassiveOpens
        # AttemptFails
        # EstabResets
        # CurrEstab
        # InSegs
        # OutSegs
        # RetransSegs
        # InErrs
        # OutRsts
        rex = r'^tcp:\s*'
        rex += r'(?P<RtoAlgorithm>\d*)\s*'
        rex += r'(?P<RtoMin>\d*)\s*'
        rex += r'(?P<RtoMax>\d*)\s*'
        rex += r'(?P<MaxConn>-?\d*)\s*'
        rex += r'(?P<ActiveOpens>\d*)\s*'
        rex += r'(?P<PassiveOpens>\d*)\s*'
        rex += r'(?P<AttemptFails>\d*)\s*'
        rex += r'(?P<EstabResets>\d*)\s*'
        rex += r'(?P<CurrEstab>\d*)\s*'
        rex += r'(?P<InSegs>\d*)\s*'
        rex += r'(?P<OutSegs>\d*)\s*'
        rex += r'(?P<RetransSegs>\d*)\s*'
        rex += r'(?P<InErrs>\d*)\s*'
        rex += r'(?P<OutRsts>\d*)\s*'
        ma = re.match(rex, line, re.IGNORECASE)
        if ma is not None:
            self.__dict__.update(ma.groupdict())

class tcp(object):
    def __init__(self, main):
        self.main=main
        main.add_argument('', '--' + self.__class__.__name__, action=argparser.ACTION_SWITCH, help='TCP traffic & connection data')

    def __call__(self):
        self.pre = self.read_tcp()
        time.sleep(.2)
        self.cur = self.read_tcp()
        # st_tcp.ActiveOpens,
#              st_tcp.PassiveOpens,
#              st_tcp.InSegs,
#              st_tcp.OutSegs,
#              st_tcp.RetransSegs
        # 两次探测间隔时间（s）
        self.interval = round((self.cur.time - self.pre.time) , 6)
        ret = {}
        ret['ActiveOpens'] = "%.2f" % S_VALUE(self.pre.ActiveOpens, self.cur.ActiveOpens, self.interval)
        ret['PassiveOpens'] = "%.2f" % S_VALUE(self.pre.PassiveOpens, self.cur.PassiveOpens, self.interval)
        ret['InSegs'] = "%.2f" % S_VALUE(self.pre.InSegs, self.cur.InSegs, self.interval)
        ret['OutSegs'] = "%.2f" % S_VALUE(self.pre.OutSegs, self.cur.OutSegs, self.interval)
        ret['RetransSegs'] = "%.2f" % S_VALUE(self.pre.RetransSegs, self.cur.RetransSegs, self.interval)
        self.main.Output.outputData({self.__class__.__name__:ret})

    def read_tcp(self):
        stat = None
        p = open(define.NET_SNMP, 'r')
        for line in p:
            ma = re.match('^tcp:\s*\d+', line, re.IGNORECASE)
            if ma is not None:
                stat = tcp_stat(line)
        p.close()
        return stat


from amuse.units import units, nbody_system
from amuse.ic.plummer import new_plummer_model
from amuse.ic.salpeter import new_powerlaw_mass_distribution
from amuse.io import write_set_to_file

Nstars = 1000
Ngas = 1000

mass = new_powerlaw_mass_distribution(
        Nstars,
        mass_max=20. | units.MSun,
        )
conv = nbody_system.nbody_to_si(
        mass.sum(),
        1. | units.parsec,
        )
p = new_plummer_model(
        Nstars,
        conv,
        )
p.mass = mass
write_set_to_file(p, "stars.hdf5", "amuse", append_to_file=False)

p = new_plummer_model(
        Ngas,
        conv,
        )
p.h_smooth = 0.1 | units.parsec
write_set_to_file(p, "gas.hdf5", "amuse", append_to_file=False)


import RPi.GPIO as GPIO
import time

GPIO.setmode(GPIO.BCM)
GPIO.setup(18, GPIO.OUT)
GPIO.setup(17, GPIO.OUT)
GPIO.setup(27, GPIO.OUT)
GPIO.setup(22, GPIO.OUT)

m1_pwm = GPIO.PWM(18, 100)
m2_pwm = GPIO.PWM(27, 100)

GPIO.output(17, GPIO.HIGH)
GPIO.output(22, GPIO.HIGH)

m1_pwm.start(50)
m2_pwm.start(50)

time.sleep(5)

m1_pwm.ChangeDutyCycle(75)
m2_pwm.ChangeDutyCycle(75)

time.sleep(5)

m1_pwm.stop()
m2_pwm.stop()

GPIO.cleanup()


import functools
import logging
from dataclasses import dataclass
from typing import Dict, Any, List, Callable, Optional, Union

from .Singleton import Singleton
from .filehandler import check_if_file_exists, to_abs_file_path, load_file

LOGGER: logging.Logger = logging.getLogger(__name__)

@dataclass
class KeyType:
    """
    Defines a config value.
        - The [name] is the config key
        - The [type] is the type to be expected
        - The required indicates if the config value is mandatory for the program. e.g. Api Key
        - The [model_from_dict] is an optional callable to parse raw data to a class.
          See the code example for more details

    ```
    @dataclass
    class User:
        username: str
        birthday: datetime

    def _user_from_raw(raw: Dict[str, Any]) -> User:
        return User(username=raw["username"],
                    birthday=datetime.strptime(raw["birthday"], "%Y.%m.%d"))

    key = KeyType(name="user", type=dict, model_from_dict=_user_from_raw)
    ```

    Would parse a config like:
    ```
    {
        "user": {
            "username": "foo"
            "birthday" "1970.1.1"
        }
    }
    ```
    To a User object
    """
    name: str
    type: Any
    required: bool = False
    model_from_dict: Optional[Callable[[Dict[str, Any]], Any]] = None

def _instance_check(func: Callable[[List[Any], Dict], Any]):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        self: ConfigLoader = args[0]
        if self is None:
            LOGGER.warning(
                f"Decorator 'instance_check' not on none-static object method applied ('{func.__name__}'). 'self' is None")
        elif not hasattr(self, "_is_instantiated"):
            LOGGER.warning(f"{type(self)} has no '_is_instantiated' attribute")
        else:
            _is_instantiated = getattr(self, "_is_instantiated")
            if not isinstance(_is_instantiated, bool):
                LOGGER.warning(
                    f"{type(self)}'s attribute '_is_instantiated' is not of type 'bool' but of type '{type(_is_instantiated)}'")
            elif not _is_instantiated:
                LOGGER.warning(f"{type(self)} is not instantiated. Called on '{type(self)}.{func.__name__}'")
                return None
        return func(*args, **kwargs)

    return wrapper

class ConfigLoader:
    """
    This class helps to load and validate config entries.
    It checks a given config file for settings and caches them.

    After initialization the `instantiate` function has to be called.
    """
    _is_instantiated: bool = False
    _key_types: List[KeyType]
    _config_file: str
    _config: Dict[str, Any] = {}

    def instantiate(self, key_types: List[KeyType], config_file: str) -> None:
        """
        Setup method
        :param key_types: a list of [KeyType] to check for.
        :param config_file: the file to look for
        :return: None
        """
        self._config_file = config_file
        self._key_types = key_types
        self._is_instantiated = True

        self.check()

    @_instance_check
    def get_value(self, key: str, default: Any = None) -> Any:
        """
        Gets a config value with key [key]
        :param key: the key of the value
        :param default: if set and no key is provided, it returns the default
        :return: the value associated with the [key]. If it is none, [default] will be returned
        """
        if key in self._config:
            return self._config[key]
        return default

    @_instance_check
    def add_keys(self, key_types: Union[KeyType, List[KeyType]]) -> None:
        """
        Adds additional [KeyType]s and refreshes the cached config
        :param key_types: the new key types to check for
        :return: None
        """
        if isinstance(key_types, KeyType):
            key_types = [key_types]
        for key_type in key_types:
            self._config[key_type.name] = key_type
        self.check()

    @_instance_check
    def remove_keys(self, key_types: Union[KeyType, List[KeyType]]) -> bool:
        """
        removes a set of keys from the config
        :param key_types: the keys to remove
        :return: bool of success. Could fail due to some unknown keys.
        It only indicates that some failed, the rest could be successful
        """
        if isinstance(key_types, KeyType):
            key_types = [key_types]
        success: bool = True
        for key_type in key_types:
            if key_type.name in self._config:
                self._config.pop(key_type.name)
            else:
                success = False
        return success

    @_instance_check
    def check(self) -> None:
        """
        Refreshes the cached config values.
        It checks the [self._config_file] for the specified keys, if the keys are of the corresponding type.
        If some required key_types fail, an exception will be raised
        :return: None
        :raises FileNotFoundError, ValueError
        """
        if not check_if_file_exists(self._config_file):
            msg: str = f"No config found at {to_abs_file_path(self._config_file)}"
            for key_type in self._key_types:
                if key_type.required:
                    LOGGER.error(msg)
                    raise FileNotFoundError(msg)
            LOGGER.warning(msg)
            return

        raw_config = load_file(self._config_file)
        if not raw_config:
            msg: str = f"Config is empty or not valid! ({to_abs_file_path(self._config_file)})"
            LOGGER.error(msg)
            raise ValueError(msg)
        if not isinstance(raw_config, dict):
            msg: str = f"Config must be dictionary! ({to_abs_file_path(self._config_file)})"
            LOGGER.error(msg)
            raise ValueError(msg)

        failure: bool = False
        for keyType in self._key_types:
            if keyType.name not in raw_config.keys():
                msg: str = f"Key '{keyType.name}' is not provided in {to_abs_file_path(self._config_file)}"
                if keyType.required:
                    LOGGER.error(msg)
                    failure = True
                else:
                    LOGGER.debug(msg)
                continue

            raw_key: Any = raw_config[keyType.name]
            if not isinstance(raw_key, keyType.type):
                msg: str = f"Key '{keyType.name}' has to be of type '{keyType.type}', not '{type(raw_key)}'"
                if keyType.required:
                    LOGGER.error(msg)
                    failure = True
                else:
                    LOGGER.warning(f"{msg}. Value gets ignored")
                continue

            if keyType.model_from_dict:
                try:
                    self._config[keyType.name] = keyType.model_from_dict(raw_key)
                except Exception:
                    msg: str = f"Could not parse {keyType.name}"
                    if keyType.required:
                        LOGGER.error(msg)
                        failure = True
                    else:
                        LOGGER.warning(f"{msg}. Ignoring key.")
                    continue
            else:
                self._config[keyType.name] = raw_key

        if failure:
            raise ValueError(f"Required config values missing. Look in the logs for details")

class ConfigLoaderSingleton(ConfigLoader, Singleton):
    """
    A [ConfigLoader] that can be globally initialized to get a reference to.
    This is just a convenience class of the [ConfigLoader], for more detail look at its documentation
    """
    pass


#!/usr/bin/env python3
import os, sys
import numpy as np
from pprint import pprint
import time
import re
import matplotlib.pyplot as plt
# import pygame
# import curses
from pynput.keyboard import Listener, Key

sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)), '..', '..'))

from utils.timeit import timeit

class VirtualMachine(object):
    def __init__(self, reg=[0, 0, 0, 0, 0, 0], debug=False):
        self.__registers = reg

        self.debug = debug

        self.__ip = None
        self.__pc = 0

        self.instructions = []

        self.ops = {
            0: self.banr,
            1: self.addr,
            2: self.eqri,
            3: self.setr,
            4: self.gtrr,
            5: self.bori,
            6: self.gtir,
            7: self.seti,
            8: self.borr,
            9: self.bani,
            10: self.eqir,
            11: self.eqrr,
            12: self.gtri,
            13: self.addi,
            14: self.muli,
            15: self.mulr,

            16: self.addr,
            17: self.addi,
            18: self.mulr,
            19: self.muli,
            20: self.banr,
            21: self.bani,
            22: self.borr,
            23: self.bori,
            24: self.setr,
            25: self.seti,
            26: self.gtir,
            27: self.gtri,
            28: self.gtrr,
            29: self.eqir,
            30: self.eqri,
            31: self.eqrr,

            'addr': self.addr,
            'addi': self.addi,
            'mulr': self.mulr,
            'muli': self.muli,
            'banr': self.banr,
            'bani': self.bani,
            'borr': self.borr,
            'bori': self.bori,
            'setr': self.setr,
            'seti': self.seti,
            'gtir': self.gtir,
            'gtri': self.gtri,
            'gtrr': self.gtrr,
            'eqir': self.eqir,
            'eqri': self.eqri,
            'eqrr': self.eqrr,
        }

    @property
    def registers(self):
        return self.__registers

    @property
    def ip(self):
        return self.__ip

    def addr(self, A, B, C):
        self.__registers[C] = self.registers[A] + self.registers[B]

    def addi(self, A, B, C):
        self.__registers[C] = self.registers[A] + B

    def mulr(self, A, B, C):
        self.__registers[C] = self.registers[A] * self.registers[B]

    def muli(self, A, B, C):
        self.__registers[C] = self.registers[A] * B

    def banr(self, A, B, C):
        self.__registers[C] = self.registers[A] & self.registers[B]

    def bani(self, A, B, C):
        self.__registers[C] = self.registers[A] & B

    def borr(self, A, B, C):
        self.__registers[C] = self.registers[A] | self.registers[B]

    def bori(self, A, B, C):
        self.__registers[C] = self.registers[A] | B

    def setr(self, A, B, C):
        self.__registers[C] = self.registers[A]

    def seti(self, A, B, C):
        self.__registers[C] = A

    def gtir(self, A, B, C):
        self.__registers[C] = 1 if A > self.registers[B] else 0

    def gtri(self, A, B, C):
        self.__registers[C] = 1 if self.registers[A] > B else 0

    def gtrr(self, A, B, C):
        self.__registers[C] = 1 if self.registers[A] > self.registers[B] else 0

    def eqir(self, A, B, C):
        self.__registers[C] = 1 if A == self.registers[B] else 0

    def eqri(self, A, B, C):
        self.__registers[C] = 1 if self.registers[A] == B else 0

    def eqrr(self, A, B, C):
        self.__registers[C] = 1 if self.registers[A] == self.registers[B] else 0

    def setip(self, A, B, C):
        last = self.registers[self.__ip if self.__ip is not None else self.__pc]
        self.__ip = A
        try:
            del self.instructions[last]
        except:
            pass

    def load(self, lines):
        ins = re.compile(r'(#?[A-Za-z]+)\s+(\d+)\s*(\d*)\s*(\d*)')
        for l in lines:
            c = ins.match(l)
            instr  = c.groups(1)[0]
            if instr == '#ip':
                A, B, C = int(c.groups(1)[1]),0,0
                self.setip(A,B,C)
            else:
                A, B, C = int(c.groups(1)[1]),int(c.groups(1)[2]),int(c.groups(1)[3])
                self.instructions.append([instr,A,B,C])

    def exec_once(self, inst=None):
        if self.debug: print(f'pc={self.__pc} ip={self.__registers[self.__ip]} {self.registers}', end=' ')
        if inst is not None:
            op, A, B, C = inst
        else:
            op, A, B, C = self.instructions[self.registers[self.__ip] if self.__ip is not None else self.__pc]
        if op == '#ip':
            self.setip(A,B,C)
            return
        if self.debug: print(f'{op} {A} {B} {C}', end=' ')
        if not op in self.ops.keys():
            raise ValueError(f'Invalid op code: {op} ({inst})')
        self.ops[op](A, B, C)
        if self.debug: print(self.registers)
        if self.__ip is not None:
            self.__registers[self.__ip] += 1
        self.__pc += 1

    def exec(self):
        while (self.registers[self.__ip] if self.__ip is not None else self.__pc) < len(self.instructions):
            self.exec_once()

class VMTests:
    @staticmethod
    def run_tests():
        VMTests.test_addr()
        VMTests.test_addi()
        VMTests.test_mulr()
        VMTests.test_muli()
        VMTests.test_banr()
        VMTests.test_bani()
        VMTests.test_borr()
        VMTests.test_bori()
        VMTests.test_setr()
        VMTests.test_seti()
        VMTests.test_gtir()
        VMTests.test_gtri()
        VMTests.test_gtrr()
        VMTests.test_eqir()
        VMTests.test_eqri()
        VMTests.test_eqrr()
        VMTests.test_s_ip()
        print('All tests executed successfully.')

    @staticmethod
    def test_addr():
        vm = VirtualMachine([1,2,3,0])
        vm.exec_once([16, 0, 1, 3])
        assert(vm.registers == [1,2,3,3])
        vm.exec_once([16, 0, 2, 3])
        assert(vm.registers == [1,2,3,4])
        vm.exec_once([16, 1, 2, 3])
        assert(vm.registers == [1,2,3,5])
        vm.exec_once([16, 2, 3, 0])
        assert(vm.registers == [8,2,3,5])
        vm.exec_once([16, 0, 0, 0])
        assert(vm.registers == [16,2,3,5])

    @staticmethod
    def test_addi():
        vm = VirtualMachine([1,2,3,0])
        vm.exec_once([17, 0, 1, 3])
        assert(vm.registers == [1,2,3,2])
        vm.exec_once([17, 0, 2, 3])
        assert(vm.registers == [1,2,3,3])
        vm.exec_once([17, 1, 2, 3])
        assert(vm.registers == [1,2,3,4])
        vm.exec_once([17, 2, 3, 0])
        assert(vm.registers == [6,2,3,4])
        vm.exec_once([17, 0, 0, 0])
        assert(vm.registers == [6,2,3,4])
        vm.exec_once([17, 0, -6, 0])
        assert(vm.registers == [0,2,3,4])

    @staticmethod
    def test_mulr():
        vm = VirtualMachine([1,2,3,0])
        vm.exec_once([18, 0, 1, 3])
        assert(vm.registers == [1,2,3,2])
        vm.exec_once([18, 0, 2, 3])
        assert(vm.registers == [1,2,3,3])
        vm.exec_once([18, 1, 2, 3])
        assert(vm.registers == [1,2,3,6])
        vm.exec_once([18, 2, 3, 0])
        assert(vm.registers == [18,2,3,6])
        vm.exec_once([18, 0, 0, 0])
        assert(vm.registers == [18*18,2,3,6])

    @staticmethod
    def test_muli():
        vm = VirtualMachine([1,2,3,0])
        vm.exec_once([19, 0, 1, 3])
        assert(vm.registers == [1,2,3,1])
        vm.exec_once([19, 0, 2, 3])
        assert(vm.registers == [1,2,3,2])
        vm.exec_once([19, 1, 2, 3])
        assert(vm.registers == [1,2,3,4])
        vm.exec_once([19, 2, 3, 0])
        assert(vm.registers == [9,2,3,4])
        vm.exec_once([19, 0, 0, 0])
        assert(vm.registers == [0,2,3,4])

    @staticmethod
    def test_banr():
        vm = VirtualMachine([1,2,3,0])
        vm.exec_once([20, 0, 1, 3])
        assert(vm.registers == [1,2,3,0])
        vm.exec_once([20, 0, 2, 3])
        assert(vm.registers == [1,2,3,1])

    @staticmethod
    def test_bani():
        vm = VirtualMachine([1,2,3,0])
        vm.exec_once([21, 0, 1, 3])
        assert(vm.registers == [1,2,3,1])
        vm.exec_once([21, 0, 2, 3])
        assert(vm.registers == [1,2,3,0])

    @staticmethod
    def test_borr():
        vm = VirtualMachine([1,2,3,0])
        vm.exec_once([22, 0, 1, 3])
        assert(vm.registers == [1,2,3,3])
        vm.exec_once([22, 0, 2, 3])
        assert(vm.registers == [1,2,3,3])

    @staticmethod
    def test_bori():
        vm = VirtualMachine([1,2,3,0])
        vm.exec_once([23, 0, 1, 3])
        assert(vm.registers == [1,2,3,1])
        vm.exec_once([23, 0, 2, 3])
        assert(vm.registers == [1,2,3,3])

    @staticmethod
    def test_setr():
        vm = VirtualMachine([1,2,3,0])
        vm.exec_once([24, 0, 1, 3])
        assert(vm.registers == [1,2,3,1])

    @staticmethod
    def test_seti():
        vm = VirtualMachine([1,2,3,0])
        vm.exec_once([25, 0, 1, 3])
        assert(vm.registers == [1,2,3,0])

    @staticmethod
    def test_gtir():
        vm = VirtualMachine([1,2,3,0])
        vm.exec_once([26, 0, 1, 3])
        assert(vm.registers == [1,2,3,0])
        vm.exec_once([26, 4, 2, 3])
        assert(vm.registers == [1,2,3,1])

    @staticmethod
    def test_gtri():
        vm = VirtualMachine([1,2,3,0])
        vm.exec_once([27, 0, 1, 3])
        assert(vm.registers == [1,2,3,0])
        vm.exec_once([27, 2, 2, 3])
        assert(vm.registers == [1,2,3,1])

    @staticmethod
    def test_gtrr():
        vm = VirtualMachine([1,2,3,0])
        vm.exec_once([28, 0, 1, 3])
        assert(vm.registers == [1,2,3,0])
        vm.exec_once([28, 1, 2, 3])
        assert(vm.registers == [1,2,3,0])
        vm.exec_once([28, 2, 3, 3])
        assert(vm.registers == [1,2,3,1])

    @staticmethod
    def test_eqir():
        vm = VirtualMachine([1,2,3,0])
        vm.exec_once([29, 0, 1, 3])
        assert(vm.registers == [1,2,3,0])
        vm.exec_once([29, 3, 2, 3])
        assert(vm.registers == [1,2,3,1])

    @staticmethod
    def test_eqri():
        vm = VirtualMachine([1,2,3,0])
        vm.exec_once([30, 0, 1, 3])
        assert(vm.registers == [1,2,3,1])
        vm.exec_once([30, 2, 2, 3])
        assert(vm.registers == [1,2,3,0])

    @staticmethod
    def test_eqrr():
        vm = VirtualMachine([1,2,3,0])
        vm.exec_once([31, 0, 1, 3])
        assert(vm.registers == [1,2,3,0])
        vm.exec_once([31, 1, 2, 3])
        assert(vm.registers == [1,2,3,0])
        vm.exec_once([31, 2, 2, 3])
        assert(vm.registers == [1,2,3,1])

    @staticmethod
    def test_s_ip():
        vm = VirtualMachine([0, 0, 0, 0])
        vm.exec_once(['#ip', 0, 0, 0])
        vm.exec_once(['seti', 2, 0, 0])
        assert(vm.registers == [3, 0, 0, 0])

def main(lines):
    vm = VirtualMachine([1, 0, 0, 0, 0, 0], debug=False)
    vm.load(lines)

    try:
        vm.exec()
    except KeyboardInterrupt:
        pass
    except:
        pass

    # history = [[],[],[],[],[]]

    # def on_press(key):
    #     if key == Key.esc:
    #         sys.exit(0)
    #     # if key == Key.enter:
    #     #     vm._VirtualMachine__registers[3] = 10551331
    #     #     vm._VirtualMachine__registers[4] = 10551331
    #     vm.exec_once()
    #     for i in range(5):
    #         history[i].append(vm.registers[i])
    #     # if vm.registers[vm.ip] == 4:
    #     #     vm._VirtualMachine__registers[3] = 10551331
    #     #     vm._VirtualMachine__registers[4] = 10551331
    #     # if vm.registers[vm.ip] == 12:
    #     #     vm._VirtualMachine__registers[1] = 10551331

    # with Listener(on_press=on_press) as listener:
    #     listener.join()

    # plt.figure()
    # for i in range(5):
    #     plt.plot(history[i])
    # plt.show()

    pprint(vm.registers)
if __name__ == '__main__':
    VMTests.run_tests()

    if len(sys.argv) <= 1:
        sys.exit('No entries given.  Answer: unknown')

    if len(sys.argv) == 2:
        with open(sys.argv[1], 'r') as f:
            lines = f.readlines()
    else:
        sys.exit('Too many args.')

    main(lines)


"""
Author: Gabriel Fernandez
GitHub: github (dot) com / gabriel80808
Date: June 6, 2018

Algorithms to be implemented
"""
import numpy as np
import tensorflow as tf
from dynamics import NNDynamicsModel
import controllers
import utils

class ModelBasedMPC:
    """Model-Based Model Predictive Control
    Algorithm: “Neural Network Dynamics for Model-Based Deep Reinforcement Learning
    with Model-Free Fine-Tuning”, Anusha Nagabandi, Gregory Kahn, Ronald S. Fearing,
    Sergey Levine. https://arxiv.org/abs/1708.02596
    """
    def __init__(self,
                 environment='Hex',
                 controller='MPC',  # Place holder for object not string
                 function_approx='feed_forward',  # Place holder for object not string
                 exp_name='mb_mpc',
                 seed=1,
                 render=True,
                 learning_rate=1e-3,
                 onpol_iters=3,  # On-policy iteration roll-outs
                 dyn_iters=60,  # Iterations for dynamics model
                 batch_size=512,
                 random_paths=10,  # Paths represents generated dynamical trajectories
                 onpol_paths=10,
                 simulated_paths=200,  # 1000
                 ep_len=200,  # 1000 simulation episode length
                 mpc_horizon=20  # 15
                 ):
        # Major Components
        self.env_name = environment
        self.con = controller
        self.func_ap = function_approx

        # Experiment meta-params
        self.exp_name = exp_name
        self.seed = seed
        self.render = render

        # Training args
        self.learning_rate = learning_rate
        self.onpol_iters = onpol_iters
        self.dyn_iters = dyn_iters
        self.batch_size = batch_size

        # Data collection
        self.random_paths = random_paths
        self.onpol_paths = onpol_paths
        self.simulated_paths = simulated_paths
        self.ep_len = ep_len
        self.mpc_horizon = mpc_horizon

        # Check/create if directory exists
        self.log_dir = utils.data_dir(self.env_name, self.exp_name)

        # Set random seed for experiments
        np.random.seed(self.seed)
        tf.set_random_seed(self.seed)

        # Get environment and cost function
        self.env = utils.get_env(self.env_name)
        self.cost_fn = utils.get_cost_fn(self.env_name)

    def set_param(self, names_value_dict):
        for name, value in names_value_dict:
            setattr(self, name, value)
        return None

    def train(self):
        """Train Model-Based MPC"""
        data = utils.run_sim(self.env,  # Generate random data to train dynamics model
                             controller=controllers.Random(self.env),
                             num_paths=self.random_paths,
                             horizon=self.ep_len,
                             render=False,
                             verbose=True,
                             section="Generating Data for Dynamics with Random Controller")

        # Build dynamics model and MPC controllers.
        sess = tf.Session()

        """Need to add function appx"""
        dyn_model = NNDynamicsModel(env=self.env,
                                    func_name=self.func_ap,
                                    normalization=utils.compute_normalization(data),
                                    batch_size=self.batch_size,
                                    iterations=self.dyn_iters,
                                    learning_rate=self.learning_rate,
                                    sess=sess)

        mpc_controller = controllers.MPC(env=self.env,
                                         dyn_model=dyn_model,
                                         horizon=self.mpc_horizon,
                                         cost_fn=self.cost_fn,
                                         num_simulated_paths=self.simulated_paths)
        # Build tf session
        sess.__enter__()
        tf.global_variables_initializer().run()

        # Iterate: Run policy, aggregate, & refit dynamics
        for itr in range(self.onpol_iters):
            dyn_model.fit(data)
            print('On Policy Iteration: ', itr)
            paths = utils.run_sim(self.env,
                                  mpc_controller,
                                  self.onpol_paths,
                                  horizon=self.ep_len,
                                  render=True,
                                  verbose=True,
                                  section="MPC Trajectories")
            data = np.concatenate((data, paths))

            returns = [np.sum(path["reward"]) for path in paths]
            costs = [utils.path_cost(self.cost_fn, path) for path in paths]
            utils.log_stats(self.log_dir, itr, {"Costs": costs, "Returns": returns})

        return returns, costs