import ROOT as root
import numpy as np
import time
import uncertainties.unumpy as unp
from uncertainties import ufloat
from uncertainties.unumpy import nominal_values as noms
from uncertainties.unumpy import std_devs as stds
from uncertainties import correlated_values
from array import array
import os

 ########################## Data of .C file, insert data underneath ######################################
#!!!!!!!!!!!!!!!! ATTENTION !!!!!!!!!!!!!!!!!!!! At 13 mm I had to use the not zoomed version of .C, becuase the other do not include any data !!!!!!!!!!!!!!!!!!!!!!!!

qMap_Ag_C0_V0 = root.TProfile2D("qMap_Ag_C0_V0","qMap_Ag_C0 (V0)",52,0,52,80,0,80,0,0);
qMap_Ag_C0_V0.SetBinEntries(2344,19);
qMap_Ag_C0_V0.SetBinEntries(2345,14183);
qMap_Ag_C0_V0.SetBinEntries(2346,14333);
qMap_Ag_C0_V0.SetBinEntries(2347,2);
qMap_Ag_C0_V0.SetBinEntries(2398,16046);
qMap_Ag_C0_V0.SetBinEntries(2399,17020);
qMap_Ag_C0_V0.SetBinEntries(2400,14729);
qMap_Ag_C0_V0.SetBinEntries(2401,14369);
qMap_Ag_C0_V0.SetBinEntries(2451,14419);
qMap_Ag_C0_V0.SetBinEntries(2452,32449);
qMap_Ag_C0_V0.SetBinEntries(2453,32476);
qMap_Ag_C0_V0.SetBinEntries(2454,25302);
qMap_Ag_C0_V0.SetBinEntries(2455,14472);
qMap_Ag_C0_V0.SetBinEntries(2505,14523);
qMap_Ag_C0_V0.SetBinEntries(2506,44794);
qMap_Ag_C0_V0.SetBinEntries(2507,51170);
qMap_Ag_C0_V0.SetBinEntries(2508,33516);
qMap_Ag_C0_V0.SetBinEntries(2509,20073);
qMap_Ag_C0_V0.SetBinEntries(2559,14075);
qMap_Ag_C0_V0.SetBinEntries(2560,48868);
qMap_Ag_C0_V0.SetBinEntries(2561,58641);
qMap_Ag_C0_V0.SetBinEntries(2562,42812);
qMap_Ag_C0_V0.SetBinEntries(2563,17004);
qMap_Ag_C0_V0.SetBinEntries(2613,14522);
qMap_Ag_C0_V0.SetBinEntries(2614,51382);
qMap_Ag_C0_V0.SetBinEntries(2615,50565);
qMap_Ag_C0_V0.SetBinEntries(2616,44027);
qMap_Ag_C0_V0.SetBinEntries(2617,21044);
qMap_Ag_C0_V0.SetBinEntries(2618,26);
qMap_Ag_C0_V0.SetBinEntries(2668,39093);
qMap_Ag_C0_V0.SetBinEntries(2669,39409);
qMap_Ag_C0_V0.SetBinEntries(2670,30846);
qMap_Ag_C0_V0.SetBinEntries(2671,18157);
qMap_Ag_C0_V0.SetBinEntries(2723,17396);
qMap_Ag_C0_V0.SetBinEntries(2724,14476);
qMap_Ag_C0_V0.SetBinContent(2344,1727);
qMap_Ag_C0_V0.SetBinContent(2345,1541144);
qMap_Ag_C0_V0.SetBinContent(2346,1448396);
qMap_Ag_C0_V0.SetBinContent(2347,188);
qMap_Ag_C0_V0.SetBinContent(2398,1930663);
qMap_Ag_C0_V0.SetBinContent(2399,2260151);
qMap_Ag_C0_V0.SetBinContent(2400,1813706);
qMap_Ag_C0_V0.SetBinContent(2401,1458013);
qMap_Ag_C0_V0.SetBinContent(2451,1170875);
qMap_Ag_C0_V0.SetBinContent(2452,3727473);
qMap_Ag_C0_V0.SetBinContent(2453,4640774);
qMap_Ag_C0_V0.SetBinContent(2454,3218252);
qMap_Ag_C0_V0.SetBinContent(2455,1685736);
qMap_Ag_C0_V0.SetBinContent(2505,1776388);
qMap_Ag_C0_V0.SetBinContent(2506,5142724);
qMap_Ag_C0_V0.SetBinContent(2507,6040486);
qMap_Ag_C0_V0.SetBinContent(2508,4818365);
qMap_Ag_C0_V0.SetBinContent(2509,2239111);
qMap_Ag_C0_V0.SetBinContent(2559,1975378);
qMap_Ag_C0_V0.SetBinContent(2560,5531553);
qMap_Ag_C0_V0.SetBinContent(2561,8887634);
qMap_Ag_C0_V0.SetBinContent(2562,1.672008e+07);
qMap_Ag_C0_V0.SetBinContent(2563,2151277);
qMap_Ag_C0_V0.SetBinContent(2613,1603932);
qMap_Ag_C0_V0.SetBinContent(2614,5852932);
qMap_Ag_C0_V0.SetBinContent(2615,5665718);
qMap_Ag_C0_V0.SetBinContent(2616,5689668);
qMap_Ag_C0_V0.SetBinContent(2617,2508568);
qMap_Ag_C0_V0.SetBinContent(2618,2136);
qMap_Ag_C0_V0.SetBinContent(2668,4186165);
qMap_Ag_C0_V0.SetBinContent(2669,4672891);
qMap_Ag_C0_V0.SetBinContent(2670,3811861);
qMap_Ag_C0_V0.SetBinContent(2671,2255823);
qMap_Ag_C0_V0.SetBinContent(2723,2049516);
qMap_Ag_C0_V0.SetBinContent(2724,1700067);
qMap_Ag_C0_V0.SetBinError(2344,402.8734);
qMap_Ag_C0_V0.SetBinError(2345,12979.8);
qMap_Ag_C0_V0.SetBinError(2346,12166.45);
qMap_Ag_C0_V0.SetBinError(2347,134.4024);
qMap_Ag_C0_V0.SetBinError(2398,93836.64);
qMap_Ag_C0_V0.SetBinError(2399,161228.9);
qMap_Ag_C0_V0.SetBinError(2400,114305.9);
qMap_Ag_C0_V0.SetBinError(2401,12252.54);
qMap_Ag_C0_V0.SetBinError(2451,9901.228);
qMap_Ag_C0_V0.SetBinError(2452,147939.1);
qMap_Ag_C0_V0.SetBinError(2453,293709.8);
qMap_Ag_C0_V0.SetBinError(2454,218026.8);
qMap_Ag_C0_V0.SetBinError(2455,14092.93);
qMap_Ag_C0_V0.SetBinError(2505,14969.22);
qMap_Ag_C0_V0.SetBinError(2506,26160.92);
qMap_Ag_C0_V0.SetBinError(2507,198427.4);
qMap_Ag_C0_V0.SetBinError(2508,246180.4);
qMap_Ag_C0_V0.SetBinError(2509,16509.44);
qMap_Ag_C0_V0.SetBinError(2559,16717.33);
qMap_Ag_C0_V0.SetBinError(2560,27341.63);
qMap_Ag_C0_V0.SetBinError(2561,366223.7);
qMap_Ag_C0_V0.SetBinError(2562,886649);
qMap_Ag_C0_V0.SetBinError(2563,16870.58);
qMap_Ag_C0_V0.SetBinError(2613,13409.44);
qMap_Ag_C0_V0.SetBinError(2614,29248.24);
qMap_Ag_C0_V0.SetBinError(2615,96545.96);
qMap_Ag_C0_V0.SetBinError(2616,208872.5);
qMap_Ag_C0_V0.SetBinError(2617,17967.6);
qMap_Ag_C0_V0.SetBinError(2618,423.1099);
qMap_Ag_C0_V0.SetBinError(2668,69887.67);
qMap_Ag_C0_V0.SetBinError(2669,148421.3);
qMap_Ag_C0_V0.SetBinError(2670,161830.9);
qMap_Ag_C0_V0.SetBinError(2671,17107.6);
qMap_Ag_C0_V0.SetBinError(2723,16283.99);
qMap_Ag_C0_V0.SetBinError(2724,14283.94);
qMap_Ag_C0_V0.SetMinimum(0);
qMap_Ag_C0_V0.SetEntries(896238);
qMap_Ag_C0_V0.SetStats(0);
qMap_Ag_C0_V0.SetContour(20);
qMap_Ag_C0_V0.SetContourLevel(0,0);
qMap_Ag_C0_V0.SetContourLevel(1,19.52733);
qMap_Ag_C0_V0.SetContourLevel(2,39.05465);
qMap_Ag_C0_V0.SetContourLevel(3,58.58198);
qMap_Ag_C0_V0.SetContourLevel(4,78.10931);
qMap_Ag_C0_V0.SetContourLevel(5,97.63663);
qMap_Ag_C0_V0.SetContourLevel(6,117.164);
qMap_Ag_C0_V0.SetContourLevel(7,136.6913);
qMap_Ag_C0_V0.SetContourLevel(8,156.2186);
qMap_Ag_C0_V0.SetContourLevel(9,175.7459);
qMap_Ag_C0_V0.SetContourLevel(10,195.2733);
qMap_Ag_C0_V0.SetContourLevel(11,214.8006);
qMap_Ag_C0_V0.SetContourLevel(12,234.3279);
qMap_Ag_C0_V0.SetContourLevel(13,253.8552);
qMap_Ag_C0_V0.SetContourLevel(14,273.3826);
qMap_Ag_C0_V0.SetContourLevel(15,292.9099);
qMap_Ag_C0_V0.SetContourLevel(16,312.4372);
qMap_Ag_C0_V0.SetContourLevel(17,331.9645);
qMap_Ag_C0_V0.SetContourLevel(18,351.4919);
qMap_Ag_C0_V0.SetContourLevel(19,371.0192);

ci = root.TColor.GetColor("#000099");
qMap_Ag_C0_V0.SetLineColor(ci);
qMap_Ag_C0_V0.GetXaxis().SetTitle("col");
qMap_Ag_C0_V0.GetXaxis().SetRange(1,52);
qMap_Ag_C0_V0.GetXaxis().SetNdivisions(508);
qMap_Ag_C0_V0.GetXaxis().SetLabelFont(42);
qMap_Ag_C0_V0.GetXaxis().SetLabelSize(0.05);
qMap_Ag_C0_V0.GetXaxis().SetTitleSize(0.05);
qMap_Ag_C0_V0.GetXaxis().SetTitleOffset(1.1);
qMap_Ag_C0_V0.GetXaxis().SetTitleFont(42);
qMap_Ag_C0_V0.GetYaxis().SetTitle("row");
qMap_Ag_C0_V0.GetYaxis().SetRange(1,80);
qMap_Ag_C0_V0.GetYaxis().SetLabelFont(42);
qMap_Ag_C0_V0.GetYaxis().SetLabelSize(0.05);
qMap_Ag_C0_V0.GetYaxis().SetTitleSize(0.05);
qMap_Ag_C0_V0.GetYaxis().SetTitleOffset(1.1);
qMap_Ag_C0_V0.GetYaxis().SetTitleFont(42);
qMap_Ag_C0_V0.GetZaxis().SetLabelFont(42);
qMap_Ag_C0_V0.GetZaxis().SetLabelSize(0.035);
qMap_Ag_C0_V0.GetZaxis().SetTitleSize(0.035);
qMap_Ag_C0_V0.GetZaxis().SetTitleFont(42);

