import asyncio
import logging
from contextlib import contextmanager
from datetime import datetime, timedelta
from typing import Optional, Iterator

import time

from pyrogram import Client
from pyrogram.errors import RpcMcgetFail
from pyrogram.handlers.handler import Handler
from typing_extensions import Final

from tgintegration.containers.response import InvalidResponseError, Response
from .awaitableaction import AwaitableAction

SLEEP_DURATION: Final[float] = 0.15

class ResponseCollectorClient(Client):
    def __init__(self, *args, global_action_delay=None, **kwargs):
        super().__init__(*args, **kwargs)
        self.logger = logging.getLogger(self.__class__.__name__)
        self.global_action_delay = global_action_delay
        self._last_response = None

    async def _wait_global(self):
        if self.global_action_delay and self._last_response:
            # Sleep for as long as the global delay prescribes
            sleep = self.global_action_delay - (time.time() - self._last_response.started)
            if sleep > 0:
                await asyncio.sleep(sleep)

    @contextmanager
    def add_handler_nolock(self, handler: Handler) -> Iterator[None]:
        # TODO: find next empty group
        group = -99
        if group not in self.dispatcher.groups:
            self.dispatcher.groups[group] = []
        self.dispatcher.groups[group].append(handler)
        yield
        self.remove_handler(handler, group)

    @contextmanager
    def collect(self, handler: Handler) -> Iterator[Response]:
        pass  # TODO

    async def act_await_response(self, action: AwaitableAction, raise_=True) -> Optional[Response]:

        await self._wait_global()

        response = Response(self, action)

        async def collect(_, message):
            # noinspection PyProtectedMember
            response._add_message(message)

        # Add handler to empty group manually, as the Pyrogram lock seems to never release when trying
        # to use .add_handler
        group = -99  # TODO: find new, empty group
        if group not in self.dispatcher.groups:
            self.dispatcher.groups[group] = []
        handler = MessageHandler(callback=collect, filters=action.filters)
        self.dispatcher.groups[group].append(handler)

        try:
            # Start timer
            response.started = time.time()

            # Execute the action
            response.action_result = await action.func(*action.args, **action.kwargs)

            # Calculate maximum wait time
            timeout_end = datetime.now() + timedelta(seconds=action.max_wait)

            # Wait for the first reply
            while response.empty:

                if time.time() - response.started > 5:
                    self.logger.debug("No response received yet after 5 seconds")

                if datetime.now() > timeout_end:
                    msg = "Aborting because no response was received after {} seconds.".format(
                        action.max_wait
                    )

                    if raise_:
                        raise InvalidResponseError(msg)
                    else:
                        self.logger.debug(msg)
                        return response

                await asyncio.sleep(SLEEP_DURATION)

            # A response was received
            if action.consecutive_wait:
                # Wait for more consecutive messages from the peer
                consecutive_delta = timedelta(seconds=action.consecutive_wait)

                while True:
                    now = datetime.now()

                    if action.num_expected is not None:
                        # User has set explicit number of messages to await
                        if response.num_messages < action.num_expected:
                            # Less messages than expected (so far)
                            if now > timeout_end:
                                # Timed out
                                msg = (
                                    "Expected {} messages but only received {} "
                                    "after waiting {} seconds.".format(
                                        action.num_expected,
                                        response.num_messages,
                                        action.max_wait,
                                    )
                                )

                                if raise_:  # TODO: should this really be toggleable? raise always?
                                    raise InvalidResponseError(msg)
                                else:
                                    self.logger.debug(msg)
                                    return None
                            # else: continue

                        elif response.num_messages > action.num_expected:
                            # More messages than expected
                            msg = "Expected {} messages but received {}.".format(
                                action.num_expected, response.num_messages
                            )

                            if raise_:  # TODO: should this really be toggleable? raise always?
                                raise InvalidResponseError(msg)
                            else:
                                self.logger.debug(msg)
                                return None
                        else:
                            self._last_response = response
                            return response
                    else:
                        # User has not provided an expected number of messages
                        if (
                            now > (response.last_message_timestamp + consecutive_delta)
                            or now > timeout_end
                        ):
                            self._last_response = response
                            return response

                    await asyncio.sleep(SLEEP_DURATION)

            self._last_response = response
            return response

        except RpcMcgetFail as e:
            self.logger.warning(e)
            await asyncio.sleep(60)  # Internal Telegram error
        finally:
            # Remove the one-off handler for this action
            self.remove_handler(handler, group)


# Copyright (c) 2011, Dorian Scholz, TU Darmstadt
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
#   * Redistributions of source code must retain the above copyright
#     notice, this list of conditions and the following disclaimer.
#   * Redistributions in binary form must reproduce the above
#     copyright notice, this list of conditions and the following
#     disclaimer in the documentation and/or other materials provided
#     with the distribution.
#   * Neither the name of the TU Darmstadt nor the names of its
#     contributors may be used to endorse or promote products derived
#     from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

from __future__ import division
import os
import rospkg

from python_qt_binding import loadUi
from python_qt_binding.QtCore import Qt, QTimer, qWarning, Slot
from python_qt_binding.QtGui import QAction, QMenu, QWidget

import rospy
from rostopic import get_topic_class
from tf.transformations import quaternion_matrix, quaternion_about_axis

from OpenGL.GL import glBegin, glColor3f, glEnd, glLineWidth, glMultMatrixf, glTranslatef, glVertex3f, GL_LINES, GL_QUADS
from .gl_widget import GLWidget

# main class inherits from the ui window class
class PoseViewWidget(QWidget):

    def __init__(self, plugin):
        super(PoseViewWidget, self).__init__()
        rp = rospkg.RosPack()
        ui_file = os.path.join(rp.get_path('rqt_pose_view'), 'resource', 'PoseViewWidget.ui')
        loadUi(ui_file, self)
        self._plugin = plugin

        self._position = (0.0, 0.0, 0.0)
        self._orientation = quaternion_about_axis(0.0, (1.0, 0.0, 0.0))
        self._topic_name = None
        self._subscriber = None

        # create GL view
        self._gl_view = GLWidget()
        self._gl_view.setAcceptDrops(True)

        # backup and replace original paint method
        self._gl_view.paintGL_original = self._gl_view.paintGL
        self._gl_view.paintGL = self._gl_view_paintGL

        # backup and replace original mouse release method
        self._gl_view.mouseReleaseEvent_original = self._gl_view.mouseReleaseEvent
        self._gl_view.mouseReleaseEvent = self._gl_view_mouseReleaseEvent

        # add GL view to widget layout
        self.layout().addWidget(self._gl_view)

        # init and start update timer with 40ms (25fps)
        self._update_timer = QTimer(self)
        self._update_timer.timeout.connect(self.update_timeout)
        self._update_timer.start(40)

    def save_settings(self, plugin_settings, instance_settings):
        view_matrix_string = repr(self._gl_view.get_view_matrix())
        instance_settings.set_value('view_matrix', view_matrix_string)

    def restore_settings(self, plugin_settings, instance_settings):
        view_matrix_string = instance_settings.value('view_matrix')
        try:
            view_matrix = eval(view_matrix_string)
        except Exception:
            view_matrix = None

        if view_matrix is not None:
            self._gl_view.set_view_matrix(view_matrix)
        else:
            self._set_default_view()

    def _set_default_view(self):
        self._gl_view.makeCurrent()
        self._gl_view.reset_view()
        self._gl_view.rotate((0, 0, 1), 45)
        self._gl_view.rotate((1, 0, 0), -45)
        self._gl_view.translate((0, -3, -15))

    def message_callback(self, message):
        self._position = (message.position.x, message.position.y, message.position.z)
        self._orientation = (message.orientation.x, message.orientation.y, message.orientation.z, message.orientation.w)

    def update_timeout(self):
        self._gl_view.makeCurrent()
        self._gl_view.updateGL()

    def _gl_view_paintGL(self):
        self._gl_view.paintGL_original()
        self._paintGLGrid()
        self._paintGLCoorsystem()
        self._paintGLBox()

    def _paintGLBox(self):
        self._position = (2.0, 2.0, 2.0)  # Set fixed translation for now
        glTranslatef(*self._position)     # Translate Box

        matrix = quaternion_matrix(self._orientation)  # convert quaternion to translation matrix
        glMultMatrixf(matrix)             # Rotate Box

        glBegin(GL_QUADS)                 # Start Drawing The Box

        glColor3f(0.0, 1.0, 0.0)
        glVertex3f(1.0, 1.0, -1.0)        # Top Right Of The Quad (Top)
        glVertex3f(-1.0, 1.0, -1.0)       # Top Left Of The Quad (Top)
        glVertex3f(-1.0, 1.0, 1.0)        # Bottom Left Of The Quad (Top)
        glVertex3f(1.0, 1.0, 1.0)         # Bottom Right Of The Quad (Top)

        glColor3f(0.5, 1.0, 0.5)
        glVertex3f(1.0, -1.0, 1.0)        # Top Right Of The Quad (Bottom)
        glVertex3f(-1.0, -1.0, 1.0)       # Top Left Of The Quad (Bottom)
        glVertex3f(-1.0, -1.0, -1.0)      # Bottom Left Of The Quad (Bottom)
        glVertex3f(1.0, -1.0, -1.0)       # Bottom Right Of The Quad (Bottom)

        glColor3f(0.0, 0.0, 1.0)
        glVertex3f(1.0, 1.0, 1.0)         # Top Right Of The Quad (Front)
        glVertex3f(-1.0, 1.0, 1.0)        # Top Left Of The Quad (Front)
        glVertex3f(-1.0, -1.0, 1.0)       # Bottom Left Of The Quad (Front)
        glVertex3f(1.0, -1.0, 1.0)        # Bottom Right Of The Quad (Front)

        glColor3f(0.5, 0.5, 1.0)
        glVertex3f(1.0, -1.0, -1.0)       # Bottom Left Of The Quad (Back)
        glVertex3f(-1.0, -1.0, -1.0)      # Bottom Right Of The Quad (Back)
        glVertex3f(-1.0, 1.0, -1.0)       # Top Right Of The Quad (Back)
        glVertex3f(1.0, 1.0, -1.0)        # Top Left Of The Quad (Back)

        glColor3f(1.0, 0.5, 0.5)
        glVertex3f(-1.0, 1.0, 1.0)        # Top Right Of The Quad (Left)
        glVertex3f(-1.0, 1.0, -1.0)       # Top Left Of The Quad (Left)
        glVertex3f(-1.0, -1.0, -1.0)      # Bottom Left Of The Quad (Left)
        glVertex3f(-1.0, -1.0, 1.0)       # Bottom Right Of The Quad (Left)

        glColor3f(1.0, 0.0, 0.0)
        glVertex3f(1.0, 1.0, -1.0)        # Top Right Of The Quad (Right)
        glVertex3f(1.0, 1.0, 1.0)         # Top Left Of The Quad (Right)
        glVertex3f(1.0, -1.0, 1.0)        # Bottom Left Of The Quad (Right)
        glVertex3f(1.0, -1.0, -1.0)       # Bottom Right Of The Quad (Right)
        glEnd()                           # Done Drawing The Quad

    def _paintGLGrid(self):
        resolutionMillimeters = 1
        griddedAreaSize = 100

        glLineWidth(1.0)

        glBegin(GL_LINES)

        glColor3f(1.0, 1.0, 1.0)

        glVertex3f(griddedAreaSize, 0, 0)
        glVertex3f(-griddedAreaSize, 0, 0)
        glVertex3f(0, griddedAreaSize, 0)
        glVertex3f(0, -griddedAreaSize, 0)

        numOfLines = int(griddedAreaSize / resolutionMillimeters)

        for i in range(numOfLines):
            glVertex3f(resolutionMillimeters * i, -griddedAreaSize, 0)
            glVertex3f(resolutionMillimeters * i, griddedAreaSize, 0)
            glVertex3f(griddedAreaSize, resolutionMillimeters * i, 0)
            glVertex3f(-griddedAreaSize, resolutionMillimeters * i, 0)

            glVertex3f(resolutionMillimeters * (-i), -griddedAreaSize, 0)
            glVertex3f(resolutionMillimeters * (-i), griddedAreaSize, 0)
            glVertex3f(griddedAreaSize, resolutionMillimeters * (-i), 0)
            glVertex3f(-griddedAreaSize, resolutionMillimeters * (-i), 0)

        glEnd()

    def _paintGLCoorsystem(self):
        glLineWidth(10.0)

        glBegin(GL_LINES)

        glColor3f(1.0, 0.0, 0.0)
        glVertex3f(0.0, 0.0, 0.0)
        glVertex3f(1.0, 0.0, 0.0)

        glColor3f(0.0, 1.0, 0.0)
        glVertex3f(0.0, 0.0, 0.0)
        glVertex3f(0.0, 1.0, 0.0)

        glColor3f(0.0, 0.0, 1.0)
        glVertex3f(0.0, 0.0, 0.0)
        glVertex3f(0.0, 0.0, 1.0)

        glEnd()

    def _gl_view_mouseReleaseEvent(self, event):
        if event.button() == Qt.RightButton:
            menu = QMenu(self._gl_view)
            action = QAction(self._gl_view.tr("Reset view"), self._gl_view)
            menu.addAction(action)
            action.triggered.connect(self._set_default_view)
            menu.exec_(self._gl_view.mapToGlobal(event.pos()))

    @Slot('QDragEnterEvent*')
    def dragEnterEvent(self, event):
        if not event.mimeData().hasText():
            if not hasattr(event.source(), 'selectedItems') or len(event.source().selectedItems()) == 0:
                qWarning('Plot.dragEnterEvent(): not hasattr(event.source(), selectedItems) or len(event.source().selectedItems()) == 0')
                return
            item = event.source().selectedItems()[0]
            ros_topic_name = item.data(0, Qt.UserRole)
            if ros_topic_name is None:
                qWarning('Plot.dragEnterEvent(): not hasattr(item, ros_topic_name_)')
                return
        # TODO: do some checks for valid topic here
        event.acceptProposedAction()

    @Slot('QDropEvent*')
    def dropEvent(self, event):
        if event.mimeData().hasText():
            topic_name = str(event.mimeData().text())
        else:
            droped_item = event.source().selectedItems()[0]
            topic_name = str(droped_item.data(0, Qt.UserRole))

        self.unregister_topic()
        self.subscribe_topic(topic_name)

    def unregister_topic(self):
        if self._subscriber:
            self._subscriber.unregister()

    def subscribe_topic(self, topic_name):
        msg_class, self._topic_name, _ = get_topic_class(topic_name)
        self._subscriber = rospy.Subscriber(self._topic_name, msg_class, self.message_callback)

    def shutdown_plugin(self):
        self.unregister_topic()


