#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# C++ version Copyright (c) 2006-2011 Erin Catto http://www.box2d.org
# Python port by Ken Lauer / http://pybox2d.googlecode.com
#
# This software is provided 'as-is', without any express or implied
# warranty.  In no event will the authors be held liable for any damages
# arising from the use of this software.
# Permission is granted to anyone to use this software for any purpose,
# including commercial applications, and to alter it and redistribute it
# freely, subject to the following restrictions:
# 1. The origin of this software must not be misrepresented; you must not
# claim that you wrote the original software. If you use this software
# in a product, an acknowledgment in the product documentation would be
# appreciated but is not required.
# 2. Altered source versions must be plainly marked as such, and must not be
# misrepresented as being the original software.
# 3. This notice may not be removed or altered from any source distribution.

from __future__ import absolute_import

__all__ = ('ABSOLUTE_TOL', 'RELATIVE_TOL',
           'ISOLATED', 'CONCAVE', 'FLAT', 'CONVEX',
           'collide_circles', 'collide_polygon_circle',
           'collide_polygons', 'collide_edge_circle', 'collide_edge_polygon',
           'EPAxis', 'EPFatEdge', 'EPProxy', 'EPCollider'
           )
__version__ = "$Revision$"
__date__ = "$Date$"
# $Source$

from .common import (Vec2, distance_squared)
from .collision_util import (max_separation, find_incident_edge, clip_segment_to_line)
from .contact_util import (ManifoldPoint, ClipVertex, Manifold)
from .settings import (MAX_FLOAT, EPSILON, MAX_MANIFOLD_POINTS, POLYGON_RADIUS, ANGULAR_SLOP)

# -- defines --
RELATIVE_TOL = 0.98
ABSOLUTE_TOL = 0.001

# Edge types
ISOLATED, CONCAVE, FLAT, CONVEX = range(4)

# -- circles --
def collide_circles(manifold, circle_a, xf_a, circle_b, xf_b):
    manifold.point_count = 0
    pa = xf_a * circle_a.position
    pb = xf_b * circle_b.position

    d = pb - pa
    dist_sqr = d.length_squared
    ra = circle_a.radius
    rb = circle_b.radius
    radius = ra + rb
    if dist_sqr > radius**2:
        return

    manifold.type = Manifold.CIRCLES
    manifold.local_point = Vec2(*circle_a.position)
    manifold.local_normal = Vec2()
    manifold.point_count = 1

    point=manifold.points[0]
    point.local_point = Vec2(*circle_b.position)
    point.set_id(0,0,0,0)

def collide_polygon_circle(manifold, polygon_a, xf_a, circle_b, xf_b):
    manifold.point_count = 0

    # Compute circle position in the frame of the polygon
    c = xf_b * circle_b.position
    c_local = xf_a.mul_t(c)

    # Find the min separating edge
    normal_index = 0
    separation = -MAX_FLOAT
    radius = polygon_a.radius + circle_b.radius

    vertices = polygon_a._vertices
    normals = polygon_a._normals

    vertex_count = len(vertices)

    for i in range(vertex_count):
        s = normals[i].dot(c_local - vertices[i])

        if s > radius:
            # Early out
            return
        if s > separation:
            separation = s
            normal_index = i

    # Vertices that subtend the incident face
    i1 = normal_index
    i2 = (i1 + 1) % vertex_count
    v1 = vertices[i1]
    v2 = vertices[i2]

    # If the center is inside the polygon...
    if separation < EPSILON:
        manifold.point_count = 1
        manifold.type = Manifold.FACE_A
        manifold.local_normal = Vec2(*normals[normal_index])
        manifold.local_point = 0.5 * (v1 + v2)
        point=manifold.points[0]
        point.local_point = Vec2(*circle_b.position)
        point.set_id(0, 0, 0, 0)
        return

    # Compute the barycentric coordinates
    u1 = (c_local - v1).dot(v2 - v1)
    u2 = (c_local - v2).dot(v1 - v2)

    if u1 < 0.0:
        if distance_squared(c_local, v1) > radius**2:
            return

        manifold.local_normal = (c_local - v1).normalized
        manifold.local_point = Vec2(*v1)
    elif u2 < 0.0:
        if distance_squared(c_local, v2) > radius**2:
            return

        manifold.local_normal = (c_local - v2).normalized
        manifold.local_point = Vec2(*v2)
    else:
        face_center = 0.5 * (v1 + v2)
        separation = (c_local - face_center).dot(normals[i1])
        if separation > radius:
            return

        manifold.local_normal = Vec2(*normals[i1])
        manifold.local_point = face_center

    manifold.point_count = 1
    manifold.type = Manifold.FACE_A
    point=manifold.points[0]
    point.local_point = Vec2(*circle_b.position)
    point.set_id(0, 0, 0, 0)

# -- polygons --
def collide_polygons(manifold, poly_a, xf_a, poly_b, xf_b):
    # Find edge normal of max separation on A - return if separating axis is found
    # Find edge normal of max separation on B - return if separation axis is found
    # Choose reference edge as min(minA, minB)
    # Find incident edge
    # Clip
    #
    # The normal points from 1 to 2

    manifold.point_count = 0
    total_radius = poly_a.radius + poly_b.radius

    separation_a, edge_a = max_separation(poly_a, xf_a, poly_b, xf_b)
    if separation_a > total_radius:
        return

    separation_b, edge_b = max_separation(poly_b, xf_b, poly_a, xf_a)
    if separation_b > total_radius:
        return

    # poly1: reference polygon
    # poly2: incident polygon
    # edge1: reference edge

    if separation_b > (RELATIVE_TOL * separation_a + ABSOLUTE_TOL):
        poly1 = poly_b
        poly2 = poly_a
        xf1 = xf_b
        xf2 = xf_a
        edge1 = edge_b
        manifold.type = Manifold.FACE_B
        flip = True
    else:
        poly1 = poly_a
        poly2 = poly_b
        xf1 = xf_a
        xf2 = xf_b
        edge1 = edge_a
        manifold.type = Manifold.FACE_A
        flip = False

    incident_edge = find_incident_edge(poly1, xf1, edge1, poly2, xf2)

    vertices1 = poly1._vertices

    iv1 = edge1
    iv2 = (edge1 + 1) % len(poly1._vertices)

    v11 = vertices1[iv1]
    v12 = vertices1[iv2]

    localtangent = (v12 - v11).normalized

    local_normal = localtangent.cross(1.0)
    planepoint = 0.5 * (v11 + v12)

    tangent = xf1._rotation * localtangent
    normal = tangent.cross(1.0)

    v11 = xf1 * v11
    v12 = xf1 * v12

    # Face offset.
    front_offset = normal.dot(v11)

    # Side offsets, extended by polytope skin thickness.
    side_offset1 = -tangent.dot(v11) + total_radius
    side_offset2 = tangent.dot(v12) + total_radius

    # Clip incident edge against extruded edge1 side edges.
    # Clip to box side 1
    clip_points1 = clip_segment_to_line(incident_edge, -tangent, side_offset1, iv1)
    if len(clip_points1) < 2:
        return

    # Clip to negative box side 1
    clip_points2 = clip_segment_to_line(clip_points1, tangent, side_offset2, iv2)
    if len(clip_points2) < 2:
        return

    # Now clip_points2 contains the clipped points.
    manifold.local_normal = local_normal
    manifold.local_point = planepoint

    point_count = 0
    for i in range(MAX_MANIFOLD_POINTS):
        separation = normal.dot(clip_points2[i].v) - front_offset

        if separation <= total_radius:
            cp = manifold.points[point_count]
            cp.local_point = xf2.mul_t(clip_points2[i].v)
            cp.set_id(*clip_points2[i].id)
            if flip:
                # Swap features
                cp.set_id(cp.index_b, cp.index_a, cp.type_b, cp.type_a)
            point_count += 1

    manifold.point_count = point_count

# -- edge --

def collide_edge_circle(manifold, edge_a, xf_a, circle_b, xf_b):
    """
    Compute contact points for edge versus circle.
    This accounts for edge connectivity.
    """
    manifold.point_count = 0

    # Compute circle in frame of edge
    Q = xf_a.mul_t(xf_b * circle_b.position)

    A = Vec2(*edge_a._vertex1)
    B = Vec2(*edge_a._vertex2)
    e = B - A

    # Barycentric coordinates
    u = e.dot(B - Q)
    v = e.dot(Q - A)

    radius = edge_a.radius + circle_b.radius

    index_b = 0
    type_b = ManifoldPoint.VERTEX

    # Region A
    if v <= 0.0:
        P = A
        d = Q - P
        dd = d.dot(d)
        if dd > radius**2:
            return

        # Is there an edge connected to A?
        if edge_a._vertex0 is not None:
            A1 = edge_a._vertex0
            B1 = A
            e1 = B1 - A1
            u1 = e1.dot(B1 - Q)

            # Is the circle in Region AB of the previous edge?
            if u1 > 0.0:
                return

        index_a = 0
        type_a = ManifoldPoint.VERTEX
        manifold.point_count = 1
        manifold.type = Manifold.CIRCLES
        manifold.local_normal = Vec2()
        manifold.local_point = P
        mp = manifold.points[0]
        mp.set_id(index_a, index_b, type_a, type_b)
        mp.local_point = Vec2(*circle_b.position)
        return

    # Region B
    if u <= 0.0:
        P = B
        d = Q - P
        dd = d.dot(d)
        if dd > radius**2:
            return

        # Is there an edge connected to B?
        if edge_a._vertex3 is not None:
            B2 = edge_a._vertex3
            A2 = B
            e2 = B2 - A2
            v2 = e2.dot(Q - A2)

            # Is the circle in Region A_b of the next edge?
            if v2 > 0.0:
                return

        index_a = 1
        type_a =ManifoldPoint.VERTEX
        manifold.point_count = 1
        manifold.type = Manifold.CIRCLES
        manifold.local_normal = Vec2()
        manifold.local_point = P

        mp = manifold.points[0]
        mp.set_id(index_a, index_b, type_a, type_b)
        mp.local_point = Vec2(*circle_b.position)
        return

    # Region A_b
    den = e.dot(e)
    assert(den > 0.0)
    P = (1.0 / den) * (u * A + v * B)
    d = Q - P
    dd = d.dot(d)
    if dd > radius**2:
        return

    n = Vec2(-e.y, e.x)
    if n.dot(Q - A) < 0.0:
        n = Vec2(-n.x, -n.y)
    n = n.normalized

    index_a = 0
    type_a = ManifoldPoint.FACE
    manifold.point_count = 1
    manifold.type = Manifold.FACE_A
    manifold.local_normal = n
    manifold.local_point = A

    mp = manifold.points[0]
    mp.set_id(index_a, index_b, type_a, type_b)
    mp.local_point = Vec2(*circle_b.position)

def collide_edge_polygon(manifold, edge_a, xf_a, polygon_b, xf_b):
    collider = EPCollider(edge_a, xf_a, polygon_b, xf_b)
    collider.collide(manifold)

# -- edge and polygon --
class EPAxis(object):
    __slots__ = ['type', 'index', 'separation']

    # Edge-polygon axis
    UNKNOWN, EDGE_A, EDGE_B = range(3)
    def __init__(self):
        self.type = EPAxis.UNKNOWN
        self.index = 0
        self.separation = 0.0

class EPFatEdge(object):
    """Edge shape plus more stuff."""
    __slots__ = ['v0', 'v1', 'v2', 'v3', 'normal']
    def __init__(self):
        pass

class EPProxy(object):
    """
    This lets us treate and edge shape and a polygon in the same
    way in the SAT collider.
    """
    __slots__ = ['vertices', 'normals', 'centroid']
    def __init__(self):
        pass