root.gStyle.SetOptTitle(0)
##################################################### Insert Data above line #############################

############################### Get name of folder #################################

name_of_folder = os.path.basename( os.getcwd() )

############################### Save Data in list #######################################

mean_value_col_list = []
mean_error_col_list = []
x_value = []
y_value = []

##############################################################################################################################

################################### Getting the mean hit value of all columns near the laserspot #############################

###############################################################################################################################

################################## Set sum area ###############################

xmin = 18
xmax = 27

ymin = 40
ymax = 52

####################################  calculating mean of each coloum ################################

for i in range(xmin,xmax): # going thru all col
    content = []
    error = []

    x_value.append(i)
    y_value.append(0)

    for j in range(ymin,ymax): # going thru all rows
        content.append( qMap_Ag_C0_V0.GetBinContent(i,j))
        error.append( qMap_Ag_C0_V0.GetBinError(i,j)) # Is this the real error

    content_bin = unp.uarray( content, error)
    mean_content_col = content_bin.mean() # mean value of each bin in the col

    # Saving values in lists
    mean_value_col_list.append( noms(mean_content_col))
    mean_error_col_list.append( stds(mean_content_col))

########################### Create errorbar plot #####################################

errorbar_plot_col = root.TGraphErrors( len(x_value), array( 'f', x_value), array( 'f', mean_value_col_list), array( 'f', y_value), array( 'f', mean_error_col_list) )

############################## Set axis label of errobar plot ##################################

errorbar_plot_col.GetXaxis().SetTitle("Col")
errorbar_plot_col.GetYaxis().SetTitle("Mean Hit / Vcal")

####################### create Canvas ##########################################

c1 = root.TCanvas("c1", "c1", 1980, 1080)
errorbar_plot_col.Fit('gaus')
errorbar_plot_col.Draw("ALP")

name_params = [ "amplitude/[MeanVcal]", "mean/[Col]", "sigma/[Col]"]

############################### Create legend ####################################
gaus_fit_col = errorbar_plot_col.GetFunction('gaus')
legend = root.TLegend(0.75,0.75,0.98,0.98)
legend.AddEntry(errorbar_plot_col,"Coloumn mean hit value","lep")
legend.AddEntry( gaus_fit_col,"Gaussian Fit","l")
legend.Draw()

############################### Save parameter and plot ###########################################

with open( f'../fit_params/{name_of_folder}_fit_parameters_col_xaxis.txt', 'w') as file:
    for i in range(0,3):
        file.write( name_params[i] + ' ' + str( gaus_fit_col.GetParameter(i) ) + ' ' + str(gaus_fit_col.GetParError(i)) + '\n')

with open( f'../sigma_col_xaxis.txt', 'a') as file:
        file.write( name_params[i] + '_' + name_of_folder + ' ' + str( gaus_fit_col.GetParameter(2) ) + ' ' + str(gaus_fit_col.GetParError(2)) + '\n')

c1.SaveAs(f'../plots/{name_of_folder}_erorbar_plot_col.pdf')

##############################################################################################################################

################################### Getting the mean hit value of all rows near the laserspot #############################

###############################################################################################################################

############################Reset lists###########################################
mean_value_col_list = []
mean_error_col_list = []
x_value = []
y_value = []

#################################### calculating mean of each row #####################################

for i in range(ymin,ymax): # going thru all rows
    content = []
    error = []

    x_value.append(i)
    y_value.append(0)

    for j in range(xmin,xmax): # going thru all col
        content.append( qMap_Ag_C0_V0.GetBinContent(j,i))
        error.append( qMap_Ag_C0_V0.GetBinError(j,i)) # Is this the real error

    content_bin = unp.uarray( content, error)
    mean_content_col = content_bin.mean() # mean value of each bin in the col

    # Saving values in lists
    mean_value_col_list.append( noms(mean_content_col))
    mean_error_col_list.append( stds(mean_content_col))

############################# Create new errorbar plot ####################################
errorbar_plot_rows = root.TGraphErrors( len(x_value), array( 'f', x_value), array( 'f', mean_value_col_list), array( 'f', y_value), array( 'f', mean_error_col_list) )

############################### create Canvas ########################################
c2 = root.TCanvas("c2", "c2", 1980, 1080);

############################## Set axis label of errobar plot ##################################

errorbar_plot_rows.GetXaxis().SetTitle("Row")
errorbar_plot_rows.GetYaxis().SetTitle("Mean Hit / Vcal")

############################### Plot fucntion and fit #############################################
errorbar_plot_rows.Fit('gaus')
errorbar_plot_rows.Draw("ALP")

##################################### create legend ################################################
gaus_fit_row = errorbar_plot_rows.GetFunction('gaus')
legend = root.TLegend(0.75,0.75,0.98,0.98)
legend.AddEntry(errorbar_plot_rows,"Mean Mean Vcal of each row","lep")
legend.AddEntry( gaus_fit_row,"Gaussian Fit","l")
legend.Draw()

########################################### saveplot and fit params ########################################
with open( f'../fit_params/{name_of_folder}_fit_parameters_row_yaxis.txt', 'w') as file:
    for i in range(0,3):
        file.write( name_params[i] + ' ' + str( gaus_fit_row.GetParameter(i) ) + ' ' + str(gaus_fit_row.GetParError(i)) + '\n')

with open( f'../sigma_row_yaxis.txt', 'a') as file:
        file.write( name_params[i] +'_' + name_of_folder + ' ' + str( gaus_fit_row.GetParameter(2) ) + ' ' + str(gaus_fit_row.GetParError(2)) + '\n')

c2.SaveAs(f'../plots/{name_of_folder}_erorbar_plot_row.pdf')


"""This module contains data structures that are robocup specific, ie. Robot, Ball,
WorldState"""

from enum import Enum
from typing import List, Optional

import numpy as np
import warnings

RobotId = int

class Robot:
    """State of a robot. Pose: [x, y, theta]. Twist: [dx, dy, dtheta]. Properties are
    to enforce that instances of this class should not be mutated."""

    __slots__ = [
        "__id", "__is_ours", "__pose", "__twist",
        "__visible", "__has_ball_sense", "__kicker_charged",
        "__kicker_healthy", "__lethal_fault"
    ]

    __id: RobotId
    __is_ours: bool
    __pose: np.ndarray
    __twist: np.ndarray
    __visible: bool
    __has_ball_sense: bool
    __kicker_charged: bool
    __kicker_healthy: bool
    __lethal_fault: bool

    def __init__(self, robot_id: RobotId, is_ours: bool, pose: np.ndarray,
                 twist: np.ndarray, visible: bool,
                 has_ball_sense: bool, kicker_charged: bool,
                 kicker_healthy: bool, lethal_fault: bool):
        """
        :param robot_id: Shell id of the robot.
        :param is_ours: Whether the robot is one of our robots
        :param pose: Pose of the Robot. [x, y, theta].
        :param twist: Twist of the robot. [dx, dy, dtheta].
        :param visible: Whether the robot is being seen by the global vision system
        :param has_ball_sense: Whether the robots ball sensor is functional
        :param kicker_charged: Whether the robots kicker capacitors are charged
        :param kicker_healthy: Whether the robots kicker is "healthy"
        :param lethal_fault: Whether the robot has experienced a fault that will prevent it from functioning
        """
        self.__id = robot_id
        self.__is_ours = is_ours
        self.__pose = pose
        self.__twist = twist
        self.__visible = visible
        self.__has_ball_sense = has_ball_sense
        self.__kicker_charged = kicker_charged
        self.__kicker_healthy = kicker_healthy
        self.__lethal_fault = lethal_fault

    def __repr__(self) -> str:
        return "Robot(id:{}, is_ours:{}, pose:{}, twist:{}, visible:{})".format(
            self.__id, self.__is_ours, self.__pose, self.__twist,
            self.__visible)

    def __eq__(self, other) -> bool:
        if isinstance(other, Robot):
            if (self.is_ours == other.is_ours and self.id == other.id):
                return True
        return False

    @property
    def id(self) -> RobotId:
        """
        :return: Id of the robot.
        """
        return self.__id

    @property
    def is_ours(self) -> bool:
        """
        :return: True if the robot is our robot.
        """
        return self.__is_ours

    @property
    def pose(self) -> np.ndarray:
        """
        :return: Pose of the robot. [x, y, theta].
        """
        if (not self.visible):
            #I could see removing this as it's a thing that may happen fairly often
            warnings.warn(
                "Attempting to retrieve robot pose from non-visible robot",
                RuntimeWarning)

        return self.__pose

    @property
    def twist(self) -> np.ndarray:
        """
        :return: Twist of the robot. [dx, dy, dtheta].
        """
        if (not self.visible):
            #I could see removing this as it's a thing that may happen fairly often
            warnings.warn(
                "Attempting to retrieve robot pose from non-visible robot",
                RuntimeWarning)

        return self.__twist

    @property
    def visible(self) -> bool:
        """
        :return: True if the robot is visible
        """
        return self.__visible

    @property
    def has_ball_sense(self) -> bool:
        """
        :return: True if this robot has functioning ball sensors
        """
        if (not self.is_ours):
            warnings.warn(
                "Attempting to retrieve ball sense status from an opposing robot",
                RuntimeWarning)

        return self.__has_ball_sense

    @property
    def kicker_charged(self) -> bool:
        """
        :return: True if the kicker capacitors are charged
        """
        if (not self.is_ours):
            warnings.warn(
                "Attempting to retrieve kicker charge status from an opposing robot",
                RuntimeWarning)
            return False

        return self.__kicker_charged

    @property
    def kicker_healthy(self) -> bool:
        """
        :return: True if the kicker is healthy
        """
        if (not self.is_ours):
            warnings.warn(
                "Attempting to retrieve kicker health status from an opposing robot",
                RuntimeWarning)
            return False

        return self.__kicker_healthy

    @property
    def lethal_fault(self) -> bool:
        """
        :return: True if the robot has encounted a fault that will prevent further play, such as an FPGA or motor fault.
        """
        if (not self.is_ours):
            warnings.warn(
                "Attempting to retrieve lethal fault information from an opposing robot",
                RuntimeWarning)

        return self.__lethal_fault

class Ball:
    """State of the ball. Properties are used to enforce that instances of this class
    should not be mutated."""

    __slots__ = ["__pos", "__vel", "__visible"]

    __pos: np.ndarray
    __vel: np.ndarray
    __visible: bool

    def __init__(self, pos: np.ndarray, vel: np.ndarray, visible: bool):
        """
        :param pos: ball position np.ndarray([x, y])
        :param vel: ball velocity np.ndarray([x, y])
        :param visible: True if the ball is visible
        """
        self.__pos = pos
        self.__vel = vel
        self.__visible = visible

    def __repr__(self) -> str:
        return "Ball(pos:{}, vel:{}, visible:{})".format(
            self.__pos, self.__vel, self.__visible)

    @property
    def pos(self) -> np.ndarray:
        """
        :return: Position of the ball. [x, y].
        """
        if (not self.visible):
            warnings.warn("Retrieved the position of a non-visible ball",
                          RuntimeWarning)

        return self.__pos

    @property
    def vel(self) -> np.ndarray:
        """
        :return: Velocity of the ball. [dx, dy].
        """
        if (not self.visible):
            warnings.warn("Retrieved the velocity of a non-visible ball",
                          RuntimeWarning)

        return self.__vel

    @property
    def visible(self) -> bool:
        """
        :return: True if the ball can currently be seen by the global vision system.

        Note that with filtering or additional sensors like whiskers we may still know where the ball is located
        """
        return self.__visible