"""
使用收盘价训练模型
训练时，使用过去一段时间的历史预测下一天的价格，然后将这个价格与真实价格
两个标量作均方误差，训练模型
"""

import tensorflow as tf
import tushare as ts
import numpy as np
import collections
import matplotlib.pyplot as plt

class LstmModel:
    def __init__(self):
        # 多长时间跨度作为上下文, 10 maybe better
        self.TIME_STEP = 10
        # lstm的隐藏单元个数, 5 maybe better
        self.NUM_HIDDEN = 5
        # 迭代次数 20 maybe better
        self.epochs = 30
        # 目前训练样本到什么位置了
        self.index = self.TIME_STEP - 1
        # 测试样本数
        self.test_sample_num = 40
        # whether print log
        self.log_print_on = False
        # whether plot datasource
        self.plot_figure_on = False
        self.is_plot_line = True

        self.buffer = collections.deque(maxlen=self.TIME_STEP)
        self._session = tf.Session()
        self.init_plot()

    def init_plot(self):
        if self.is_plot_line:
            pass
        else:
            plt.ion()
            plt.axis([0, 600, 6, 15])

    def gen_data(self):
        if self.index >= self.series_length:
            print("error")
        if self.index <= self.TIME_STEP - 1:
            for single in self.stock_example[:self.TIME_STEP]:
                self.buffer.append(single)
        else:
            self.buffer.append(self.stock_example[self.index])

        result = []
        result.append(list(self.buffer))
        return result

    def gen_target(self):
        # return np.array(datasource[index])
        # print(datasource[index + 1][np.newaxis, :])
        return self.stock_example[self.index + 1][np.newaxis, :]

    def get_data(self):
        stock = ts.get_hist_data("000001")
        close_price = stock.close.values
        # print(close_price)
        self.series_length = len(close_price)
        self.stock_example = close_price[:, np.newaxis]
        # print(self.stock_example[:40])

    def build_graph(self):
        self.data = tf.placeholder(tf.float32, [1, self.TIME_STEP, 1])
        self.target = tf.placeholder(tf.float32, [1, 1])

        # cell = tf.nn.rnn_cell.LSTMCell(NUM_HIDDEN)
        cell = tf.nn.rnn_cell.LSTMCell(self.NUM_HIDDEN)
        val, state = tf.nn.dynamic_rnn(cell, self.data, dtype=tf.float32)
        val = tf.transpose(val, [1, 0, 2])
        self.lstm_output = tf.gather(val, self.TIME_STEP - 1)
        self.lstm_output = tf.gather(self.lstm_output, 0)
        self.lstm_output = tf.gather(self.lstm_output, self.NUM_HIDDEN - 1)

        # weight = tf.Variable(tf.truncated_normal(shape=[self.NUM_HIDDEN, 1]), dtype=tf.float32)
        # bias = tf.Variable(tf.constant(10.0, shape=[1, 1]), dtype=tf.float32)
        self.weight = tf.Variable(tf.truncated_normal([1]), dtype=tf.float32)
        self.bias = tf.Variable(tf.constant(10.0, shape=[1]), dtype=tf.float32)

        self.stock_predict_price = tf.mul(self.lstm_output, self.weight) + self.bias

        self.diff = tf.reduce_sum(np.square(self.stock_predict_price - self.target))
        # diff = np.square(lstm_output - target)
        self.minimize = tf.train.AdamOptimizer().minimize(self.diff)

    def train(self):
        init_op = tf.initialize_all_variables()
        self._session.run(init_op)
        for i in range(self.epochs):
            print("epoch %i " % i)
            for day in range(self.TIME_STEP - 1, self.series_length - self.test_sample_num - 1, 1):
                self.index = day
                self._session.run(self.minimize, {self.data: self.gen_data(), self.target: self.gen_target()})
                predict_stock_price = self._session.run(
                    self.stock_predict_price, {self.data: self.gen_data(), self.target: self.gen_target()})
                diff_price = self._session.run(self.diff, {self.data: self.gen_data(), self.target: self.gen_target()})

                if self.index % 100 == 0 and self.log_print_on:
                    print("epoch is %d ,stock price is %f ,and real price is %f, diff is %f"
                          % (i, predict_stock_price, self.gen_target()[0][0], diff_price))

                if i >= self.epochs - 1 and self.plot_figure_on:
                    self.plot_scatter(day, predict_stock_price, self.gen_target()[0][0])

        # self._session.close()

    def plot_scatter(self, no_day, predict, real):
        plt.scatter(no_day, predict, c='r')
        plt.scatter(no_day, real, c='b')
        plt.pause(0.000000001)

    def plot_line(self, days, predict, real):
        plt.plot(days, predict, 'r-')
        plt.plot(days, real, 'b-')
        plt.show()

    def run(self):
        self.get_data()
        self.build_graph()
        self.train()

    def test(self):
        test_index = self.series_length - self.test_sample_num
        predict, real, days = [], [], []
        tmp_index = 1
        for day in range(test_index, self.series_length - 1, 1):
            self.index = day
            predict_stock_price = self._session.run(
                self.stock_predict_price, {self.data: self.gen_data()})
            real_stock_price = self.gen_target()[0][0]
            print("test===> predict price is %f, real price is %f" % (predict_stock_price, real_stock_price))
            days.append(tmp_index)
            predict.append(predict_stock_price)
            real.append(real_stock_price)
            tmp_index += 1
        self.plot_line(days, predict, real)

    # 进行预测时先把数据进行处理成合适的维度，input：1Dim output：3Dim
    def handle_data_format(self, data):
        data = data[:, np.newaxis]
        tmp = []
        for i in data:
            tmp.append(i)
        result = []
        result.append(tmp)
        print(result)
        return result

    # 用训练好的模型进行预测, datasource: 1Dim
    def predict(self, data):
        data = self.handle_data_format(data)
        predict_stock_price = self._session.run(
            self.stock_predict_price, {self.data: data})
        print("predict price is %f" % predict_stock_price)
        return predict_stock_price

    def destory(self):
        self._session.close()

if __name__ == "__main__":
    lstm = LstmModel()
    lstm.run()
    lstm.test()
    lstm.destory()


from tubeserv.api import api
from tubeserv.api.exceptions import InternalException

from flask import jsonify, request, g

from tubeserv.database import db
from tubeserv.database.models import Video, Session

from tubeserv.extractors import YouTubeExtractor, SoundCloudExtractor, DefaultExtractor

from threading import Thread

from random import choice

def metafetcher(video, videos):
    provider, url = video.provider, video.url
    seed = choice(videos)

    if provider == 0:
        video.title = YouTubeExtractor(url).extract_title()
    elif provider == 1:
        video.title = SoundCloudExtractor(url).extract_title()

    if seed[0] == 0:
        video.related_url = YouTubeExtractor(seed[1]).extract_related()
    else:
        sc_title = SoundCloudExtractor(seed[1]).extract_title()
        sc_yt = DefaultExtractor(sc_title).extract_video()
        video.related_url = YouTubeExtractor(sc_yt).extract_related()

    try:
        db.session.add(video)
        db.session.commit()
    except Exception as e:
        db.session.rollback()