class EPCollider(object):
    """
    This class collides and edge and a polygon, taking into account edge adjacency.

    Collide an edge and polygon. This uses the SAT and clipping to produce up to 2 contact points.
    Edge adjacency is handle to produce locally valid contact points and normals. This is intended
    to allow the polygon to slide smoothly over an edge chain.

    Algorithm
    1. Classify front-side or back-side collision with edge.
    2. Compute separation
    3. Process adjacent edges
    4. Classify adjacent edge as convex, flat, null, or concave
    5. Skip null or concave edges. Concave edges get a separate manifold.
    6. If the edge is flat, compute contact points as normal. Discard boundary points.
    7. If the edge is convex, compute it's separation.
    8. Use the minimum separation of up to three edges. If the minimum separation
       is not the primary edge, return.
    9. If the minimum separation is the primary edge, compute the contact points and return.
    """
    __slots__ = ['_xf', '_edge_a', '_proxy_a', '_proxy_b', '_radius',
                '_limit11', '_limit12', '_limit21', '_limit22']

    def __init__(self, edge_a, xf_a, polygon_b, xf_b):
        # Transform
        self._xf = xf_a.mul_t(xf_b)

        # Edge geometry
        self._edge_a = EPFatEdge()
        self._edge_a.v0 = edge_a._vertex0
        self._edge_a.v1 = edge_a._vertex1
        self._edge_a.v2 = edge_a._vertex2
        self._edge_a.v3 = edge_a._vertex3
        e = self._edge_a.v2 - self._edge_a.v1

        # Normal points outwards in CCW order.
        self._edge_a.normal = Vec2(e.y, -e.x).normalized

        # Proxy for edge
        self._proxy_a = EPProxy()
        self._proxy_a.vertices = [self._edge_a.v1, self._edge_a.v2]
        self._proxy_a.normals = [self._edge_a.normal, -self._edge_a.normal]
        self._proxy_a.centroid = 0.5 * (self._edge_a.v1 + self._edge_a.v2)

        # Proxy for polygon
        self._proxy_b = EPProxy()
        self._proxy_b.centroid = self._xf * polygon_b._centroid
        self._proxy_b.vertices = [self._xf * vertex for vertex in polygon_b._vertices]
        self._proxy_b.normals = [self._xf._rotation * normal for normal in polygon_b._normals]

        self._radius = 2.0 * POLYGON_RADIUS

        self._limit11 = Vec2()
        self._limit12 = Vec2()
        self._limit21 = Vec2()
        self._limit22 = Vec2()

    def collide(self, manifold):
        manifold.point_count = 0

        self.compute_adjacency()

        edge_axis = self.compute_edge_separation()

        # If no valid normal can be found than this edge should not collide.
        # This can happen on the middle edge of a 3-edge zig-zag chain.
        if edge_axis.type == EPAxis.UNKNOWN:
            return

        if edge_axis.separation > self._radius:
            return

        polygon_axis = self.compute_polygon_separation()
        if polygon_axis.type != EPAxis.UNKNOWN and polygon_axis.separation > self._radius:
            return

        # Use hysteresis for jitter reduction.
        if polygon_axis.type == EPAxis.UNKNOWN:
            primary_axis = edge_axis
        elif polygon_axis.separation > RELATIVE_TOL * edge_axis.separation + ABSOLUTE_TOL:
            primary_axis = polygon_axis
        else:
            primary_axis = edge_axis

        if primary_axis.type == EPAxis.EDGE_A:
            proxy1 = self._proxy_a
            proxy2 = self._proxy_b
            manifold.type = Manifold.FACE_A
        else:
            proxy1 = self._proxy_b
            proxy2 = self._proxy_a
            manifold.type = Manifold.FACE_B

        edge1 = primary_axis.index

        incident_edge = self.find_incident_edge(proxy1, primary_axis.index, proxy2)
        vertices1 = proxy1.vertices

        iv1 = edge1
        iv2 = (edge1 + 1) % len(proxy1.vertices)

        v11 = vertices1[iv1]
        v12 = vertices1[iv2]

        tangent = (v12 - v11).normalized

        normal = tangent.cross(1.0)
        plane_point = 0.5 * (v11 + v12)

        # Face offset.
        front_offset = normal.dot(v11)

        # Side offsets, extended by polytope skin thickness.
        side_offset1 = -tangent.dot(v11) + self._radius
        side_offset2 =  tangent.dot(v12) + self._radius

        # Clip incident edge against extruded edge1 side edges.
        # Clip to box side 1
        clip_points1 = clip_segment_to_line(incident_edge, -tangent, side_offset1, iv1)
        if len(clip_points1) < MAX_MANIFOLD_POINTS:
            return

        # Clip to negative box side 1
        clip_points2 = clip_segment_to_line(clip_points1, tangent, side_offset2, iv2)
        if len(clip_points2) < MAX_MANIFOLD_POINTS:
            return

        # Now clip_points2 contains the clipped points.
        if primary_axis.type == EPAxis.EDGE_A:
            manifold.local_normal = normal
            manifold.local_point = plane_point
        else:
            manifold.local_normal = self._xf._rotation.mul_t(normal)
            manifold.local_point = self._xf.mul_t(plane_point)

        point_count = 0
        for clip_point in clip_points2:
            separation = normal.dot(clip_point.v) - front_offset

            if separation <= self._radius:
                cp = manifold.points[point_count]

                if primary_axis.type == EPAxis.EDGE_A:
                    cp.local_point = self._xf.mul_t(clip_point.v)
                    cp.id = tuple(clip_point.id)
                else:
                    cp.local_point = clip_point.v
                    cp.set_id(clip_point.type_b, clip_point.type_a, clip_point.index_b, clip_point.index_a)

                point_count += 1
        manifold.point_count = point_count

    # Compute allowable normal ranges based on adjacency.
    # A normal n is allowable iff:
    # cross(n, n1) >= 0.0 and cross(n2, n) >= 0.0
    # n points from A to B (edge to polygon)
    def compute_adjacency(self):
        v0 = self._edge_a.v0
        v1 = self._edge_a.v1
        v2 = self._edge_a.v2
        v3 = self._edge_a.v3

        # Determine allowable the normal regions based on adjacency.
        # Note: it may be possible that no normal is admissable.
        center_b = self._proxy_b.centroid
        if self._edge_a.v0 is not None:
            e0 = v1 - v0
            e1 = v2 - v1
            n0 = Vec2(e0.y, -e0.x).normalized
            n1 = Vec2(e1.y, -e1.x).normalized

            convex = n0.cross(n1) >= 0.0
            front0 = n0.dot(center_b - v0) >= 0.0
            front1 = n1.dot(center_b - v1) >= 0.0

            if convex:
                if front0 or front1:
                    self._limit11 = n1
                    self._limit12 = n0
                else:
                    self._limit11 = -n1
                    self._limit12 = -n0
            else:
                if front0 and front1:
                    self._limit11 = n0
                    self._limit12 = n1
                else:
                    self._limit11 = -n0
                    self._limit12 = -n1
        else:
            self._limit11 = Vec2()
            self._limit12 = Vec2()

        if self._edge_a.v3 is not None:
            e1 = v2 - v1
            e2 = v3 - v2
            n1 = Vec2(e1.y, -e1.x).normalized
            n2 = Vec2(e2.y, -e2.x).normalized

            convex = n1.cross(n2) >= 0.0
            front1 = n1.dot(center_b - v1) >= 0.0
            front2 = n2.dot(center_b - v2) >= 0.0

            if convex:
                if front1 or front2:
                    self._limit21 = n2
                    self._limit22 = n1
                else:
                    self._limit21 = -n2
                    self._limit22 = -n1
            else:
                if front1 and front2:
                    self._limit21 = n1
                    self._limit22 = n2
                else:
                    self._limit21 = -n1
                    self._limit22 = -n2
        else:
            self._limit21 = Vec2()
            self._limit22 = Vec2()

    def compute_edge_separation(self):
        # Edge_a separation
        best_axis = EPAxis()
        best_axis.type = EPAxis.UNKNOWN
        best_axis.index = -1
        best_axis.separation = -MAX_FLOAT
        normals = [self._edge_a.normal, -self._edge_a.normal]

        for i, n in enumerate(normals):
            # Adjacency
            valid1 = (n.cross(self._limit11) >= -ANGULAR_SLOP) and (self._limit12.cross(n) >= -ANGULAR_SLOP)
            valid2 = (n.cross(self._limit21) >= -ANGULAR_SLOP) and (self._limit22.cross(n) >= -ANGULAR_SLOP)

            if not valid1 or not valid2:
                continue

            axis = EPAxis()
            axis.type = EPAxis.EDGE_A
            axis.index = i

            dots = [n.dot(vertex - self._edge_a.v1) for vertex in self._proxy_b.vertices]
            axis.separation = min(dots)

            if axis.separation > self._radius:
                return axis

            if axis.separation > best_axis.separation:
                best_axis = axis

        return best_axis

    def compute_polygon_separation(self):
        axis = EPAxis()
        axis.type = EPAxis.UNKNOWN
        axis.index = -1
        axis.separation = -MAX_FLOAT
        for i in range(len(self._proxy_b.vertices)):
            n = -self._proxy_b.normals[i]

            # Adjacency
            valid1 = (n.cross(self._limit11) >= -ANGULAR_SLOP) and (self._limit12.cross(n) >= -ANGULAR_SLOP)
            valid2 = (n.cross(self._limit21) >= -ANGULAR_SLOP) and (self._limit22.cross(n) >= -ANGULAR_SLOP)

            if not valid1 or not valid2:
                continue

            s1 = n.dot(self._proxy_b.vertices[i] - self._edge_a.v1)
            s2 = n.dot(self._proxy_b.vertices[i] - self._edge_a.v2)
            s = min(s1, s2)

            if s > self._radius:
                axis.type = EPAxis.EDGE_B
                axis.index = i
                axis.separation = s

            if s > axis.separation:
                axis.type = EPAxis.EDGE_B
                axis.index = i
                axis.separation = s

        return axis

    def find_incident_edge(self, proxy1, edge1, proxy2):
        normals1 = proxy1.normals

        vertices2 = proxy2.vertices
        normals2 = proxy2.normals

        assert(0 <= edge1 < len(normals1))

        # Get the normal of the reference edge in proxy2's frame.
        normal1 = normals1[edge1]

        # Find the incident edge on proxy2.
        dots = [normal1.dot(v) for v in normals2]
        min_dot = min(dots)
        index = dots.index(min_dot)

        # Build the clip vertices for the incident edge.
        i1 = index
        i2 = (i1 + 1) % len(vertices2)

        c0=ClipVertex(vertices2[i1], edge1, i1, ManifoldPoint.FACE, ManifoldPoint.VERTEX)
        c1=ClipVertex(vertices2[i2], edge1, i2, ManifoldPoint.FACE, ManifoldPoint.VERTEX)
        return [c0, c1]


"""
single capsule network to classify mnist digits.
"""
from collections import OrderedDict

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision.transforms as transforms
from torchvision import datasets

TRAIN_ON_GPU = torch.cuda.is_available()

if (TRAIN_ON_GPU):
    print('Training on GPU!')
else:
    print('Only CPU available')

# plt.interactive(False)
seed = 1
np.random.seed(seed)
torch.manual_seed(seed)

num_workers = 0
batch_size = 32

ip_transform = transforms.Compose(
    [
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ]
)

# get datasets
train_data = datasets.MNIST(root='data', train=True, download=True, transform=ip_transform)
test_data = datasets.MNIST(root='data', train=False, download=True, transform=ip_transform)

# data loader
train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, num_workers=num_workers)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, num_workers=num_workers)

# obtain one batch of training images
dataiter = iter(train_loader)
images, labels = dataiter.next()
images = images.numpy()

# plot the images in the batch, along with the corresponding labels
# fig = plt.figure(figsize=(25, 4))
# for idx in np.arange(batch_size):
#     ax = fig.add_subplot(2, batch_size/2, idx+1, xticks=[], yticks=[])
#     ax.imshow(np.squeeze(images[idx]), cmap='gray')
#     # print out the correct label for each image
#     # .item() gets the value contained in a Tensor
#     ax.set_title(str(labels[idx].item()))
#
# plt.show()