class GamePeriod(Enum):
    """Game period."""

    FIRST_HALF = 0
    HALF_TIME = 1
    SECOND_HALF = 2
    OVERTIME1 = 3
    OVERTIME2 = 4
    PENALTY_SHOOTOUT = 5

class GameState(Enum):
    """State of the game."""

    HALT = 0  # Robots must not move.
    STOP = 1  # Robots must stay 500mm away from the ball.
    SETUP = 2  # Robots not on starting team msut stay 500mm away from ball.
    READY = 3  # A robot on the starting team may kick the ball.
    PLAYING = 4  # Normal play.
    PENALTY_PLAYING = 5  # All robots except the striker and the goalie must stay on the opposite side of the field.

class GameRestart(Enum):
    """What kind of restart."""

    NONE = 0
    KICKOFF = 1
    DIRECT = 2
    INDIRECT = 3
    PENALTY = 4
    PLACEMENT = 5

class Field:
    """Information about the field."""

    __slots__ = [
        "__length_m", "__width_m", "__border_m", "__line_width_m",
        "__goal_width_m", "__goal_depth_m", "__goal_height_m",
        "__def_area_short_dist_m", "__def_area_long_dist_m",
        "__center_radius_m", "__center_diameter_m", "__goal_flat_m",
        "__floor_length_m", "__floor_width_m"
    ]

    __length_m: float
    __width_m: float
    __border_m: float
    __line_width_m: float
    __goal_width_m: float
    __goal_depth_m: float
    __goal_height_m: float
    __def_area_short_dist_m: float
    __def_area_long_dist_m: float
    __center_radius_m: float
    __center_diameter_m: float
    __goal_flat_m: float
    __floor_length_m: float
    __floor_width_m: float

    def __init__(self, length_m: float, width_m: float, border_m: float,
                 line_width_m: float, goal_width_m: float, goal_depth_m: float,
                 goal_height_m: float, def_area_short_dist_m: float,
                 def_area_long_dist_m: float, center_radius_m: float,
                 center_diameter_m: float, goal_flat_m: float,
                 floor_length_m: float, floor_width_m: float):
        self.__length_m = length_m
        self.__width_m = width_m
        self.__border_m = border_m
        self.__line_width_m = line_width_m
        self.__goal_width_m = goal_width_m
        self.__goal_depth_m = goal_depth_m
        self.__goal_height_m = goal_height_m
        self.__def_area_short_dist_m = def_area_short_dist_m
        self.__def_area_long_dist_m = def_area_long_dist_m
        self.__center_radius_m = center_radius_m
        self.__center_diameter_m = center_diameter_m
        self.__goal_flat_m = goal_flat_m
        self.__floor_length_m = floor_length_m
        self.__floor_width_m = floor_width_m

    @property
    def our_goal_loc(self) -> np.ndarray:
        """
        Conveniance function for getting our goal location
        :return: the location of our goal - its always (0,0)
        """
        return np.array([0.0, 0.0])

    @property
    def center_field_loc(self) -> np.ndarray:
        """
        Conveniance function for getting the center field location
        :return: the location of the center of the field
        """
        return np.array([0.0, self.length_m / 2])

    @property
    def their_goal_loc(self) -> np.ndarray:
        """
        Conveniance function for getting the opponents field location
        :return: the location of the opponents goal
        """
        return np.array([0.0, self.length_m])

    @property
    def floor_width_m(self) -> float:
        """
        :return: check on this one
        """
        return self.__floor_width_m

    @property
    def floor_length_m(self) -> float:
        """
        :return: check on this one
        """
        return self.__floor_length_m

    @property
    def goal_flat_m(self) -> float:
        """
        :return: check on this one
        """
        return self.__goal_flat_m

    @property
    def center_diameter_m(self) -> float:
        """
        :return: returns the diameter of the center of the field
        """
        return self.__center_diameter_m

    @property
    def center_radius_m(self) -> float:
        """
        :return: returns the radius of the center of the field
        """
        return self.__center_radius_m

    @property
    def def_area_long_dist_m(self) -> float:
        """
        :return: double check on this one
        """
        return self.__def_area_long_dist_m

    @property
    def def_area_short_dist_m(self) -> float:
        """
        :return: double check on this one
        """
        return self.__def_area_short_dist_m

    @property
    def border_m(self) -> float:
        """
        :return: The size of the border of the field
        """
        return self.__border_m

    @property
    def line_width_m(self) -> float:
        """
        :return: The width of the lines of the field
        """
        return self.__line_width_m

    @property
    def length_m(self) -> float:
        """
        :return: The length of the field in meters
        """
        return self.__length_m

    @property
    def width_m(self) -> float:
        """
        :return: the width of the field in meters
        """
        return self.__width_m

    @property
    def goal_width_m(self) -> float:
        """
        :return: the width of the goals in meters
        """
        return self.__goal_width_m

    @property
    def goal_depth_m(self) -> float:
        """
        :return: the depth of the goals in meters
        """
        return self.__goal_depth_m

    @property
    def goal_height_m(self) -> float:
        """
        :return: the height of the goals in meters
        """
        return self.__goal_height_m

class GameInfo:
    """State of the soccer game. Corresponds to a combination of the C++-side PlayState and MatchState"""

    __slots__ = [
        "__period",
        "__state",
        "__restart",
        "__our_restart",
        "__ball_placement",
    ]

    __period: GamePeriod
    __state: GameState
    __restart: GameRestart
    __our_restart: bool
    __ball_placement: np.array

    def __init__(self, period: GamePeriod, state: GameState,
                 restart: GameRestart, our_restart: bool,
                 ball_placement: np.array):
        self.__period = period
        self.__state = state
        self.__restart = restart
        self.__our_restart = our_restart
        self.__ball_placement = ball_placement

    @property
    def period(self) -> GamePeriod:
        """
        :return: The game period
        """
        return self.__period

    @property
    def state(self) -> GameState:
        """
        :return: The game state
        """
        return self.__state

    @property
    def restart(self) -> GameRestart:
        """
        :return: The game restart state
        """
        return self.__restart

    @property
    def our_restart(self) -> bool:
        """
        :return: True if it is our restart
        """
        if (not self.is_restart()):
            warnings.warn(
                "Retrieved if it is our restart when it is not a restart at all",
                RuntimeWarning)
            return False  #Is returning this dangerous?

        return self.__our_restart

    @property
    def their_restart(self) -> bool:
        """
        :return: True if it is their restart
        """
        if (not self.is_restart()):
            warnings.warn(
                "Retrieved if it is our restart when it is not a restart at all",
                RuntimeWarning)
            return False  #Is returning this dangerous?

        return not self.__our_restart

    def is_stopped(self) -> bool:
        """
        :return: True if play is stopped.
        """
        return self.state == GameState.STOP

    def is_ready(self) -> bool:
        """
        :return: True if the field is waiting on a team to kick the ball in a restart.
        """
        return self.state == GameState.READY

    def is_setup(self) -> bool:
        """
        :return: True if the field is setting up for a penalty kick or kickoff.
        """
        return self.state == GameState.SETUP

    def is_restart(self) -> bool:
        """
        :return: True if there is a restart.
        """
        return self.restart != GameRestart.NONE

    def is_kickoff(self) -> bool:
        """
        :return: True if the restart is a kickoff.
        """
        return self.restart == GameRestart.KICKOFF

    def is_penalty(self) -> bool:
        """
        :return: True if the restart is a penalty.
        """
        return self.restart == GameRestart.PENALTY

    def is_direct(self) -> bool:
        """
        :return: True if the restart is a direct kick.
        """
        return self.restart == GameRestart.DIRECT

    def is_indirect(self) -> bool:
        """
        :return: True if the restart is an indirect kick.
        """
        return self.restart == GameRestart.INDIRECT

    def is_free_placement(self) -> bool:
        """
        :return: True if the restart is free placement.
        """
        return self.restart == GameRestart.PLACEMENT

    def ball_placement(self) -> Optional[np.ndarray]:
        """
        :return: True if the restart is free placement.
        """
        return self.__ball_placement if self.is_free_placement() else None

class WorldState:
    """Current state of the world."""

    __slots__ = [
        "__our_robots", "__their_robots", "__ball", "__game_info", "__field",
        "__goalie_id"
    ]

    __our_robots: List[Robot]
    __their_robots: List[Robot]
    __ball: Ball
    __game_info: GameInfo
    __field: Field
    __goalie_id: int

    def __init__(self, our_robots: List[Robot], their_robots: List[Robot],
                 ball: Ball, game_info: GameInfo, field: Field,
                 goalie_id: int):
        self.__our_robots = our_robots
        self.__their_robots = their_robots
        self.__ball = ball
        self.__game_info = game_info
        self.__field = field
        self.__goalie_id = goalie_id

    @property
    def robots(self) -> List[Robot]:
        """
        :return: A list of all robots (created by merging our_robots with their_robots)
        """
        return self.our_robots + self.their_robots

    @property
    def our_robots(self) -> List[Robot]:
        """
        :return: A list of our robots
        """
        return self.__our_robots

    @property
    def their_robots(self) -> List[Robot]:
        """
        :return: A list of their robots
        """
        return self.__their_robots

    @property
    def ball(self) -> Ball:
        """
        :return: The ball
        """
        return self.__ball

    @property
    def game_info(self) -> GameInfo:
        """
        :return: The GameInfo object
        """
        return self.__game_info

    @property
    def field(self) -> Field:
        """
        :return: The Field object
        """
        return self.__field

    @property
    def goalie_id(self) -> int:
        """
        :return: The Field object
        """
        return self.__goalie_id

    #def get_visible_robots(self) -> List[Robot]:
    #    pass

    #def get_our_visible_robots(self) -> List[Robot]:
    #    pass

    #def get_their_visible_robots(self) -> List[Robot]:
    #    pass


#!/usr/bin/python
# -*- coding: utf-8 -*-
# http://elecrow.com/
# original author: https://github.com/glennklockwood/raspberrypi

import random
import time
import spi
import numpy as np

### Initialize an SPI connection using BCM-mode pins 21, 20, and 16
max7219 = spi.SPI(clk=21, cs=20, mosi=16, miso=None)

### Zero out all registers
for cmd in range(16):
    packet = cmd << 8
    max7219.put(packet,12)

### Enable all columns via the scan limit register
max7219.put(int("101100000111",2),12)

### Disable shutdown mode
max7219.put(int("110000000001",2),12)

### Define the values for each column that gives us a heart shape
heart_shape = np.array( [
    [ 0, 0, 0, 0, 0, 0, 0, 0 ],
    [ 0, 1, 1, 0, 0, 1, 1, 0 ],
    [ 1, 1, 1, 1, 1, 1, 1, 1 ],
    [ 1, 1, 1, 1, 1, 1, 1, 1 ],
    [ 0, 1, 1, 1, 1, 1, 1, 0 ],
    [ 0, 0, 1, 1, 1, 1, 0, 0 ],
    [ 0, 0, 0, 1, 1, 0, 0, 0 ],
    [ 0, 0, 0, 0, 0, 0, 0, 0 ],
    ] )