@api.route("/sessions/<uuid:session>/videos/", methods=["POST"], defaults={"video":""})
@api.route("/sessions/<uuid:session>/videos/<int:video>", methods=["GET", "PUT", "PATCH", "DELETE"])
def video(**kwargs):
    session = getattr(g, "session", None)
    video = getattr(g, "video", None)
    json = getattr(g, "json", None)
    res = None

    metafetcher_thread = None

    if video:
        if request.method in ["PUT","PATCH"]:
            video.state = json.get("state", video.state)
        elif request.method in ["DELETE"]:
            res = video.to_dict()
            db.session.delete(video)
    else:
        video = Video(session_id=session.id,
                client_id=json.get("client_id"))
        video.url = json.get("url")
        video.provider = json.get("provider", 0)

        videos = [(v.provider, v.url) for v in
                Video.query.filter(Video.session_id == session.id).all()]
        videos.append((video.provider, video.url))

        metafetcher_thread = Thread(target=metafetcher, args=(video, videos))

        db.session.add(video)

    try:
        db.session.commit()

        if metafetcher_thread is not None:
            db.session.refresh(video)
            metafetcher_thread.start()

        res = video.to_dict()
    except Exception as e:
        db.session.rollback()

        raise InternalException()

    return jsonify(res)


import uuid

from boto.exception import SWFResponseError
from boto.swf.layer1_decisions import Layer1Decisions

from flowy.utils import logger

INPUT_SIZE = RESULT_SIZE = 32768
REASON_SIZE = 256

class SWFActivityDecision(object):
    def __init__(self, layer1, token):
        self.layer1 = layer1
        self.token = token

    def heartbeat(self):
        try:
            self.layer1.record_activity_task_heartbeat(task_token=str(self.token))
        except SWFResponseError:
            logger.exception('Error while sending the heartbeat:')
            return False
        return True

    def fail(self, reason):
        try:
            self.layer1.respond_activity_task_failed(
                reason=str(reason)[:256], task_token=str(self.token))
        except SWFResponseError:
            logger.exception('Error while failing the activity:')
            return False
        return True

    def flush(self):
        self.fail("Cannot flush activities.")

    def restart(self, input_data):
        self.fail("Cannot restart activities.")

    def finish(self, result):
        result = str(result)
        if len(result) > RESULT_SIZE:
            self.fail("Result too large: %s/%s" % (len(result), RESULT_SIZE))
        try:
            self.layer1.respond_activity_task_completed(
                result=result, task_token=str(self.token))
        except SWFResponseError:
            logger.exception('Error while finishing the activity:')
            return False
        return True

class SWFWorkflowDecision(object):
    def __init__(self, layer1, token, name, version, task_list,
                 decision_duration, workflow_duration, tags, child_policy):
        self.layer1 = layer1
        self.token = token
        self.task_list = task_list
        self.decision_duration = decision_duration
        self.workflow_duration = workflow_duration
        self.tags = tags
        self.child_policy = child_policy
        self.decisions = Layer1Decisions()
        self.closed = False

    def fail(self, reason):
        """Fail the workflow and flush.

        Any other decisions queued are cleared.
        The reason is truncated if too large.
        """
        decisions = self.decisions = Layer1Decisions()
        decisions.fail_workflow_execution(reason=str(reason)[:REASON_SIZE])
        self.flush()

    def flush(self):
        """Flush the decisions; no other decisions can be sent after that."""
        if self.closed:
            return
        self.closed = True
        try:
            self.layer1.respond_decision_task_completed(
                task_token=str(self.token), decisions=self.decisions._data)
        except SWFResponseError:
            logger.exception('Error while sending the decisions:')
            # ignore the error and let the decision timeout and retry

    def restart(self, input_data):
        """Restart the workflow and flush.

        Any other decisions queued are cleared.
        """
        decisions = self.decisions = Layer1Decisions()
        input_data = str(input_data)
        if len(input_data) > INPUT_SIZE:
            self.fail("Restart input too large: %s/%s" % (len(input_data), INPUT_SIZE))
        else:
            decisions.continue_as_new_workflow_execution(
                start_to_close_timeout=self.decision_duration,
                execution_start_to_close_timeout=self.workflow_duration,
                task_list=self.task_list,
                input=input_data,
                tag_list=self.tags,
                child_policy=self.child_policy)
        self.flush()

    def finish(self, result):
        """Finish the workflow execution and flush.

        Any other decisions queued are cleared.
        """
        decisions = self.decisions = Layer1Decisions()
        result = str(result)
        if len(result) > RESULT_SIZE:
            self.fail("Result too large: %s/%s" % (len(result), RESULT_SIZE))
        else:
            decisions.complete_workflow_execution(result)
            self.flush()

    def schedule_timer(self, call_key, delay):
        """Schedule a timer. This is used to delay execution of tasks."""
        self.decisions.start_timer(timer_id=timer_key(call_key),
                                   start_to_fire_timeout=str(delay))

    def schedule_activity(self, call_key, name, version, input_data, task_list,
                          heartbeat, schedule_to_close, schedule_to_start,
                          start_to_close):
        """Schedule an activity execution."""
        input_data = str(input_data)
        if len(input_data) > INPUT_SIZE:
            self.fail("Activity input too large: %s/%s" % (len(input_data), INPUT_SIZE))
        self.decisions.schedule_activity_task(
            call_key, name, version,
            heartbeat_timeout=heartbeat,
            schedule_to_close_timeout=schedule_to_close,
            schedule_to_start_timeout=schedule_to_start,
            start_to_close_timeout=start_to_close,
            task_list=task_list,
            input=input_data)

    def schedule_workflow(self, call_key, name, version, input_data, task_list,
                          workflow_duration, decision_duration, child_policy):
        """Schedule a workflow execution."""
        input_data = str(input_data)
        if len(input_data) > INPUT_SIZE:
            self.fail("Workflow input too large: %s/%s" % (len(input_data), INPUT_SIZE))
        call_key = '%s:%s' % (uuid.uuid4(), call_key)
        self.decisions.start_child_workflow_execution(
            name, version, call_key,
            task_start_to_close_timeout=decision_duration,
            execution_start_to_close_timeout=workflow_duration,
            task_list=task_list,
            input=input_data,
            child_policy=child_policy)

class SWFWorkflowTaskDecision(object):
    def __init__(self, decision, execution_history, proxy_factory, rate_limit):
        self.decision = decision
        self.execution_history = execution_history
        self.proxy_factory = proxy_factory
        self.rate_limit = rate_limit

    def fail(self, reason):
        self.decision.fail(reason)

    def schedule(self, call_number, retry_number, delay, input_data):
        if not self.rate_limit.consume():
            return
        tk = task_key(self.proxy_factory.identity, call_number, retry_number)
        if delay > 0:
            if self.execution_history.is_timer_ready(tk):
                self._schedule(tk, input_data)
            elif not self.execution_history.is_timer_running(tk):
                self.decision.schedule_timer(tk, delay)
        else:
            self._schedule(tk, input_data)

    def _schedule(self, task_key, input_data):
        self.decision.schedule_workflow(
            task_key, self.proxy_factory.name, self.proxy_factory.version, input_data,
            self.proxy_factory.task_list, self.proxy_factory.workflow_duration,
            self.proxy_factory.decision_duration, self.proxy_factory.child_policy)

class SWFActivityTaskDecision(SWFWorkflowTaskDecision):
    def _schedule(self, task_key, input_data):
        self.decision.schedule_activity(
            task_key, self.proxy_factory.name, self.proxy_factory.version, input_data,
            self.proxy_factory.task_list, self.proxy_factory.heartbeat,
            self.proxy_factory.schedule_to_close, self.proxy_factory.schedule_to_start,
            self.proxy_factory.start_to_close)

def timer_key(call_key):
    return '%s:t' % call_key

def task_key(identity, call_number, retry_number):
    return '%s-%s-%s' % (identity, call_number, retry_number)


from __future__ import absolute_import
from __future__ import unicode_literals

from django.apps import apps
from django.conf import settings
from django.db import migrations

from corehq.sql_db.operations import RawSQLMigration

def get_view_migrations():
    sql_views = [
        'awc_location_months.sql',
        'agg_awc_monthly.sql',
        'agg_ccs_record_monthly.sql',
        'agg_child_health_monthly.sql',
        'daily_attendance.sql',
        'agg_awc_daily.sql',
        'child_health_monthly.sql',
        'disha_indicators.sql',
        'ccs_record_monthly_view.sql',
        'agg_ls_monthly.sql',
        'service_delivery_monthly.sql',
        'aww_incentive_report_monthly.sql',
    ]
    migrator = RawSQLMigration(('custom', 'icds_reports', 'migrations', 'sql_templates', 'database_views'))
    operations = []
    for view in sql_views:
        operations.append(migrator.get_migration(view))
    return operations

def _citus_composite_key_sql(model_cls):
    pkey_name = '{}_pkey'.format(model_cls._meta.db_table)
    fields = list(model_cls._meta.unique_together)[0]
    columns = [model_cls._meta.get_field(field).column for field in fields]
    index = '{}_{}_uniq'.format(model_cls._meta.db_table, '_'.join(columns))
    sql = """
        CREATE UNIQUE INDEX {index} on "{table}" ({cols});
        ALTER TABLE "{table}" DROP CONSTRAINT {pkey_name};
        ALTER TABLE "{table}" ADD CONSTRAINT {pkey_name} PRIMARY KEY USING INDEX {index};
    """.format(
        table=model_cls._meta.db_table,
        pkey_name=pkey_name,
        index=index,
        cols=','.join(columns)
    )
    reverse_sql = """
        ALTER TABLE "{table}" DROP CONSTRAINT IF EXISTS {index},
        DROP CONSTRAINT IF EXISTS {pkey_name};
        ALTER TABLE "{table}" ADD CONSTRAINT {pkey_name} PRIMARY KEY ({pkey});
        DROP INDEX IF EXISTS {index};
    """.format(
        table=model_cls._meta.db_table,
        pkey_name=pkey_name,
        pkey=model_cls._meta.pk.name,
        index=index,
    )
    if getattr(settings, 'UNIT_TESTING', False):
        return sql, reverse_sql
    else:
        return migrations.RunSQL.noop, reverse_sql

def get_composite_primary_key_migrations(models_to_update):
    operations = []
    for model_name in models_to_update:
        model = apps.get_model('icds_reports', model_name)
        sql, reverse_sql = _citus_composite_key_sql(model)
        operations.append(migrations.RunSQL(
            sql,
            reverse_sql,
            state_operations=[
                migrations.AlterUniqueTogether(
                    name=model_name,
                    unique_together=model._meta.unique_together,
                ),
            ]
        ))
    return operations


import numpy as np
from keras.preprocessing.image import load_img, img_to_array

def create_dataset(fold, names, base_dir):
    dataset = np.ndarray(shape=(len(names), 50, 160, 3), dtype = np.float32)
    i = 0
    for f in names:
        img = load_img(base_dir + fold + "/" + f)
        x = img_to_array(img)
        x = (x - 128) / 128
        dataset[i] = x
        i += 1
    dataset = dataset[:,:,:,0]
    dataset = np.expand_dims(dataset, axis = 3)
    return dataset