def squash(input_tensor):
    '''Squashes an input Tensor so it has a magnitude between 0-1.
       param input_tensor: a stack of capsule inputs, s_j
       return: a stack of normalized, capsule output vectors, v_j
    '''
    squared_norm = (input_tensor ** 2).sum(dim=-1, keepdim=True)
    scale = squared_norm / (1 + squared_norm)  # normalization coeff
    output_tensor = scale * input_tensor / torch.sqrt(squared_norm)
    return output_tensor

def transpose_softmax(input_tensor, dim=1):
    # transpose input
    transposed_input = input_tensor.transpose(dim, len(input_tensor.size()) - 1)
    # calculate softmax
    softmaxed_output = F.softmax(transposed_input.contiguous().view(-1, transposed_input.size(-1)), dim=-1)
    # un-transpose result
    return softmaxed_output.view(*transposed_input.size()).transpose(dim, len(input_tensor.size()) - 1)

# Initial Convolutino later
class ConvolutionLayer(nn.Module):

    def __init__(self, in_channels=1, out_channels=256):
        super(ConvolutionLayer, self).__init__()
        # define convolution kernel
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=9, stride=1, padding=0)

    def forward(self, x):
        features = F.relu(self.conv(x))
        if TRAIN_ON_GPU:
            features = features.cuda()
        return features

# scalar ip and vector op for primary capsule
class PrimaryCapsule(nn.Module):

    def __init__(self, num_capsules=8, in_channels=256, out_channels=32, kernel_size=9, stride=2, padding=0):
        super(PrimaryCapsule, self).__init__()
        # create a list of convnet for capsules
        self.capsules = nn.ModuleList \
                (
                [
                    nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size,
                              stride=stride, padding=padding)
                    for _ in range(num_capsules)
                ]
            )

    def forward(self, x):
        batch_size = x.size(0)
        # get output and reshape them
        u = [capsule(x).view(batch_size, 32 * 6 * 6, 1) for capsule in self.capsules]
        # concatinate outputs
        u = torch.cat(u, dim=-1)  # TOLEARN: Understand why dim =-1 and why caat in such manner
        # squash for non linearity
        u_squash = squash(u)
        return u_squash

# Dynamic Routing
def dynamic_routing(b_ij, u_hat, routing_iteration=3):
    '''
    :param b_ij: var that will yield coupling coeff; initial log probabilities that capsule i should be coupled to capsule j
    :param u_hat: input, weighted capsule vectors, W u; op of last layer after the inverse transform.
    :param routing_iteration: no. of times the dynamic routing should happen. by paper 3-5 times.
    :return: final vector op of the this layer. last layer was u already

    last layer op = u.
    after transform = W u = u_hat, note : this operation is not done in this function. (like the inverse graphics step to remove positional variance or whatever)

    final op of last layer = vj.
    vj = squash(sj)
    sj = sum(cij.u_hat)
    cij = softmax(b ij, over j)
    b ij <- bij + agreement
    agreement = dot product u_hat i.e. transformed vec and v
    '''
    for iteration in range(routing_iteration):
        # softmax calculation of coupling coefficients, c_ij
        c_ij = transpose_softmax(b_ij, dim=2)  # TOLEARN: understand why transpose softmax

        # calculating total capsule inputs, s_j = sum(c_ij*u_hat)
        s_j = (c_ij * u_hat).sum(dim=2, keepdim=True)

        # squash
        v_j = squash(s_j)

        if iteration < routing_iteration - 1:
            # get agreement
            a_ij = (u_hat * v_j).sum(dim=-1, keepdim=True)

            # update b_ij
            b_ij += a_ij
    return v_j

# capsules for digits
class DigitCapsule(nn.Module):
    def __init__(self, num_capsule=10, num_cap_prev=8, cap_dim=16, cap_dim_prev=32 * 6 * 6):
        '''
        'inverse graphics' will happen here. u_hat = W*u, supposed to remove variance from representation

        3rd network and 2nd capsule network
        :param num_capsule: number of capsule in current layer. (= no. of class if this is last layer)
        :param num_cap_prev: number of capsules in previous later
        :param cap_dim: elements in current capsule layer = 16
        :param cap_dim_prev: elements in previous layer capsule (32*6*6)
        '''
        super(DigitCapsule, self).__init__()
        self.num_capsules = num_capsule
        self.num_capsules_prev = num_cap_prev
        self.cap_dim = cap_dim
        self.cap_dim_prev = cap_dim_prev
        # TOLEARN: why the shape of W as given below ? how des 4d mat mul happen ?
        self.W = nn.Parameter(torch.randn(num_capsule, cap_dim_prev, num_cap_prev, cap_dim))

    def forward(self, u):
        # TOLEARN: understand whats happening with the matrices below
        u = u[None, :, :, None, :]
        W = self.W[:, None, :, :, :]
        u_hat = torch.matmul(u, W)
        # TOLEARN understand the dimensionality here
        b_ij = torch.zeros(*u_hat.size())
        if TRAIN_ON_GPU:
            b_ij = b_ij.cuda()
        v_j = dynamic_routing(b_ij, u_hat, routing_iteration=3)
        return v_j

# Decoding
class Decoder(nn.Module):
    def __init__(self, input_capsule_dim=16, input_capsules=10, hidden_dim=512):
        super(Decoder, self).__init__()
        input_dim = input_capsule_dim * input_capsules

        self.linear_decoder = nn.Sequential(OrderedDict([
            ('fc1', nn.Linear(input_dim, hidden_dim)),
            ('relu1', nn.ReLU(inplace=True)),
            ('fc2', nn.Linear(hidden_dim, hidden_dim * 2)),
            ('relu2', nn.ReLU(inplace=True)),
            ('fc3', nn.Linear(hidden_dim * 2, 28 * 28 * 1)),
            ('Sigmoid', nn.Sigmoid())
        ]))

    def forward(self, x):
        # probability of each class based on the magnitude of the vector
        class_prob = (x ** 2).sum(dim=-1) ** 0.5
        class_pred = F.softmax(class_prob, dim=-1)

        # Find capsule with max length, i.e. prediction
        _, max_length_idx = class_pred.max(dim=1)

        # identity matrix
        sparse_matrix = torch.eye(10)  # 10 classes
        if TRAIN_ON_GPU:
            sparse_matrix = sparse_matrix.cuda()

        # fancy way to get a vector with 1hot representation of the prediction
        y = sparse_matrix.index_select(dim=0, index=max_length_idx.data)

        # reconstructed pixels
        # so that we reconstruct only the predicted class, only have those in the reconstruction matrix
        x = x * y[:, :, None]
        flattened_x = x.view(x.size(0), -1)
        reconstructions = self.linear_decoder(flattened_x)

        # return reconstructions, class_pred
        return reconstructions, y

# Final capsule class that usess all the above class.
class CapsuleNetwork(nn.Module):
    def __init__(self):
        super(CapsuleNetwork, self).__init__()
        self.convLayer = ConvolutionLayer()
        self.primaryCapsle = PrimaryCapsule()
        self.digitCapsule = DigitCapsule()
        self.decoder = Decoder()

    def forward(self, images):
        primary_cap_op = self.primaryCapsle(self.convLayer(images))
        caps_out = self.digitCapsule(primary_cap_op).squeeze().transpose(0, 1)
        reconstruction, y = self.decoder(caps_out)
        return caps_out, reconstruction, y

capsnet = CapsuleNetwork()
print(capsnet)

# Loss for Capsule;
class CapsuleLoss(nn.Module):
    def __init__(self):
        super(CapsuleLoss, self).__init__()
        self.reconstructionLoss = nn.MSELoss(size_average=False)

    def forward(self, x, labels, images, reconstructions):
        batch_size = x.size(0)

        # TOTRY: maybe skip this, pass logits in above decoder fn instead of y
        v_c = torch.sqrt((x ** 2).sum(dim=2, keepdim=True))

        # get correct and incorrect loss
        left = F.relu(0.9 - v_c).view(batch_size, -1)
        right = F.relu(v_c - 0.1).view(batch_size, -1)

        margin_loss = labels * left + (1. - labels) * right
        margin_loss = margin_loss.sum()

        # get reconstruction loss
        images = images.view(reconstructions.size()[0], -1)
        reconstruction_loss = self.reconstructionLoss(reconstructions, images)

        # return a weighted sum
        return (margin_loss + 0.0005 * reconstruction_loss) / images.size(0)

criterion = CapsuleLoss()
optimizer = optim.Adam(capsnet.parameters())

if TRAIN_ON_GPU:
    capsnet = capsnet.cuda()

# train
def train(capsule_net, criterion, optimizer, n_epochs=1, print_every=1000):
    losses = []
    for epoch in range(1, n_epochs + 1):
        train_loss = 0.0
        capsule_net.train()

        for batch_i, (images, target) in enumerate(train_loader):

            # reshape and get target class
            target = torch.eye(10).index_select(dim=0, index=target)

            if TRAIN_ON_GPU:
                images, target = images.cuda(), target.cuda()

            # zero out gradients
            optimizer.zero_grad()
            # get model outputs
            caps_output, reconstructions, y = capsule_net(images)
            # calculate loss
            loss = criterion(caps_output, target, images, reconstructions)
            # perform backpropagation and optimization
            loss.backward()
            optimizer.step()

            train_loss += loss.item()  # accumulated training loss

            # print and record training stats
            if batch_i != 0 and batch_i % print_every == 0:
                avg_train_loss = train_loss / print_every
                losses.append(avg_train_loss)
                print('Epoch: {} \tTraining Loss: {:.8f}'.format(epoch, avg_train_loss))
                train_loss = 0  # reset accumulated training loss

    return losses

n_epochs = 3
losses = train(capsnet, criterion, optimizer, n_epochs=n_epochs, print_every=100)

plt.plot(losses)
plt.title("Training Loss")
plt.show()

def test(capsule_net, test_loader):
    '''Prints out test statistics for a given capsule net.
       param capsule_net: trained capsule network
       param test_loader: test dataloader
       return: returns last batch of test image data and corresponding reconstructions
       '''
    class_correct = list(0. for i in range(10))
    class_total = list(0. for i in range(10))

    test_loss = 0  # loss tracking

    capsule_net.eval()  # eval mode

    for batch_i, (images, target) in enumerate(test_loader):
        target = torch.eye(10).index_select(dim=0, index=target)

        batch_size = images.size(0)

        if TRAIN_ON_GPU:
            images, target = images.cuda(), target.cuda()

        # forward pass: compute predicted outputs by passing inputs to the model
        caps_output, reconstructions, y = capsule_net(images)
        # calculate the loss
        loss = criterion(caps_output, target, images, reconstructions)
        # update average test loss
        test_loss += loss.item()
        # convert output probabilities to predicted class
        _, pred = torch.max(y.data.cpu(), 1)
        _, target_shape = torch.max(target.data.cpu(), 1)

        # compare predictions to true label
        correct = np.squeeze(pred.eq(target_shape.data.view_as(pred)))
        # calculate test accuracy for each object class
        for i in range(batch_size):
            label = target_shape.data[i]
            class_correct[label] += correct[i].item()
            class_total[label] += 1

    # avg test loss
    avg_test_loss = test_loss / len(test_loader)
    print('Test Loss: {:.8f}\n'.format(avg_test_loss))

    for i in range(10):
        if class_total[i] > 0:
            print('Test Accuracy of %5s: %2d%% (%2d/%2d)' % (
                str(i), 100 * class_correct[i] / class_total[i],
                np.sum(class_correct[i]), np.sum(class_total[i])))
        else:
            # print('Test Accuracy of %5s: N/A (no training examples)' % (classes[i]))
            pass

    print('\nTest Accuracy (Overall): %2d%% (%2d/%2d)' % (
        100. * np.sum(class_correct) / np.sum(class_total),
        np.sum(class_correct), np.sum(class_total)))

    # return last batch of capsule vectors, images, reconstructions
    return caps_output, images, reconstructions

caps_output, images, reconstructions = test(capsnet, test_loader)