### Set each column according to our heart shape matrix
for col in range(8):
    cmd = (col+1) << 8
    values = 0
    ### Without flipud, the heart would be upside-down since the LED matrix
    ### defines 0,0 as the bottom left corner while numpy puts it at the top
    ### left.
    for i in np.flipud(heart_shape[:,col]):
        values <<= 1
        if i != 0:
            values |= 1
    max7219.put(cmd|values, 12)

### MAX7219 supports intensity up to 31, but that's REALLY bright
MAX_INTENSITY = 6
input_value = 0.0
while True:
    ### This math just makes the intensity pulse like a sine wave
    input_value += np.pi / MAX_INTENSITY
    value = int((np.sin(input_value)+1.0) * MAX_INTENSITY)

    ### 1010 is the LED intensity register
    cmd = int("1010",2) << 8
    max7219.put(cmd|value, 12)
    time.sleep(0.1)


import os
import argparse
import glob
import cv2
from utils import face_utils
from utils import cv_utils
import face_recognition
from PIL import Image
import torchvision.transforms as transforms
import torch
import pickle
import numpy as np
from models.models import ModelsFactory
from options.test_options import TestOptions

class MorphFacesInTheWild:
    def __init__(self, opt):
        self._opt = opt
        self._model = ModelsFactory.get_by_name(self._opt.model, self._opt)
        self._model.set_eval()
        self._transform = transforms.Compose([transforms.ToTensor(),
                                              transforms.Normalize(mean=[0.5, 0.5, 0.5],
                                                                   std=[0.5, 0.5, 0.5])
                                              ])

    def morph_file(self, img_path, expresion):
        img = cv_utils.read_cv2_img(img_path)
        morphed_img = self._img_morph(img, expresion)
        output_name = '%s_out.png' % os.path.basename(img_path)
        self._save_img(morphed_img, output_name)

    def _img_morph(self, img, expresion):
        bbs = face_recognition.face_locations(img)
        if len(bbs) > 0:
            y, right, bottom, x = bbs[0]
            bb = x, y, (right - x), (bottom - y)
            face = face_utils.crop_face_with_bb(img, bb)
            face = face_utils.resize_face(face)
        else:
            face = face_utils.resize_face(img)

        morphed_face = self._morph_face(face, expresion)

        return morphed_face

    def _morph_face(self, face, expresion):
        face = torch.unsqueeze(self._transform(Image.fromarray(face)), 0)
        expresion = torch.unsqueeze(torch.from_numpy(expresion/5.0), 0)
        test_batch = {'real_img': face, 'real_cond': expresion, 'desired_cond': expresion, 'sample_id': torch.FloatTensor(), 'real_img_path': []}
        self._model.set_input(test_batch)
        imgs, _ = self._model.forward(keep_data_for_visuals=False, return_estimates=True)
        return imgs['concat']

    def _save_img(self, img, filename):
        filepath = os.path.join(self._opt.output_dir, filename)
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        cv2.imwrite(filepath, img)

def main():
    opt = TestOptions().parse()
    if not os.path.isdir(opt.output_dir):
        os.makedirs(opt.output_dir)

    morph = MorphFacesInTheWild(opt)

    image_path = opt.input_path
    expression = np.random.uniform(0, 1, opt.cond_nc)
    morph.morph_file(image_path, expression)

if __name__ == '__main__':
    main()


 #!/usr/bin/env python

import rospy
'''
EKF slam implementation for MUR driverless

Author: Jack McRobbie
'''
import numpy as np
import rospy

class EkfSlam:

    def __init__(self,sensors):
        self.data = sensors

    def doSlam(self):
        while(True):
            mean = np.mean(self.data.cameraCloud)
            rospy.loginfo(mean)


#https://www.influxdata.com/blog/getting-started-python-influxdb/

from time import sleep
from datetime import *
from influxdb import InfluxDBClient
import numpy as np
import random

#Conectarse a la case de datos
client = InfluxDBClient('localhost', 8086, 'root', 'root', 'Health')
#Crear base de datos
client.create_database('Health')

times = datetime.utcnow();
for i in range(7000000):
    nivelAzucar = round(random.uniform(10,200),2)
    ritmoCardiaco = round(random.uniform(50,210),2)
    calorias = round(random.uniform(0,2000),2)
    times = times + timedelta(minutes = 1)
    sensores = [{'measurement': 'sensor',
                    'tags':{'tipo':'medidorInsulina'},
                    'time':times.isoformat(),
                    'fields': {'nivelAzucar':nivelAzucar}},
                    {'measurement': 'sensor',
                    'tags':{'tipo':'reloj'},
                    'time':times.isoformat(),
                    'fields': {'ritmoCardiaco':ritmoCardiaco}},
                    {'measurement': 'sensor',
                    'tags':{'tipo':'reloj'},
                    'time':times.isoformat(),
                    'fields': {'calorias':ritmoCardiaco}}]
    client.write_points(sensores)

for i in range(7000000):
    caloriasIngeridas = round(random.uniform(10,1000),2)
    caloriasQuemadas = round(random.uniform(0,2000),2)
    times = times + timedelta(minutes = 1)
    calorias = [{'measurement': 'porcentajeCalorico',
                    'tags':{'tipo':'reloj'},
                    'time':times.isoformat(),
                    'fields': {'caloriasIngeridas':caloriasIngeridas}},
                    {'measurement': 'porcentajeCalorico',
                    'tags':{'tipo':'reloj'},
                    'time':times.isoformat(),
                    'fields': {'caloriasQuemadas':caloriasQuemadas}}]
    client.write_points(calorias)

for i in range(7000000):
    pasosTotales = round(random.uniform(0,1000),2)
    times = times + timedelta(minutes = 1)
    pasos = [{'measurement': 'pasos',
                    'tags':{'tipo':'reloj'},
                    'time':times.isoformat(),
                    'fields': {'pasosTotales':pasosTotales}}]
    client.write_points(pasos)


# -*- coding: utf-8 -*-
# Copyright (c) 2018 João Martins
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import numpy as np

class SparseGenerator:
    def __init__(self):
        self._current_batch = 0
        self._data = []
        self.grid_indexes = []
        self.charges = []
        self.masses = []
        self.aa_one_hot = []
        self._hold = True

    def load_data(self, protein_feature, grid_feature):
        for index, prot in enumerate(protein_feature):
            loader = np.load(prot)
            data = dict(zip((k for k in loader), (loader[k] for k in loader)))
            grid_loader = np.load(grid_feature[index])
            grid_data = dict(zip((k for k in grid_loader), (grid_loader[k] for k in grid_loader)))
            for dex, value in enumerate(grid_data['indices']):
                aa_one_hot = data['aa_one_hot'][dex]
                indexer = grid_data['selector'][dex][grid_data['selector'][dex] >= 0]
                grid_indexes = grid_data['indices'][dex][:len(indexer)]
                charges = data['charges'][indexer]
                masses = data['masses'][indexer]
                self.grid_indexes.append(grid_indexes)
                self.charges.append(charges)
                self.masses.append(masses)
                self.aa_one_hot.append(aa_one_hot)

    def next(self, batch_size):
        self._current_batch += 1
        if self._current_batch * batch_size >= len(self.grid_indexes):
            self._hold = False
            start, end = ((self._current_batch - 1) * batch_size, len(self.grid_indexes)) # FIXME: Change to wrap around
            start = -batch_size
            indices = [np.insert(x, 0, ind, axis=1) for ind, x in enumerate(self.grid_indexes[start:end])]
            indices = np.array([np.concatenate((np.insert(x, 4, 0, axis=1),
                                                np.insert(x, 4, 1, axis=1))) for ind, x in enumerate(indices)])
            values = np.array([np.concatenate((self.masses[start:end][index], self.charges[start:end][index]))
                               for index, val in enumerate(self.masses[start:end])])
            shape = np.array([end-start, 24, 76, 151, 2])
            hots = np.array(self.aa_one_hot[start:end])
            return np.concatenate(indices), np.concatenate(values).ravel(), shape, hots
        else:
            start, end = ((self._current_batch-1) * batch_size, self._current_batch * batch_size)
            indices = [np.insert(x, 0, ind, axis=1) for ind, x in enumerate(self.grid_indexes[start:end])]
            indices = np.array([np.concatenate((np.insert(x, 4, 0, axis=1),
                                np.insert(x, 4, 1, axis=1))) for ind, x in enumerate(indices)])
            values = np.array([np.concatenate((self.masses[start:end][index], self.charges[start:end][index]))
                               for index, val in enumerate(self.masses[start:end])])
            shape = np.array([end-start, 24, 76, 151, 2])
            hots = np.array(self.aa_one_hot[start:end])
            return np.concatenate(indices), np.concatenate(values).ravel(), shape, hots

    def __repr__(self):
        return str(len(self.grid_indexes))

    def hold(self):
        return self._hold

    def restart(self):
        self._current_batch = 0
        self._hold = True

    def grid_shape(self):
        return [24, 76, 151, 2]  # FIXME: ADD CALL TO CLASS INSTANTIATION

    def infer(self, residue_index):
        zeroth_index = residue_index - 1
        # zeroth_index = residue_index
        indices = np.insert(self.grid_indexes[zeroth_index], 0, 0, axis=1)
        # indices = np.insert(indices, 4, self.masses[zeroth_index], axis=1)
        indices = np.array([(np.insert(x, 4, 0), np.insert(x, 4, 1)) for ind, x in enumerate(indices)])
        values = np.array([np.concatenate((self.masses[zeroth_index][index], self.charges[zeroth_index][index]))
                           for index, val in enumerate(self.masses[zeroth_index])])
        # values = np.array([len(self.masses[zeroth_index]) * [1]])
        hots = np.array(self.aa_one_hot[zeroth_index])
        return np.concatenate(indices), np.concatenate(values).ravel(), hots


from simtk.openmm import app
import simtk.openmm as mm
from simtk import unit
from numpy import sqrt