def create_labels(names):
    labels = [np.zeros((len(names),10)) for i in range(5)]
    i = 0
    for n in range(len(names)):
        seq = names[n][0:5]
        for digit in range(len(seq)):
            label = int(seq[digit])
            labels[digit][n][label] = 1
        i += 1
    return labels


import shutil
import os

import torch

class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.val = None
        self.avg = None
        self.sum = None
        self.count = None
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def save_IPN_checkpoint(model_I, model_P, ckpt_dir, epoch=None):
    os.makedirs(ckpt_dir, exist_ok=True)
    if epoch is None:
        dict_network = model_I.state_dict()
        torch.save(dict_network, os.path.join(ckpt_dir, f'I.pth'))
        dict_network = model_P.state_dict()
        torch.save(dict_network, os.path.join(ckpt_dir, f'P.pth'))
    else:
        dict_network = model_I.state_dict()
        torch.save(dict_network, os.path.join(ckpt_dir, f'I_{epoch}.pth'))
        dict_network = model_P.state_dict()
        torch.save(dict_network, os.path.join(ckpt_dir, f'P_{epoch}.pth'))


# vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4
import serial,sys,time
import argparse

#https://github.com/espressif/esptool/wiki/ESP32-Boot-Mode-Selection
if 1 and __name__=="__main__":
    parser = argparse.ArgumentParser(description='''
To reset the esp32:
python send_byte.py -d 0
''')
    parser.add_argument("-s","--send", help="bytes to send in hex format eg 0101 for two bytes")
    parser.add_argument("-d","--dtr",help="sets value 0 or 1 to DTR line in esp32 its GPIO0")
    parser.add_argument("-r","--rts",help="sets value 0 or 1 to RTS line in esp32 its EN")
    parser.add_argument("-b","--baudrate",help="sets baudrate default 512000",default="460800")
    parser.add_argument("-u","--usbdevice",help="sets usbdevice default /dev/ttyUSB1",default='/dev/ttyUSB1')
    args = parser.parse_args()
#import pdb;pdb.set_trace()
ser = serial.Serial(args.usbdevice,baudrate=int(args.baudrate), rtscts=True, dsrdtr=True)
if args.send is not None:
    tosend = b'%c'%int(args.send,16)
    print('sending ...',tosend)
    ser.write(tosend)
    #print(ser.read(1))
    ser.flush()
if  args.rts is not None:
    ser.rts=int(args.rts)
if  args.dtr is not None:
    ser.dtr=int(args.dtr)

time.sleep(1)
#if ser.inWaiting()>0:
#    print(ser.read())
ser.close()


# Imports.
import numpy as np
from multiprocessing import Pool
import matplotlib.pyplot as plt
import numpy.random as npr
import pygame as pg

from SwingyMonkey import SwingyMonkey

class Random(object):
    '''
    This agent jumps randomly.
    '''

    def __init__(self):
        self.last_state = None
        self.last_action = None
        self.last_reward = None

    def reset(self):
        self.last_state = None
        self.last_action = None
        self.last_reward = None

    def action_callback(self, state):
        '''
        Implement this function to learn things and take actions.
        Return 0 if you don't want to jump and 1 if you do.
        '''
        # You might do some learning here
        # based on the current state and the last state.
        # You'll need to select and action and return it.
        # Return 0 to swing and 1 to jump.
        new_action = npr.rand() < 0.1
        new_state = state
        self.last_action = new_action
        self.last_state = new_state
        return self.last_action

    def reward_callback(self, reward):
        '''This gets called so you can see what reward you get.'''
        self.last_reward = reward

class QLearner(object):

    def __init__(self):
        # Constants - Same each game
        self.LEARNING_RATE = 0.8
        self.DISCOUNT = 1
        self.EPSILON = 0.1
        self.A = [0, 1]
        self.SCREEN_WIDTH = 600
        self.SCREEN_HEIGHT = 400
        # Constant - Randomly initialized at start of each game
        self.GRAVITY = None
        self.PROP_REWARD_STATES_LENGTH = 5
        # Q-learning
        self.last_state_full = None
        self.last_state_discrete = None
        self.last_action = None
        self.last_reward = None
        self.q = {}
        # FIFO stack of (s, a) pairs to propagate reward back to
        self.prop_reward_states = []

    def reset(self):
        self.last_state_full = None
        self.last_state_discrete = None
        self.last_action = None
        self.last_reward = None
        self.GRAVITY = None

    def greedy_epsilon(self, potential_actions):
        num = np.random.uniform(0, 1)
        print ("num", num)
        if num < self.EPSILON:
            # Pick random action with probability self.EPSILON
            print ("pot act", potential_actions)
            return np.random.choice(range(len(potential_actions)))
        return np.argmax(potential_actions)

    def discrete_state(self, state):
        tree_dist = state['tree']['dist']
        tree_top = state['tree']['top']
        tree_bot = state['tree']['bot']
        monkey_vel = state['monkey']['vel']
        monkey_top = state['monkey']['top']
        monkey_bot = state['monkey']['bot']

        #
        # Feature extraction
        #
        above_tree_dist = monkey_top - tree_top

        #
        # Transformations
        #

        # Bin all continuous pixel ranges
        dist_to_tree_trans = np.floor_divide(tree_dist, 50)
        if dist_to_tree_trans < 0:
            dist_to_tree_trans = 0
        dist_above_tree_trans = np.floor_divide(above_tree_dist, 50)
        # Take gravity effects on monkey velocity into account
        vel_trans = monkey_vel - (
            self.GRAVITY if self.GRAVITY is not None else 0)
        # Bin velocity
        motion_trans = int(vel_trans / 2)
        new_state = (dist_to_tree_trans, dist_above_tree_trans, motion_trans)
        return new_state

    def calc_gravity(self, current_state, last_state):
        last_monkey_vel = last_state['monkey']['vel']
        current_monkey_vel = current_state['monkey']['vel']
        self.GRAVITY = last_monkey_vel - current_monkey_vel

    def q_lookup(self, s, a):
        if (s, a) not in self.q:
            return 0
        else:
            return self.q[(s, a)]

    def q_update(self, s, a, r, s_prime):
        print("&&&&&&&&&&.   " + str(r))

        q_idx = (s, a)
        self.prop_reward_states.insert(0, q_idx)

        if q_idx not in self.q:
            self.q[q_idx] = 0
            print(
                "New state (" + str(s) + ", " + str(a) + ") = " +
                str(self.q[q_idx]))
        else:
            print(
                "Already seen (" + str(s) + ", " + str(a) + ") = " +
                str(self.q[q_idx]))

        # Q-learning algorithm
        if r < 0:
            print(": R < 0")
            # If non-zero reward, propagate back
            # to all states in self.prop_reward_states
            _s_prime = s_prime
            # Work backwards to update
            # If we have (s1, a1) -> (s2, a2) -> (s3, a3) -> (s4, a4)
            # Then set Q(s3, a3) based on (s4, 0/1),
            # Q(s2, a2) based on Q(s3, 0/1), etc.
            for (_s, _a) in self.prop_reward_states:
                max_q = np.max(
                    [self.q_lookup(_s_prime, 0), self.q_lookup(_s_prime, 1)])
                self.q[(_s, _a)] = (
                    1 - self.LEARNING_RATE) * self.q_lookup(
                    _s, _a) + self.LEARNING_RATE * (
                    r * 100 + self.DISCOUNT * max_q)
                _s_prime = _s
            self.prop_reward_states = []
        else:
            # No penalty, so give slight reward
            max_q = np.max(
                [self.q_lookup(s_prime, 0), self.q_lookup(s_prime, 1)])
            self.q[(s, a)] = (
                1 - self.LEARNING_RATE) * self.q_lookup(
                s, a) + self.LEARNING_RATE * (1 + self.DISCOUNT * max_q)
            if len(self.prop_reward_states) > self.PROP_REWARD_STATES_LENGTH:
                del self.prop_reward_states[-1]

    def action_callback(self, state):
        '''
        Implement this function to learn things and take actions.
        Return 0 if you don't want to jump and 1 if you do.
        '''
        # You might do some learning here
        # based on the current state and the last state.
        # You'll need to select an action and return it.
        # Return 0 to swing and 1 to jump.

        # At state s (self.last_state), we took action
        # a (self.last_action), got reward r (self.last_reward),
        # and now are in state s' and need to decide a'
        s_prime = self.discrete_state(state)

        if self.last_action is None:
            # first action, default to not doing anything to measure gravity
            a_prime = 0
        else:
            # Calculate gravity based on monkey position
            if self.GRAVITY is None:
                self.calc_gravity(state, self.last_state_full)
            # This new action a' will not be the very first action we take
            # Thus, we can safely check previous state/action (s/a)
            a = self.last_action
            s = self.last_state_discrete
            r = self.last_reward
            self.q_update(s, a, r, s_prime)
            a_prime = 0 if self.q_lookup(
                s_prime, 0) >= self.q_lookup(s_prime, 1) else 1
            print(
                "(s, a, r, s', a'): " + str(s) + ", " +
                str(a) + ", " + str(r) + ", " + str(s_prime) +
                ", " + str(a_prime))
            print(
                str(self.q_lookup(s_prime, 0)) + " v. " +
                str(self.q_lookup(s_prime, 1)))

        # Save current s_prime/a_prime as last s/a
        self.last_action = a_prime
        self.last_state_discrete = s_prime
        self.last_state_full = state
        return a_prime

    def reward_callback(self, reward):
        '''This gets called so you can see what reward you get.'''
        self.last_reward = reward