def display_images(images, reconstructions):
    '''Plot one row of original MNIST images and another row (below)
       of their reconstructions.'''
    # convert to numpy images
    images = images.data.cpu().numpy()
    reconstructions = reconstructions.view(-1, 1, 28, 28)
    reconstructions = reconstructions.data.cpu().numpy()

    # plot the first ten input images and then reconstructed images
    fig, axes = plt.subplots(nrows=2, ncols=10, sharex=True, sharey=True, figsize=(26, 5))

    # input images on top row, reconstructions on bottom
    for images, row in zip([images, reconstructions], axes):
        for img, ax in zip(images, row):
            ax.imshow(np.squeeze(img), cmap='gray')
            ax.get_xaxis().set_visible(False)
            ax.get_yaxis().set_visible(False)

display_images(images, reconstructions)


import responsavel_financeiro
import boto3
from moto import mock_s3

@mock_s3
def test_put():
    conn = boto3.resource('s3', region_name='us-east-1')
    conn.create_bucket(Bucket='mybucket')
    body = conn.Object('mybucket', 'steve').get()['Body'].read().decode("utf-8")

    responsavel_financeiro.get_dados_responsavel_financeiro('mybucket', {'file_key':'file_key'})

    assert body == 'file_key'

@mock_s3
def test_empty_key_set_on_existing_key():
    s3 = boto3.resource("s3")
    client = boto3.client("s3")
    s3.create_bucket(Bucket="foobar")

    key = s3.Object("foobar", "the-key")
    key.put(Body=b"some content")

    resp = responsavel_financeiro.get_dados_responsavel_financeiro(Bucket="foobar", Key="the-key")
    resp["Body"].read().should.equal(b"some content")

    key.put(Body=b"")

    resp = client.get_object(Bucket="foobar", Key="the-key")
    resp.should.have.key("ContentLength").equal(0)
    resp["Body"].read().should.equal(b"")


# Dada uma placa quadrada 1x1 m^2 e as temperaturas de fronteiras
# calcule a temperatura nesta placa e visualize a distribuição
# usando Gauss-Seidel e um grid nxn

# Aluno: Leonardo Meireles Murtha Oliveira
# NºUSP: 4182085

import GaussSeidel as gsd
import numpy as np
import matplotlib.pyplot as plt

# Retorna a matriz de coeficientes e o vetor B
def createLinearSystem(grid, n):
    m = (n-2)**2
    A = 4*np.identity(m, float) # Vezes quatro pois a variavel é a média entre 4
    b = np.zeros(m, float)

    for i in range(1, n-1):
        for j in range(1,n-1):
            g = 0

            if(grid[i,j-1] == 3.14):
                A[(i-1)*(n-2) + (j-1)][(i-1)*(n-2) + (j-2)] = -1
                #print('{} = {} {}'.format(((i-1)*(n-2) + (j-1)),i-1, (i-1)*(n-2) + (j-2)))
            else: g = g + grid[i, j-1]

            if(grid[i-1,j] == 3.14):
                A[(i-1)*(n-2) + (j-1)][(i-2)*(n-2) + (j-1)] = -1
                #print('{} = {} {}'.format(((i-1)*(n-2) + (j-1)),i-1, (i-2)*(n-2) + (j-1)))
            else: g = g + grid[i-1, j]

            if(grid[i,j+1] == 3.14):
                A[(i-1)*(n-2) + (j-1)][(i-1)*(n-2) + j] = -1
                #print('{} = {} {}'.format(((i-1)*(n-2) + (j-1)),i-1, (i-1)*(n-2) + j))
            else: g = g + grid[i, j+1]

            if(grid[i+1,j] == 3.14):
                A[(i-1)*(n-2) + (j-1)][i*(n-2) + ((j)-1)] = -1
                #print('{} = {} {}'.format(((i-1)*(n-2) + (j-1)),i-1, i*(n-2) + ((j)-1)))
            else: g = g + grid[i+1, j]

            b[(i-1)*(n-2) + (j-1)] = g

    return A, b

# Inicializa o grid com as temp. fronteiras
def initGrid(dR, n):
    grid = np.full((n, n), float(3.14),float)

    print('>>Inicializando o grid.')
    grid[0,0] = (dR[0] + dR[1])/2
    grid[n-1,0] = (dR[1] + dR[2])/2
    grid[n-1,n-1] = (dR[2] + dR[3])/2
    grid[0, n-1] = (dR[3] + dR[0])/2

    for i in range(1,n-1):
        grid[i,0] = dR[0] # Parte esquerda
        grid[0,i] = dR[1] # Parte do topo
        grid[i,n-1] = dR[2] # Parte direita
        grid[n-1,i] = dR[3] # Parde de baixo

    return grid

# Atualiza o grid com a solução
def updateGrid(grid, xf, n):
    for i in range(1, n-1):
        for j in range(1,n-1):
            grid[i][j] = xf[(i-1)*(n-2) + (j-1)]

def main():

    print('>>>A coleta de fronteira é sentido horário começando pelo lado esquerdo.')
    print('>>>Lembrando que a dimensão do grid é n gera (n-2)^2 variaveis.')
    n = int(input('>>>Digite a dimensão do grid.\n'))
    dR = np.zeros(4) # Vetor de temperaturas das fronteiras

    for i in range(4):
        dR[i] = float(input('>>>Digite a temp da fronteira [{}].\n'.format(i+1)))

    grid = initGrid(dR, n) # Inicia o grid com as temperatura de fronteiras em equilibrio
    A, b = createLinearSystem(grid, n)  # Cria a matriz a partir das equações
    xi = gsd.initialX(A, b) # Pega o chute inicial
    t = 20 # Número de iterações do Gauss-Seidel
    xf = gsd.solve(A, b, xi, t) # Resolve o sistema
    updateGrid(grid, xf, n) # Atualiza o grid com a solução aproximada

    # Parte de visualização
    plt.title('Heat Distribution')
    plt.xlabel('Coluna do grid')
    plt.ylabel('Linha do grid')
    plt.imshow(grid, cmap='plasma', interpolation='nearest')
    plt.colorbar()
    plt.show()

if __name__ == "__main__":
    main()


import asyncio, aiohttp, async_timeout, lxml, json, datetime as dt, itertools, copy, pickle, math, scipy
from bs4 import BeautifulSoup, SoupStrainer
from time import strptime
from statistics import mean
from scipy.stats.mstats import gmean

class Team:
	def __init__(self, teamid, longname, shortname=None, isd1=False, conference=None):
		self.teamid = teamid
		self.longname = longname
		self.shortname = shortname
		self.isd1 = isd1
		self.conference = conference
		self.games = set()

	def winslosses(self):
		sortedlist = [[],[]]
		for game in self.games:
			if game.winner() is self:
				sortedlist[0].append(game)
			elif game.loser() is self:
				sortedlist[1].append(game)
		return sortedlist

	def opponent(self, game):
		if game.hometeam is self:
			return game.awayteam
		elif game.awayteam is self:
			return game.hometeam
		else:
			print(self.longname+' did not play in game '+str(game.gameid)+'!')

	def D1winslosses(self):
		winslosses = self.winslosses()
		return [[win for win in winslosses[0] if self.opponent(win).isd1], [loss for loss in winslosses[1] if self.opponent(loss).isd1]]

	def record(self):
		winslosses = self.winslosses()
		return str(len(winslosses[0]))+'-'+str(len(winslosses[1]))

	def D1record(self):
		D1winslosses = self.D1winslosses()
		return str(len(D1winslosses[0]))+'-'+str(len(D1winslosses[1]))

	def winpct(self):
		winslosses = self.winslosses()
		wins = winslosses[0]
		losses = winslosses[1]
		if not wins and not losses:
			return 0.
		else:
			return len(wins)/(len(wins)+len(losses))

	def D1winpct(self):
		D1winslosses = self.D1winslosses()
		D1wins = D1winslosses[0]
		D1losses = D1winslosses[1]
		if not D1wins and not D1losses:
			return 0.
		else:
			return len(D1wins)/(len(D1wins)+len(D1losses))

	def D1winpctwithoutopponent(self, opponent):
		D1winslosses = self.D1winslosses()
		D1winswithoutopponent = [win for win in D1winslosses[0] if self.opponent(win) is not opponent]
		D1losseswithoutopponent = [loss for loss in D1winslosses[1] if self.opponent(loss) is not opponent]
		if not D1winswithoutopponent and not D1losseswithoutopponent:
			return 0.
		else:
			return len(D1winswithoutopponent)/(len(D1winswithoutopponent)+len(D1losseswithoutopponent))

	def opponents(self):
		return {self.opponent(game) for game in self.games}

	def D1opponents(self):
		return {self.opponent(game) for game in self.games if self.opponent(game).isd1}

	def OWP(self):
		return mean([opponent.D1winpctwithoutopponent(self) for opponent in self.D1opponents()])

	def OOWP(self):
		return mean([opponent.OWP() for opponent in self.D1opponents()])

	def RPI(self):
		return 0.25*self.D1winpct()+0.5*self.OWP()+0.25*self.OOWP()

class Game:
	def __init__(self, gameid, time, hometeam, homescore, awayteam, awayscore, ots=0):
		self.gameid = gameid
		self.time = time
		self.hometeam = hometeam
		self.homescore = homescore
		self.awayteam = awayteam
		self.awayscore = awayscore
		self.ots = ots

	def winner(self):
		if self.homescore > self.awayscore:
			return self.hometeam
		elif self.awayscore > self.homescore:
			return self.awayteam

	def loser(self):
		if self.homescore > self.awayscore:
			return self.awayteam
		elif self.awayscore > self.homescore:
			return self.hometeam

	def pointtotal(self):
		return self.homescore+self.awayscore

	def pointdifferential(self):
		return abs(self.homescore-self.awayscore)

def savedata(teams, games, savefile):
	saveteams = [[team.teamid, team.longname, team.shortname, team.isd1, team.conference, [game.gameid for game in team.games]] for team in teams.values()]
	savegames = [[game.gameid, game.time, game.hometeam.teamid, game.homescore, game.awayteam.teamid, game.awayscore, game.ots] for game in games.values()]
	pickle.dump([saveteams, savegames], open(savefile,'wb'))

def recalldata(savefile):
	teamsgames = pickle.load(open(savefile,'rb'))
	recallteams = {team[0] : Team(team[0], team[1], shortname=team[2], isd1=team[3], conference=team[4]) for team in teamsgames[0]}
	recallgames = {game[0] : Game(game[0], game[1], recallteams[game[2]], game[3], recallteams[game[4]], game[5], ots=game[6]) for game in teamsgames[1]}
	for game in recallgames.values():
		game.hometeam.games.add(game)
		game.awayteam.games.add(game)
	return recallteams, recallgames

async def getD1teams(session):
	teams = dict()
	async with session.get('https://www.espn.com/womens-college-basketball/teams') as resp:
		html = await resp.text()
		teamsoup = BeautifulSoup(html, 'lxml', parse_only=SoupStrainer(text=lambda string: string.startswith("window['__espnfitt__']")))
	columns = json.loads(teamsoup.text[23:-1])['page']['content']['leagueTeams']['columns']
	for column in columns:
		groups = column['groups']
		for group in groups:
			conference = group['nm']
			teamsjson = group['tms']
			for teamjson in teamsjson:
				teamid = int(teamjson['id'])
				longname = teamjson['n']
				team = Team(teamid, longname, isd1=True, conference=conference)
				teams[teamid] = team
	return teams

async def dictyieldvalues(dictionary):
	for i in dictionary.values():
		yield i

async def getD1teamschedulejson(teamid, session):
	#print('getting '+str(teamid))
	schedulejson = None
	while not schedulejson:
		async with session.get('https://www.espn.com/womens-college-basketball/team/schedule/_/id/'+str(teamid)) as resp:
			try:
				with async_timeout.timeout(5.0):
					html = await asyncio.wait_for(resp.text(), timeout=5.0)
			except asyncio.TimeoutError:
				#print('still waiting for '+str(teamid)+', restarting')
				continue
			try:
				schedulesoup = BeautifulSoup(html, 'lxml', parse_only=SoupStrainer(text=lambda string: string.startswith("window['__espnfitt__']")))
				schedulejson = json.loads(schedulesoup.text[23:-1])['page']['content']['scheduleData']
			except json.decoder.JSONDecodeError:
				#print('problem with '+str(teamid))
				pass
	#print('done getting '+str(teamid))
	return teamid, schedulejson