class VSiteTester:
    """Test the OPLS and virtual site implementation in openmm"""

    def __init__(self, pdb_file, xml_file):
        self.pdb = pdb_file
        self.xml = xml_file
        self.system = None
        self.opls = True
        self.openmm_system()

    def openmm_system(self):
        """Initialise the OpenMM system we will use to evaluate the energies."""

        # Load the initial coords into the system and initialise
        pdb = app.PDBFile(self.pdb)
        forcefield = app.ForceField(self.xml)
        modeller = app.Modeller(pdb.topology, pdb.positions)  # set the initial positions from the pdb
        # now we need to check if there are extra sites and add them to the system
        print(list(modeller.topology.bonds()))
        modeller.addExtraParticles(forcefield)
        # now we can make sure that the particles have been added
        app.PDBFile.writeFile(modeller.topology, modeller.positions, open('modeller.pdb', 'w+'))
        # before setting up the system we need to make sure the modeller.topology includes a parent-site bond

        # for atom in modeller.topology.atoms():
        #     print(atom)
        # atoms = [atom for atom in modeller.topology.atoms()]
        # modeller.topology.addBond(atoms[3], atoms[11])
        for bond in modeller.topology.bonds():
            print(bond)
        data = forcefield._SystemData()
        print(data.atoms)
        print(data.impropers)

        exit()

        self.system = forcefield.createSystem(modeller.topology, nonbondedMethod=app.NoCutoff, constraints=None)

        if self.opls:
            print('using opls rules')
            self.opls_lj()

        exit()

        temperature = 298.15 * unit.kelvin
        integrator = mm.LangevinIntegrator(temperature, 5 / unit.picoseconds, 0.001 * unit.picoseconds)

        self.simulation = app.Simulation(modeller.topology, self.system, integrator)
        self.simulation.context.setPositions(modeller.positions)

    def opls_lj(self, excep_pairs=None, normal_pairs=None):
        """This function changes the standard OpenMM combination rules to use OPLS, execp and normal pairs are only
        required if their are virtual sites in the molecule.
        """

        # Get the system information from the openmm system
        forces = {self.system.getForce(index).__class__.__name__: self.system.getForce(index) for index in
                  range(self.system.getNumForces())}
        # Use the nondonded_force to get the same rules
        nonbonded_force = forces['NonbondedForce']
        lorentz = mm.CustomNonbondedForce(
            'epsilon*((sigma/r)^12-(sigma/r)^6); sigma=sqrt(sigma1*sigma2); epsilon=sqrt(epsilon1*epsilon2)*4.0')
        lorentz.setNonbondedMethod(nonbonded_force.getNonbondedMethod())
        lorentz.addPerParticleParameter('sigma')
        lorentz.addPerParticleParameter('epsilon')
        lorentz.setCutoffDistance(nonbonded_force.getCutoffDistance())
        self.system.addForce(lorentz)

        l_j_set = {}
        # For each particle, collect the particle parameters into the lj method
        for index in range(nonbonded_force.getNumParticles()):
            charge, sigma, epsilon = nonbonded_force.getParticleParameters(index)
            # collect the lj params for the custom method
            l_j_set[index] = (sigma, epsilon)
            lorentz.addParticle([sigma, epsilon])
            # set the standard nonbonded lj parameters to 0 so they are not double counted
            nonbonded_force.setParticleParameters(index, charge, sigma, epsilon * 0)

        for i in range(nonbonded_force.getNumExceptions()):
            (p1, p2, q, sig, eps) = nonbonded_force.getExceptionParameters(i)
            print(p1, p2, q, sig, eps)
            # ALL THE 12,13 and 14 interactions are EXCLUDED FROM CUSTOM NONBONDED FORCE
            lorentz.addExclusion(p1, p2)
            if eps._value != 0.0:
                sig14 = sqrt(l_j_set[p1][0] * l_j_set[p2][0])
                eps14 = sqrt(l_j_set[p1][1] * l_j_set[p2][1])
                nonbonded_force.setExceptionParameters(i, p1, p2, q, sig14, eps14)
            # # If there is a virtual site in the molecule we have to change the exceptions and pairs lists
            # # Old method which needs updating
            # if excep_pairs:
            #     for x in range(len(excep_pairs)):  # scale 14 interactions
            #         if p1 == excep_pairs[x, 0] and p2 == excep_pairs[x, 1] or p2 == excep_pairs[x, 0] and p1 == \
            #                 excep_pairs[x, 1]:
            #             charge1, sigma1, epsilon1 = nonbonded_force.getParticleParameters(p1)
            #             charge2, sigma2, epsilon2 = nonbonded_force.getParticleParameters(p2)
            #             q = charge1 * charge2 * 0.5
            #             sig14 = sqrt(sigma1 * sigma2) * 0.5
            #             eps = sqrt(epsilon1 * epsilon2) * 0.5
            #             nonbonded_force.setExceptionParameters(i, p1, p2, q, sig14, eps)
            #
            # if normal_pairs:
            #     for x in range(len(normal_pairs)):
            #         if p1 == normal_pairs[x, 0] and p2 == normal_pairs[x, 1] or p2 == normal_pairs[x, 0] and p1 == \
            #                 normal_pairs[x, 1]:
            #             charge1, sigma1, epsilon1 = nonbonded_force.getParticleParameters(p1)
            #             charge2, sigma2, epsilon2 = nonbonded_force.getParticleParameters(p2)
            #             q = charge1 * charge2
            #             sig14 = sqrt(sigma1 * sigma2)
            #             eps = sqrt(epsilon1 * epsilon2)
            #             nonbonded_force.setExceptionParameters(i, p1, p2, q, sig14, eps)

        return self.system

def main():
    tester = VSiteTester('pyridine.pdb', 'pyridine.xml')

if __name__ == '__main__':
    main()


import tempfile
import logging
import os
import random
import types
import warnings

from pathlib import Path
from typing import Any, Dict, Optional

import ray
from ray import train
from ray.train._internal.accelerator import Accelerator
from ray.train.constants import PYTORCH_PROFILER_KEY
from torch.optim import Optimizer
from ray.train._internal.session import get_accelerator, set_accelerator
from ray.util.annotations import PublicAPI, Deprecated

import numpy as np

import torch
from torch.cuda.amp import autocast, GradScaler
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import (
    DistributedSampler,
    DataLoader,
    IterableDataset,
    SequentialSampler,
    RandomSampler,
)

try:
    from torch.profiler import profile
except ImportError:
    profile = None

logger = logging.getLogger(__name__)

@PublicAPI(stability="beta")
def get_device() -> torch.device:
    """Gets the correct torch device to use for training."""
    return get_accelerator(_TorchAccelerator).get_device()

@PublicAPI(stability="beta")
def prepare_model(
    model: torch.nn.Module,
    move_to_device: bool = True,
    wrap_ddp: bool = True,
    ddp_kwargs: Optional[Dict[str, Any]] = None,
) -> torch.nn.Module:
    """Prepares the model for distributed execution.

    This allows you to use the same exact code regardless of number of
    workers or the device type being used (CPU, GPU).

    Args:
        model (torch.nn.Module): A torch model to prepare.
        move_to_device: Whether to move the model to the correct
            device. If set to False, the model needs to manually be moved
            to the correct device.
        wrap_ddp: Whether to wrap models in
            ``DistributedDataParallel``.
        ddp_kwargs (Dict[str, Any]): Args to pass into
            ``DistributedDataParallel`` initialization if ``wrap_ddp`` is
            set to True.
    """
    return get_accelerator(_TorchAccelerator).prepare_model(
        model,
        move_to_device=move_to_device,
        wrap_ddp=wrap_ddp,
        ddp_kwargs=ddp_kwargs,
    )

@PublicAPI(stability="beta")
def prepare_data_loader(
    data_loader: torch.utils.data.DataLoader,
    add_dist_sampler: bool = True,
    move_to_device: bool = True,
    auto_transfer: bool = True,
) -> torch.utils.data.DataLoader:
    """Prepares DataLoader for distributed execution.

    This allows you to use the same exact code regardless of number of
    workers or the device type being used (CPU, GPU).

    Args:
        data_loader (torch.utils.data.DataLoader): The DataLoader to
            prepare.
        add_dist_sampler: Whether to add a DistributedSampler to
            the provided DataLoader.
        move_to_device: If set, automatically move the data
            returned by the data loader to the correct device.
        auto_transfer: If set and device is GPU, another CUDA stream
            is created to automatically copy data from host (CPU) memory
            to device (GPU) memory (the default CUDA stream still runs the
            training procedure). If device is CPU, it will be disabled
            regardless of the setting. This configuration will be ignored
            if ``move_to_device`` is False.
    """
    return get_accelerator(_TorchAccelerator).prepare_data_loader(
        data_loader,
        add_dist_sampler=add_dist_sampler,
        move_to_device=move_to_device,
        auto_transfer=auto_transfer,
    )

@PublicAPI(stability="beta")
def accelerate(amp: bool = False) -> None:
    """Enables training optimizations.

    Arguments:
        amp: If true, perform training with automatic mixed precision.
            Otherwise, use full precision.

    .. warning:: ``train.torch.accelerate`` cannot be called more than once, and it
       must be called before any other ``train.torch`` utility function.
    """
    try:
        set_accelerator(_TorchAccelerator(amp=amp))
    except RuntimeError:
        raise RuntimeError(
            "An accelerator has already been set. Make sure "
            "`train.torch.accelerate()` is not called multiple times, and is called "
            "before any of the prepare methods."
        )

@PublicAPI(stability="beta")
def prepare_optimizer(optimizer: torch.optim.Optimizer) -> torch.optim.Optimizer:
    """Wraps optimizer to support automatic mixed precision.

    Args:
        optimizer (torch.optim.Optimizer): The DataLoader to prepare.

    Returns:
        A wrapped optimizer.
    """
    return get_accelerator(_TorchAccelerator).prepare_optimizer(optimizer)

@PublicAPI(stability="beta")
def backward(tensor: torch.Tensor) -> None:
    """Computes the gradient of the specified tensor w.r.t. graph leaves.

    Args:
        tensor (torch.Tensor): Tensor of which the derivative will be computed.
    """
    get_accelerator(_TorchAccelerator).backward(tensor)

@PublicAPI(stability="beta")
def enable_reproducibility(seed: int = 0) -> None:
    """Limits sources of nondeterministic behavior.

    This function:

        * Seeds PyTorch, Python, and NumPy.
        * Disables CUDA convolution benchmarking.
        * Configures PyTorch to use determinstic algorithms.
        * Seeds workers spawned for multi-process data loading.

    Args:
        seed: The number to seed libraries and data workers with.

    .. warning:: ``train.torch.enable_reproducibility()`` can't guarantee
        completely reproducible results across executions. To learn more, read
        the `PyTorch notes on randomness
        <https://pytorch.org/docs/stable/notes/randomness.html>`_.
    """
    get_accelerator(_TorchAccelerator).enable_reproducibility(seed)

@Deprecated
class TorchWorkerProfiler:
    """Utility class for running PyTorch Profiler on a Train worker.

    Args:
        trace_dir (Optional[str]): The directory to store traces on the
           worker node. If ``None``, this will use a default temporary dir.
    """

    WORKER_TRACE_DIR_NAME = "pytorch_profiler_worker_traces"

    def __init__(self, trace_dir: Optional[str] = None):
        warnings.warn(
            "The `ray.train.torch.TorchWorkerProfiler` API is deprecated in Ray 2.0",
            DeprecationWarning,
            stacklevel=2,
        )
        if profile is None:
            raise ImportError(
                "Torch Profiler requires torch>=1.8.1. "
                "Run `pip install 'torch>=1.8.1'` to use TorchWorkerProfiler."
            )

        trace_dir = trace_dir or Path(tempfile.gettempdir()).joinpath(
            self.WORKER_TRACE_DIR_NAME
        )
        self.trace_dir = Path(trace_dir)
        self.trace_dir.mkdir(parents=True, exist_ok=True)
        # Accumulated traces.
        self.profiler_trace_filenames = []

    def trace_handler(self, p: profile):
        """A stateful PyTorch Profiler trace handler.

        This will the export chrome trace to a file on disk.

        These exported traces can then be fetched by calling
        ``get_and_clear_profile_traces``.

        Args:
            p: A PyTorch Profiler profile.
        """
        trace_filename = f"worker_{train.world_rank()}_epoch_{p.step_num}.pt.trace.json"
        trace_path = self.trace_dir.joinpath(trace_filename)

        logger.debug(f"Writing worker trace to {trace_path}.")
        p.export_chrome_trace(str(trace_path))
        self.profiler_trace_filenames.append(trace_filename)

    def get_and_clear_profile_traces(self):
        """Reads unread Profiler traces from this worker.

        Returns:
            The traces in a format consumable by
            ``TorchTensorboardProfilerCallback``.
        """

        def get_trace(filename):
            trace_path = self.trace_dir.joinpath(filename)
            return trace_path.read_text()

        traces = [
            (trace_filename, get_trace(trace_filename))
            for trace_filename in self.profiler_trace_filenames
        ]

        self.profiler_trace_files = []
        return {PYTORCH_PROFILER_KEY: traces}