class OldQLearner(object):

    def __init__(self):
        # Constants - Same each game
        self.LEARNING_RATE = 0.8
        self.DISCOUNT = 0.90
        self.EPSILON = 0.1
        self.A = [0, 1]
        self.SCREEN_WIDTH = 600
        self.SCREEN_HEIGHT = 400
        self.NEAR_GROUND = np.floor_divide(100, 50)
        # Constant - Randomly initialized at start of each game
        self.gravity = None
        # Q-learning
        self.last_state_full = None
        self.last_state_discrete = None
        self.last_action = None
        self.last_reward = None
        self.q = {}

    def reset(self):
        self.last_state_full = None
        self.last_state_discrete = None
        self.last_action = None
        self.last_reward = None
        self.gravity = None

    def greedy_epsilon(self, potential_actions):
        num = np.random.uniform(0, 1)
        print ("num", num)
        if num < self.EPSILON:
            # Pick random action with probability self.EPSILON
            # print ("pot act", potential_actions)
            return np.random.choice(range(len(potential_actions)))
        return np.argmax(potential_actions)

    def discrete_state(self, state):
        tree_dist = state['tree']['dist']
        tree_top = state['tree']['top']
        tree_bot = state['tree']['bot']
        monkey_vel = state['monkey']['vel']
        monkey_top = state['monkey']['top']
        monkey_bot = state['monkey']['bot']

        #
        # Feature extraction
        #
        tree_center = (tree_top - tree_bot) / 2 + tree_bot
        monkey_center = (monkey_top - monkey_bot) / 2 + monkey_bot

        dist_to_tree_center = np.sqrt(
            (tree_dist) ** 2 + (tree_center - monkey_center) ** 2)
        above_tree_center = (monkey_center > tree_center)
        # above_tree_center = (monkey_top < tree_top) and (monkey_bot > tree_bot)
        # above_tree_center = True
        dist_to_ground = monkey_center
        # dist_to_tree_top = np.sqrt(tree_dist**2 + (tree_top - monkey_top)**2)
        # dist_to_tree_bot = np.sqrt(tree_dist**2 + (tree_bot - monkey_bot)**2)
        # dist_to_tree = np.array([dist_to_tree_bot, dist_to_tree_top])

        #
        # Transformations
        #

        # Bin all continuous pixel ranges into groups of 50 (e.g. 400-449 -> 4)
        dist_to_tree_center_trans = np.floor_divide(dist_to_tree_center, 25)
        if dist_to_tree_center_trans < 0:
            dist_to_tree_center_trans = 0
        dist_to_ground_trans = np.floor_divide(dist_to_ground, 25)
        dist_to_tree = np.floor_divide(tree_dist, 25)
        # Take gravity effects on monkey velocity into account
        vel_trans = monkey_vel - (self.gravity if self.gravity is not None else 0)
        # Bin velocity into groups of 10
        motion_trans = int(vel_trans/10)
        new_state = (above_tree_center, dist_to_ground_trans, dist_to_tree, motion_trans)
        return new_state

    def calc_gravity(self, current_state, last_state):
        last_monkey_vel = last_state['monkey']['vel']
        current_monkey_vel = current_state['monkey']['vel']
        self.gravity = last_monkey_vel - current_monkey_vel

    def q_lookup(self, s, a):

        # if (s, a) not in self.q:
        #     # Return closest Q entry if near bottom of screen
        #     if s[1] <= self.NEAR_GROUND:
        #         return self.q[]
        #     # Return closest Q entry if tree center is within 50
        #    return 0

        # Action
        closest_a = None
        if a in self.q:
            closest_a = a
        else:
            # Not in Q
            return 0

        # Above tree center
        closest_above_tree_center = None
        if s[0] in self.q[a]:
            closest_above_tree_center = s[0]
        else:
            # Not in Q
            if (not s[0]) in self.q[a]:
                closest_above_tree_center = not s[0]
            else:
                return 0

        # Dist to ground
        closest_dist_to_ground = None
        if s[1] in self.q[a][closest_above_tree_center]:
            closest_dist_to_ground = s[1]
        else:
            # Not in Q
            # Find closest dist to ground
            for dist, q_val in self.q[a][closest_above_tree_center].items():
                if (closest_dist_to_ground is None or np.abs(closest_dist_to_ground - s[2]) > np.abs(int(dist) - s[3])):
                    closest_dist_to_ground = int(dist)
            if closest_dist_to_ground is None:
                return 0

        # Dist to tree
        closest_dist_to_tree = None
        if s[2] in self.q[a][closest_above_tree_center][closest_dist_to_ground]:
            closest_dist_to_tree = s[2]
        else:
            # Not in Q
            # Find closest dist to tree
            for dist, q_val in self.q[a][closest_above_tree_center][closest_dist_to_ground].items():
                if closest_dist_to_tree is None or np.abs(closest_dist_to_tree - s[2]) > np.abs(int(dist) - s[3]):
                    closest_dist_to_tree = int(dist)
            if closest_dist_to_tree is None:
                return 0

        # Motion
        closest_motion = None
        if s[3] in self.q[a][closest_above_tree_center][closest_dist_to_ground][closest_dist_to_tree]:
            closest_motion = s[3]
        else:
            # Not in Q
            # Find closest motion
            for motion, q_val in self.q[a][closest_above_tree_center][closest_dist_to_ground][closest_dist_to_tree].items():
                if closest_motion is None or np.abs(closest_motion - s[3]) > np.abs(int(motion) - s[3]):
                    closest_motion = int(motion)
            if closest_motion is None:
                return 0

        return self.q[a][closest_above_tree_center][closest_dist_to_ground][closest_dist_to_tree][closest_motion]

    def q_update(self, s, a, r, s_prime):
        try:
            # print(
            #     "Already seen (" + str(s) + ", " + str(a) + ") = " +
            #     str(self.q[a][s[0]][s[1]][s[2]][s[3]]))
            # Q-learning algorithm

            max_q = np.max(
                [self.q_lookup(s_prime, 0), self.q_lookup(s_prime, 1)])
            self.q[a][s[0]][s[1]][s[2]][s[3]] = (
                1 - self.LEARNING_RATE) * self.q_lookup(
                s, a) + self.LEARNING_RATE * (r + self.DISCOUNT * max_q)
        except:
            # Haven't seen this state/action pair before in Q
            if a not in self.q:
                self.q[a] = {}
            if s[0] not in self.q[a]:
                self.q[a][s[0]] = {}
            if s[1] not in self.q[a][s[0]]:
                self.q[a][s[0]][s[1]] = {}
            if s[2] not in self.q[a][s[0]][s[1]]:
                self.q[a][s[0]][s[1]][s[2]] = {}
            if s[3] not in self.q[a][s[0]][s[1]][s[2]]:
                self.q[a][s[0]][s[1]][s[2]][s[3]] = 0
            # print(
            #     "New state (" + str(s) + ", " + str(a) + ") = " +
            #     str(self.q[a][s[0]][s[1]][s[2]][s[3]]))

    def action_callback(self, state):
        '''
        Implement this function to learn things and take actions.
        Return 0 if you don't want to jump and 1 if you do.
        '''
        # You might do some learning here
        # based on the current state and the last state.
        # You'll need to select an action and return it.
        # Return 0 to swing and 1 to jump.
        # print("action call", state)

        # At state s (self.last_state), we took action a (self.last_action), got reward r (self.last_reward), and now are in state s' and need to decide a'
        s_prime = self.discrete_state(state)

        if self.last_action is None:
            # first action, default to not doing anything to measure gravity
            a_prime = 0

        else:
            # Calculate gravity based on monkey position
            if self.gravity is None:
                self.calc_gravity(state, self.last_state_full)
            # This new action a' will not be the very first action we take
            # Thus, we can safely check previous state/action (s/a)
            a = self.last_action
            s = self.last_state_discrete
            r = self.last_reward
            self.q_update(s, a, r, s_prime)

            swing = self.q_lookup(s_prime, 0)
            jump = self.q_lookup(s_prime, 1)

            a_prime = np.argmax([swing, jump])

            # print(
            #     "(s, a, r, s', a'): " + str(s) + ", " + str(a) + ", " +
            #     str(r) + ", " + str(s_prime) + ", " + str(a_prime))
            # print(
            #     str(self.q_lookup(s_prime, 0)) + " v. " +
            #     str(self.q_lookup(s_prime, 1)))

            # Get new policy from Q
            a_prime = self.greedy_epsilon([swing, jump])

        # Save current s_prime/a_prime as last s/a
        self.last_action = a_prime
        self.last_state_discrete = s_prime
        self.last_state_full = state
        return a_prime

    def reward_callback(self, reward):
        '''This gets called so you can see what reward you get.'''
        self.last_reward = reward

class DhruvQLearner(object):

    def __init__(self, discount=0.90, epsilon=0.01):
        # Constants - Same each game
        self.BINS = 50
        self.VBINS = 5

        self.DISCOUNT = discount
        self.EPSILON = epsilon
        self.SCREEN_WIDTH = 600
        self.SCREEN_HEIGHT = 400
        # self.gravity = None

        self.Q = np.zeros(
            (2, int(self.SCREEN_WIDTH / self.BINS) + 1,
             int(self.SCREEN_HEIGHT / self.BINS) + 1, self.VBINS))
        self.k = np.zeros(
            (2, int(self.SCREEN_WIDTH / self.BINS) + 1,
             int(self.SCREEN_HEIGHT / self.BINS) + 1, self.VBINS))

        self.iters = 0
        self.mem = [0, 0]
        self.epoch = 1

        self.last_state = None
        self.last_action = None
        self.last_reward = None

    def reset(self):
        self.last_state = None
        self.last_action = None
        self.last_reward = None
        self.epoch += 1

    def action_callback(self, state):
        '''
        Implement this function to learn things and take actions.
        Return 0 if you don't want to jump and 1 if you do.
        '''
        # You might do some learning here
        # based on the current state and the last state.
        # You'll need to select an action and return it.
        # Return 0 to swing and 1 to jump.

        self.iters += 1

        tree_dist = int(state['tree']['dist'] / self.BINS)

        if tree_dist < 0:
            tree_dist = 0

        tree_bot = int(
            (state['tree']['top'] - state['monkey']['top'] + 0) / self.BINS)

        monkey_vel = int((state['monkey']['vel']) / 20)

        if np.abs(monkey_vel) > 2:
            monkey_vel = np.sign(monkey_vel) * 2
        if npr.rand() < 0.5:
            aprime = 0
        else:
            aprime = 1

        if not self.last_action == None:

            last_tree_dist = int(self.last_state['tree']['dist'] / self.BINS)
            last_tree_bot = int(
                (self.last_state['tree']['top'] -
                    self.last_state['monkey']['top']) / self.BINS)
            last_monkey_vel = int(self.last_state['monkey']['vel'] / 20)

            if np.abs(last_monkey_vel) > 2:
                last_monkey_vel = np.sign(last_monkey_vel) * 2

            qmax = np.max(self.Q[:, tree_dist, tree_bot, monkey_vel])

            if self.Q[1][tree_dist, tree_bot, monkey_vel] > self.Q[0][tree_dist, tree_bot, monkey_vel]:
                aprime = 1
            else:
                aprime = 0

            if self.k[aprime][tree_dist, tree_bot, monkey_vel] > 0:
                # reduce epsilon if you have a high Q reward
                eps = self.EPSILON / self.k[aprime][tree_dist, tree_bot, monkey_vel]
            else:
                eps = self.EPSILON

            r = npr.rand()

            if r < eps:
                if r < 0.5:
                    aprime = 0
                else:
                    aprime = 1

            alph = 1 / self.k[self.last_action][last_tree_dist, last_tree_bot, last_monkey_vel] # dynamically set learning rate
            self.Q[self.last_action][last_tree_dist, last_tree_bot, last_monkey_vel] += alph*(self.last_reward + self.DISCOUNT * qmax - self.Q[self.last_action][last_tree_dist, last_tree_bot, last_monkey_vel])

        self.mem[0] = state['monkey']['top']

        self.last_action = aprime
        self.last_state = state
        self.k[aprime][tree_dist, tree_bot, monkey_vel] += 1  # ERROR HERE
        return aprime

    def reward_callback(self, reward):
        '''This gets called so you can see what reward you get.'''

        self.last_reward = reward