def processschedulejson(teams, team, schedulejson):
	games = dict()
	team.shortname = schedulejson['team']['abbrev']
	gamesjson = list(itertools.chain.from_iterable([schedule['events']['post'] for schedule in schedulejson['teamSchedule']]))
	for gamejson in gamesjson:
		gameid = int(gamejson['time']['link'].split('=')[1])
		datetime = dt.datetime.strptime(gamejson['date']['date'], '%Y-%m-%dT%H:%MZ').replace(tzinfo=dt.timezone.utc)

		opponentjson = gamejson['opponent']
		opponentid = int(opponentjson['id'])
		if opponentid not in teams:
			try:
				opponentshortname = opponentjson['abbrev']
			except KeyError:
				opponentshortname = None
			opponent = Team(opponentid, opponentjson['displayName'], shortname=opponentshortname)
			teams[opponentid] = opponent
		else:
			opponent = teams[opponentid]
		vsat = opponentjson['homeAwaySymbol']
		neutralsite = opponentjson['neutralSite']

		result = gamejson['result']
		teamscore = int(result['currentTeamScore'])
		opponentscore = int(result['opponentTeamScore'])

		try:
			overtime = result['overtime']
			if overtime == 'OT':
				ots = 1
			else:
				ots = int(overtime[:-2])
		except KeyError:
			ots = 0

		if vsat == 'vs':
			if neutralsite:
				gameteams = sorted([[team, teamscore], [opponent, opponentscore]], key=lambda x: x[0].teamid)
				hometeam = gameteams[0][0]
				homescore = gameteams[0][1]
				awayteam = gameteams[1][0]
				awayscore = gameteams[1][1]
				game = Game(gameid, datetime, hometeam, homescore, awayteam, awayscore, ots=ots)
				games[gameid] = game
			else:
				hometeam = team
				homescore = teamscore
				awayteam = opponent
				awayscore = opponentscore
				game = Game(gameid, datetime, hometeam, homescore, awayteam, awayscore, ots=ots)
				games[gameid] = game
		elif vsat == '@':
			hometeam = opponent
			homescore = opponentscore
			awayteam = team
			awayscore = teamscore
			game = Game(gameid, datetime, hometeam, homescore, awayteam, awayscore, ots=ots)
			games[gameid] = game
	return games

async def getD1teamsgames(savefile=None):
	async with aiohttp.ClientSession() as session:
		teams = await getD1teams(session)
		schedulejsons = await asyncio.gather(*[getD1teamschedulejson(team.teamid, session) async for team in dictyieldvalues(teams)])
	games = dict()
	for teamid, schedulejson in schedulejsons:
		games.update(processschedulejson(teams, teams[teamid], schedulejson))
	for game in games.values():
		game.hometeam.games.add(game)
		game.awayteam.games.add(game)
	return teams, games

#KRACH rankings
def cleanexistingties(teamsin, gamesin):
	teams = {team.teamid: Team(team.teamid, team.longname, shortname=team.shortname, isd1=team.isd1, conference=team.conference) for team in teamsin.values()}
	games = {game.gameid: Game(game.gameid, game.time, teams[game.hometeam.teamid], game.homescore, teams[game.awayteam.teamid], game.awayscore) for game in gamesin.values() if game.winner()}
	for game in list(games.values()):
		game.hometeam.games.add(game)
		game.awayteam.games.add(game)
	return teams, games

def cleanbeforetime(time, teamsin, gamesin):
	teams = {team.teamid: Team(team.teamid, team.longname, shortname=team.shortname, isd1=team.isd1, conference=team.conference) for team in teamsin.values()}
	games = {game.gameid: Game(game.gameid, game.time, teams[game.hometeam.teamid], game.homescore, teams[game.awayteam.teamid], game.awayscore) for game in gamesin.values() if game.time >= time}
	for game in list(games.values()):
		game.hometeam.games.add(game)
		game.awayteam.games.add(game)
	return teams, games

def cleanaftertime(time, teamsin, gamesin):
	teams = {team.teamid: Team(team.teamid, team.longname, shortname=team.shortname, isd1=team.isd1, conference=team.conference) for team in teamsin.values()}
	games = {game.gameid: Game(game.gameid, game.time, teams[game.hometeam.teamid], game.homescore, teams[game.awayteam.teamid], game.awayscore) for game in gamesin.values() if game.time < time}
	for game in list(games.values()):
		game.hometeam.games.add(game)
		game.awayteam.games.add(game)
	return teams, games

def numplayed(team1, team2):
	return len(team1.games.intersection(team2.games))

def sos(krachratings, team):
	return gmean([krachratings[team.opponent(game).teamid] for game in team.games if team.opponent(game).isd1])

def oocsos(krachratings, team):
	return gmean([krachratings[team.opponent(game).teamid] for game in team.games if team.opponent(game).isd1 and team.opponent(game).conference != team.conference])

'''def sos(krachratings, team):
	return sum(krachratings[opponent.teamid]*numplayed(team, opponent)/(krachratings[opponent.teamid]+krachratings[team.teamid]) for opponent in team.D1opponents())/sum(numplayed(team, opponent)/(krachratings[opponent.teamid]+krachratings[team.teamid]) for opponent in team.D1opponents())'''

def conferencestrength(krachratings, teams, conference):
	return gmean([krachratings[team.teamid] for team in teams.values() if team.conference == conference])

def victorypoints(team, alpha):
	D1winslosses = team.D1winslosses()
	return sum([1/(1+math.exp(-game.pointdifferential()/(alpha*game.pointtotal()))) for game in D1winslosses[0]]+[1/(1+math.exp(game.pointdifferential()/(alpha*game.pointtotal()))) for game in D1winslosses[1]])

def rrwp(krachratings, team):
	teamrating = krachratings[team.teamid]
	return sum([teamrating/(teamrating+krachratings[opponentid]) for opponentid in krachratings if opponentid != team.teamid])/(len(krachratings)-1)

def rrwpkrach(krachratings, rating):
	return sum([rating/(rating+opprating) for opprating in krachratings.values()])/len(krachratings)

def krachadj(krachratings):
	return scipy.optimize.fsolve(lambda rating: rrwpkrach(krachratings, rating)-0.5, 100)[0]

def calckrachratings(teamsin, gamesin, vpalpha=5, goaldelta=1e-10, time=None, sincetime=None, savefile=None, calcteamsos=True):
	teams, games = cleanexistingties(teamsin, gamesin)
	if time:
		teams, games = cleanaftertime(time, teams, games)
	if sincetime:
		teams, games = cleanbeforetime(sincetime, teams, games)
	D1teams = {team.teamid: team for team in teams.values() if team.isd1}

	krachratings = {teamid: 100. for teamid, team in D1teams.items() if team.isd1}

	iterations = 0
	alphaadj = vpalpha/mean([game.pointtotal() for game in games.values() if game.hometeam.isd1 and game.awayteam.isd1])
	victorypoints_dict = dict()
	D1opponents_dict = dict()
	numplayed_dict = dict()
	for team in D1teams.values():
			victorypoints_dict[team] = victorypoints(team, alphaadj)
			D1opponents = team.D1opponents()
			D1opponents_dict[team] = D1opponents
			for opponent in D1opponents:
				numplayed_dict[frozenset({team, opponent})] = numplayed(team, opponent)
	while True:
		print('Iteration '+str(iterations+1))
		newkrachratings = dict(krachratings)
		delta = 0.

		for team in D1teams.values():
			newkrachratings[team.teamid] = victorypoints_dict[team]/sum(numplayed_dict[frozenset({team, opponent})]/(krachratings[team.teamid]+krachratings[opponent.teamid]) for opponent in D1opponents_dict[team])
			delta += abs(newkrachratings[team.teamid]-krachratings[team.teamid])

		krachratings = dict(newkrachratings)

		print(delta)
		if delta < goaldelta*gmean(list(krachratings.values())):
			adj = krachadj(krachratings)
			krachratings = {k: v*100/adj for k, v in krachratings.items()}
			break

		iterations += 1
	if calcteamsos:
		teamsos = {team.teamid: sos(krachratings, team) for team in D1teams.values()}
	if savefile:
		if calcteamsos:
			pickle.dump([krachratings, teamsos], open(savefile,'wb'))
		else:
			pickle.dump(krachratings, open(savefile,'wb'))
	if calcteamsos:
		return krachratings, teamsos
	else:
		return krachratings

def calckrachratingsold(teamsin, gamesin, goaldelta=1e-10, time=None, savefile=None, calcteamsos=True):
	teams, games = cleanexistingties(teamsin, gamesin)
	if time:
		teams, games = cleanaftertime(time, teams, games)
	D1teams = {team.teamid: team for team in teams.values() if team.isd1}

	if [team for team in D1teams.values() if (not team.D1winslosses()[0] or not team.D1winslosses()[1])]:
		needbogusteam = True
	else:
		needbogusteam = False
	krachratings = {teamid: 100. for teamid, team in D1teams.items() if team.isd1}
	if needbogusteam:
		krachratings[-1] = 100.

	iterations = 0
	numwins_dict = dict()
	D1opponents_dict = dict()
	numplayed_dict = dict()
	for team in D1teams.values():
		numwins_dict[team] = len(team.D1winslosses()[0])
		D1opponents = team.D1opponents()
		D1opponents_dict[team] = D1opponents
		for opponent in D1opponents:
			numplayed_dict[frozenset({team, opponent})] = numplayed(team, opponent)
	while True:
		print('Iteration '+str(iterations+1))
		newkrachratings = dict(krachratings)
		delta = 0.

		if not needbogusteam:
			for team in D1teams.values():
				newkrachratings[team.teamid] = numwins_dict[team]/sum(numplayed_dict[frozenset({team, opponent})]/(krachratings[team.teamid]+krachratings[opponent.teamid]) for opponent in D1opponents_dict[team])
				delta += abs(newkrachratings[team.teamid]-krachratings[team.teamid])
		if needbogusteam:
			for team in D1teams.values():
				newkrachratings[team.teamid] = (numwins_dict[team]+0.5)/(1/(krachratings[team.teamid]+krachratings[-1])+sum(numplayed_dict[frozenset({team, opponent})]/(krachratings[team.teamid]+krachratings[opponent.teamid]) for opponent in D1opponents_dict[team]))
				delta += abs(newkrachratings[team.teamid]-krachratings[team.teamid])
			newkrachratings[-1] = (len(D1teams)/2)/sum(1/(krachratings[team.teamid]+krachratings[-1]) for team in D1teams.values())

		krachratings = dict(newkrachratings)

		print(delta)
		if delta < goaldelta:
			krachratings = {k: v*100./krachratings[-1] for k, v in krachratings.items()}
			break

		iterations += 1
	if calcteamsos:
		teamsos = {team.teamid: sos(krachratings, team) for team in D1teams.values()}
	if savefile:
		if calcteamsos:
			pickle.dump([krachratings, teamsos], open(savefile,'wb'))
		else:
			pickle.dump(krachratings, open(savefile,'wb'))
	if calcteamsos:
		return krachratings, teamsos
	else:
		return krachratings


'''
A function module
'''

import pygame
import os
import random
import asset
from asset import Block
from player import Player
from enemy import Portal, Imp, FatImp
from bullet import Bullet