class _TorchAccelerator(Accelerator):
    """A utility that implements methods to accelerate PyTorch training.

    Arguments:
        amp: If true, perform training with automatic mixed precision.
            Otherwise, use full precision.
    """

    def __init__(self, amp: bool = False):
        self.amp_is_enabled = amp
        self.scaler = GradScaler() if amp else None
        self._seed = None

    def prepare_model(
        self,
        model: torch.nn.Module,
        move_to_device: bool = True,
        wrap_ddp: bool = True,
        ddp_kwargs: Optional[Dict[str, Any]] = None,
    ) -> torch.nn.Module:
        """Prepares the model for distributed execution.

        This allows you to use the same exact code regardless of number of
        workers or the device type being used (CPU, GPU).

        Args:
            model (torch.nn.Module): A torch model to prepare.
            move_to_device: Whether to move the model to the correct
                device. If set to False, the model needs to manually be moved
                to the correct device.
            wrap_ddp: Whether to wrap models in
                ``DistributedDataParallel``.
            ddp_kwargs (Dict[str, Any]): Args to pass into
                ``DistributedDataParallel`` initialization if ``wrap_ddp`` is
                set to True.
        """
        ddp_kwargs = ddp_kwargs or {}

        rank = train.local_rank()

        device = self.get_device()

        if torch.cuda.is_available():
            torch.cuda.set_device(device)

        if move_to_device:
            logger.info(f"Moving model to device: {device}")
            model = model.to(device)

        def model_get_state(self):
            # `__getstate__` is an special method that informs pickle which attributes
            # to serialize. This custom implementation ensures that the wrapped forward
            # method and custom `__getstate__` method aren't serialized.
            if hasattr(self, "_original_get_state"):
                state = self._original_get_state()
                state["__getstate__"] = state["_original_get_state"]
                del state["_original_get_state"]
            else:
                # If model does not have a `__getstate__` already defined, use default
                # implementation.
                state = self.__dict__.copy()
                del state["__getstate__"]
            state["forward"] = state["_unwrapped_forward"]
            del state["_unwrapped_forward"]

            return state

        if self.amp_is_enabled:
            # Pickle cannot serialize the wrapped forward method. As a workaround,
            # define a custom `__getstate__` method that unwraps the forward method.
            model._unwrapped_forward = model.forward
            model.forward = autocast()(model.forward)

            # TODO(amogkam): Replace below logic with a generic "unpack model" method.
            # Replacing the `model.forward` method makes the model no longer
            # serializable. When serializing the model, we have to override the
            # `__getstate__` method to set back the original forward method.
            if hasattr(model, "__getstate__"):
                model._original_get_state = model.__getstate__
            # `__getstate__` must be a bound method rather than an callable attribute.
            # See https://stackoverflow.com/questions/972/adding-a-method-to-an-existing-object-instance.  # noqa: E501
            model.__getstate__ = types.MethodType(model_get_state, model)

        if wrap_ddp and train.world_size() > 1:
            logger.info("Wrapping provided model in DDP.")
            if torch.cuda.is_available():
                model = DistributedDataParallel(
                    model, device_ids=[rank], output_device=rank, **ddp_kwargs
                )
            else:
                model = DistributedDataParallel(model, **ddp_kwargs)

        return model

    def prepare_data_loader(
        self,
        data_loader: torch.utils.data.DataLoader,
        add_dist_sampler: bool = True,
        move_to_device: bool = True,
        auto_transfer: bool = True,
    ) -> torch.utils.data.DataLoader:
        """Prepares DataLoader for distributed execution.

        This allows you to use the same exact code regardless of number of
        workers or the device type being used (CPU, GPU).

        Args:
            data_loader (torch.utils.data.DataLoader): The DataLoader to
                prepare.
            add_dist_sampler: Whether to add a DistributedSampler to
                the provided DataLoader.
            move_to_device: If set, automatically move the data
                returned by the data loader to the correct device.
            auto_transfer: If set and device is GPU, another CUDA stream
                is created to automatically copy data from host (CPU) memory
                to device (GPU) memory (the default CUDA stream still runs the
                training procedure). If device is CPU, it will be disabled
                regardless of the setting. This configuration will be ignored
                if ``move_to_device`` is False.
        """

        # Only add Distributed Sampler if the following conditions hold:
        # 1. More than one training worker is being used.
        # 2. A DistributedSampler has not already been added by the user.
        # 3. The dataset is not an IterableDataset. Samplers do not worker with
        # IterableDatasets.
        if (
            train.world_size() > 1
            and not isinstance(data_loader.sampler, DistributedSampler)
            and not (
                hasattr(data_loader, "dataset")
                and isinstance(data_loader.dataset, IterableDataset)
            )
            and add_dist_sampler
        ):

            def with_sampler(loader):
                # Automatically set the DistributedSampler

                # If you're using a sampler, the DataLoader shuffle flag must be set to
                # False. Shuffling is instead determined by the shuffle argument passed
                # to the DistributedSampler constructor.

                # If no sampler is passed to the DataLoader constructor, Torch
                # constructs a default sampler. The default sampler is a RandomSampler
                # if shuffling is enabled and a SequentialSampler otherwise. DataLoader
                # does not have a shuffle attribute, so we instead identify whether
                # shuffling is enabled by checking the default sampler type.
                shuffle = not isinstance(loader.sampler, SequentialSampler)

                def seeded_worker_init_fn(worker_init_fn):
                    def wrapper(worker_id):
                        worker_seed = torch.initial_seed() % 2 ** 32
                        np.random.seed(worker_seed)
                        random.seed(worker_seed)
                        worker_init_fn(worker_id)

                    return wrapper

                worker_init_fn = loader.worker_init_fn
                generator = loader.generator
                if self._seed is not None:
                    worker_init_fn = seeded_worker_init_fn(loader.worker_init_fn)
                    generator = torch.Generator()
                    generator.manual_seed(self._seed)

                using_default_sampler = isinstance(
                    loader.sampler, (SequentialSampler, RandomSampler)
                )
                if not using_default_sampler and train.world_rank() == 0:
                    logger.warn(
                        f"The {loader.sampler.__class__.__name__} will be overwritten "
                        "with a DistributedSampler. You can disable this by setting "
                        "`with_sampler` to False in `prepare_data_loader`."
                    )

                data_loader_args = {
                    "dataset": loader.dataset,
                    "batch_size": loader.batch_size,
                    "shuffle": False,
                    "num_workers": loader.num_workers,
                    "collate_fn": loader.collate_fn,
                    "pin_memory": loader.pin_memory,
                    "drop_last": loader.drop_last,
                    "timeout": loader.timeout,
                    "worker_init_fn": worker_init_fn,
                    "generator": generator,
                    "sampler": DistributedSampler(loader.dataset, shuffle=shuffle),
                }
                return DataLoader(**data_loader_args)

            data_loader = with_sampler(data_loader)

        if move_to_device:
            device = self.get_device()
            data_loader = _WrappedDataLoader(data_loader, device, auto_transfer)

        return data_loader

    def get_device(self) -> torch.device:
        """Gets the correct torch device to use for training."""
        if torch.cuda.is_available():
            gpu_ids = ray.get_gpu_ids()
            if len(gpu_ids) > 0:
                device_id = gpu_ids[0]
            else:
                # If called on the driver or outside of Ray Train, return the
                # 0th device.
                device_id = 0
            device = torch.device(f"cuda:{device_id}")
        else:
            device = torch.device("cpu")

        return device

    def prepare_optimizer(self, optimizer: Optimizer) -> Optimizer:
        """Wraps optimizer to support automatic mixed precision.

        Args:
            optimizer (torch.optim.Optimizer): The DataLoader to prepare.

        Returns:
            A wrapped optimizer.
        """
        return _WrappedOptimizer(optimizer, scaler=self.scaler)

    def backward(self, tensor: torch.Tensor) -> None:
        """Computes the gradient of the specified tensor w.r.t. graph leaves.

        Args:
            tensor (torch.Tensor): Tensor of which the derivative will be computed.
        """
        if self.amp_is_enabled:
            self.scaler.scale(tensor).backward()
        else:
            tensor.backward()

    def enable_reproducibility(self, seed: int = 0) -> None:
        """Limits sources of nondeterministic behavior."""
        self._seed = seed

        torch.manual_seed(seed)
        random.seed(seed)
        np.random.seed(seed)

        torch.use_deterministic_algorithms(True)
        torch.backends.cudnn.benchmark = False

        # If you want to use deterministic algorithms with CUDA, then you need to set
        # the CUBLAS_WORKSPACE_CONFIG environment variable; otherwise, Torch errors.
        # See https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility.
        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

class _WrappedDataLoader(DataLoader):
    def __init__(
        self, base_dataloader: DataLoader, device: torch.device, auto_transfer: bool
    ):

        self.__dict__.update(getattr(base_dataloader, "__dict__", {}))
        self._dataloader = base_dataloader
        self.dataloader_iter = None
        self.device = device
        # disable auto transfer (host->device) if cpu is used
        self._auto_transfer = auto_transfer if device.type == "cuda" else False
        # create a new CUDA stream to move data from host to device concurrently
        self._memcpy_stream = (
            torch.cuda.Stream()
            if device.type == "cuda" and self._auto_transfer
            else None
        )
        self.next_batch = None

    def _move_to_device(self, item):
        if item is None:
            return None

        def try_move_device(i):
            try:
                i = i.to(self.device, non_blocking=self._auto_transfer)
            except AttributeError:
                logger.debug(f"Item {i} cannot be moved to device " f"{self.device}.")
            return i

        with torch.cuda.stream(self._memcpy_stream):
            return tuple(try_move_device(i) for i in item)

    def _wait_for_batch(self, item):
        if self._memcpy_stream is None:
            return
        # Reference:
        # https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html
        # The training stream (current) needs to wait until
        # the memory copy stream finishes.
        curr_stream = torch.cuda.current_stream()
        curr_stream.wait_stream(self._memcpy_stream)
        # When a tensor is used by CUDA streams different from
        # its original allocator, we need to call ``record_stream``
        # to inform the allocator of all these streams. Otherwise,
        # the tensor might be freed once it is no longer used by
        # the creator stream.
        for i in item:
            # The Pytorch DataLoader has no restrictions on what is outputted for
            # each batch. We should only ``record_stream`` if the item has the
            # ability to do so.
            try:
                i.record_stream(curr_stream)
            except AttributeError:
                pass

    def __len__(self):
        return len(self._dataloader)

    def _prefetch_next_batch(self):
        next_batch = next(self.dataloader_iter, None)
        self.next_batch = self._move_to_device(next_batch)

    def __iter__(self):
        self.dataloader_iter = iter(self._dataloader)
        self._prefetch_next_batch()
        return self

    def __next__(self):
        next_batch = self.next_batch
        if next_batch is None:
            raise StopIteration
        self._wait_for_batch(next_batch)
        self._prefetch_next_batch()
        return next_batch