def run_games(learner, hist, iters=100, t_len=100):

    '''
    Driver function to simulate learning by
    having the agent play a sequence of games.
    '''
    for ii in range(iters):
        # Make a new monkey object.
        swing = SwingyMonkey(  # Don't play sounds.
            sound=False,
            text="Epoch %d" % (ii),  # Display the epoch on screen.
            tick_length=t_len,          # Make game ticks super fast.
            action_callback=learner.action_callback,
            reward_callback=learner.reward_callback)
        # Loop until you hit something.
        while swing.game_loop():
            pass
        # Save score history.
        hist.append(swing.score)
        # print(swing.score)
        # Reset the state of the learner.
        # print("---- Q ----")
        # print(learner.Q)
        # print("\n\n")
        # print("---- NEW GAME ----")
        learner.reset()
    pg.quit()
    return

if __name__ == '__main__':

    TEST = False

    if TEST:
        ITERS = 500
        TRIALS = 9
        epsilon_to_test = [0.01, 0.03, 0.05, 0.07, 0.09]
        discount_to_test = [0.5, 0.6, 0.7, 0.8, 0.9]

        def epsilon_test(epsilon):
            for trial in range(TRIALS):
                # Select agent.
                agent = DhruvQLearner(epsilon=epsilon)

                # Empty list to save history.
                hist = []

                # Run games.
                run_games(agent, hist, ITERS, 1)

                # Save history.
                print("epsilon - " + str(epsilon))
                print(hist)
                np.save(
                    'hist-epsilon-%f-%i' % (epsilon, trial + 10),
                    np.array(hist))

        def discount_test(discount):
            for trial in TRIALS:
                # Select agent.
                agent = DhruvQLearner(discount=discount, epsilon=0.09)

                # Empty list to save history.
                hist = []

                # Run games.
                run_games(agent, hist, ITERS, 1)
                print("discount - " + str(discount))
                print(hist)

                np.save(
                    'hist-discount-%f-%i' % (discount, trial), np.array(hist))

        epsilon_test(float(input("epsilon:")))

        # for idx, discount in enumerate(discount_to_test):
        #     pool.map(discount_test, discount)
    else:
        # Run normally with optimal parameters
        agent = DhruvQLearner(epsilon=0.09, discount=0.7)
        hist = []
        run_games(agent, hist, 500, 1)
        print('Optimal')
        print(hist)
        np.save('hist-optimal', np.array(hist))


import networkx as nx
import matplotlib.pyplot as plot

A=nx.DiGraph()
A.add_edge('Pc1','Pc1')
A.add_edge('Pc1','Pc2')
A.add_edge('Pc1','Pc3')
A.add_edge('Pc1','Pc4')
A.add_edge('Pc1','Pc5')
A.add_edge('Pc1','Pc6')
A.add_edge('Pc1','Pc7')
nodes_reflex = {'Pc1'}
nodes_no_reflex = {'Pc2','Pc3','Pc4','Pc5','Pc6','Pc6','Pc7'}
posicion=nx.spring_layout(A)
nx.draw_networkx_nodes(A,posicion,nodelist=nodes_reflex,
                       node_size=500, node_color= 'r')
nx.draw_networkx_nodes(A,posicion,nodelist=nodes_no_reflex,
                       node_size=500, node_color= 'y')
nx.draw_networkx_edges(A,posicion, width=2)
nx.draw_networkx_labels(A,posicion, font_size=11,font_family='arial')

plot.axis('off')
plot.savefig("Graf6.eps")
plot.show(A)


import xml.etree.ElementTree as ET
tree = ET.parse('../cddata.xml')
root = tree.getroot()

CDCheapest = 0
CDExpensive = 0

for cd in root.iter('CD'):
    cdPrice = float(cd.find('PRICE').text)

    # Checks if it's cheaper than CDCheapest
    if CDCheapest == 0:
        CDCheapest = cd
    else:
        CDCheapestPrice = float(CDCheapest.find('PRICE').text)
        if cdPrice < CDCheapestPrice:
            CDCheapest = cd

    # Checks if it's more expensive than CDExpensive
    if CDExpensive == 0:
        CDExpensive = cd
    else:
        CDExpensivePrice = float(CDExpensive.find('PRICE').text)
        if cdPrice > CDExpensivePrice:
            CDExpensive = cd

print("The cheapes CD was", CDCheapest.find('TITLE').text, "by", CDCheapest.find('ARTIST').text, "at", CDCheapest.find('PRICE').text)
print("The most expensive CD was", CDExpensive.find('TITLE').text, "by", CDExpensive.find('ARTIST').text, "at", CDExpensive.find('PRICE').text)

Kveldsarrangement/21.02 - Intro til datainnhenting/XML/Fasit/Oppgave3.py


import numpy as np
import matplotlib.pyplot as plt

import subprocess
from tempfile import mkstemp

from sklearn.tree import export_graphviz

from IPython.core.display import HTML

def plot_surface(clf, X, y, n_steps=250, subplot=None, show=True,
                 ylim=None, xlim=None):
    if subplot is None:
        fig = plt.figure()
    else:
        plt.subplot(*subplot)

    if xlim is None:
        xlim = X[:, 0].min(), X[:, 0].max()
    if ylim is None:
        ylim = X[:, 1].min(), X[:, 1].max()
    xx, yy = np.meshgrid(np.linspace(xlim[0], xlim[1], n_steps),
                         np.linspace(ylim[0], ylim[1], n_steps))

    z = clf.predict(np.c_[xx.ravel(), yy.ravel()])[:, 1]

    z = z.reshape(xx.shape)
    plt.contourf(xx, yy, z, alpha=0.8, cmap=plt.cm.RdBu_r)
    plt.scatter(X[:, 0], X[:, 1], c=y)
    plt.xlim(*xlim)
    plt.ylim(*ylim)

    if show:
        plt.show()


"""The gen-sample command."""
import logging
import os
import shutil
import sys

from ..runway_command import RunwayCommand
from ...tfenv import get_latest_tf_version

LOGGER = logging.getLogger('runway')

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))

def generate_sample_module(module_dir):
    """Generate skeleton sample module."""
    if os.path.isdir(module_dir):
        LOGGER.error("Error generating sample module -- directory %s "
                     "already exists!",
                     module_dir)
        sys.exit(1)
    os.mkdir(module_dir)

def generate_sample_sls_module(env_root, module_dir=None):
    """Generate skeleton Serverless sample module."""
    if module_dir is None:
        module_dir = os.path.join(env_root, 'sampleapp.sls')
    generate_sample_module(module_dir)
    for i in ['config-dev-us-east-1.json', 'handler.py', 'package.json',
              'serverless.yml']:
        shutil.copyfile(
            os.path.join(ROOT,
                         'templates',
                         'serverless',
                         i),
            os.path.join(module_dir, i),
        )
    LOGGER.info("Sample Serverless module created at %s",
                module_dir)

def generate_sample_sls_tsc_module(env_root, module_dir=None):
    """Generate skeleton Serverless TypeScript sample module."""
    if module_dir is None:
        module_dir = os.path.join(env_root, 'sampleapp.sls')
    generate_sample_module(module_dir)
    for i in ['package.json', 'serverless.yml', 'tsconfig.json',
              'webpack.config.js']:
        shutil.copyfile(
            os.path.join(ROOT,
                         'templates',
                         'sls-tsc',
                         i),
            os.path.join(module_dir, i),
        )
    os.mkdir(os.path.join(module_dir, 'src'))
    for i in ['handler.spec.ts', 'handler.ts']:
        shutil.copyfile(
            os.path.join(ROOT,
                         'templates',
                         'sls-tsc',
                         'src',
                         i),
            os.path.join(module_dir, 'src', i),
        )
    LOGGER.info("Sample Serverless TypeScript module created at %s",
                module_dir)

def generate_sample_cdk_tsc_module(env_root, module_dir=None):
    """Generate skeleton CDK TS sample module."""
    if module_dir is None:
        module_dir = os.path.join(env_root, 'sampleapp.cdk')
    generate_sample_module(module_dir)
    for i in ['.npmignore', 'cdk.json', 'package.json', 'runway.module.yml',
              'tsconfig.json', 'README.md']:
        shutil.copyfile(
            os.path.join(ROOT,
                         'templates',
                         'cdk-tsc',
                         i),
            os.path.join(module_dir, i),
        )
    for i in [['bin', 'sample.ts'], ['lib', 'sample-stack.ts']]:
        os.mkdir(os.path.join(module_dir, i[0]))
        shutil.copyfile(
            os.path.join(ROOT,
                         'templates',
                         'cdk-tsc',
                         i[0],
                         i[1]),
            os.path.join(module_dir, i[0], i[1]),
        )
    with open(os.path.join(module_dir, '.gitignore'), 'w') as stream:
        stream.write('*.js\n')
        stream.write('*.d.ts\n')
        stream.write('node_modules\n')
    LOGGER.info("Sample CDK module created at %s", module_dir)
    LOGGER.info('To finish its setup, change to the %s directory and execute '
                '"npm install" to generate its lockfile.', module_dir)

def generate_sample_cdk_cs_module(env_root, module_dir=None):
    """Generate skeleton CDK C# sample module."""
    if module_dir is None:
        module_dir = os.path.join(env_root, 'sampleapp.cdk')
    generate_sample_module(module_dir)
    for i in ['add-project.hook.d.ts', 'cdk.json', 'package.json',
              'runway.module.yml', 'README.md']:
        shutil.copyfile(
            os.path.join(ROOT,
                         'templates',
                         'cdk-csharp',
                         i),
            os.path.join(module_dir, i),
        )
    shutil.copyfile(
        os.path.join(ROOT,
                     'templates',
                     'cdk-csharp',
                     'dot_gitignore'),
        os.path.join(module_dir, '.gitignore'),
    )
    os.mkdir(os.path.join(module_dir, 'src'))
    shutil.copyfile(
        os.path.join(ROOT,
                     'templates',
                     'cdk-csharp',
                     'src',
                     'HelloCdk.sln'),
        os.path.join(module_dir, 'src', 'HelloCdk.sln'),
    )
    os.mkdir(os.path.join(module_dir, 'src', 'HelloCdk'))
    for i in ['HelloCdk.csproj', 'HelloConstruct.cs', 'HelloStack.cs',
              'Program.cs']:
        shutil.copyfile(
            os.path.join(ROOT,
                         'templates',
                         'cdk-csharp',
                         'src',
                         'HelloCdk',
                         i),
            os.path.join(module_dir, 'src', 'HelloCdk', i),
        )
    LOGGER.info("Sample C# CDK module created at %s", module_dir)
    LOGGER.info('To finish its setup, change to the %s directory and execute '
                '"npm install" to generate its lockfile.', module_dir)

def generate_sample_cdk_py_module(env_root, module_dir=None):
    """Generate skeleton CDK python sample module."""
    if module_dir is None:
        module_dir = os.path.join(env_root, 'sampleapp.cdk')
    generate_sample_module(module_dir)
    for i in ['app.py', 'cdk.json', 'lambda-index.py', 'package.json',
              'runway.module.yml', 'Pipfile']:
        shutil.copyfile(
            os.path.join(ROOT,
                         'templates',
                         'cdk-py',
                         i),
            os.path.join(module_dir, i),
        )
    with open(os.path.join(module_dir, '.gitignore'), 'w') as stream:
        stream.write('node_modules')
    LOGGER.info("Sample CDK module created at %s", module_dir)
    LOGGER.info('To finish its setup, change to the %s directory and execute '
                '"npm install" and "pipenv update -d --three" to generate its '
                'lockfiles.', module_dir)