#wrapper function
def run_game(width, height, fps, starting_scene):
    '''
    wrapper function that runs the whole game
    all global variables should be defined here
    '''
    os.environ['SDL_VIDEO_CENTERED'] = '1'
    pygame.init()
    window = pygame.display.set_mode((width, height))
    pygame.display.set_caption("Imps & Assasin")
    clock = pygame.time.Clock()
    font = pygame.font.SysFont("calibri", 30)
    font_small = pygame.font.SysFont("calibri", 25)

    #instantiate sprite objects
    assasin = Player(asset.NINJA_R[0])
    #add sprites to groups
    asset.PLAYER_GROUP.add(assasin)

    active_scene = starting_scene

    #GAME LOOP###
    while active_scene != None:
        key = pygame.key.get_pressed()

        # Event filtering
        filtered_events = []
        for event in pygame.event.get():
            quit_attempt = False
            if event.type == pygame.QUIT:
                quit_attempt = True
            elif event.type == pygame.KEYDOWN:
                alt_pressed = key[pygame.K_LALT] or \
                              key[pygame.K_RALT]
                if event.key == pygame.K_ESCAPE:
                    quit_attempt = True
                elif event.key == pygame.K_F4 and alt_pressed:
                    quit_attempt = True

            if quit_attempt:
                active_scene.Terminate()
            else:
                filtered_events.append(event)

        #pass variables to scene methods
        active_scene.ProcessInput(filtered_events, key, assasin)
        active_scene.Update(assasin)
        active_scene.Render(window, font, font_small)

        active_scene = active_scene.next

        pygame.display.flip()

        clock.tick(fps)