class _WrappedOptimizer(Optimizer):
    def __init__(self, optimizer: Optimizer, scaler: Optional[GradScaler] = None):
        self.optimizer = optimizer
        self.scaler = scaler

    @property
    def state(self):
        return self.optimizer.state

    @state.setter
    def state(self, state):
        self.optimizer.state = state

    @property
    def param_groups(self):
        return self.optimizer.param_groups

    @param_groups.setter
    def param_groups(self, param_groups):
        self.optimizer.param_groups = param_groups

    @property
    def defaults(self):
        return self.optimizer.defaults

    @defaults.setter
    def defaults(self, defaults):
        self.optimizer.defaults = defaults

    def add_param_group(self, param_group):
        self.optimizer.add_param_group(param_group)

    def load_state_dict(self, state_dict):
        self.optimizer.load_state_dict(state_dict)

    def state_dict(self):
        return self.optimizer.state_dict()

    def zero_grad(self):
        self.optimizer.zero_grad()

    def step(self, closure=None):
        if self.scaler is not None:
            self.scaler.step(self.optimizer, closure)
            self.scaler.update()
        else:
            self.optimizer.step(closure)


#!/usr/bin/env python

# Light each LED in sequence, and repeat.

#import fastopc, time
import numpy as np
import mido
import time

class Board:
    def __init__(self):
        self.knobs      = np.zeros(128)
        self.pitchwheel = 0.0
        self.notes      = np.zeros(128)
        self.velocities = np.zeros(128)
    def update(self, msg):
        print(msg)
        if msg.type == 'note_on':
            self.notes[msg.note] = 1
            self.velocities[msg.note] = msg.velocity
        elif msg.type == 'note_off':
            self.notes[msg.note] = 0
        elif msg.type == 'pitchwheel':
            self.pitchwheel = msg.pitch
        elif msg.type == 'control_change':
            self.knobs[msg.control] = msg.value

class midiSequence:
    def __init__(self):
        self.msgList = []
        self.tList   = []
    def update(self, msg):
        if self.tList == 0: self.t0 = time.time
        self.tList.append(time.time()-self.t0)
        self.msgList.append(msg)
    def getMsgs(self, msg):
        msgReturnList = []
        t = time.time()
        for n in range(len(self.tList)):
            if t<tList[n]:
                msgReturnList.append(msgList[n])