def generate_sample_cfn_module(env_root, module_dir=None):
    """Generate skeleton CloudFormation sample module."""
    if module_dir is None:
        module_dir = os.path.join(env_root, 'sampleapp.cfn')
    generate_sample_module(module_dir)
    for i in ['stacks.yaml', 'dev-us-east-1.env']:
        shutil.copyfile(
            os.path.join(ROOT,
                         'templates',
                         'cfn',
                         i),
            os.path.join(module_dir, i)
        )
    os.mkdir(os.path.join(module_dir, 'templates'))
    shutil.copyfile(
        os.path.join(ROOT,
                     'templates',
                     'cfn',
                     'templates',
                     'tf_state.yml'),
        os.path.join(module_dir, 'templates', 'tf_state.yml')
    )
    LOGGER.info("Sample CloudFormation module created at %s",
                module_dir)

def generate_sample_stacker_module(env_root, module_dir=None):
    """Generate skeleton Stacker sample module."""
    if module_dir is None:
        module_dir = os.path.join(env_root,
                                  'runway-sample-tfstate.cfn')
    generate_sample_module(module_dir)
    for i in ['stacks.yaml', 'dev-us-east-1.env']:
        shutil.copyfile(
            os.path.join(ROOT,
                         'templates',
                         'stacker',
                         i),
            os.path.join(module_dir, i)
        )
    os.mkdir(os.path.join(module_dir, 'tfstate_blueprints'))
    for i in ['__init__.py', 'tf_state.py']:
        shutil.copyfile(
            os.path.join(ROOT,
                         'templates',
                         'stacker',
                         'tfstate_blueprints',
                         i),
            os.path.join(module_dir, 'tfstate_blueprints', i)
        )
    os.chmod(  # make blueprint executable
        os.path.join(module_dir, 'tfstate_blueprints', 'tf_state.py'),
        os.stat(os.path.join(module_dir,
                             'tfstate_blueprints',
                             'tf_state.py')).st_mode | 0o0111
    )
    LOGGER.info("Sample Stacker module created at %s",
                module_dir)

def generate_sample_tf_module(env_root, module_dir=None):
    """Generate skeleton Terraform sample module."""
    if module_dir is None:
        module_dir = os.path.join(env_root, 'sampleapp.tf')
    generate_sample_module(module_dir)
    for i in ['backend-us-east-1.tfvars', 'dev-us-east-1.tfvars', 'main.tf']:
        shutil.copyfile(
            os.path.join(ROOT,
                         'templates',
                         'terraform',
                         i),
            os.path.join(module_dir, i),
        )
    tf_ver_template = os.path.join(ROOT,
                                   'templates',
                                   'terraform',
                                   '.terraform-version')
    if os.path.isfile(tf_ver_template):
        shutil.copyfile(
            tf_ver_template,
            os.path.join(module_dir, '.terraform-version'),
        )
    else:  # running directly from git
        latest_tf_ver = get_latest_tf_version()
        with open(os.path.join(module_dir,
                               '.terraform-version'), 'w') as stream:
            stream.write(latest_tf_ver)

    LOGGER.info("Sample Terraform app created at %s",
                module_dir)

class GenSample(RunwayCommand):
    """Extend Base with execute to run the module generators."""

    def execute(self):
        """Run selected module generator."""
        if self._cli_arguments['cfn']:
            generate_sample_cfn_module(self.env_root)
        elif self._cli_arguments['sls']:
            generate_sample_sls_module(self.env_root)
        elif self._cli_arguments['sls-tsc']:
            generate_sample_sls_tsc_module(self.env_root)
        elif self._cli_arguments['stacker']:
            generate_sample_stacker_module(self.env_root)
        elif self._cli_arguments['tf']:
            generate_sample_tf_module(self.env_root)
        elif self._cli_arguments['cdk-tsc']:
            generate_sample_cdk_tsc_module(self.env_root)
        elif self._cli_arguments['cdk-py']:
            generate_sample_cdk_py_module(self.env_root)
        elif self._cli_arguments['cdk-csharp']:
            generate_sample_cdk_cs_module(self.env_root)


"""
VPfit

Fit Voigt Profiles using MCMC. Uses bayesian model selection to pick the appropriate number
of profiles for a given absorption.

The main class containing the fitting functionality is `VPfit`.

"""

import matplotlib.pyplot as plt

import datetime

import numpy as np
import pymc as mc

# Voigt modules
# from scipy.special import wofz
from astropy.modeling.models import Voigt1D

# peak finding modules
from scipy.signal import find_peaks_cwt
from scipy.signal import savgol_filter

from physics import *

#TODO: change the way that "Models are instantiated" (need to understand that first)
import warnings
warnings.filterwarnings("ignore", message="Instantiating a Model object directly is deprecated. We recommend passing variables directly to the Model subclass.")