#functions passed to player methods
def identify_death(player, screen, orientation):
    if player.health <= 0:
        if orientation == 0:
            screen.blit(asset.NINJA_DEATH_L[player.deathblocker // 3], (player.rect.x, player.rect.y))
        else:
            screen.blit(asset.NINJA_DEATH_R[player.deathblocker // 3], (player.rect.x, player.rect.y))

        if player.deathblocker < 6:
            player.deathblocker += 1

def identify_jump_attack(player, screen, orientation):
    if player.health > 0 and player.attack == True and player.shoot_bullet == False and player.jump == True:
        if orientation == 0:
            screen.blit(asset.NINJA_ATTACK_L[player.meleblocker // 3], (player.rect.x-40, player.rect.y))
        else:
            screen.blit(asset.NINJA_ATTACK_R[player.meleblocker // 3], (player.rect.x, player.rect.y))

        if player.jumpblocker < 6:
            player.jumpblocker += 1

def identify_jump_throw(player, screen, orientation):
    if player.health > 0 and player.attack == False and player.shoot_bullet == True and player.jump == True:
        if orientation == 0:
            screen.blit(asset.NINJA_SHOOT_L[player.throwblocker // 3], (player.rect.x, player.rect.y))
        else:
            screen.blit(asset.NINJA_SHOOT_R[player.throwblocker // 3], (player.rect.x, player.rect.y))

        if player.throwblocker < 6:
            player.throwblocker += 1

def identify_jump(player, screen, orientation):
    if player.health > 0 and player.attack == False and player.shoot_bullet == False and player.jump == True:
        if orientation == 0:
            screen.blit(asset.NINJA_JUMP_L[player.jumpblocker // 3], (player.rect.x, player.rect.y))
        else:
            screen.blit(asset.NINJA_JUMP_R[player.jumpblocker // 3], (player.rect.x, player.rect.y))

        if player.jumpblocker < 6:
            player.jumpblocker += 1

def identify_attack(player, screen, orientation):
    if player.health > 0 and player.attack == True and player.shoot_bullet == False and player.jump == False:
        if orientation == 0:
            screen.blit(asset.NINJA_ATTACK_L[player.meleblocker // 3], (player.rect.x-40, player.rect.y))
        else:
            screen.blit(asset.NINJA_ATTACK_R[player.meleblocker // 3], (player.rect.x, player.rect.y))

        if player.jumpblocker < 6:
            player.jumpblocker += 1

def identify_throw(player, screen, orientation):
    if player.health > 0 and player.attack == False and player.shoot_bullet == True and player.jump == False:
        if orientation == 0:
            screen.blit(asset.NINJA_SHOOT_L[player.throwblocker // 3], (player.rect.x, player.rect.y))
        else:
            screen.blit(asset.NINJA_SHOOT_R[player.meleblocker // 3], (player.rect.x, player.rect.y))

        if player.throwblocker < 6:
            player.throwblocker += 1

def identify_walk(player, screen, orientation):
    if player.health > 0 and player.attack == False and player.shoot_bullet == False and player.jump == False:
        if orientation == 0:
            screen.blit(asset.NINJA_L[player.walkblocker // 3], (player.rect.x, player.rect.y))
        else:
            screen.blit(asset.NINJA_R[player.walkblocker // 3], (player.rect.x, player.rect.y))

    if player.walkblocker + 1 < 9:
        player.walkblocker += 1
    else:
        player.walkblocker = 0

#functions passed to enemy methods
def enemy_walk(enemy):
    dice_roll_1 = (random.random()*100)

    if dice_roll_1 < enemy.walk_dynamics:
        enemy.vel *= -1

    if enemy.vel > 0:
        if enemy.rect.x + enemy.vel < enemy.path[1]:
            enemy.rect.x += enemy.vel
        else:
            enemy.vel *= -1
            enemy.walkcount = 0
    else:
        if enemy.rect.x + enemy.vel > enemy.path[0]:
            enemy.rect.x += enemy.vel
        else:
            enemy.vel *= -1
            enemy.walkcount = 0

def enemy_jump(enemy):
    dice_roll_2 = (random.random()*100)

    if enemy.jump == False:
        if dice_roll_2 < enemy.jump_dynamics:
            enemy.jump = True
    else:
        if enemy.jumpcount >= -10:
            neg = 1
            if enemy.jumpcount < 0:
                neg = -1
            enemy.rect.y -= (enemy.jumpcount ** 2) * enemy.jump_factor * neg
            enemy.jumpcount -= 1
        else:
            enemy.rect.y = enemy.y_cap
            enemy.jump = False
            enemy.jumpcount = 10

#functions passed to scene methods
def movement_manager(player):
    #player sprite movement mechanics
    keys = pygame.key.get_pressed()

    if keys[pygame.K_LEFT] and player.rect.x > 0:
        player.rect.x -= player.vel
        player.right = False
        player.left = True
    elif keys[pygame.K_RIGHT] and player.rect.x < asset.WINDOW_W - player.rect.width:
        player.rect.x += player.vel
        player.right = True
        player.left = False
    else:
        player.walkblocker = 0
    #jumping mechanics

def jump_manager(player):
    keys = pygame.key.get_pressed()
    if player.jump == False:
        if keys[pygame.K_UP]:
            player.jump = True
    else:
        if player.jumpcount >= -12:
            neg = 1
            if player.jumpcount < 0:
                neg = -1
            player.rect.y -= (player.jumpcount ** 2) * 0.35 * neg
            player.jumpcount -= 1

        else:
            player.jump = False
            player.jumpcount = player.jumpcount_cap
            player.rect.y = player.rect_y_cap

def shooting_manager(player):
    #player sprite attack controls mechanics
    keys = pygame.key.get_pressed()

    #shooting bullets mechanics
    if player.shoot_bullet == False:
        if keys[pygame.K_x]:
            player.shoot_bullet = True
    else:
        if player.left == True:
            shuriken = Bullet(asset.BULLET_IMG,
                              player.rect.x, player.rect.y+45, 3, -1)
            asset.BULLET_GROUP.add(shuriken)
            player.shoot_bullet = False

        else:
            shuriken = Bullet(asset.BULLET_IMG,
                              player.rect.x+25, player.rect.y+45, 3, 1)
            asset.BULLET_GROUP.add(shuriken)
            player.shoot_bullet = False

def attack_manager(player):
    #player sprite attack controls mechanics
    keys = pygame.key.get_pressed()

    #mele attack switch
    if player.attack == False:
        if keys[pygame.K_SPACE]:
            player.attack = True
            print(player.attack)
    else:
        player.attack = False

def redraw_manager(screen, bg, font, player_container, bullet_container, enemy_container):
    #blit background image
    screen.blit(bg, (0, -190))
    #render and blit text
    text = font.render("Assasin HP", 1, (250, 250, 250), )
    screen.blit(text, (13, 25))
    #blit static health bar (foreground)
    screen.blit(asset.H_BAR_IMG, (10, 50))

    #draw health bar
    for player in player_container:
        if player.health > 0: #just to make sure no leftover green bar is blitted
            health_factor = 8 * (30 - player.health)
            pygame.draw.rect(screen, (20, 89, 12), (15, 55, 240 - health_factor, 20))
    #update projectile data and draw
    for bullet in bullet_container:
        bullet.bullet_update()
        asset.BULLET_GROUP.draw(screen)
    #update enemy data and draw
    for enemy in enemy_container:
        enemy.update()
        enemy.draw(screen)
        #draw health bars for enemies
        health_factor = (60/enemy.health_cap) * (enemy.health_cap - enemy.health)
        pygame.draw.rect(screen, (100, 0, 0), (enemy.rect.x, enemy.rect.y - 30, 60, 8))
        pygame.draw.rect(screen, (35, 89, 12), (enemy.rect.x, enemy.rect.y - 30, 60 - health_factor, 8))
    #draw player using sprite container
    for player in player_container:
        player.draw(screen)

    pygame.display.flip()

def collision_manager(player_container, bullet_container, enemy_container, block_container):
    '''
    this function gathers all collision methods from instantiated objects
    and executes them
    '''
    #player mele attack
    for player in player_container:
        player.mele_enemy_collision(player, enemy_container)
    #bullet collisions
    for elem in bullet_container:
        elem.bullet_enemy_collision(bullet_container, enemy_container)
        elem.bullet_block_collision(bullet_container, block_container)
    #enemy mele attack
    for enemy in enemy_container:
        for player in player_container:
            enemy.enemy_player_collision(player, enemy, enemy.damage)

def death_manager(*args):
    for arg in args:
        for elem in arg:
            if elem.health <= 0:
                arg.remove(elem)
                print("Sprite killed")

def enemy_generator(impno, fatno, acidno, mamano, impcap, fatcap, acidcap, mamacap):
    seed_x = random.randint(200,700)
    portal = Portal(asset.PORTAL[0], seed_x, 25)
    asset.ENEMY_GROUP.add(portal)

    if impno > impcap:
        for n in range(0, impcap):
            imp = Imp(asset.IMP_R[0], asset.WINDOW_W, seed_x, 5)
            asset.ENEMY_GROUP.add(imp)
    else:
        for n in range(0, impno):
            imp = Imp(asset.IMP_R[0], asset.WINDOW_W, seed_x, 5)
            asset.ENEMY_GROUP.add(imp)

    if fatno > fatcap:
        for n in range(0, fatcap):
            fat_imp = FatImp(asset.IMP_R[0], asset.WINDOW_W, seed_x, 15)
            asset.ENEMY_GROUP.add(fat_imp)
    else:
        for n in range(0, fatno):
            fat_imp = FatImp(asset.FAT_IMP_R[0], asset.WINDOW_W, seed_x, 15)
            asset.ENEMY_GROUP.add(fat_imp)

    #add acid imp and mama imp

def transition_message(screen, font, smallfont):
    text1 = font.render("You banished all imps from this land!", 1, (250, 250, 250), )
    text2 = font.render("Hit enter when you are ready to move to next map",
    1, (250, 250, 250))
    screen.blit(text1, (200, 250))
    screen.blit(text2, (130, 285))

def game_over_message(screen, font, smallfont):
    text3 = font.render("YOU DIED", 1, (250, 250, 250))
    text4 = smallfont.render("Hit enter to restart", 1, (250, 250, 250))
    screen.blit(text3, (350, 250))
    screen.blit(text4, (320, 290))


#!/usr/bin/env python

# Avery Tan altan 1392212
# CMPUT366 A5

"""
  Author: Adam White, Matthew Schlegel, Mohammad M. Ajallooeian, Sina Ghiassian
  Purpose: Skeleton code for Monte Carlo Exploring Starts Control Agent
           for use on A3 of Reinforcement learning course University of Alberta Fall 2017

"""

from utils import rand_in_range, rand_un
import numpy as np
import pickle
import operator
import random

pi = None
last_action=None
alpha = 0.1
epsilon = 0.1
currX=None
currY=None
Q=None
gamma = 0.95
model = None
old_state = None
random.seed(a=1000)
N=5
theta = 0.00066
maxP = None
PQueue=None

prev_obs_states= None
prev_obs_actions_in_states= None

def calculate_poss_moves(xx,yy):
    """
    calculates possible moves given current x and y coordinates
    returns a list containing all valid legal moves
    """

    global currX , currY
    action_space=[(1,0),(-1,0),(0,1),(0,-1)]
    legal_moves=list()

    #filtering out illegal moves:
    for i in action_space:
        if i[0]+currX>=0 and i[0]+currX<=8 and i[1]+currY>=0 and i[1]+currY<=5:
            legal_moves.append(i)

    return legal_moves

def agent_init():
    """
    Hint: Initializes Q, where we store the Q(s,a). Q is a 2d array such that Q[x][y][a] where x and y are
    intergers and represent a particular location and a is an action and Q[x][y][a] is the state-action value for
    that particular action taken in that particular state
    Returns: nothing
    """
    global Q,model, prev_obs_states,prev_obs_actions_in_states,pi,PQueue
    prev_obs_states=set()
    prev_obs_actions_in_states=dict()
    Q=list()
    for i in range(9):
        Q.append([{(1,0):0,(-1,0):0,(0,1):0,(0,-1):0},
            {(1,0):0,(-1,0):0,(0,1):0,(0,-1):0},
            {(1,0):0,(-1,0):0,(0,1):0,(0,-1):0},
            {(1,0):0,(-1,0):0,(0,1):0,(0,-1):0},
            {(1,0):0,(-1,0):0,(0,1):0,(0,-1):0},
            {(1,0):0,(-1,0):0,(0,1):0,(0,-1):0}])
    model=dict()
    random.seed(a=1000)
    pi=dict()
    PQueue = dict()

def record_StateAndAction(os,la):
    '''
    function adds a state and action to 'memory'
    '''

    global prev_obs_states,prev_obs_actions_in_states
    #adding to list of seen state
    prev_obs_states.add(os)
    if os not in prev_obs_actions_in_states:
        prev_obs_actions_in_states[os]=list()
        prev_obs_actions_in_states[os].append(la)
    else:
        prev_obs_actions_in_states[os].append(la)

    return

def get_max_a_Q(currX,currY):
    '''
    returns the action with the largest estimated curr Q value as a tuple (x,y)
    '''
    global Q
    poss_moves=[(1,0),(-1,0),(0,-1),(0,1)]
    qlist=dict()
    for i in poss_moves:
        qlist[i]=Q[currX][currY][i]
    max_q=qlist[max(qlist, key=qlist.get)]

    #choose randomly all optimal legal moves
    equal_q_list=list()
    for i in poss_moves:
        if Q[currX][currY][i] >= max_q:
            equal_q_list.append(i)
    action_index=rand_in_range(len(equal_q_list))
    action = equal_q_list[action_index]

    return action

def agent_start(state):
    """
    Arguments: state: tuple (x,y) representing coordinates on gridworld
    Returns: action: tuple (x,y) representing an action
    """
    global currX,currY,Q,last_action,model,pi, old_state

    currX=state[0]
    currY=state[1]

    equal_q_list=list()
    action_space = [(1,0),(-1,0),(0,1),(0,-1)]

    #getting the highest q
    qlist=dict()
    for i in action_space:
        qlist[i]=Q[currX][currY][i]
    max_q=qlist[max(qlist, key=qlist.get)]

    #going through possible actions again to get all estimated optimal actions
    for i in action_space:
        if qlist[i]==max_q:
            equal_q_list.append(i)

    #randomly choose an action
    action_index=rand_in_range(len(equal_q_list))
    action = equal_q_list[action_index]

    old_state = (currX,currY)

    #adding to list of seen state
    record_StateAndAction(old_state,action)

    last_action=action
    return action

def agent_step(reward, state): # returns NumPy array, reward: floating point, this_observation: NumPy array
    """
    Arguments: reward: floting point, state: tuple (x,y) representing coordinates of agent
    Returns: action: tuple (x,y) representing move agent can make
    """
    global currX,currY,Q, epsilon,alpha,last_action,old_state,gamma,n,N,pi,theta,PQueue

    # simple worldmap implementation to visualize what is happening
    # worldmap = list()
    # obstructions = [(2,1),(2,2),(2,3),(5,4),(7,0),(7,1),(7,2)]
    # for u in range(9):
    #     worldmap.append(['.','.','.','.','.','.'])
    # worldmap[8][0]='G'
    # worldmap[currX][currY]='s'
    # for j in obstructions:
    #     worldmap[j[0]][j[1]]='*'
    # for i in worldmap:
    #     print i
    # print currX,currY
    # print '\n\n'

    action = None
    currX=state[0]
    currY=state[1]

    #update the model
    model[(old_state,last_action)]=(reward,state)

    action = get_max_a_Q(currX,currY)
    #LEARN!
    Q[old_state[0]][old_state[1]][last_action]=Q[old_state[0]][old_state[1]][last_action]+alpha*(reward+gamma*Q[currX][currY][action]-Q[old_state[0]][old_state[1]][last_action])

    #Calculate the priority P
    max_a=get_max_a_Q(currX,currY)
    priority=abs(reward+gamma*Q[currX][currY][max_a]-Q[old_state[0]][old_state[1]][last_action])
    if priority>theta:
        #ensure that if a state action pair already is in the queue, that this new priority is greater.
        if (old_state,last_action) in PQueue:
            if priority > PQueue[(old_state,last_action)]:
                PQueue[(old_state,last_action)]=priority
        else:
            PQueue[(old_state,last_action)]=priority

    #doing the planning update:
    if N != 0:
        n=random.randint(1,N)
        count_n=0
        while count_n<n and PQueue:
            count_n+=1
            S,A = max(PQueue, key=PQueue.get)
            maxP=PQueue[max(PQueue, key=PQueue.get)]
            del PQueue[(S,A)] #pop

            re,next_state = model[(S,A)]
            caction = get_max_a_Q(next_state[0],next_state[1])
            Q[S[0]][S[1]][A]+=alpha*(re+gamma*Q[next_state[0]][next_state[1]][caction]-Q[S[0]][S[1]][A])
            r_sp=None
            for i in model:
                if model[i][1]==S: #for every state leading to this state
                    r_sp = model[i][0]
                    preceding_state = i[0]
                    preceding_action = i[1]
                    max_a=get_max_a_Q(S[0],S[1])
                    priority = abs(r_sp+gamma*Q[S[0]][S[1]][max_a]-Q[preceding_state[0]][preceding_state[1]][preceding_action])
                    if priority > theta:
                        if (preceding_state,preceding_action) in PQueue:
                            if priority > PQueue[(preceding_state,preceding_action)]:
                                PQueue[(preceding_state,preceding_action)]=priority
                        else:
                            PQueue[(preceding_state,preceding_action)]=priority

    #now action selection for the next direct RL
    rndm=rand_un()
    if rndm < epsilon: #explore!
        poss_moves= [(1,0),(-1,0),(0,1),(0,-1)]
        action_index=rand_in_range(len(poss_moves))
        action = poss_moves[action_index]

    else: #exploit!
        action = get_max_a_Q(currX,currY)
        pass

    old_state = (currX,currY)
    last_action=action
    #adding to list of seen state
    record_StateAndAction(old_state,last_action)
    return action

def agent_end(reward):
    """
    Arguments: reward: floating point
    Returns: Nothing
    """
    # do learning and update pi
    global Q,old_state,last_action,alpha,N

    model[(old_state,last_action)]=(reward,(8,0))

    #LEARN!
    Q[old_state[0]][old_state[1]][last_action]=Q[old_state[0]][old_state[1]][last_action]+alpha*(reward-Q[old_state[0]][old_state[1]][last_action])

    #Calculate the priority P
    max_a=get_max_a_Q(currX,currY)
    priority=abs(reward-Q[old_state[0]][old_state[1]][last_action])
    if priority>theta:
        if (old_state,last_action) in PQueue:
            if priority > PQueue[(old_state,last_action)]:
                PQueue[(old_state,last_action)]=priority
        else:
            PQueue[(old_state,last_action)]=priority

    #doing the planning update:
    if N != 0:
        n=random.randint(1,N)
        count_n=0
        while count_n<n and PQueue:
            count_n+=1
            S,A = max(PQueue, key=PQueue.get)
            maxP=PQueue[max(PQueue, key=PQueue.get)]
            del PQueue[(S,A)]

            re,next_state = model[(S,A)]
            caction = get_max_a_Q(next_state[0],next_state[1])
            Q[S[0]][S[1]][A]+=alpha*(re+gamma*Q[next_state[0]][next_state[1]][caction]-Q[S[0]][S[1]][A])
            r_sp=None
            for i in model:
                if model[i][1]==S:
                    r_sp = model[i][0]
                    preceding_state = i[0]
                    preceding_action = i[1]
                    max_a=get_max_a_Q(S[0],S[1])
                    priority = abs(r_sp+gamma*Q[S[0]][S[1]][max_a]-Q[preceding_state[0]][preceding_state[1]][preceding_action])
                    if priority > theta:
                        if (preceding_state,preceding_action) in PQueue:
                            if priority > PQueue[(preceding_state,preceding_action)]:
                                PQueue[(preceding_state,preceding_action)]=priority
                        else:
                            PQueue[(preceding_state,preceding_action)]=priority

    return

def agent_cleanup():
    """
    This function is not used
    """
    # clean up

    return

def agent_message(in_message): # returns string, in_message: string
    global Q
    """
    Arguments: in_message: string
    returns: The value function as a string.
    This function is complete. You do not need to add code here.
    """
    # should not need to modify this function. Modify at your own risk
    if (in_message == 'ValueFunction'):
        return pickle.dumps(np.max(Q, axis=1), protocol=0)
    else:
        return "I don't know what to return!!"


import re
import time
import pytest
from unittest.mock import patch

from iredis.utils import timer, strip_quote_args
from iredis.commands import split_command_args, split_unknown_args
from iredis.utils import command_syntax
from iredis.style import STYLE
from iredis.exceptions import InvalidArguments, AmbiguousCommand
from iredis.commands import commands_summary
from prompt_toolkit import print_formatted_text

def test_timer():
    with patch("iredis.utils.logger") as mock_logger:
        timer("foo")
        time.sleep(0.1)
        timer("bar")
        mock_logger.debug.assert_called()
        args, kwargs = mock_logger.debug.call_args
        matched = re.match(r"\[timer (\d)\] (0\.\d+) -> bar", args[0])

        assert matched.group(1) == str(3)
        assert 0.1 <= float(matched.group(2)) <= 0.2

        # --- test again ---
        timer("foo")
        time.sleep(0.2)
        timer("bar")
        mock_logger.debug.assert_called()
        args, kwargs = mock_logger.debug.call_args
        matched = re.match(r"\[timer (\d)\] (0\.\d+) -> bar", args[0])

        assert matched.group(1) == str(5)
        assert 0.2 <= float(matched.group(2)) <= 0.3

@pytest.mark.parametrize(
    "test_input,expected",
    [
        ("hello world", ["hello", "world"]),
        ("'hello world'", ["hello world"]),
        ('''hello"world"''', ["helloworld"]),
        (r'''hello\"world"''', [r"hello\world"]),
        (r'"\\"', [r"\\"]),
        (r"\\", [r"\\"]),
        (r"\abcd ef", [r"\abcd", "ef"]),
        # quotes in quotes
        (r""" 'hello"world' """, ['hello"world']),
        (r""" "hello'world" """, ["hello'world"]),
        (r""" 'hello\'world'""", ["hello'world"]),
        (r""" "hello\"world" """, ['hello"world']),
        (r"''", [""]),  # set foo "" is a legal command
        (r'""', [""]),  # set foo "" is a legal command
        (r"\\", ["\\\\"]),  # backslash are legal
        ("\\hello\\", ["\\hello\\"]),  # backslash are legal
    ],
)
def test_stripe_quote_escape_in_quote(test_input, expected):
    assert list(strip_quote_args(test_input)) == expected

@pytest.mark.parametrize(
    "command,expected,args",
    [
        ("GET a", "GET", ["a"]),
        ("cluster info", "cluster info", []),
        ("getbit foo 17", "getbit", ["foo", "17"]),
        ("command ", "command", []),
        (" command count  ", "command count", []),
        (" command  count  ", "command  count", []),  # command with multi space
        (" command  count    ' hello   world'", "command  count", [" hello   world"]),
        ("set foo 'hello   world'", "set", ["foo", "hello   world"]),
    ],
)
def test_split_commands(command, expected, args):
    parsed_command, parsed_args = split_command_args(command)
    assert expected == parsed_command
    assert args == parsed_args

def test_split_commands_fail_on_unknown_command():
    with pytest.raises(InvalidArguments):
        split_command_args("FOO BAR")

@pytest.mark.parametrize(
    "command",
    ["command in", "command   in", "Command   in", "COMMAND     in"],
)
def test_split_commands_fail_on_partially_input(command):
    with pytest.raises(AmbiguousCommand):
        split_command_args(command)

def test_split_commands_fail_on_unfinished_command():
    with pytest.raises(InvalidArguments):
        split_command_args("setn")

def test_render_bottom_with_command_json():
    for command, info in commands_summary.items():
        print_formatted_text(command_syntax(command, info), style=STYLE)

@pytest.mark.parametrize(
    "raw,command,args",
    [
        ("abc 123", "abc", ["123"]),
        ("abc", "abc", []),
        ("abc foo bar", "abc", ["foo", "bar"]),
        ("abc 'foo bar'", "abc", ["foo bar"]),
        ('abc "foo bar"', "abc", ["foo bar"]),
        ('abc "foo bar" 3 hello', "abc", ["foo bar", "3", "hello"]),
        ('abc "foo \nbar"', "abc", ["foo \nbar"]),
    ],
)
def test_split_unknown_commands(raw, command, args):
    parsed_command, parsed_args = split_unknown_args(raw)
    assert command == parsed_command
    assert args == parsed_args


#!/usr/bin/env python
# -*- coding: utf-8 -*-

#############################################################################
# This script contains all test functions for some helper functions
# Eventu<lly, every helper function has to have its test function
# DIsclaimer, not all test functions have "assert" statements. The funtionality
# for some of them dependo on what is succesfully printed on the screen as an
# image.

def test_window_3masks():
    '''
    Draws mask of for a screen, so that the smaller screen is insede the big
    one like the scheme: the center (C) at the top of the small screen is in
    the center of the big one and that the big screen fills all the resolution
    in the x-axis.

                        #######################
                        #                     #
                        #                     #
                        #                     #
                        #     #####C#####     #
                        #     #         #     #
                        #     #         #     #
                        #     #         #     #
                        #######################

    Screen size of my laptop in psychopy = 1536,864
    '''

    import psychopy
    from psychopy import visual,core,logging,event, gui, monitors
    from psychopy.visual.windowwarp import Warper # perspective correction
    from helper import window_3masks

    # Windows creatinon
    _width = 400
    _height = 400
    _monitor = 'testMonitor'

    win = visual.Window(fullscr = False, monitor=_monitor,
                                size = [_width,_height], viewScale = [1,1],
                                screen = 0,
                                color=[-1,-1,-1],useFBO = True,allowGUI=False,
                                viewOri = 0.0)

    win_mask_ls  = window_3masks(win=win,_monitor = 'testMonitor')

    # Stimulus creation #
    circle_stim = visual.Circle(win=win, radius=400, units='pix',
                        fillColor=[1,1,1], edges = 128)

    # Stimulus drawing
    circle_stim.draw()
    win.flip()
    win_mask_ls[0].flip()
    win_mask_ls[1].flip()
    win_mask_ls[2].flip()

    core.wait(3)

    win.close()
    win_mask_ls[0].close()
    win_mask_ls[1].close()
    win_mask_ls[2].close()
    '''
    # Printing useful info
    print(f"The used monitor '{win.monitor.name}' has a resolution of: {win.monitor.getSizePix()} pixels")
    print(f"Main screnn located at: {win.pos} pixels")
    print(type(win.size[0]))
    print(win.size[1])
    '''

def test_window_4masks(_width=400,_height= 400,_xpos=568,_ypos=232):

    '''
    NOT FINISHED, check TODO

    Draws mask of for a screen, so that the smaller screen is insede the big
    one like the scheme: the center (C) at the top of the small screen is in
    the center of the big one and that the big screen fills all the resolution
    in the x-axis.

                        #######################
                        #                     #
                        #                     #
                        #     #####C#####     #
                        #     #         #     #
                        #     #         #     #
                        #     #         #     #
                        #     ###########     #
                        #                     #
                        #######################

    Screen size of my laptop in psychopy = 1536,864
    '''

    import psychopy
    from psychopy import visual,core,logging,event, gui, monitors
    from psychopy.visual.windowwarp import Warper # perspective correction
    from helper import max_horizontal_angle

    ########################## Windows creation ############################

    win = visual.Window(fullscr = False, monitor='testMonitor',
                                size = [_width,_height], viewScale = [1,1],
                                screen = 0, pos = [_xpos,_ypos],
                                color=[-1,-1,-1],useFBO = True,allowGUI=False,
                                viewOri = 0.0)

    scr_width = win.scrWidthCM
    scr_distance = win.scrDistCM
    maxhorang = max_horizontal_angle(scr_width, scr_distance)

    # TO DO: midified the aboved code to move the small screen up 15 deg and add a 4th mask
    '''
    win_mask_sizes =[[_width,_height/2],[_width/4,_height],[_width/4,_height]]
    win_mask_positions =[[_xpos,_ypos],[_xpos,_ypos],[_xpos+(3*(_width/4)),_ypos]]

    win_mask_0 = visual.Window(fullscr = False, monitor='testMonitor',
                                size = win_mask_sizes[0], viewScale = [1,1],
                                screen = 0, pos = win_mask_positions[0],
                                color=[-1,-1,-1],useFBO = True,allowGUI=False,
                                viewOri = 0.0)

    win_mask_1 = visual.Window(fullscr = False, monitor='testMonitor',
                                size = win_mask_sizes[1], viewScale = [1,1],
                                screen = 0, pos = win_mask_positions[1],
                                color=[-1,-1,-1],useFBO = True,allowGUI=False,
                                viewOri = 0.0)

    win_mask_2 = visual.Window(fullscr = False, monitor='testMonitor',
                                size = win_mask_sizes[2], viewScale = [1,1],
                                screen = 0, pos = win_mask_positions[2],
                                color=[-1,-1,-1],useFBO = True,allowGUI=False,
                                viewOri = 0.0)

    ########################## Stimulus creation ############################

    circle_stim = visual.Circle(win=win, radius=400, units='pix',fillColor=[1,1,1], edges = 128)

     ########################## Stimulus drawing ############################
    circle_stim.draw()
    win.flip()
    win_mask_0.flip()
    win_mask_1.flip()
    win_mask_2.flip()

    core.wait(5)

    win.close()
    win_mask_0.close()
    win_mask_1.close()
    win_mask_2.close()

    ########################### Printing useful info ########################
    print(f"The used monitor '{win.monitor.name}' has a resolution of: {win.monitor.getSizePix()} pixels")
    print(f"Main screnn located at: {win.pos} pixels")
    '''


import numpy as np
from pandas import period_range
from scipy.optimize import root
import matplotlib.pyplot as plt
from matplotlib.offsetbox import AnchoredText
from matplotlib import animation

def solicitar_requerimientos(requirements):
    datos_exportados=[]
    #Recorre la lista de requerimientos para una función
    for n in requirements:
        datum=float(input(f"Por favor introduzca el {n} "))
        datos_exportados.append(datum)
    return datos_exportados

def calcular_masa_total_minima_launch(check):
    #Datos Default
    # mu = 0.172 #Coeficiente de masa total
    # e = 0.08 #Coeficiente de masa estructural
    # delta_v = 7000.0 #Delta v requerido (m/s)
    # m_c = 500.0 #Masa de carga (kg)
    # I_sp = 311.0 #Impulso especifico (s)
    # g0 = 9.81 #Aceleracion gravitacional en la superficie terrestre (m/s^2)
    # m_p = 321.21 #Flujo masico (kg/s)

    if check ==1: #Caso de usar datos propios
        requerimientos=["Coeficiente de Masa Total (μ)",
        "Coeficiente de Masa estructural(e)",
        "ΔV",
        "Payload(Masa de carga)",
        "Impulso Específico",
        "Aceleración Gravitacional (m/s^2)",
        "Flujo Másico (kg/s)"
        ]
        datos_solicitados=solicitar_requerimientos(requerimientos)
        datos_necesarios = {
        "mu":datos_solicitados[0],
        "e":datos_solicitados[1],
        "delta_v":datos_solicitados[2],
        "m_c":datos_solicitados[3],
        "I_sp":datos_solicitados[4],
        "g0":datos_solicitados[5],
        "m_p":datos_solicitados[6]
        }

    elif check ==2: #Caso de usar los datos default
        datos_necesarios = {
        "mu": 0.172,
        "e":0.08,
        "delta_v":7000,
        "m_c":500,
        "I_sp":311,
        "g0":9.81,
        "m_p":321.21
        }
    def funcion_optimizable(x): #Definicion de la funcion que se va a optimizar.
        f = datos_necesarios["g0"]*datos_necesarios["I_sp"]*np.log(x/(datos_necesarios["e"]*(x-datos_necesarios["m_c"])+datos_necesarios["m_c"]))-(datos_necesarios["g0"]*x*(1-datos_necesarios["mu"])/datos_necesarios["m_p"])-datos_necesarios["delta_v"]
        return f

    m_0 = root(fun=funcion_optimizable,x0=10000) #Se usa un estimado inicial de m_0 = 10000 kg.
    return m_0.x[0]#Se retorna el valor de interés

def calcular_posiciones_orbit(e,a,nombre):
    """
    Estructura de datos ej:
    resultado = [punto en x, puntos en y, apogeo, perigeo, coodenadas apo, coordenadas peri]
    """
    resultado = [] #
    x=[]
    y=[]
    perio=843840923843290
    apo=0
    coordenadas_apo={"x":0,"y":0}
    coordenadas_perio={"x":0,"y":0}
    coleccion = np.linspace(0,2*3.1415927585,500)
    for cuenta,phi in enumerate(coleccion):
        r = (a*(1-e**2))/(1+e*np.cos(phi))
        #r = (1*(1-e**2))/(1+e*np.cos(phi))
        x.append(r*np.cos(phi))
        y.append(r*np.sin(phi))
        distancia= (x[cuenta]**2 + y[cuenta]**2)**(1/2)
        if distancia>apo:
            apo = distancia
            coordenadas_apo["x"]=x[cuenta]
            coordenadas_apo["y"]=y[cuenta]
        elif distancia<perio:
            perio=distancia
            coordenadas_perio["x"]=x[cuenta]
            coordenadas_perio["y"]=y[cuenta]
    resultado.append(x)
    resultado.append(y)
    resultado.append(apo)
    resultado.append(perio)
    resultado.append(coordenadas_apo)
    resultado.append(coordenadas_perio)
    resultado.append(nombre)
    return resultado

def dibujar_orbita(databas):
    fig = plt.figure()
    ax=fig.add_subplot()
    ax.set_title("Órbita solicitada")
    ax.set_xlabel("Eje X en UA")
    ax.set_ylabel("Eje y en UA")
    ax.set_facecolor('black')
    ax.plot(0,0,marker=f"$✦$",markersize=4,label="Centro de Masa",color="yellow")
    for n in databas:
        database = n
        linea,=ax.plot(database[0],database[1],'-',label = f"Trayectoria de la Orbita {database[6]}")
        ax.legend(loc='lower right',prop={'size': 6})
    point,=ax.plot(0,0,marker=8,color="red",label="Cuerpo en Orbita")
    def update_point(n, x, y,point):
        point.set_data(np.array([x[n], y[n]]))
        return point
    ani=animation.FuncAnimation(fig, update_point, fargs=(database[0], database[1], point),interval=1)
    writergif = animation.PillowWriter(fps=30)
    ani.save("orbita.gif",writer=writergif)