class Pixels():
    def __init__(self, nStrips, lStrip, minDisplay, upSampleFactor=1):
        self.upSampleFactor=upSampleFactor
        self.nStrips    = nStrips
        self.lStrip     = lStrip
        self.nLed       = nStrips*lStrip
        self.minDisplay = minDisplay
        self.array      = np.zeros([self.nLed, 3])
    def update(self, arrayNew, alphaRise, alphaDecay):
        if arrayNew.shape == (self.lStrip, 3):
            arrayNew = np.tile(arrayNew, (self.nStrips, 1))
            alpha = arrayNew - self.array
            alpha[alpha > 0.0 ] = alphaRise
            alpha[alpha <= 0.0] = alphaDecay
            self.array = alpha*arrayNew + (1.0-alpha)*self.array
        elif arrayNew.shape == (self.nLed,   3):
            alpha = arrayNew - self.array
            alpha[alpha > 0.0 ] = alphaRise
            alpha[alpha <= 0.0] = alphaDecay
            self.array = alpha*arrayNew + (1.0-alpha)*self.array
        elif arrayNew.shape == (self.nStrips, self.lStrip, 3):
            arrayNew = arrayNew.reshape(self.nLed, 3)
            alpha    = arrayNew - self.array
            alpha[alpha > 0.0 ] = alphaRise
            alpha[alpha <= 0.0] = alphaDecay
            self.array = alpha*arrayNew + (1.0-alpha)*self.array
        elif arrayNew.shape == (self.lStrip//2,   3):
            arrayNewFlip = np.flip(arrayNew, axis=0)
            arrayNewDouble = np.concatenate((arrayNewFlip,arrayNew), axis=0)
            arrayNew = np.tile(arrayNewDouble, (self.nStrips, 1))
            alpha = arrayNew - self.array
            alpha[alpha > 0.0 ] = alphaRise
            alpha[alpha <= 0.0] = alphaDecay
            self.array = alpha*arrayNew + (1.0-alpha)*self.array
        elif arrayNew.shape == (self.nLed//2,   3):
            arrayNewFull = np.zeros_like(self.array)
            for i in range(self.nStrips):
                arrayNewFlip = np.flip(arrayNew, axis=0)
                arrayNewDouble = np.concatenate((arrayNewFlip,arrayNew), axis=0)
                arrayNewFull[(i-1)*self.lStrip:(i)*self.lStrip] = arrayNewDouble
            alpha = arrayNew - self.array
            alpha[alpha > 0.0 ] = alphaRise
            alpha[alpha <= 0.0] = alphaDecay
            self.array = alpha*arrayNewFull + (1.0-alpha)*self.array
    def updateSimple(self, arrayNew):
        if arrayNew.shape == (self.lStrip, 3):
            arrayNew = np.tile(arrayNew, (self.nStrips, 1))
            self.array = arrayNew
        elif arrayNew.shape == (self.nLed,   3):
            self.array = arrayNew
    def getArrayForDisplay(self):
        if self.upSampleFactor==1:
            returnArray = self.array
        else:
            returnArray = np.mean(self.array.reshape(-1,self.upSampleFactor), axis=0)
        if self.minDisplay != 0: returnArray[returnArray < self.minDisplay] = 0
        return returnArray

class PixelsHalfStrips():
    def __init__(self, nStrips, lStrip, minDisplay):
        self.nStrips    = nStrips
        self.lStrip     = lStrip
        self.nLed       = nStrips*lStrip
        self.minDisplay = minDisplay
        self.array      = np.zeros([self.nLed, 3])
    def update(self, arrayNew, alphaRise, alphaDecay):
        if arrayNew.shape == (self.lStrip, 3):
            arrayNew = np.tile(arrayNew, (self.nStrips, 1))
            alpha = arrayNew - self.array
            alpha[alpha > 0.0 ] = alphaRise
            alpha[alpha <= 0.0] = alphaDecay
            self.array = alpha*arrayNew + (1.0-alpha)*self.array
        elif arrayNew.shape == (self.nLed,   3):
            alpha = arrayNew - self.array
            alpha[alpha > 0.0 ] = alphaRise
            alpha[alpha <= 0.0] = alphaDecay
            self.array = alpha*arrayNew + (1.0-alpha)*self.array
        elif arrayNew.shape == (self.lStrip//2,   3):
            arrayNewFlip = np.flip(arrayNew, axis=0)
            arrayNewDouble = np.concatenate((arrayNewFlip,arrayNew), axis=0)
            arrayNew = np.tile(arrayNewDouble, (self.nStrips, 1))
            alpha = arrayNew - self.array
            alpha[alpha > 0.0 ] = alphaRise
            alpha[alpha <= 0.0] = alphaDecay
            self.array = alpha*arrayNew + (1.0-alpha)*self.array
        elif arrayNew.shape == (self.nLed//2,   3):
            arrayNewFull = np.zeros_like(self.array)
            for i in range(self.nStrips):
                arrayNewFlip = np.flip(arrayNew, axis=0)
                arrayNewDouble = np.concatenate((arrayNewFlip,arrayNew), axis=0)
                arrayNewFull[(i-1)*self.lStrip:(i)*self.lStrip] = arrayNewDouble
            alpha = arrayNew - self.array
            alpha[alpha > 0.0 ] = alphaRise
            alpha[alpha <= 0.0] = alphaDecay
            self.array = alpha*arrayNewFull + (1.0-alpha)*self.array
    def updateSimple(self, arrayNew):
        if arrayNew.shape == (self.lStrip, 3):
            arrayNew = np.tile(arrayNew, (self.nStrips, 1))
            self.array = arrayNew
        elif arrayNew.shape == (self.nLed,   3):
            self.array = arrayNew
    def getArrayForDisplay(self):
        returnArray = self.array
        if self.minDisplay != 0: returnArray[returnArray < self.minDisplay] = 0
        return returnArray

#
class BassPixels():
    'neck then body then strap'
    def __init__(self, lNeck,lBody, lStrap, minDisplay):
        self.lNeck  = lNeck
        self.lBody  = lBody
        self.lStrap = lStrap
        self.nLed = lNeck*2 + lBody + lStrap
        self.minDisplay = minDisplay
        self.array      = np.zeros([self.nLed, 3])
    def updateNeck(self, arrayNew, alphaRise, alphaDecay):
        alpha = arrayNew - self.array[0:self.lNeck]
        alpha[alpha > 0.0 ] = alphaRise
        alpha[alpha <= 0.0] = alphaDecay
        self.array[0:self.lNeck] = alpha*arrayNew + (1.0-alpha)*self.array[0:self.lNeck]
        self.array[self.lNeck:2*self.lNeck] = alpha*arrayNew + (1.0-alpha)*self.array[self.lNeck:2*self.lNeck]
    def updateNeck1(self, arrayNew, alphaRise, alphaDecay):
        alpha = arrayNew - self.array[0:self.lNeck]
        alpha[alpha > 0.0 ] = alphaRise
        alpha[alpha <= 0.0] = alphaDecay
        self.array[0:self.lNeck] = alpha*arrayNew + (1.0-alpha)*self.array[0:self.lNeck]
    def updateNeck2(self, arrayNew, alphaRise, alphaDecay):
        alpha = arrayNew - self.array[self.lNeck:2*self.lNeck]
        alpha[alpha > 0.0 ] = alphaRise
        alpha[alpha <= 0.0] = alphaDecay
        self.array[self.lNeck:2*self.lNeck] = alpha*arrayNew + (1.0-alpha)*self.array[self.lNeck:2*self.lNeck]
    def updateBody(self, arrayNew, alphaRise, alphaDecay):
        alpha = arrayNew - self.array[self.lNeck*2:self.lNeck*2+self.lBody]
        alpha[alpha > 0.0 ] = alphaRise
        alpha[alpha <= 0.0] = alphaDecay
        self.array[self.lNeck*2:self.lNeck*2+self.lBody] = alpha*arrayNew + (1.0-alpha)*self.array[self.lNeck*2:self.lNeck*2+self.lBody]
    def updateStrap(self, arrayNew, alphaRise, alphaDecay):
        alpha = arrayNew - self.array[self.lNeck*2+self.lBody:]
        alpha[alpha > 0.0 ] = alphaRise
        alpha[alpha <= 0.0] = alphaDecay
        self.array[self.lNeck*2+self.lBody:] = alpha*arrayNew + (1.0-alpha)*self.array[self.lNeck*2+self.lBody:]
    def getArrayForDisplay(self):
        returnArray = self.array
        if self.minDisplay != 0: returnArray[returnArray < self.minDisplay] = 0
        return returnArray

class ExpFilter:
    """Temporal exponential smoothing filter
    """
    def __init__(self, val=0.0, alpha_decay=0.5, alpha_rise=0.5):
        """Small rise / decay factors = more smoothing"""
        self.alpha_decay = alpha_decay
        self.alpha_rise = alpha_rise
        self.value = val
    def update(self, value):
        if isinstance(self.value, (list, np.ndarray, tuple)):
            alpha = value - self.value
            alpha[alpha > 0.0] = self.alpha_rise
            alpha[alpha <= 0.0] = self.alpha_decay
        else:
            alpha = self.alpha_rise if value > self.value else self.alpha_decay
        self.value = alpha * value + (1.0 - alpha) * self.value

class Bouncer :
    def __init__(self, n, speed, color, nPixels):
        self.n = n
        self.speed = speed
        self.color=color/np.amax(color)
        self.color=self.color*255.0
        for i in range(0,3): self.color[i]=int(self.color[i])
        self.nPixels = nPixels
        self.locInt = 10
        self.locFloat = float(self.locInt)
        self.outArray = np.zeros([nPixels, 3])
        self.outZeros = np.zeros_like(self.outArray)
        for i in range(self.locInt-self.n, self.locInt+self.n+1):
            if i == self.locInt:
                self.outArray[i] = self.color
            else:
                self.outArray[i] = self.color
    def update(self):
        self.locFloat = self.locFloat + self.speed
        if int(self.locFloat) != self.locInt:
            self.locInt = int(self.locFloat)
            self.outArray = np.roll(self.outArray, int(np.sign(self.speed)), axis=0)
        if self.locInt == self.n or self.locInt == self.nPixels-1-self.n:
            self.speed = -self.speed
    def getFullOutArray(self):
        returnArray = self.outArray
        return returnArray.astype(int)

'''
def getColorWheel(nTot):
    colorWheel = np.zeros([nTot,3])
    nTot3 = nTot//3
    for n in range(nTot):
        if n <= nTot3:
            colorWheel[n,0] = 1.0 - float(n)/float(nTot3)
            colorWheel[n,1] = 0.0 + float(n)/float(nTot3)
            colorWheel[n,2] = 0.0
        elif nTot3 < n <= 2*nTot3:
            colorWheel[n,0] = 0.0
            colorWheel[n,1] = 1.0 - float(n-nTot3)/float(nTot3)
            colorWheel[n,2] = 0.0 + float(n-nTot3)/float(nTot3)
        elif 2*nTot3 < n <= nTot:
            colorWheel[n,0] = 0.0 + float(n-2*nTot3)/float(nTot3)
            colorWheel[n,1] = 0.0
            colorWheel[n,2] = 1.0 - float(n-2*nTot3)/float(nTot3)
    return colorWheel
'''

def getColorWheel(nTot):
    colorWheel = np.zeros([nTot,3])
    nTot3 = nTot//3
    for n in range(nTot):
        if n <= nTot3:
            colorWheel[n,0] = 1.0 - np.power(float(n)/float(nTot3), 1)
            colorWheel[n,1] = 0.0 + np.power(float(n)/float(nTot3), 1)
            colorWheel[n,2] = 0.0
        elif nTot3 < n <= 2*nTot3:
            colorWheel[n,0] = 0.0
            colorWheel[n,1] = 1.0 - np.power(float(n-nTot3)/float(nTot3), 1)
            colorWheel[n,2] = 0.0 + np.power(float(n-nTot3)/float(nTot3), 1)
        elif 2*nTot3 < n <= nTot:
            colorWheel[n,0] = 0.0 + np.power(float(n-2*nTot3)/float(nTot3), 1)
            colorWheel[n,1] = 0.0
            colorWheel[n,2] = 1.0 - np.power(float(n-2*nTot3)/float(nTot3), 1)
    return colorWheel

def getColorWheelRedBlue(nTot):
    colorWheel = np.zeros([nTot,3])
    for n in range(nTot):
            colorWheel[n,0] = 0.0 + np.power(float(n)/float(nTot), 0.25)
            colorWheel[n,1] = 0.0
            colorWheel[n,2] = 1.0 - np.power(float(n)/float(nTot), 4)
    return colorWheel

def clDisplay(powerNorm1):
    maxDisplay     = 100.0
    maxOutFactor   = 5
    segLen         = int(maxDisplay / maxOutFactor)
    powerNorm      = powerNorm1 * maxDisplay/maxOutFactor
    string1        = 'O' * int(min(powerNorm, maxDisplay))
    string2        = '-' * int(max(maxDisplay-powerNorm,0))
    finalString    = string1 + string2
    charList = list(finalString+'|')
    for i in range(maxOutFactor): charList[i*segLen]='|'
    finalString = ''.join(charList)
    print(finalString)


import numpy as np
import time
import cv2
import math
import imutils

from scipy.spatial import distance as dist
import numpy as np
import imutils
import cv2
import os
# base path to YOLO directory
MODEL_PATH = "yolo-coco"

# initialize minimum probability to filter weak detections along with
# the threshold when applying non-maxima suppression
MIN_CONF = 0.3
NMS_THRESH = 0.3

# boolean indicating if NVIDIA CUDA GPU should be used
USE_GPU = False

# define the minimum safe distance (in pixels) that two people can be
# from each other
MIN_DISTANCE = 50

# load the COCO class labels our YOLO model was trained on
labelsPath = os.path.sep.join([MODEL_PATH, "obj.names"])
LABELS = open(labelsPath).read().strip().split("\n")

print(LABELS)

print(len(LABELS))

# derive the paths to the YOLO weights and model configuration
weightsPath = os.path.sep.join([MODEL_PATH, "yolov4.weights"])
configPath = os.path.sep.join([MODEL_PATH, "yolov4.cfg"])

# load our YOLO object detector trained on COCO dataset (80 classes)
print("[INFO] loading YOLO from disk...")
net = cv2.dnn.readNetFromDarknet(configPath, weightsPath)

# check if we are going to use GPU
if USE_GPU:
    # set CUDA as the preferable backend and target
    print("Setting preferable backend and target to CUDA...")
    net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
    net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)

# determine only the *output* layer names that we need from YOLO
ln = net.getLayerNames()
ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]

# check if we are going to use GPU
if USE_GPU:
    # set CUDA as the preferable backend and target
    print("Setting preferable backend and target to CUDA...")
    net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
    net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)

# determine only the *output* layer names that we need from YOLO

print("Accessing video stream...")
#upload the video file you want to check social distancing for- below.

class Video_Web(object):
    def __init__(self):
        self.video = cv2.VideoCapture(0)
        fps = self.video.get(cv2.CAP_PROP_FPS)
        print("FPS of the current video: ",fps)

        num_frames = self.video.get(cv2.CAP_PROP_FRAME_COUNT)
        print("Number of frames in the video: ",num_frames)

        writer = None
    def __del__(self):
        self.video.release()
    def get_frame(self):
        def detect_people(frame, net, ln, personIdx=0):
    # grab the dimensions of the frame and  initialize the list of
    # results
            (H, W) = frame.shape[:2]
            results = []

    # construct a blob from the input frame and then perform a forward
    # pass of the YOLO object detector, giving us our bounding boxes
    # and associated probabilities
            blob = cv2.dnn.blobFromImage(frame, 1 / 255.0, (416, 416), swapRB=True, crop=False)
            net.setInput(blob)
            layerOutputs = net.forward(ln)

    # initialize our lists of detected bounding boxes, centroids, and
    # confidences, respectively
            boxes = []
            centroids = []
            confidences = []

    # loop over each of the layer outputs
            for output in layerOutputs:
        # loop over each of the detections
                for detection in output:
            # extract the class ID and confidence (i.e., probability)
            # of the current object detection
                    scores = detection[5:]
                    classID = np.argmax(scores)
                    confidence = scores[classID]

            # filter detections by (1) ensuring that the object
            # detected was a person and (2) that the minimum
            # confidence is met
                    if classID == personIdx and confidence > MIN_CONF:
                # scale the bounding box coordinates back relative to
                # the size of the image, keeping in mind that YOLO
                # actually returns the center (x, y)-coordinates of
                # the bounding box followed by the boxes' width and
                # height
                        box = detection[0:4] * np.array([W, H, W, H])
                        (centerX, centerY, width, height) = box.astype("int")

                # use the center (x, y)-coordinates to derive the top
                # and and left corner of the bounding box
                        x = int(centerX - (width / 2))
                        y = int(centerY - (height / 2))

                # update our list of bounding box coordinates,
                # centroids, and confidences
                        boxes.append([x, y, int(width), int(height)])
                        centroids.append((centerX, centerY))
                        confidences.append(float(confidence))

    # apply non-maxima suppression to suppress weak, overlapping
    # bounding boxes
            idxs = cv2.dnn.NMSBoxes(boxes, confidences, MIN_CONF, NMS_THRESH)

    # ensure at least one detection exists
            if len(idxs) > 0:
        # loop over the indexes we are keeping
                    for i in idxs.flatten():
            # extract the bounding box coordinates
                        (x, y) = (boxes[i][0], boxes[i][1])
                        (w, h) = (boxes[i][2], boxes[i][3])

            # update our results list to consist of the person
            # prediction probability, bounding box coordinates,
            # and the centroid
                        r = (confidences[i], (x, y, x + w, y + h), centroids[i])
                        results.append(r)
    # return the list of results
            return results

    # read the next frame from the file
        grabbed, frame = self.video.read()

    # resize the frame and then detect people (and only people) in it
        frame = imutils.resize(frame, width=700)
        results = detect_people(frame, net, ln, personIdx=LABELS.index("person"))

    # initialize the set of indexes that violate the minimum social
    # distance
        violate = set()

    # ensure there are *at least* two people detections (required in
    # order to compute our pairwise distance maps)
        if len(results) >= 2:
        # extract all centroids from the results and compute the
        # Euclidean distances between all pairs of the centroids
                centroids = np.array([r[2] for r in results])
                D = dist.cdist(centroids, centroids, metric="euclidean")

        # loop over the upper triangular of the distance matrix
                for i in range(0, D.shape[0]):
                    for j in range(i + 1, D.shape[1]):
                # check to see if the distance between any two
                # centroid pairs is less than the configured number
                # of pixels
                        if D[i, j] < MIN_DISTANCE:
                    # update our violation set with the indexes of
                    # the centroid pairs
                            violate.add(i)
                            violate.add(j)

    # loop over the results
        for (i, (prob, bbox, centroid)) in enumerate(results):
        # extract the bounding box and centroid coordinates, then
        # initialize the color of the annotation
                (startX, startY, endX, endY) = bbox
                (cX, cY) = centroid
                color = (0, 255, 0)

        # if the index pair exists within the violation set, then
        # update the color
                if i in violate:
                    color = (0, 0, 255)

        # draw (1) a bounding box around the person and (2) the
        # centroid coordinates of the person,
                cv2.rectangle(frame, (startX, startY), (endX, endY), color, 2)
                cv2.circle(frame, (cX, cY), 5, color, 1)

    # draw the total number of social distancing violations on the
    # output frame
        text = "Social Distancing Violations: {}".format(len(violate))
        cv2.putText(frame, text, (10, frame.shape[0] - 25), cv2.FONT_HERSHEY_SIMPLEX, 0.85, (0, 0, 255), 3)

        # ret,frame=self.video.read()
        grabbed,jpg = cv2.imencode('.jpg',frame)
        return jpg.tobytes()