class VPfit():

    def __init__(self, noise=None):
        """
        Initialise noise variable
        """
        self.std_deviation = 1./(mc.Uniform("sd", 0, 1) if noise is None else noise)**2
        self.verbose = False

    @staticmethod
    def GaussFunction(x, amplitude, centroid, sigma):
        """
        Gaussian.

        Args:
            x (numpy array): wavelength array
            amplitude (float)
            centroid (float): must be between the limits of wavelength_array
            sigma (float)
        """
        return amplitude * np.exp(-0.5 * ((x - centroid) / sigma) ** 2)

    @staticmethod
    def VoigtFunction(x, centroid, amplitude, L_fwhm, G_fwhm):
        """
        Return the Voigt line shape at x with Lorentzian component HWHM gamma
        and Gaussian component HWHM alpha.

        Source: http://scipython.com/book/chapter-8-scipy/examples/the-voigt-profile/

        Args:
            x (numpy array): wavelength array
            centroid (float): center of profile
            alpha (float): Gaussian HWHM
            gamma (float): Lorentzian HWHM
        """

        #sigma = alpha / np.sqrt(2 * np.log(2))
        #return np.real(wofz((x - centroid + 1j*gamma)/sigma/np.sqrt(2))) / sigma / np.sqrt(2*np.pi)

        v = Voigt1D(x_0=centroid, amplitude_L=amplitude, fwhm_L=L_fwhm, fwhm_G=G_fwhm)
        return v(x)

    @staticmethod
    def GaussianWidth(G_fwhm):
        """
        Find the std deviation of the Gaussian associated with a Voigt profile, from the
        Full Width Half Maximum (FWHM)

        Args:
            G_fwhm (float): Full Width Half Maximum (FWHM)
        """
        return G_fwhm / (2.* np.sqrt(2.*np.log(2.)))

    def trunc_normal(self, name, mu, tau, a, b):
        """
        Take a sample from a truncated Gaussian distribution with PyMC Normal.
        Args:
            name (string): name of the component
            mu (float): mean of the distribution
            tau (float): precision of the distribution, which corresponds to 1/sigma**2 (tau > 0).
            a (float): lower bound of the distribution
            b (float): upper bound of the distribution
        """
        for i in range(1000):
            x = pymc.Normal(name, mu, tau)
            if (a < x.value) and (b > x.value):
                break
        if (a > x.value) or (b < x.value):
            raise ValueError('Computational time exceeded.')
        return x

    @staticmethod
    def Chisquared(observed, expected, noise):
        """
        Calculate the chi-squared goodness of fit

        Args:
            observed (array)
            expected (array): same shape as observed
        """
        return sum( ((observed - expected) /  noise)**2 )

    @staticmethod
    def ReducedChisquared(observed, expected, noise, freedom):
        """
        Calculate the reduced chih-squared goodness of fit

        Args:
            observed (array)
            expected (array): same shape as observed
            freedom (int): degrees of freedom
        """
        return VPfit.Chisquared(observed, expected, noise) / freedom

    def plot(self, wavelength_array, flux_array, clouds=None, n=1, onesigmaerror = 0.02,
            start_pix=None, end_pix=None, filename=None):
        """
        Plot the fitted absorption profile

        Args:
            wavelength_array (numpy array):
            flux_array (numpy array): original flux array, same length as wavelength_array
            filename (string): for saving the plot
	    clouds (pandas dataframe): dataframe containing details on each absorption feature
            n (int): number of *fitted* absorption profiles
            onesigmaerror (float): noise on profile plot
        """

        if not start_pix:
            start_pix = 0
        if not end_pix:
            end_pix = len(wavelength_array)

        f, (ax1, ax2, ax3) = plt.subplots(3, sharex=True, sharey=False, figsize=(10,10))

        ax1.plot(wavelength_array, (flux_array - self.total.value) / onesigmaerror)
        ax1.hlines(1, wavelength_array[0], wavelength_array[-1], color='red', linestyles='-')
        ax1.hlines(-1, wavelength_array[0], wavelength_array[-1], color='red', linestyles='-')
        ax1.hlines(3, wavelength_array[0], wavelength_array[-1], color='red', linestyles='--')
        ax1.hlines(-3, wavelength_array[0], wavelength_array[-1], color='red', linestyles='--')

        ax2.plot(wavelength_array, flux_array, color='black', linewidth=1.0)

        if clouds is not None:
            for c in range(len(clouds)):
                if c==0:
                    ax2.plot(wavelength_array, Tau2flux(clouds.ix[c]['tau'][start_pix:end_pix]),
                             color="red", label="Actual", lw=1.5)
                else:
                    ax2.plot(wavelength_array, Tau2flux(clouds.ix[c]['tau'][start_pix:end_pix]),
                             color="red", lw=1.5)

        for c in range(n):
            if c==0:
                ax2.plot(wavelength_array, Tau2flux(self.estimated_profiles[c].value),
                         color="green", label="Fit")
            else:
                ax2.plot(wavelength_array, Tau2flux(self.estimated_profiles[c].value),
                         color="green")

        ax2.legend()

        ax3.plot(wavelength_array, flux_array, label="Measured")
        ax3.plot(wavelength_array, self.total.value, color='green', label="Fit", linewidth=2.0)
        ax3.legend()

        f.subplots_adjust(hspace=0)

        if hasattr(self, 'fit_time'): ax1.set_title("Fit time: " + self.fit_time)
        ax1.set_ylabel("Residuals")
        ax2.set_ylabel("Normalised Flux")
        ax3.set_ylabel("Normalised Flux")
        ax3.set_xlabel("$ \lambda (\AA)$")

        if filename:
            plt.savefig(filename)
        else:
            plt.show()

    def find_local_minima(self, f_array, window=101):
        """
        Find the local minima of an absorption profile.

        Args:
            f_array: flux array
            window: smoothing window, pixels
        Returns:
            indices of local minima in flux_array
        """

        # smooth flux profile
        smoothed_flux = savgol_filter(f_array, window, 1)

        return find_peaks_cwt(smoothed_flux * -1, np.array([window/3]))

    def initialise_components(self, frequency_array, n, sigma_max):
        """
        Initialise each fitted component of the model in optical depth space. Each component
        consists of three variables, height, centroid and sigma. These variables are encapsulated
        in a deterministic profile variable. The variables are stored in a dictionary, `estimated_variables`,
        and the profiles in a list, `estimated_profiles`.

        Args:
            frequency_array (numpy array)
            n (int): number of components
            sigma_max (float): maximum permitted range of fitted sigma values
        """

        self.estimated_variables = {}
        self.estimated_profiles = []

        for component in range(n):

            self.estimated_variables[component] = {}

            @mc.stochastic(name='xexp_%d' % component)
            def xexp(value=0.5):
                if value < 0:
                    return -np.inf
                else:
                    return np.log(value * np.exp(-value))

            self.estimated_variables[component]['amplitude'] = xexp
            #self.estimated_variables[component]['height'] = mc.Uniform("est_height_%d" % component, 0, 5)

            self.estimated_variables[component]['centroid'] = mc.Uniform("est_centroid_%d" % component,
                                                                         frequency_array[0], frequency_array[-1])

            self.estimated_variables[component]['sigma'] = mc.Uniform("est_sigma_%d" % component, 0, sigma_max)

            @mc.deterministic(name='component_%d' % component, trace = True)
            def profile(x=frequency_array,
                        centroid=self.estimated_variables[component]['centroid'],
                        sigma=self.estimated_variables[component]['sigma'],
                        height=self.estimated_variables[component]['amplitude']):

                return self.GaussFunction(x, height, centroid, sigma )

            self.estimated_profiles.append(profile)

    def initialise_voigt_profiles(self, frequency_array, n, sigma_max, local_minima=[]):
        """
        Args:
            frequency_array (numpy array)
            n (int): number of components
            sigma_max (float): maximum permitted range of fitted sigma values
        """

        #if n < len(local_minima):
        #    raise ValueError("Less profiles than number of minima.")

        self.estimated_variables = {}
        self.estimated_profiles = []

        for component in range(n):

            self.estimated_variables[component] = {}

            @mc.stochastic(name='xexp_%d' % component)
            def xexp(value=0.5):
                if value < 0:
                    return -np.inf
                else:
                    return np.log(value * np.exp(-value))

            self.estimated_variables[component]['amplitude'] = xexp

            self.estimated_variables[component]['centroid'] = mc.Uniform("est_centroid_%d" % component,
                                                                            frequency_array[0], frequency_array[-1])

            self.estimated_variables[component]['L_fwhm'] = mc.Uniform("est_L_%d" % component, 0, sigma_max)
            self.estimated_variables[component]['G_fwhm'] = mc.Uniform("est_G_%d" % component, 0, sigma_max)

            @mc.deterministic(name='component_%d' % component, trace = True)
            def profile(x=frequency_array,
                        centroid=self.estimated_variables[component]['centroid'],
                        amplitude=self.estimated_variables[component]['amplitude'],
                        L=self.estimated_variables[component]['L_fwhm'],
                        G=self.estimated_variables[component]['G_fwhm']):
                return self.VoigtFunction(x, centroid, amplitude, L, G)

            self.estimated_profiles.append(profile)

    def initialise_model(self, frequency, flux, n, local_minima=[], voigt=False):
        """
        Initialise deterministic model of all absorption features, in normalised flux.

        Args:
            frequency (numpy array)
            flux (numpy array): flux values at each wavelength
            n (int): number of absorption profiles to fit
        """

        self.sigma_max = (frequency[-1] - frequency[0]) / 2.

        # always reinitialise profiles, otherwise starts sampling from previously calculated parameter values.
        if(voigt):
            if self.verbose:
                print("Initialising Voigt profile components.")
            self.fwhm_max = self.sigma_max * 2*np.sqrt(2*np.log(2.))
            self.initialise_voigt_profiles(frequency, n, self.fwhm_max, local_minima=local_minima)
        else:
            if self.verbose:
                print("Initialising Gaussian profile components.")
            self.initialise_components(frequency, n, self.sigma_max)

        # deterministic variable for the full profile, given in terms of normalised flux
        @mc.deterministic(name='profile', trace=False)
        def total(profile_sum=self.estimated_profiles):
            return Tau2flux(sum(profile_sum))

        self.total = total

        # represent full profile as a normal distribution
        self.profile = mc.Normal("obs", self.total, self.std_deviation, value=flux, observed=True)

        # create model with parameters of all profiles to be fitted
        # TODO: fix the deprecation error (this is the only case where a "Model" is instantiated):
        # /home/jacobc/anaconda2/envs/vamp27pymc237/lib/python2.7/site-packages/pymc/MCMC.py:81:
        # UserWarning: Instantiating a Model object directly is deprecated. We recommend passing variables directly to the Model subclass.  warnings.warn(message)

        self.model = mc.Model([self.estimated_variables[x][y] for x in self.estimated_variables \
                            for y in self.estimated_variables[x]])# + [std_deviation])

    def map_estimate(self, iterations=2000):
        """
        Compute the Maximum A Posteriori estimates for the initialised model
        """

        self.map = mc.MAP(self.model)
        self.map.fit(iterlim=iterations, tol=1e-3)

    def mcmc_fit(self, iterations=15000, burnin=100, thinning=15, step_method=mc.AdaptiveMetropolis):
        """
        MCMC fit of `n` absorption profiles to a given spectrum

        Args:
            iterations (int): How long to sample the MCMC for
            burnin (int): How many iterations to ignore from the start of the chain
            thinning (int): How many iterations to ignore each tally
            step_method (PyMC steo method object)
        """

        try:
            getattr(self, 'map')
        except AttributeError:
            print("\nWARNING: MAP estimate not provided. \nIt is recommended to compute this \
            in advance of running the MCMC so as to start the sampling with good initial values.")

        # create MCMC object
        self.mcmc = mc.MCMC(self.model)

        # change step method
        if step_method != mc.AdaptiveMetropolis:
            print("Changing step method for each parameter.")
            for index, item in self.estimated_variables.items():
                self.mcmc.use_step_method(step_method, item['sigma'])
                self.mcmc.use_step_method(step_method, item['centroid'])
                self.mcmc.use_step_method(step_method, item['height'])
        else:
            print("Using Adaptive Metropolis step method for each parameter.")

        # fit the model
        starttime=datetime.datetime.now()
        self.mcmc.sample(iter=iterations, burn=burnin, thin=thinning)
        self.fit_time = str(datetime.datetime.now() - starttime)
        print("\nTook:", self.fit_time, " to finish.")

    def find_bic(self, frequency_array, flux_array, n, noise_array, freedom, voigt=False,
                iterations=3000, thin=15, burn=300, thorough=False):
        """
        Initialise the Voigt model and run the MCMC fitting for a particular number of
        regions and return the Bayesian Information Criterion. Used to identify the
        most appropriate number of profles in the region.

        Args:
            frequency_array (numpy array)
            flux_array (numpy array)
            n (int): number of minima to fit in the region
            noise_array (numpy array)
            freedom (int): number of degrees of freedom
            voigt (Boolean): switch to fit as Voigt profile or Gaussian
            iterations, thin, burn (ints): MCMC parameters
        """

        self.bic_array = []
        self.red_chi_array = []
        for i in range(3):
            self.initialise_model(frequency_array, flux_array, n, voigt=voigt)
            self.map = mc.MAP(self.model)
            self.mcmc = mc.MCMC(self.model)
            if thorough:
                self.map.fit(iterlim=iterations, tol=1e-3)
                self.mcmc.sample(iter=iterations, burn=burn, thin=thin, progress_bar=False)
                self.map.fit(iterlim=iterations, tol=1e-3)
            self.mcmc.sample(iter=iterations, burn=burn, thin=thin, progress_bar=False)
            self.map.fit(iterlim=iterations, tol=1e-3)
            self.bic_array.append(self.map.BIC)
            self.red_chi_array.append(self.ReducedChisquared(flux_array, self.total.value, noise_array, freedom))
        return

    def chain_covariance(self, n, voigt=False):
        """
        Find the covariance matrix of the MCMC chain variables
        Args:
            n (int): the number of lines in the region
            voigt (boolean): switch to use Voigt parameters
        Returns:
            cov (numpy array): the covariance matrix
            cov =   [ var(amplitude),           cov(amplitude, sigma),  cov(amplitude, centroid) ]
                    [ cov(sigma, amplitude),    var(sigma),             cov(sigma, centroid)     ]
                    [ cov(centroid, amplitude), cov(centroid, sigma),   var(centroid)            ]

        """
        cov = np.zeros((n, 3, 3))
        for i in range(n):
            amp_samples = self.mcmc.trace('xexp_'+str(i))[:]
            if not voigt:
                sigma_samples = self.mcmc.trace('est_sigma_'+str(i))[:]
            elif voigt:
                gfwhm_samples = self.mcmc.trace('est_G_'+str(i))[:]
                sigma_samples = self.GaussianWidth(gfwhm_samples)
            c_samples = self.mcmc.trace('est_centroid_'+str(i))[:]

            cov[i] = np.cov(np.array((amp_samples, sigma_samples, c_samples)))
        return cov


from matplotlib import pyplot as plt
import math

def frange(start, final, increment):

    result=[]
    while start<=final:
        result.append(start)
        start=increment+start
    return result

def draw_graph(x,y):

    plt.plot(x,y)
    plt.xlabel('x-coord')
    plt.ylabel('y-coord')
    plt.title('Herecy!!!')

def draw_trajectory(u, theta):
    theta=math.radians(theta)
    g=9.8
    t_flight=2*u*math.sin(theta)/g
    print(t_flight)
    intervals=frange(0,t_flight,0.01)
    x=[]
    y=[]
    for t in intervals:
        x.append(u*math.cos(theta)*t)
        y_temp=u*math.sin(theta)*t-0.5*g*t*t
        if y_temp>=0:
            y.append(y_temp)
        else:
            y.append(0.0)
    draw_graph(x,y)

draw_trajectory(5,30)
plt.show()


import yaml
import torch
from model_training.dataset import BratsDataset
from model_training.unet_vae.trainer import TrainerVAE
from model_training.unet_addnet.model import UNetAddNet
from model_training.unet_vae.vae import VAE

with open("train.yaml", "r") as config_file:
    config = yaml.full_load(config_file)

train_ds = BratsDataset(config["train"]["path"])
val_ds = BratsDataset(config["val"]["path"])

train_dl = torch.utils.data.DataLoader(train_ds, batch_size=config['batch_size'], shuffle=True)
val_dl = torch.utils.data.DataLoader(val_ds, batch_size=config['batch_size'], shuffle=True)

if torch.cuda.is_available():
    print("Train on GPU")
    device = torch.device("cuda:0")
else:
    print("Train on CPU")
    device = torch.device("cpu")

model = UNetAddNet(in_channels=4, out_channels=4)
vae = VAE()

trainer = TrainerVAE(model, vae, config, train_dl, val_dl, device)
trainer.train()