from typing import Set, Optional, Literal, Type, TypeVar, Union, Generic
from dataclasses import dataclass
import logging

from angr.knowledge_plugins.variables.variable_manager import VariableManagerInternal
from angr.sim_variable import SimTemporaryVariable
from angr.sim_variable import SimMemoryVariable
from angr.sim_variable import SimStackVariable
from angr.sim_variable import SimRegisterVariable
from angr.misc.ux import once

from ...engines.light import SpOffset
from ...code_location import CodeLocation, ExternalCodeLocation
from .atoms import Atom, MemoryLocation, Register, Tmp, AtomKind, atom_kind_mapping
from .tag import Tag
from ...sim_variable import SimVariable

log = logging.getLogger(__name__)

@dataclass
class DefinitionMatchPredicate:
    """
    A dataclass indicating several facts which much all must match in order for a definition to match. Largely an
    internal class; don't worry about this.
    """

    kind: Union[AtomKind, Type[Atom], None] = None
    bbl_addr: Optional[int] = None
    ins_addr: Optional[int] = None
    variable: Optional[SimVariable] = None
    variable_manager: Union[VariableManagerInternal, None, Literal[False]] = None
    stack_offset: Optional[int] = None
    reg_name: Optional[Union[str, int]] = None
    heap_offset: Optional[int] = None
    global_addr: Optional[int] = None
    tmp_idx: Optional[int] = None
    const_val: Optional[int] = None
    extern: Optional[bool] = None

    @staticmethod
    def construct(predicate: Optional["DefinitionMatchPredicate"] = None, **kwargs) -> "DefinitionMatchPredicate":
        if predicate is None:
            predicate = DefinitionMatchPredicate(**kwargs)
            predicate.normalize()
        return predicate

    def normalize(self):
        if self.variable is not None:
            if isinstance(self.variable, SimRegisterVariable):
                self.reg_name = self.variable.reg
            elif isinstance(self.variable, SimStackVariable):
                if self.variable.base != "bp":
                    log.warning("Cannot match against variables with %s base (need bp)", self.variable.base)
                else:
                    self.stack_offset = self.variable.offset
            elif isinstance(self.variable, SimMemoryVariable):
                # TODO region
                if isinstance(self.variable.addr, int):
                    self.global_addr = self.variable.addr
                else:
                    log.warning(
                        "Cannot match against memory variable with %s addr (need int)",
                        type(self.variable.addr).__name__,
                    )
            elif isinstance(self.variable, SimTemporaryVariable):
                self.tmp_idx = self.variable.tmp_id
            else:
                log.warning(
                    "Cannot match against definition to %s (need reg, stack, mem, or tmp)", type(self.variable).__name__
                )

        if self.reg_name is not None:
            self.kind = AtomKind.REGISTER
        elif self.stack_offset is not None or self.heap_offset is not None or self.global_addr is not None:
            self.kind = AtomKind.MEMORY
        elif self.const_val is not None:
            self.kind = AtomKind.CONSTANT
        elif self.tmp_idx is not None:
            self.kind = AtomKind.TMP

    def matches(self, defn: "Definition") -> bool:
        if self.variable is not None:
            if self.variable_manager is False:
                pass
            elif self.variable_manager is not None:
                if not self.variable_manager.is_variable_used_at(
                    self.variable, (defn.codeloc.bbl_addr, defn.codeloc.stmt_idx)
                ):
                    return False
            elif once("definition_matches_no_variable_manager"):
                log.warning(
                    "Cannot match definitions to variables on the basis of locations without a variable manager."
                )
                log.warning("Pass variable_manager=False to acknowledge this explicitly.")
        if self.bbl_addr is not None and defn.codeloc.block_addr != self.bbl_addr:
            return False
        if self.ins_addr is not None and defn.codeloc.ins_addr != self.ins_addr:
            return False
        if self.extern is not None and isinstance(defn.codeloc, ExternalCodeLocation) != self.extern:
            return False

        if self.kind is not None:
            if not isinstance(self.kind, type):
                self.kind = atom_kind_mapping[self.kind]
            if not isinstance(defn.atom, self.kind):
                return False

        if isinstance(defn.atom, Register):
            if self.reg_name is not None:
                if isinstance(self.reg_name, int):
                    if not defn.atom.reg_offset <= self.reg_name < defn.atom.reg_offset + defn.atom.size:
                        return False
                elif isinstance(self.reg_name, str):
                    if defn.atom.arch is not None:
                        if self.reg_name != defn.atom.name:
                            return False
                    else:
                        log.warning(
                            "Attempting to match by register name against a definition which does not have an arch"
                        )
                        return False
                else:
                    raise TypeError(self.reg_name)
        elif isinstance(defn.atom, MemoryLocation):
            if self.stack_offset is not None:
                if (
                    not isinstance(defn.atom.addr, SpOffset)
                    or defn.atom.addr.base != "sp"  # TODO???????
                    or defn.atom.addr.offset != self.stack_offset
                ):
                    return False
        elif isinstance(defn.atom, Tmp):
            if self.tmp_idx is not None and self.tmp_idx != defn.atom.tmp_idx:
                return False

        return True

A = TypeVar("A", bound=Atom)

class Definition(Generic[A]):
    """
    An atom definition.

    :ivar atom:     The atom being defined.
    :ivar codeloc:  Where this definition is created in the original binary code.
    :ivar dummy:    Tell whether the definition should be considered dummy or not. During simplification by AILment,
                    definitions marked as dummy will not be removed.
    :ivar tags:     A set of tags containing information about the definition gathered during analyses.
    """

    __slots__ = (
        "atom",
        "codeloc",
        "dummy",
        "tags",
        "_hash",
    )

    def __init__(self, atom: A, codeloc: CodeLocation, dummy: bool = False, tags: Optional[Set[Tag]] = None):
        self.atom: A = atom
        self.codeloc: CodeLocation = codeloc
        self.dummy: bool = dummy
        self.tags = tags or set()
        self._hash = None

    def __eq__(self, other):
        return self.atom == other.atom and self.codeloc == other.codeloc

    def __repr__(self):
        if not self.tags:
            return "<Definition {{Atom:{}, Codeloc:{}}}{}>".format(
                self.atom, self.codeloc, "" if not self.dummy else "dummy"
            )
        else:
            return "<Definition {{Tags:{}, Atom:{}, Codeloc:{}}}{}>".format(
                repr(self.tags), self.atom, self.codeloc, "" if not self.dummy else " dummy"
            )

    def __str__(self):
        pretty_tags = "\n".join([str(tag) for tag in self.tags])
        return f"Definition:\n" f"Atom: {self.atom}\n" f"CodeLoc: {self.codeloc}\n" f"Tags: {pretty_tags}"

    def __hash__(self):
        if self._hash is None:
            self._hash = hash((self.atom, self.codeloc))
        return self._hash

    @property
    def offset(self) -> int:
        if isinstance(self.atom, Register):
            return self.atom.reg_offset
        elif isinstance(self.atom, MemoryLocation):
            if isinstance(self.atom.addr, SpOffset):
                return self.atom.addr.offset
            else:
                return self.atom.addr
        else:
            raise ValueError("Unsupported operation offset on %s." % type(self.atom))

    @property
    def size(self) -> int:
        if isinstance(self.atom, Register):
            return self.atom.size
        elif isinstance(self.atom, MemoryLocation):
            return self.atom.bits // 8
        else:
            raise ValueError("Unsupported operation size on %s." % type(self.atom))

    def matches(self, **kwargs) -> bool:
        """
        Return whether this definition has certain characteristics.

        """
        return DefinitionMatchPredicate.construct(**kwargs).matches(self)


import numpy as np
import ase.io
from nice.blocks import *
from nice.utilities import *
from nice.rascal_coefficients import process_structures
import copy
from rascal.representations import SphericalInvariants

def get_nice_powerspectrum():
    return StandardSequence([
        StandardBlock(None, None, None,
                      ThresholdExpansioner(num_expand=None, mode='invariants'),
                      None, None)
    ],
                            initial_scaler=None)

def get_nice_ps_kernel(structures, hypers):

    all_species = get_all_species(structures)
    coefficients = get_spherical_expansion(structures, hypers, all_species, split_by_central_specie=False,
                                           show_progress = False)
    nice = get_nice_powerspectrum()
    nice.fit(coefficients)
    nice_ps = nice.transform(coefficients, return_only_invariants = True)[2]
    return nice_ps.dot(nice_ps.T)

def get_rascal_ps_kernel(structures, hypers):
    structures = process_structures(structures)
    soap = SphericalInvariants(**hypers)
    librascal_ps = soap.transform(structures).get_features(soap)
    return librascal_ps.dot(librascal_ps.T)

def test_powerspectrum_kernels(epsilon = 1e-10):
    structures = ase.io.read('../reference_configurations/methane_100.extxyz', index = ':')
    HYPERS = {
        'interaction_cutoff': 6.3,
        'max_radial': 5,
        'max_angular': 5,
        'gaussian_sigma_type': 'Constant',
        'gaussian_sigma_constant': 0.3,
        'cutoff_smooth_width': 0.3,
        'radial_basis': 'GTO',

    }

    HYPERS_PS = copy.deepcopy(HYPERS)
    HYPERS_PS['normalize'] = False
    HYPERS_PS['soap_type'] = 'PowerSpectrum'

    nice_kernel = get_nice_ps_kernel(structures, HYPERS)
    rascal_kernel = get_rascal_ps_kernel(structures, HYPERS_PS)

    nice_kernel = np.reshape(nice_kernel, [-1])
    rascal_kernel = np.reshape(rascal_kernel, [-1])

    mask = rascal_kernel > epsilon
    nice_kernel = nice_kernel[mask]
    rascal_kernel = rascal_kernel[mask]

    ratios = nice_kernel / rascal_kernel
    discrepancy = (np.max(ratios) - np.min(ratios)) / np.mean(ratios)
    assert discrepancy < epsilon


import abc
from copy import deepcopy
from inspect import isclass
from typing import Any, Dict, List, Optional, Set, Type, Union

FlatDict = Dict[str, Any]
Config = Dict[str, Any]

class ConfigManager:

    def __init__(
        self,
        init_cfg: Config,
        components: Optional[List[Union[Type["ConfigComponent"], "ConfigComponent"]]] = None
    ):
        self.init_cfg = init_cfg
        self._components: List[Union[Type["ConfigComponent"], "ConfigComponent"]] = []

        if components is not None:
            for component in components:
                self.add_component(component)

    def add_component(self, component: Union[Type["ConfigComponent"], "ConfigComponent"]) -> None:
        if any(c.name() == component.name() for c in self._components):
            raise ValueError(f"Duplicate '{component.name}' ConfigComponent")

        self._components.append(component)
        component.manager = self

    @property
    def components(self) -> List[Union[Type["ConfigComponent"], "ConfigComponent"]]:
        return self._components

    def configure_config(self) -> Config:
        cfg = deepcopy(self.init_cfg)

        for component in self.components:
            if isclass(component):
                component_obj = component()
            else:
                component_obj = component

            if component_obj.is_enabled(cfg):
                cfg.update(component_obj.configure(cfg))

        return cfg

class ConfigComponent(abc.ABC):

    def __init__(self):
        self._manager: Optional["ConfigManager"] = None

    @property
    def manager(self) -> "ConfigManager":
        if self._manager is None:
            raise RuntimeError("")

        return self._manager

    @manager.setter
    def manager(self, manager: "ConfigManager") -> None:
        self._manager = manager

    def is_enabled(self, cfg: Config) -> bool:
        return True

    @classmethod
    def name(cls) -> str:
        return cls.__name__

    @abc.abstractmethod
    def configure(self, cfg: Config) -> FlatDict:
        pass

class Component1(ConfigComponent):

    def configure(self, cfg) -> FlatDict:
        return {"foo": "bar"}

class Component2(ConfigComponent):

    def configure(self, cfg) -> FlatDict:
        return {cfg["foo"]: "baz"}

if __name__ == "__main__":
    c1 = Component1()
    cm = ConfigManager(init_cfg={}, components=[c1, Component2])
    cfg = cm.configure_config()
    print(cfg)


from sklearn.svm import OneClassSVM
from matplotlib import pyplot as plt
import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler

# ## Load data

data_url = "C:/Users/Mahee/Documents/Group proj/Data/1128_1419_fx_tenor_curves_data.csv"
data = pd.read_csv(data_url)

# #### Make column for log returns of spot rate

data['Log Returns'] = np.log(data['Spot Rate'].shift(1)/data['Spot Rate'])

currency = 'GBP'
filteredData = data[data['Currency'] == currency]
dates = filteredData['Date']
spot_rates = filteredData['Spot Rate']
spot_log_ret = filteredData['Log Returns']

# resets index no. so starts from 0 (when picking a currency index does not start from 0)
dates=dates.reset_index(drop=True)
spot_rates=spot_rates.reset_index(drop=True)
spot_log_ret=spot_log_ret.reset_index(drop=True)

# ## One class SVM algorithm

def ocsvm( spot_rate_df, nu, kernel, gamma):
    spot_rates_flat = spot_rate_df.values.reshape(-1,1) # flatten data
    scaler = StandardScaler()
    spot_rates_scaled = scaler.fit_transform(spot_rates_flat) # scale data
    scaled_df = pd.DataFrame(spot_rates_scaled)

    model = OneClassSVM(nu=nu, kernel=kernel, gamma=gamma)
    model.fit(scaled_df) # train data

    # Create a column for the anomaly
    spot_rate_df['anomaly'] = pd.Series(model.predict(scaled_df)).map({1: False, -1: True})
    return spot_rate_df

def plotAnomaly(df_anomaly, dates, spot_rates):
    plt.plot(dates, spot_rates)
    plt.xlabel("Dates")
    plt.ylabel("Spot Rate")
    plt.title( "Spot rate for GBP")
    anomalyDates=dates.loc[df_anomaly['anomaly']] # locates all dates where anomaly is true
    anomalySpots = spot_rates.loc[df_anomaly['anomaly']] # locates all spot rates where anomaly is true
    plt.scatter(anomalyDates, anomalySpots, color='red')
    plt.show()

# ## Results

df_anomaly_2 = ocsvm(spot_rates.copy(), 0.01, 'rbf', 0.01)
plotAnomaly(df_anomaly_2, dates, spot_rates)

# #### Log return results

df_anomaly_3 = ocsvm(spot_log_ret.copy(), 0.001, 'rbf', 0.01)
plotAnomaly(df_anomaly_3, dates, spot_rates)
plt.plot(dates[1:], spot_log_ret[1:])
plt.scatter(dates[1:].loc[df_anomaly_3['anomaly'][1:]], spot_log_ret[1:].loc[df_anomaly_3['anomaly'][1:]], color='red')
plt.show()


from django.core.management.base import BaseCommand
from retail_api.models import Product, Store, Staff
import csv

class Command(BaseCommand):

    def handle(self, *args, **kwargs):
        with open('staff.csv') as file:
            csv_obj = csv.reader(file)

            next(csv_obj)

            for line in csv_obj:
                store_id = int(line[2])
                this_store = Store.objects.get(id=store_id)
                Staff.objects.create( first_name=line[0], last_name=line[1], store=this_store, active=True)

        self.stdout.write("all staff have bee added")


from __future__ import print_function

import os
import sys
import logging

class AIDbg(object):

    def __init__(self, kernal, user_uri, realm, verbose):

        self.kernal   = kernal
        self.user_uri = user_uri
        self.realm    = realm
        self.verbose  = verbose

        if self.verbose:
            logging.getLogger().setLevel(logging.DEBUG)
        else:
            logging.getLogger().setLevel(logging.INFO)

    def print_help(self):
        print (":h          help")
        print (":c <skills> compile <skills>")
        print (":m          show memory / context")
        print (":t          %s prolog tracing" % ('disable' if self.run_trace else 'enable'))
        print (":v          verbose logging %s" % ('off' if self.verbose else 'on'))
        print (":q          quit")

    def process_command(self, line):

        if line == ":h":
            self.print_help()

        # FIXME: we'd have to force-reload the python module
        elif line[:2] == ":c":

            parts = line.split(' ')
            if len(parts) < 2:
                logging.error('?usage')
                self.print_help()
                return

            self.kernal.compile_skill_multi (parts[1:])

        elif line == ":m":

            print ("ctx.user   = %s" % self.ctx.user)
            print ("ctx.realm  = %s" % self.ctx.realm)
            print ("ctx.lang   = %s" % self.ctx.lang)
            memd = self.kernal.mem_dump(self.ctx.realm)
            for k, v, score in memd:
                print(u'MEM(%-8s):    %-20s: %s (%f)' % (self.ctx.realm, k, v, score))
            memd = self.kernal.mem_dump(self.ctx.user)
            for k, v, score in memd:
                print(u'MEM(%-8s):    %-20s: %s (%f)' % (self.ctx.user, k, v, score))

        elif line == ":t":
            self.run_trace = not self.run_trace

        elif line == ":v":
            self.verbose = not self.verbose
            if self.verbose:
                logging.getLogger().setLevel(logging.DEBUG)
            else:
                logging.getLogger().setLevel(logging.INFO)

        else:
            logging.error("? command error ('%s')" % line)
            self.print_help()

    def run(self):

        self.ctx       = self.kernal.create_context(user=self.user_uri, realm=self.realm)
        self.run_trace = False

        while True:

            line = raw_input ('dbg (:h for help)> ')

            if not line:
                continue
            if len(line)<1:
                continue

            if line == ':q':
                break

            if line[0] == ':':
                self.process_command(line)
                continue

            out, score, action = self.kernal.process_input(self.ctx, line, run_trace=self.run_trace)

            if action:
                logging.info(u'RESP: [%6.1f] %s | action: %s' % (score, out, unicode(action)))
            else:
                logging.info(u'RESP: [%6.1f] %s ' % (score, out))


#!/usr/bin/env python3

import os
from wiktionaryparser import WiktionaryParser
import simplejson as json
import spacy

# Init
wikparser = WiktionaryParser()
nlp = spacy.load("en_core_web_sm")
trad = {
        "NOUN":"noun",
        "ADJ":"adjective",
        "VERB":"verb",
        }

words_to_fetch = [
        ("apple", "NOUN"),
        ("dragon", "NOUN"),
        ("lion", "NOUN"),
        ("house", "NOUN"),
        ("town", "NOUN"),
        ("plane", "NOUN"),
        ("cloud", "NOUN"),
        ("weather", "NOUN"),
        ("pirate", "NOUN"),
        ("yeti", "NOUN"),
        ("truck", "NOUN"),
        ("star", "NOUN"),
        ("dinosaur", "NOUN"),
        ("fairy", "NOUN"),
        ("ogre", "NOUN"),
        ("forest", "NOUN"),
        ("elf", "NOUN"),
        ("orc", "NOUN"),
        ("demon", "NOUN"),
        ("earth", "NOUN"),
        ("ground", "NOUN"),
        ("table", "NOUN"),
        ("chair", "NOUN"),
        ("insect", "NOUN"),
        ("beach", "NOUN"),
        ("ocean", "NOUN"),
        ("river", "NOUN"),
        ("sea", "NOUN"),
        ("cactus", "NOUN"),
        ("baobab", "NOUN"),
        ("spider", "NOUN"),
        ("volcano", "NOUN"),
        ("robot", "NOUN"),
        ("heater", "NOUN"),
        ("cave", "NOUN"),
        ("torch", "NOUN"),
        ("tank", "NOUN"),
        ("lion", "NOUN"),
        ("unicorn", "NOUN"),
        ("skeleton", "NOUN"),
        ("monster", "NOUN"),
        ("goose", "NOUN"),
        ("univers", "NOUN"),
        ("world", "NOUN"),
        ("cake", "NOUN"),
        ("hero", "NOUN"),
        ("castle", "NOUN"),
        ("dungeon", "NOUN"),
        ("ant", "NOUN"),
        ("carpet", "NOUN"),
        ("hell", "NOUN"),
        ("jungle", "NOUN"),
        ("ghost", "NOUN"),
        ("vampire", "NOUN"),
        ("burning", "ADJ"),
        ("static", "ADJ"),
        ("vegetal", "ADJ"),
        ("exit", "ADJ"),
        ("exit", "NOUN"),
        ("exit", "VERB"),
        ]
out_dict = {}

if os.path.exists('dict.json'):
    print("Reading dict.json")
    with open('dict.json', 'r') as f:
        out_dict = json.load(f)
    print("Updating words_to_fetch")
    for w in out_dict:
        tw = eval(w)
        if tw in words_to_fetch:
            words_to_fetch.remove(tw)
        # Language processing
        doc = nlp(out_dict[w])
        for token in doc:
            if token.pos_ in trad:
                child = (token.lemma_, token.pos_)
                if not str(child) in out_dict and not child in words_to_fetch:
                    words_to_fetch.append(child)

def fetch(word, part):
    # Requesting word on the dictionary
    try:
        wikword = wikparser.fetch(word)
    except:
        print("Couldn't fetch word :", word)
        return

    # Finding first definition
    definition = None
    for eth in wikword:
        definitions = eth["definitions"]
        for d in definitions:
            if d['partOfSpeech'] == trad[part]:
                definition = d['text'][1]
                break
        if definition:
            break

    if not definition:
        print("Getting first definition", (word, part))
        for eth in wikword:
            definitions = eth["definitions"]
            for d in definitions:
                definition = d['text'][1]
                break
            if definition:
                break

    if definition:
        out_dict[str((word, part))] = definition

        # Language processing
        doc = nlp(definition)
        for token in doc:
            if token.pos_ in trad:
                child = (token.lemma_, token.pos_)
                if not str(child) in out_dict and not child in words_to_fetch:
                    words_to_fetch.append(child)
    else:
        print("Couldn't find definition", (word, part))

print("Fetching : ", words_to_fetch)
for word, part in words_to_fetch:
    fetch(word, part)
    if len(out_dict) % 100 == 0:
        print("Saving : {}".format(len(out_dict)))
        with open('dict.json', 'w') as f:
            json.dump(out_dict, f, sort_keys=True, indent=4)

print("Saving : {}".format(len(out_dict)))
with open('dict.json', 'w') as f:
    json.dump(out_dict, f, sort_keys=True, indent=4)


import pygame,pygame.mixer
from copy import deepcopy
from game_edit import edit,get_tuple,get_bool
from random import random
from math import sqrt,hypot,pi

def load_config():
    """ Permite obtener datos de el archivo config
    """
    global Width,Height,color_palabras,color_particulas1,color_particulas2,FPS,dt,color_fondo,Kv_max,Kv_min,Kr_max,Kr_min,Kd,n,merge,bound,color_medio,Max_Points
    Width = get_tuple("Data\config.txt",0)[0]
    Height = get_tuple("Data\config.txt",0)[1]
    FPS = get_tuple("Data\config.txt",1)[0]
    dt = get_tuple("Data\config.txt",2)[0]
    color_particulas1 = get_tuple("Data\config.txt",3)
    color_particulas2 = get_tuple("Data\config.txt",4)
    color_palabras = get_tuple("Data\config.txt",5)
    color_fondo = get_tuple("Data\config.txt",6)
    Kv_max = get_tuple("Data\config.txt",7)[0]
    Kv_min = get_tuple("Data\config.txt",8)[0]
    Kr_max = get_tuple("Data\config.txt",9)[0]
    Kr_min = get_tuple("Data\config.txt",10)[0]
    Kd = get_tuple("Data\config.txt",11)[0]
    n = get_tuple("Data\config.txt",12)[0]
    merge = get_bool("Data\config.txt",13)
    bound = get_bool("Data\config.txt",14)
    Max_Points = get_tuple("Data\config.txt",15)[0]
    u = []
    for i in range(3):
        u.append(int((color_particulas1[i]+color_particulas2[i])//2))
    color_medio = tuple(u)
    print("\n**Datos cargados exitosamente**")

def out(pos,r):
    """ Recibe una posicion
    y un radio para determinar si un circulo esta afuera del cuadrado
    """
    if pos[0]-r > L or pos[1]-r > L:
        return True
    if pos[0]+r < 0 or pos[1]+r < 0:
        return True
    return False

def to_display(pos,L):
    """ Recibe una posicion y el largo de un cuadrado(logico) y retorna la posicion en donde se tiene que renderizar
    """
    return [int(Margen_x +(pos[0]*Largo_caja/L)),int(Margen_y+(pos[1]*Largo_caja/L))]
def to_logic(pos,L):
    """ Recive una posicion de el display y lo pasa a un punto del cuadrado(logico)
    """
    return [((L*(pos[0]-Margen_x))/Largo_caja),((L*(pos[1]-Margen_y))/Largo_caja)]

def norm_coef(coef):
    """ Acota el coeficiente de restitucion, para que no sean colisiones super elasticas
    """
    if coef > 1:
        return 1
    elif coef < 0:
        return 0
    return coef

def general_d(pos1,pos2):
    """ Distancia entre dos puntos
    """
    u = 0
    for i in range(2):
        u+= (pos1[i] - pos2[i])**2
    return sqrt(u)

def d(P1,P2):
    """ Distancia entre dos particulas
    """
    return general_d(P1.pos,P2.pos)

def atraccion(P1,P2,G):
    """ Atraccion gravitatoria entre dos particulas con una constante G
    retorna aceleracion que experimenta P1
    """
    for i in range(2):
        delta = P2.pos[i] - P1.pos[i]
        P1.a[i] += (P2.m*G*delta)/(r**3)# r = d(P1,P2), ya calculado (para reducir cantidad de calculos)

def merge_particulas(P1,P2,particulas):
    """ Permite que dos particulas se unan en una
    """
    P3_color = []
    for i in range(3):
        P3_col = (P1.color[i] + P2.color[i])//2
        P3_color.append(int(P3_col))
    if P1.m**2/P1.r < P2.m**2/P2.r:
        P1,P2 = P2,P1
    P1d = ((P1.m/P1.r))
    P1.coef = (P1.coef + P2.coef)/2
    for i in range(2):
        P1.v[i] = ((P1.v[i]*P1.m) + (P2.v[i]*P2.m)) / (P1.m + P2.m)
    P1.m += P2.m
    P1.r = P1.m/P1d
    P1.color = P3_color
    particulas.remove(P2)

def gravedad_caja(P,g):
    """ Gravedad de la caja sobre una particula
    """
    for i in range(2):
        P.a[i] += g[i]

def colisionar(P1,P2):
    """ Colision entre dos particulas, puede ser inelastica o elastica
    y ademas evita que las particulas se queden pegadas
    """
    C = (P1.coef + P2.coef) / 2
    E = r-(P1.r + P2.r)# r = d(P1,P2), ya calculado
    for i in range(2):
        vf1 = (C*P2.m*(P2.v[i]-P1.v[i]) + P1.m * P1.v[i] + P2.m * P2.v[i]) / (P1.m + P2.m)
        P2.v[i] += (P1.m/P2.m) * (P1.v[i]-vf1)
        P1.v[i] = vf1
        collision_pos = ((P1.pos[i] * P2.r) + (P2.pos[i] * P1.r)) / (P2.r + P1.r)
        error = E * (P1.pos[i] - collision_pos)/(P1.r) # para evitar que queden atrapados juntos
        if P1.m < P2.m:
            P1.pos[i] -= 2*error
        else:
            P2.pos[i] += 2*error

def toggle(boolean):
    """ invierte boleano
    """
    if boolean:
        return False
    else:
        return True

def neg():
    """ permite que la direccion sea aleatoria
    """
    u = random()
    if u < 0.5:
        return -1
    return 1

def message_to_screen(msg,color,x,y):
    """ Recibe un mensaje con un color y lo pone en la pantalla en una posicion indicada
    """
    screen_text = font.render(msg, True, color)
    display.blit(screen_text, (x,y))

def pause():
    """ Pausa simulacion
    Sonido Incluido
    """
    message_to_screen("Pausa :3",color_palabras,Width//2,int((Height//2)+Ly))
    pygame.display.update()
    sound = pygame.mixer.Sound("Data\smb_pause.wav")
    pygame.mixer.music.pause()
    pygame.mixer.Sound.play(sound)
    pause = True
    while pause:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                simulacion = False
                pygame.quit()
                quit()
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_p:
                    pause = False
                    pygame.mixer.music.unpause()

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

class Particula:
    """ Particula que tiene como atributos
    posicion, Velocidad, Aceleracion, Masa, Radio, Color, Coeficiente de restitucion y
    Cola

    Se puede saber si una posicion esta dentro de ella,
    Puede rebotar si se le da una pared,
    Se desplaza y acelera
    Puede colisionar y unirse con otra particula
    Es atraida a todas las particulas
    y por ultimo puede tener una propia cola
    """
    def __init__(self,pos,velocidad,masa,radio,coef,color):
        self.pos = pos
        self.v = velocidad
        self.m = masa
        self.r = abs(radio)
        self.coef = norm_coef(coef)
        self.color = color
        self.cola = [to_display(deepcopy(self.pos),L)]

    def adentro(self,pos):
        """ Determina si una posicion esta adentro
        """
        u = 0
        for i in range(2):
            u += (self.pos[i] - pos[i])**2
        if u <= self.r**2:
            return True
        return False

    def get_cola(self):
        """ Obtiene cola, con un numero maximo de puntos
        """
        nueva_parte = to_display(deepcopy(self.pos),L)
        self.cola.insert(0,nueva_parte)
        while len(self.cola) > Max_Points:
            last = self.cola[-1]
            self.cola.remove(last)

    def desplazar(self,dt):
        """ Se desplaza en intervalos dt
        """
        for i in range(2):
            self.pos[i] += self.v[i]*dt

    def acelerar(self,dt):
        """ Acelera en intervalos dt
        """
        for i in range(2):
            self.v[i] += self.a[i]*dt

    def rebotar_pared(self,L):
        """ Rebota en una pared de largo L
        y ademas devuelve la particula si se encuentra
        fuera de el contenedor, cada colision es ponderada por
        el coeficiente de restitucion
        """
        for i in range(2):
            if self.pos[i]-self.r <= 0:
                self.v[i] *= -self.coef
                d = -self.pos[i]+self.r
                self.pos[i] += d
            elif self.pos[i]+self.r >= L:
                self.v[i] *= -self.coef
                d = L-(self.pos[i]+self.r)
                self.pos[i] += d

    def comprobar_colision(self,particulas,merge):
        """ Comprueba si la particula self, colisiono con alguna
        particula diferente de ella en una lista, y dependiendo de la configuracion
        puede colisionar o unirse
        """
        for particula in particulas:
            if particula is not self:
                global r
                r = d(self,particula)
                if r <= self.r + particula.r:
                    if not merge:
                        colisionar(self,particula)
                    else:
                        merge_particulas(self,particula,particulas)

    def atraccion_neta(self,particulas,G,g):
        """ Asigna una aceleracion neta a la particula
        que depende de la atraccion gravitatoria y la gravedad de la caja
        """
        self.a = [0,0]
        gravedad_caja(self,g)
        if G != 0:
            for particula in particulas:
                if particula is not self:
                    global r
                    r = d(self,particula)
                    if r > self.r*2 + particula.r*2:
                        atraccion(self,particula,G)

class Caja:
    """ Caja que contiene consigo objetos particulas que interaccionan entre ellas y la caja.
    Tiene como atributos ademas: un largo, una aceleracion de gravedad como si se tratase de la tierra
    una constante gravitatoria, y puede controlar cosas como los tipos de colisiones, puede permitir
    que hayan colas y campos vectoriales, o que existan o no limites en la caja.

    """
    def __init__(self,L,g = [0,0],G = 0,particulas = [], merge= True, bound = False,colas = False):
        self.L = L
        self.g = g
        self.G = G
        self.particulas = particulas
        self.vector_field = False
        self.merge = merge
        self.bound = bound
        self.colas = colas

    def generar(self, n, velocidad, radio, densidad, coef):
        """ Genera n particulas de manera aleaotira dependiendo de los limites que se le dan
        """
        v_max = velocidad[0]
        v_min = velocidad[1]
        r_max = radio[0]
        r_min = radio[1]
        for i in range(n):
            pos = []
            v = []
            r = random()*(r_max-r_min)+ r_min
            m = r * densidad
            for k in range(2):
                v_eje = ((random()*(v_max-v_min))+v_min)*neg()
                pos_eje = random()*self.L
                v.append(v_eje)
                pos.append(pos_eje)
            u = Particula(pos,v,m,r,coef,color_particulas1)
            self.particulas.append(u)

    def agregar(self,pos,velocidad,masa,radio,coef,color):
        """ Agrega una particula con determinadas caracteristicas
        """
        p = Particula(pos,velocidad,masa,radio,coef,color)
        self.particulas.append(p)

    def update_logic(self,dt):
        """ Actualiza todos los parametros logicos de las particulas dentro de la caja
        """
        for particula in self.particulas:
            particula.comprobar_colision(self.particulas,self.merge)
            if self.bound:
                particula.rebotar_pared(self.L)
            particula.desplazar(dt)
            particula.atraccion_neta(self.particulas,self.G,self.g)
            particula.acelerar(dt)

    def select(self,pos):
        """ Retorna una particula selecionada en una posicion
        """
        for particula in self.particulas:
            if particula.adentro(pos) == True:
                return particula

    def agarrar(self,pos,dv):
        """ Permite acelerar particulas de una posicion pos con un dv
        """
        for particula in self.particulas:
            if particula.adentro(pos) == True:
                for i in range(2):
                    particula.v[i] += dv[i]

    def detener(self,pos):
        """ Detiene una particula de una posicion determinada
        """
        for particula in self.particulas:
            if particula.adentro(pos) == True:
                for i in range(2):
                    particula.v[i] = 0

    def eliminar(self,pos):
        """ Elimina una particula en una determinada posicion
        """
        for particula in self.particulas:
            if particula.adentro(pos) == True:
                self.particulas.remove(particula)

    def probe(self,pos):
        """ Permite calcular como una particula de prueba aceleraria en un punto determinado
        """
        a = [0,0]
        try:
            for particula in self.particulas:
                r = general_d(pos,particula.pos)
                for i in range(2):
                    delta = particula.pos[i] - pos[i]
                    a[i] += (particula.m*self.G*delta)/(r**3) + self.g[i]
            return a
        except ZeroDivisionError:
            return "infinte","infinte"

    def cambiar_coef(self,d):
        """ Permite cambiar el coeficiente de restitucion de todas las particulas
        """
        for particula in self.particulas:
            particula.coef += d
            if particula.coef > 1:
                particula.coef = 1
            elif particula.coef < 0:
                particula.coef = 0

    def rapidez_promedio(self):
        """ Retorna rapidez promedio
        """
        v = 0
        i = 0
        for particula in self.particulas:
            u = 0
            for i in range(2):
                u += particula.v[i]**2
            v+= sqrt(u)
            i += 1
        try:
            return v/i
        except:
            return 0

    def get_info(self,particula):
        """ Muestra en pantalla algunos parametros de una particula selecionada
        """
        x1 = Margen_x
        try:
            pos_plot ="Pos:"+"(" + str("{:.1e}".format(round(particula.pos[0]))) + "," + str("{:.2e}".format(particula.pos[1]))+")"
            v_plot = "V:"+"(" + str("{:.1e}".format(round(particula.v[0]))) + "," + str("{:.2e}".format(particula.v[1]))+")"
            a_plot = "A:"+"("+str("{:.1e}".format(round(particula.a[0]))) + "," + str("{:.2e}".format(particula.a[1]))+")"
            m_plot = "M:"+str("{:.1e}".format(round(particula.m)))
            r_plot = "R:"+str("{:.1e}".format(round(particula.r)))
            x2 = x1 + len(pos_plot)*Kx
            x3 = x2 + len(v_plot)*Kx
            x4 = x3 + len(a_plot)*Kx
            x5 = x4 + len(m_plot)*Kx
            message_to_screen(pos_plot,color_palabras,x1,Margen_y//10)
            message_to_screen(v_plot,color_palabras,x2,Margen_y//10)
            message_to_screen(a_plot,color_palabras,x3,Margen_y//10)
            message_to_screen(m_plot,color_palabras,x4,Margen_y//10)
            message_to_screen(r_plot,color_palabras,x5,Margen_y//10)
            p = to_display((particula.pos[0],particula.pos[1]),L)
            r = int(particula.r*Largo_caja/(self.L*4))
            pygame.draw.circle(display,color_particulas2,p,r,r//2)

        except:
            pos_plot = "Pos: --"
            v_plot = "V: --"
            a_plot = "A: --"
            m_plot = "M: --"
            r_plot = "R: --"
            x2 = x1 + len(pos_plot)*Kx
            x3 = x2 + len(v_plot)*Kx
            x4 = x3 + len(a_plot)*Kx
            x5 = x4 + len(m_plot)*Kx
            message_to_screen(pos_plot,color_palabras,x1,Margen_y//10)
            message_to_screen(v_plot,color_palabras,x2,Margen_y//10)
            message_to_screen(a_plot,color_palabras,x3,Margen_y//10)
            message_to_screen(m_plot,color_palabras,x4,Margen_y//10)
            message_to_screen(r_plot,color_palabras,x5,Margen_y//10)

    def get_vectorPoints(self):
        """ Muestra en pantalla un campo vectorial de la suma neta de las fuerzas
        """
        if (self.G != 0 and self.vector_field) or (self.g !=[0,0] and self.vector_field):
            r = self.L/(n*2)
            u = 4
            v = 2
            for i in range(n):
                for k in range(n):
                    point = (int(r+(2*r*i)),int(r+(2*r*k)))
                    a = self.probe(point)
                    norm_a = sqrt(a[0]**2 + a[1]**2)
                    k = norm_a/Largo_caja*10
                    if k>=0.6:
                        k=0.6
                    try:
                        delta_x = k*a[0]*r/norm_a
                        delta_y = k*a[1]*r/norm_a
                        end_point = to_display((point[0]+delta_x,point[1]+delta_y),self.L)
                        beg_point = to_display((point[0]-delta_x,point[1]-delta_y),self.L)
                        x1 = (-a[0]*u - a[1]*v*k)/norm_a + end_point[0]
                        y1 = (-a[1]*u + a[0]*v*k)/norm_a + end_point[1]
                        x2 = (-a[0]*u + a[1]*v*k)/norm_a + end_point[0]
                        y2 = (-a[1]*u - a[0]*v*k)/norm_a + end_point[1]
                        p1 = (x1,y1)
                        p2 = (x2,y2)
                        pygame.draw.line(display,(0,0,0),end_point,beg_point)
                        pygame.draw.polygon(display, (0,0,0),(end_point,p1,p2))
                    except:
                        pass

    def sim_start(self,dt,t = -1):
        """ Parte simulacion dentro de la caja
        """
        dt_inicial = deepcopy(dt)
        G_inicial = deepcopy(self.G)
        selected_particula = self.particulas[0]
        probe_pos = [self.L/2,self.L/2]

        #Parametros Default iniciales
        troll = False
        press = False
        Up = False
        Down = False
        Right = False
        Left = False
        G_up = False
        G_down = False
        coef_up = False
        coef_down = False
        dt_up = False
        dt_down = False
        rojos = 0

        while t != 0:
            FpsA = clock.get_fps()
            if FpsA == 0:
                FpsA = 0.01
            N = len(self.particulas)
            if N == 0:
                coefEn = 1
            else:
                coefEn = self.particulas[0].coef
#------------------------------------------------------------------------------------------------------// Manejo Eventos
            for event in pygame.event.get():
                posicion_mouse = pygame.mouse.get_pos()
                pos = to_logic(posicion_mouse,self.L) #Posicion normalizada acuerdo a la caja
                rel = pygame.mouse.get_rel()
                if event.type == pygame.QUIT:
                    pygame.quit()
                    quit()

                if event.type == pygame.MOUSEBUTTONDOWN:
                    key = pygame.mouse.get_pressed()
                    if key == (1,0,0):
                        selected_particula = self.select(pos)
                        Vx = 0
                        Vy = 0
                        press = True
                    if key == (0,0,1) or key == (1,0,1):
                        self.detener(pos)
                if press == True:
                    Vx = rel[0]
                    Vy = rel[1]
                    v = [Vx,Vy]
                    self.agarrar(pos,v)

                if event.type == pygame.MOUSEBUTTONUP:
                    press = False

                if event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_k:
                        m = []
                        r = []
                        for particula in self.particulas:
                            m.append(particula.m)
                            r.append(particula.r)
                        try:
                            self.agregar(pos,[0,0],16*max(m), min(r),coefEn,color_particulas2)
                            rojos += 1
                        except:
                            pass
                    if event.key == pygame.K_p:
                        pause()
                    if event.key == pygame.K_UP:
                        Up = True
                    if event.key == pygame.K_DOWN:
                        Down = True
                    if event.key == pygame.K_RIGHT:
                        Right = True
                    if event.key == pygame.K_LEFT:
                        Left = True
                    if event.key == pygame.K_w:
                        coef_up = True
                    if event.key == pygame.K_q:
                        coef_down = True
                    if event.key == pygame.K_x:
                        dt_up = True
                    elif event.key == pygame.K_z:
                        dt_down = True
                    if event.key == pygame.K_s:
                        G_up = True
                    elif event.key == pygame.K_a:
                        G_down = True
                    if event.key == pygame.K_r:
                        self.g = [0,0]
                        dt = dt_inicial
                        self.G = G_inicial
                        for particula in self.particulas:
                            particula.coef = 1
                    if event.key == pygame.K_y:
                        self.G = 0
                    if event.key == pygame.K_t:
                        self.generar(1,((self.L*Kv_max),(self.L*Kv_min)),((self.L*Kr_max),(self.L*Kr_min)),Kd,coefEn)
                    if event.key == pygame.K_BACKSPACE:
                        self.eliminar(pos)
                    if event.key == pygame.K_h:
                        probe_pos = pos
                    if event.key == pygame.K_v:
                        self.vector_field = toggle(self.vector_field)
                    if event.key == pygame.K_m:
                        self.merge = toggle(self.merge)
                    if event.key == pygame.K_b:
                        self.bound = toggle(self.bound)
                    if event.key == pygame.K_c:
                        self.colas = toggle(self.colas)
                    if event.key == pygame.K_f:
                        dt = 0

                if event.type == pygame.KEYUP:
                    if event.key == pygame.K_UP:
                        Up = False
                    elif event.key == pygame.K_DOWN:
                        Down = False
                    if event.key == pygame.K_RIGHT:
                        Right = False
                    elif event.key == pygame.K_LEFT:
                        Left = False

                    if event.key == pygame.K_w:
                        coef_up = False
                    elif event.key == pygame.K_q:
                        coef_down = False

                    if event.key == pygame.K_x:
                        dt_up = False
                    elif event.key == pygame.K_z:
                        dt_down = False

                    if event.key == pygame.K_s:
                        G_up = False
                    elif event.key == pygame.K_a:
                        G_down = False

            if Up:
                self.g[1] -=5/FpsA
            if Down:
                self.g[1] +=5/FpsA
            if Right:
                self.g[0] +=5/FpsA
            if Left:
                self.g[0] -=5/FpsA

            if coef_up and coefEn<0.9999:
                self.cambiar_coef(0.45/FpsA)
            if coef_down and coefEn>0.0001:
                self.cambiar_coef(-0.45/FpsA)

            if dt_up:
                dt += 0.002/FpsA
            if dt_down:
                dt -= 0.002/FpsA
                if dt <= 0:
                    dt = 0
            if G_up:
                self.G += 50/FpsA
            if G_down:
                self.G -= 50/FpsA
#------------------------------------------------------------------------------------------------------// Manejo Eventos
#------------------------------------------------------------------------------------------------------// Display
            display.fill(color_particulas1)
            pygame.draw.rect(display,color_fondo,(Margen_x,Margen_y,Largo_caja,Largo_caja))
            for particula in self.particulas:
                if not out(particula.pos,particula.r):
                    p = to_display((particula.pos[0],particula.pos[1]),self.L)
                    r = int(particula.r*Largo_caja/self.L)
                    pygame.draw.circle(display,particula.color,p,r)
                    if self.colas:
                        particula.get_cola()
                        pygame.draw.lines(display,color_medio,False,particula.cola)
                    else:
                        particula.cola = [deepcopy(to_display(particula.pos,L))]

            try:
                v_promedio = self.rapidez_promedio()
                g_in_display = "g :("+str(round(self.g[0]*10)/10)+", "+str(round(self.g[1]*10)/10)+")"
                N_in_display = "NºBolas :"+str(N)
                V_in_display = "Rapidez media:" +str(round(v_promedio))
                Largo_in_display = "Largo caja:" +str(self.L)
                FPS_in_display = "FPS: "+str(round(FpsA))
                dt_in_display ="dt: " + str("{:.2e}".format(dt))
                G_in_display = "G: "+str("{:.2e}".format(self.G))
                probe_acceleration = self.probe(probe_pos)
                Probe_in_display ="Probe A:"+"(" + str("{:.2e}".format(round(probe_acceleration[0]))) + ", " + str("{:.2e}".format(probe_acceleration[0]))+")"
                coef_in_display = "coef :"+str(round(self.particulas[0].coef*100)/100)
            except:
                pass
            try:
                message_to_screen(g_in_display,color_palabras,Margen_x//10,Margen_y)
                message_to_screen(N_in_display,color_palabras,Largo_caja+((11*Margen_x)//10),Margen_y)
                message_to_screen(V_in_display,color_palabras,Margen_x//10,Height - Margen_y + Ly)
                message_to_screen(Largo_in_display,color_palabras,Largo_caja+((11*Margen_x)//10),Height - Margen_y + Ly)
                message_to_screen(FPS_in_display,color_palabras,Largo_caja+((11*Margen_x)//10),Height - Margen_y)
                message_to_screen(dt_in_display,color_palabras,Largo_caja+((11*Margen_x)//10),Margen_y + Ly)
                message_to_screen(G_in_display,color_palabras,Largo_caja+((11*Margen_x)//10),Margen_y + 2*Ly)
                message_to_screen(coef_in_display,color_palabras,Margen_x//10,Margen_y + Ly)
                message_to_screen(Probe_in_display,color_palabras,(Largo_caja/2)+Margen_x-(len(Probe_in_display)-10)*Kx,Height - (Margen_y/2))
                if self.merge:
                    m = "Merge"
                else:
                    m = "Collide"
                message_to_screen(m,color_palabras,Margen_x//10,Margen_y + Ly*4)
                if self.bound:
                    b = "Con Limte"
                else:
                    b = "Sin Limite"
                message_to_screen(b,color_palabras,Margen_x//10,Margen_y + Ly*5)
            except:
                pass
#------------------------------------------------------------------------------------------------------// Display

            self.update_logic(dt)
            self.get_info(selected_particula)
            self.get_vectorPoints()

            if (rojos >= 4) and not troll:
                troll = True
                pygame.mixer.music.play(-1)
            if troll:
                message_to_screen(":D",color_particulas2,Width//2,Height//2)

            pygame.display.update()
            clock.tick(FPS)
            t -= 1

        print("\n**Simulacion Terminada**")
        pygame.quit()
#---------------------------------------------------------------------------------------------------------------------------------------------------

print("*Simulacion de Colisiones*\n")
Q = input("(R->Cambiar config) / (Enter->continuar): ")
if Q.upper() == "R":
    config_loop = True
    while config_loop:
        print("\n(1)Resolucion\n(2)FPS_lock\n(3)dt\n(4)Colores\n(5)Parametros de particulas ingresadas despues de partir\n(6)Numero de puntos de prueba para el campo vectorial\n(7)Permitir Merge\n(8)Rebote con las paredes\n(9)Numero de puntos en las colas")
        Q = input("\nIngrese configuracion a cambiar, exit para salir: ")
        if Q == "":
            break
        if Q.upper() == "EXIT":
            break
        elif int(Q) == 1:
            print("->Resolucion Actual:",get_tuple("Data\config.txt",0))
            I = input("Cambiar? ")
            if I.upper() == "SI":
                X = input("->Ingrese resolucion de la forma (A,B): ")
                edit("Data\config.txt",0,X)
        elif int(Q) == 2:
            print("->FPS Actual:",get_tuple("Data\config.txt",1)[0])
            I = input("Cambiar? ")
            if I.upper() == "SI":
                X = input("->Ingrese FPS_lock: ")
                edit("Data\config.txt",1,X)
        elif int(Q) == 3:
            print("->dt Actual:",get_tuple("Data\config.txt",2)[0])
            I = input("Cambiar? ")
            if I.upper() == "SI":
                X = input("->Ingrese dt: ")
                edit("Data\config.txt",2,X)
        elif int(Q) == 4:
            print("\nColores(RGB):\n(1)Particula1:",get_tuple("Data\config.txt",3),"\n(2)Particula2:",get_tuple("Data\config.txt",4),"\n(3)Palabras:",get_tuple("Data\config.txt",5),"\n(4)Fondo:",get_tuple("Data\config.txt",6))
            I = input("\nNºColor a cambiar(Enter para pasar): ")
            if I == "":
                pass
            elif int(I) in range(1,5):
                col = input("->Ingrese color forma(R,G,B): ")
                edit("Data\config.txt",int(I)+2,col)
        elif int(Q) == 5:
            print("\nParametros:\n(1)Velocidad_max/Largo_caja:",get_tuple("Data\config.txt",7)[0],"\n(2)Velocidad_min/Largo_caja:",get_tuple("Data\config.txt",8)[0])
            print("(3)Radio_max/Largo_caja:",get_tuple("Data\config.txt",9)[0],"\n(4)Radio_min/Largo_caja:",get_tuple("Data\config.txt",10)[0],"\n(5)Densidad/Largo_caja:",get_tuple("Data\config.txt",11)[0])
            I = input("\nNºParametro a cambiar(Enter para pasar): ")
            if I == "":
                pass
            elif int(I) in range(1,6):
                col = input("->Ingrese Nuevo Parametro: ")
                edit("Data\config.txt",int(I)+6,col)
        elif int(Q) == 6:
            print("->n Actual:",get_tuple("Data\config.txt",12)[0])
            I = input("Cambiar? ")
            if I.upper() == "SI":
                X = input("->Ingrese n: ")
                edit("Data\config.txt",12,X)
        elif int(Q) == 7:
            print("->Merge:",get_bool("Data\config.txt",13))
            I = input("Cambiar? ")
            if I.upper() == "SI":
                X = input("->Ingrese 1 para True o 0 para False: ")
                edit("Data\config.txt",13,X)
        elif int(Q) == 8:
            print("->Bound:",get_bool("Data\config.txt",14))
            I = input("Cambiar? ")
            if I.upper() == "SI":
                X = input("->Ingrese 1 para True o 0 para False: ")
                edit("Data\config.txt",14,X)
        elif int(Q) == 9:
            print("->Puntos maximos:",get_tuple("Data\config.txt",15)[0])
            I = input("Cambiar? ")
            if I.upper() == "SI":
                X = input("->Ingrese nuevo valor: ")
                edit("Data\config.txt",15,X)
        else:
            pass

load_config()
fontl = min(Width,Height)//30
maxLen = 16
Largo_caja = min(Width,Height)//1.1
if (Width-Largo_caja)/2 < maxLen * fontl/2.5:
    fontl = int((2.5*(Width - Largo_caja))/(2*maxLen))
Kx = fontl/2.5
Ky = fontl/1.818
Ly = Ky * 1.6
Margen_x = (Width - Largo_caja)//2
Margen_y = (Height - Largo_caja)//2
pygame.init()
pygame.mixer.init()

pregunta = True
while pregunta:
    pregunta = False
    Q = input("\nQuieres ocupar los parametros default?(recomendado): ")
    if Q.upper() == "NO":
        try:
            L = float(input("\n->Ingrese el largo de el contenedor cuadrado: "))
            N = int(input("->Ingrese la cantidad de pelotas a generar: "))
            gy = float(input("->Ingrese la gravedad en el eje vertical: "))
            gx = float(input("->Ingrese la gravedad en el eje horizontal: "))
            G = float(input("->Ingrese Constante atraccion gravitatoria(0 si no hay): "))
            Vmax = float(input("->Ingrese la velocidad maxima permitida: "))
            Vmin = float(input("->Ingrese la velocidad minima permitida: "))
            fit = False
            while not fit:
                fit = True
                Rmax = float(input("->Ingrese el radio maximo permitido: "))
                if 4*Rmax >= L:
                    print("\nEso definitivamente no cabe en el contenedor de largo",L)
                    fit = False
                elif N*pi*Rmax**2 > 1.5*L**2:
                    print("\nPuede que se vea muy apretada la cosa")
                    Q = input("Continuar? ")
                    if Q.upper() == "NO":
                        fit = False
            Rmin = float(input("->Ingrese el radio minimo permitido: "))
            Densidad = float(input("->Ingrese la densidad de cada pelota: "))
            c = float(input("->Ingrese el coeficiente de restitucion : "))
            t = int(input("->Ingrese el tiempo que dura la simulacion(-1 si es infinito): "))
            V = (Vmax,Vmin)
            r = (Rmax,Rmin)
            g = [gx,gy]
        except:
            print("\nError")
            pregunta = True
    else:
        N = 6
        V = (50,30)
        r = (8,2)
        L = 500
        Densidad = 1
        c = 0.98
        t = -1
        g = [0,0]
        G = 15000
    if pregunta == False:
        print("**------------------------------------**")
        print("Entonces tendriamos:\n")
        print("El largo como",L)
        print("\nEl numero de particulas como",N)
        print("\nLa constante gravitatoria como",G)
        print("\nLa velocidad como",V)
        print("\nEl radio como",r)
        print("\nLa densidad como",Densidad)
        print("\nEl coeficientes de restitucion como",c)
        print("\nEl tiempo de simulacion es",t)
        print("\ndt :",dt)
        Q = input("\nEstas bien con esta configuracion? ")
        if Q.upper() == "NO":
            pregunta = True
            print("\nOkay, te preguntaremos de nuevo")
            input("\n(Enter para continuar)")
        else:
            input("\n(Enter para empezar simulacion)")
            print("\n******************************\n")
            font = pygame.font.SysFont(None, fontl)
            pygame.display.set_caption("Colisiones")
            clock = pygame.time.Clock()
            pygame.mixer.music.load("Data\Troll.mp3")
            display = pygame.display.set_mode((Width,Height))
            contenedor = Caja(L,g,G,[],merge,bound)
            contenedor.generar(N,V,r,Densidad,c)
            contenedor.sim_start(dt,t)


from typing import Sequence, Tuple, List
from tradebot.models.plotter import Plotter
from tradebot.models.ema import Ema

from bokeh.plotting import figure, ColumnDataSource
from bokeh.models import CDSView, HoverTool, BooleanFilter
from bokeh.palettes import Category20
from bokeh.models.formatters import NumeralTickFormatter
import pandas as pd
import talib as ta

RED = Category20[7][6]
GREEN = Category20[5][4]
VBAR_WIDTH = 0.4
DATETIME_FORMAT = "%b %d %H:%M:%S"

class CandleStickPlotter(Plotter):
    def __init__(self,  title, plot_width, plot_height, tools,  toolbar_location="above"):
        super().__init__(title, plot_width, plot_height,
                         tools,  toolbar_location)
        self.graph = None

    def __set_x_axis_labels(self, source, timeframe):
        # p.xaxis.major_label_overrides = {
        #     # i + int(data[x_axis][0]): date.strftime(DATETIME_FORMAT) for i, date in enumerate(pd.to_datetime(data["Date"]))
        #     i: date.strftime(DATETIME_FORMAT) for i, date in enumerate(pd.to_datetime(data["Date"]))
        # }
        lablels = {}
        for i, date in enumerate(pd.to_datetime(source.data["Date"])):
            # print(i)
            lablels[i] = ""
            # print("{} {} ".format("*", timeframe[-1]))
            if(timeframe.endswith("m")):
                lablels[i] = date.strftime("%M")

            elif(timeframe.endswith("h")):

                lablels[i] = date.strftime("%H")

            elif(timeframe.endswith("d")):
                lablels[i] = date.strftime("%d")

            elif(timeframe.endswith("w")):
                lablels[i] = date.strftime("%d")

            elif(timeframe.endswith("M")):
                lablels[i] = date.strftime("%-m %B")

        self.graph.xaxis.major_label_overrides = lablels

    def __add_ema(self, source, x_axis, ema: Ema):

        _ema = ta.EMA(source.data['Close'], timeperiod=ema.timeperiod)
        self.graph.line(x=source.data[x_axis], y=_ema,
                        line_color=ema.linecolor,
                        legend_label="Ema " + str(ema.timeperiod),
                        line_width=ema.linewidth)

    def plot(self, source, timeframe, emas: List[Ema]) -> figure:
        x_axis = "index"
        self.graph = figure(

            plot_width=self.plot_width,
            plot_height=self.plot_height,
            tools=self.tools,
            title=self.title,
            toolbar_location=self.toolbar_location,
        )

        for ema in emas:
            self.__add_ema(source, x_axis, ema)

        inc = source.data["Close"] > source.data["Open"]
        dec = source.data["Open"] > source.data["Close"]

        view_inc = CDSView(source=source, filters=[
            BooleanFilter(inc)])
        view_dec = CDSView(source=source, filters=[
            BooleanFilter(dec)])

        # map dataframe indices to date strings and use as label overrides

        self.__set_x_axis_labels(source, timeframe)

        self.graph.xaxis.bounds = (
            source.data[x_axis][0],  source.data[x_axis][-1])
        self.graph.segment(
            x0=x_axis,
            x1=x_axis,
            y0="Low",
            y1="High",
            color=RED,
            source=source,
            view=view_inc,
        )
        self.graph.segment(
            x0=x_axis,
            x1=x_axis,
            y0="Low",
            y1="High",
            color=GREEN,
            source=source,
            view=view_dec,
        )

        self.graph.vbar(
            x=x_axis,
            width=VBAR_WIDTH,
            top="Open",
            bottom="Close",
            fill_color=GREEN,
            line_color=GREEN,
            source=source,
            view=view_inc,
            name="price",
        )
        self.graph.vbar(
            x=x_axis,
            width=VBAR_WIDTH,
            top="Open",
            bottom="Close",
            fill_color=RED,
            line_color=RED,
            source=source,
            view=view_dec,
            name="price",
        )

        self.graph.yaxis.formatter = NumeralTickFormatter(format="$ 0,0[.]000")
        self.graph.x_range.range_padding = 0.05
        self.graph.xaxis.ticker.desired_num_ticks = 40
        self.graph.xaxis.major_label_orientation = 3.14 / 4

        # Select specific tool for the plot
        price_hover = self.graph.select(dict(type=HoverTool))

        # Choose, which glyphs are active by glyph name
        price_hover.names = ["price"]
        # Creating tooltips
        # ("Datetime", "@Date{%Y-%m-%d}"),
        price_hover.tooltips = [
            ("Datetime", "@Date{%b %d %H:%M:%S}"),
            ("Open", "@Open{$0,0.00}"),
            ("High", "@High{$0,0.00}"),
            ("Low", "@Low{$0,0.00}"),
            ("Close", "@Close{$0,0.00}"),
            ("Volume", "@Volume{($ 0.00 a)}"),
        ]
        price_hover.formatters = {"@Date": "datetime"}
        price_hover.mode = "vline"

        self.graph.legend.location = "top_left"
        self.graph.legend.border_line_alpha = 0
        self.graph.legend.background_fill_alpha = 0
        self.graph.legend.click_policy = "mute"

        #self.graph.grid.grid_line_alpha = 0.3
        self.graph.xaxis.axis_label = 'Date'
        self.graph.yaxis.axis_label = 'Price'
        return self.graph


#!/usr/bin/env python3

import rospy
from nav_msgs.msg import OccupancyGrid, MapMetaData, Path
from geometry_msgs.msg import Twist, Pose2D, PoseStamped
from std_msgs.msg import String
from std_msgs.msg import Int32
from asl_turtlebot.msg import DetectedObject
from visualization_msgs.msg import Marker
import tf
import numpy as np
from numpy import linalg
from utils.utils import wrapToPi
from utils.grids import StochOccupancyGrid2D
from planners import AStar, compute_smoothed_traj
import scipy.interpolate
import matplotlib.pyplot as plt
from controllers import PoseController, TrajectoryTracker, HeadingController
from enum import Enum
import random

from dynamic_reconfigure.server import Server
from asl_turtlebot.cfg import NavigatorConfig

# state machine modes, not all implemented
class Mode(Enum):
    IDLE = 0
    ALIGN = 1
    TRACK = 2
    PARK = 3
    TELEOP = 4
    PICKUP = 5
    STOP = 6

class Navigator:
    """
    This node handles point to point turtlebot motion, avoiding obstacles.
    It is the sole node that should publish to cmd_vel
    """

    def __init__(self):
        rospy.init_node("turtlebot_navigator", anonymous=True)
        self.mode = Mode.IDLE

        # initial state
        self.x_start = None
        self.y_start = None
        self.theta_start = None

        # current state
        self.x = 0.0
        self.y = 0.0
        self.theta = 0.0

        # detected animals
        self.detected_animals = {}

        # goal state
        self.x_g = 3.4
        self.y_g = 0.4
        self.theta_g = 3.14

        self.th_init = 0.0

        # map parameters
        self.map_width = 0
        self.map_height = 0
        self.map_resolution = 0
        self.map_origin = [0, 0]
        self.map_probs = []
        self.occupancy = None
        self.occupancy_updated = False

        # plan parameters
        self.plan_resolution = 0.1
        self.plan_horizon = 4

        # time when we started following the plan
        self.current_plan_start_time = rospy.get_rostime()
        self.current_plan_duration = 0
        self.plan_start = [0.0, 0.0]

        # Robot limits
        self.v_max = 0.3  # maximum velocity
        self.om_max = 1  # maximum angular velocity

        self.v_des = 0.2  # desired cruising velocity
        self.theta_start_thresh = 0.05  # threshold in theta to start moving forward when path-following
        self.start_pos_thresh = (
            0.2  # threshold to be far enough into the plan to recompute it
        )

        # threshold at which navigator switches from trajectory to pose control
        self.near_thresh = 0.2
        self.at_thresh = 0.05
        self.at_thresh_theta = 0.2

        # trajectory smoothing
        self.spline_alpha = 0
        self.spline_deg = 3  # cubic spline
        self.traj_dt = 0.1

        # trajectory tracking controller parameters
        self.kpx = 5
        self.kpy = 5
        self.kdx = 5
        self.kdy = 5

        # rescue parameters
        self.min_pickup_time = 3 #sec
        self.pickup_start = None
        self.rescue_phase = False
        self.num_rescue_targets = 0
        self.rescue_targets = {} # animals to rescue (key = animal, value = position)
        self.rescue_target = None # current animal to rescue

        # stop sign params
        self.stop_sign_start = None
        self.stop_time = 5
        self.last_stop_sign = rospy.get_rostime()

        # teleop parameters
        self.V_teleop = 0.0
        self.om_teleop = 0.0

        # heading controller parameters
        self.kp_th = 1.0

        self.traj_controller = TrajectoryTracker(
            self.kpx, self.kpy, self.kdx, self.kdy, self.v_max, self.om_max
        )
        self.pose_controller = PoseController(
            0.0, 0.0, 0.0, self.v_max, self.om_max
        )
        self.heading_controller = HeadingController(self.kp_th, self.om_max)

        self.nav_planned_path_pub = rospy.Publisher(
            "/planned_path", Path, queue_size=10
        )
        self.nav_smoothed_path_pub = rospy.Publisher(
            "/cmd_smoothed_path", Path, queue_size=10
        )
        self.nav_smoothed_path_rej_pub = rospy.Publisher(
            "/cmd_smoothed_path_rejected", Path, queue_size=10
        )
        self.nav_vel_pub = rospy.Publisher("/cmd_vel", Twist, queue_size=10)

        self.animal_noise_pub = rospy.Publisher('/animal_noise', String, queue_size=10)

        self.marker_pub = rospy.Publisher('/marker_topic', Marker, queue_size=10)

        self.arrived_pub = rospy.Publisher('/arrived', String, queue_size=10)

        self.trans_listener = tf.TransformListener()

        self.cfg_srv = Server(NavigatorConfig, self.dyn_cfg_callback)

        rospy.Subscriber("/map", OccupancyGrid, self.map_callback)
        rospy.Subscriber("/map_metadata", MapMetaData, self.map_md_callback)
        rospy.Subscriber("/cmd_nav", Pose2D, self.cmd_nav_callback)
        rospy.Subscriber("/cmd_vel_teleop", Twist, self.cmd_vel_teleop_callback)
        rospy.Subscriber("/animal_unique", String, self.animal_rescue_callback)
        rospy.Subscriber('/detector/stop_sign', DetectedObject, self.stop_sign_callback)
        rospy.Subscriber('/num_rescues', Int32, self.num_rescues_callback)

        # Detector subscribers
        #rospy.Subscriber("/detector/bird", DetectedObject, self.bird_callback)
        rospy.Subscriber("/detector/cat", DetectedObject, self.cat_callback)
        rospy.Subscriber("/detector/dog", DetectedObject, self.dog_callback)
        #rospy.Subscriber("/detector/horse", DetectedObject, self.horse_callback)
        #rospy.Subscriber("/detector/sheep", DetectedObject, self.sheep_callback)
        rospy.Subscriber("/detector/cow", DetectedObject, self.cow_callback)
        rospy.Subscriber("/detector/elephant", DetectedObject, self.elephant_callback)
        #rospy.Subscriber("/detector/bear", DetectedObject, self.bear_callback)
        rospy.Subscriber("/detector/zebra", DetectedObject, self.zebra_callback)
        rospy.Subscriber("/detector/giraffe", DetectedObject, self.giraffe_callback)

        msg = String()
        msg.data = "arrived"
        self.arrived_pub.publish(msg)

        rospy.sleep(10)
        self.replan()
        print("finished init")

    def dyn_cfg_callback(self, config, level):
        rospy.loginfo(
            "Reconfigure Request: k1:{k1}, k2:{k2}, k3:{k3}".format(**config)
        )
        self.pose_controller.k1 = config["k1"]
        self.pose_controller.k2 = config["k2"]
        self.pose_controller.k3 = config["k3"]
        return config

    def cmd_nav_callback(self, data):
        """
        loads in goal if different from current goal, and replans
        """
        rospy.loginfo("New goal received")
        if (
            data.x != self.x_g
            or data.y != self.y_g
            or data.theta != self.theta_g
        ):
            rospy.logdebug(f"New command nav received:\n{data}")
            self.x_g = data.x
            self.y_g = data.y
            self.theta_g = data.theta
            self.replan()

    def cmd_vel_teleop_callback(self, data):
        self.V_teleop = data.linear.x
        self.om_teleop = data.angular.z
        self.mode = Mode.TELEOP

    def num_rescues_callback(self, msg):
        self.num_rescue_targets = msg.data

    def animal_rescue_callback(self, msg): # msg.data contains animal name (string)
        """
        populates rescue_targets dictionary with animals published from command line
        """
        if not self.rescue_phase:
            animal_location = self.detected_animals[msg.data]
            self.rescue_targets[msg.data] = (animal_location['x'], animal_location['y'], animal_location['th'])
            rospy.loginfo("ADDING %s to RESCUE TARGETS", msg.data)
            rospy.loginfo("Number of rescue targets: %s", self.num_rescue_targets)
            rospy.loginfo("Length of rescue_targets: %s", len(self.rescue_targets))
            if len(self.rescue_targets) == self.num_rescue_targets:
                self.init_rescue()

    def stop_sign_callback(self, msg):
        if rospy.get_rostime() - self.last_stop_sign > rospy.Duration.from_sec(60):
            self.init_stop_sign()

    def map_md_callback(self, msg):
        """
        receives maps meta data and stores it
        """
        self.map_width = msg.width
        self.map_height = msg.height
        self.map_resolution = msg.resolution
        self.map_origin = (msg.origin.position.x, msg.origin.position.y)

    def map_callback(self, msg):
        """
        receives new map info and updates the map
        """
        self.map_probs = msg.data
        # if we've received the map metadata and have a way to update it:
        if (
            self.map_width > 0
            and self.map_height > 0
            and len(self.map_probs) > 0
        ):
            self.occupancy = StochOccupancyGrid2D(
                self.map_resolution,
                self.map_width,
                self.map_height,
                self.map_origin[0],
                self.map_origin[1],
                7,
                self.map_probs,
            )
            if self.x_g is not None:
                # if we have a goal to plan to, replan
                rospy.loginfo("replanning because of new map")
                self.replan()  # new map, need to replan

    def shutdown_callback(self):
        """
        publishes zero velocities upon rospy shutdown
        """
        cmd_vel = Twist()
        cmd_vel.linear.x = 0.0
        cmd_vel.angular.z = 0.0
        self.nav_vel_pub.publish(cmd_vel)

    def detected_animal(self, animal, object_msg, marker_id):
        if not self.rescue_phase:
            if self.occupancy.is_free((self.x, self.y)):
                if animal not in self.detected_animals or self.detected_animals[animal]['conf'] < object_msg.confidence:
                    self.detected_animals[animal] = {'x': self.x, 'y': self.y, 'th': self.theta, 'conf': object_msg.confidence}
                    rospy.loginfo("Detected a %s. Currently detected animals: %s", animal, self.detected_animals.keys())
                    self.publish_marker(marker_id)

    def publish_marker(self, marker_id):
        marker = Marker()
        marker.header.frame_id = "base_footprint"
        #marker.frame_locked = False
        marker.header.stamp = rospy.Time()
        marker.id = marker_id
        rospy.loginfo("MARKER_ID: %d", marker.id)
        marker.type = 2 # sphere
        marker.pose.position.x = 0
        marker.pose.position.y = 0
        marker.pose.position.z = 0
        marker.pose.orientation.x = 0.0
        marker.pose.orientation.y = 0.0
        marker.pose.orientation.z = 0.0
        marker.pose.orientation.w = 0.0
        marker.scale.x = 0.2
        marker.scale.y = 0.2
        marker.scale.z = 0.2
        marker.color.a = 1.0
        marker.color.r = 0.0
        marker.color.g = 1.0
        marker.color.b = 0.0
        self.marker_pub.publish(marker)
        rospy.loginfo('Published marker!')

    def bird_callback(self, object_msg):
       self.detected_animal('bird', object_msg)

    def cat_callback(self, object_msg):
        self.detected_animal('cat', object_msg, 0)
        # publish meow
        meow_msg = String()
        meow_msg.data = "MEOW"
        self.animal_noise_pub.publish(meow_msg)

    def dog_callback(self, object_msg):
        self.detected_animal('dog', object_msg, 1)
        # publish woof
        woof_msg = String()
        woof_msg.data = "WOOF"
        self.animal_noise_pub.publish(woof_msg)

    def horse_callback(self, object_msg):
        self.detected_animal('horse', object_msg)

    def sheep_callback(self, object_msg):
        self.detected_animal('sheep', object_msg)

    def cow_callback(self, object_msg):
        self.detected_animal('cow', object_msg, 2)

    def elephant_callback(self, object_msg):
        self.detected_animal('elephant', object_msg, 3)

    def bear_callback(self, object_msg):
        self.detected_animal('bear', object_msg)

    def zebra_callback(self, object_msg):
        self.detected_animal('zebra', object_msg, 4)

    def giraffe_callback(self, object_msg):
        self.detected_animal('giraffe', object_msg, 5)

    def near_goal(self):
        """
        returns whether the robot is close enough in position to the goal to
        start using the pose controller
        """
        return (
            linalg.norm(np.array([self.x - self.x_g, self.y - self.y_g]))
            < self.near_thresh
        )

    def at_goal(self):
        """
        returns whether the robot has reached the goal position with enough
        accuracy to return to idle state
        """
        return (
            linalg.norm(np.array([self.x - self.x_g, self.y - self.y_g]))
            < self.at_thresh
            and abs(wrapToPi(self.theta - self.theta_g)) < self.at_thresh_theta
        )

    def aligned(self):
        """
        returns whether robot is aligned with starting direction of path
        (enough to switch to tracking controller)
        """
        return (
            abs(wrapToPi(self.theta - self.th_init)) < self.theta_start_thresh
        )

    def close_to_plan_start(self):
        return (
            abs(self.x - self.plan_start[0]) < self.start_pos_thresh
            and abs(self.y - self.plan_start[1]) < self.start_pos_thresh
        )

    def snap_to_grid(self, x):
        return (
            self.plan_resolution * round(x[0] / self.plan_resolution),
            self.plan_resolution * round(x[1] / self.plan_resolution),
        )

    def init_rescue(self):
        self.rescue_phase = True
        rospy.loginfo("RESCUE PHASE")
        self.TSP()
        rospy.loginfo("Rescue targets sorted: %s", self.rescue_targets)
        self.rescue_target = self.sorted_rescue_targets[0]
        #self.rescue_target = random.choice(list(self.rescue_targets.keys())) # choose first rescue target randomly
        rospy.loginfo("NEXT RESCUE TARGET: %s", self.rescue_target)
        self.x_g, self.y_g, self.theta_g = self.rescue_targets[self.rescue_target]
        self.replan()

    def init_stop_sign(self):
        self.stop_sign_start = rospy.get_rostime()
        self.last_stop_sign = rospy.get_rostime()
        self.switch_mode(Mode.STOP)

    def has_stopped(self):
        """ checks if robot stops at stop sign for 3 sec """

        return rospy.get_rostime() - self.stop_sign_start > rospy.Duration.from_sec(self.stop_time)

    def TSP(self):
        rospy.loginfo("Running TSP")

        coordinates = {0: (self.x_start, self.y_start, self.theta_start)}
        num_to_animal = {}
        num = 1
        for key, value in self.rescue_targets.items():
            num_to_animal[num] = key
            coordinates[num] = (value[0], value[1], value[2])
            num += 1

        rospy.loginfo("Coordinates: %s", coordinates)
        rospy.loginfo("Num to animal: %s", num_to_animal)

        if len(self.rescue_targets) == 1 or len(self.rescue_targets) == 2:
            sorted_rescue_targets = list(self.rescue_targets.keys())
            self.sorted_rescue_targets = sorted_rescue_targets
            return
        elif len(self.rescue_targets) == 3:
            combinations = [[0,1,2,3,0], [0,1,3,2,0], [0,2,1,3,0]]
        elif len(self.rescue_targets) == 4:
            combinations = [[0,1,2,3,4,0],
                            [0,1,2,4,3,0],
                            [0,1,3,2,4,0],
                            [0,1,3,4,2,0],
                            [0,1,4,2,3,0],
                            [0,1,4,3,2,0],
                            [0,2,1,3,4,0],
                            [0,2,1,4,3,0],
                            [0,2,3,1,4,0],
                            [0,2,4,1,3,0],
                            [0,3,1,2,4,0]]

        min_length = np.inf
        min_combination = combinations[0]
        for combination in combinations:
            length_planned_path = 0
            for i in range(len(combination)-1):
                x_init = self.snap_to_grid((coordinates[combination[i]][0], coordinates[combination[i]][1]))
                x_goal = self.snap_to_grid((coordinates[combination[i+1]][0], coordinates[combination[i+1]][1]))
                problem = AStar(self.snap_to_grid((0, 0)), self.snap_to_grid((self.plan_horizon, self.plan_horizon)), x_init, x_goal, self.occupancy, self.plan_resolution)

                success = problem.solve()
                if not success:
                    rospy.loginfo("A* failed in TSP")
                    length_planned_path += np.inf
                else:
                    rospy.loginfo("A* succeded in TSP")
                    planned_path = problem.path
                    length_planned_path += len(planned_path)

            rospy.loginfo("Length of path: %s", length_planned_path)
            if length_planned_path < min_length:
                min_length = length_planned_path
                min_combination = combination

        rospy.loginfo("Shortest combination: %s", min_combination)

        sorted_rescue_targets = []
        for j in range(1,len(min_combination)-1):
            animal_num = min_combination[j]
            sorted_rescue_targets.append(num_to_animal[animal_num])

        self.sorted_rescue_targets = sorted_rescue_targets

    def has_picked_up(self):
        """ checks if animal has been picked up (robot has stopped for 3-5s) """
        if self.mode == Mode.PICKUP and rospy.get_rostime() - self.pickup_start >= rospy.Duration.from_sec(self.min_pickup_time):
            rospy.loginfo("STOPPED > 3 SEC")
            return True
        else:
            return False

    def switch_mode(self, new_mode):
        rospy.loginfo("Switching from %s -> %s", self.mode, new_mode)
        self.mode = new_mode

    def publish_planned_path(self, path, publisher):
        # publish planned plan for visualization
        path_msg = Path()
        path_msg.header.frame_id = "map"
        for state in path:
            pose_st = PoseStamped()
            pose_st.pose.position.x = state[0]
            pose_st.pose.position.y = state[1]
            pose_st.pose.orientation.w = 1
            pose_st.header.frame_id = "map"
            path_msg.poses.append(pose_st)
        publisher.publish(path_msg)

    def publish_smoothed_path(self, traj, publisher):
        # publish planned plan for visualization
        path_msg = Path()
        path_msg.header.frame_id = "map"
        for i in range(traj.shape[0]):
            pose_st = PoseStamped()
            pose_st.pose.position.x = traj[i, 0]
            pose_st.pose.position.y = traj[i, 1]
            pose_st.pose.orientation.w = 1
            pose_st.header.frame_id = "map"
            path_msg.poses.append(pose_st)
        publisher.publish(path_msg)

    def publish_control(self):
        """
        Runs appropriate controller depending on the mode. Assumes all controllers
        are all properly set up / with the correct goals loaded
        """
        t = self.get_current_plan_time()

        if self.mode == Mode.PARK:
            V, om = self.pose_controller.compute_control(
                self.x, self.y, self.theta, t
            )
        elif self.mode == Mode.TRACK:
            V, om = self.traj_controller.compute_control(
                self.x, self.y, self.theta, t
            )
        elif self.mode == Mode.ALIGN:
            V, om = self.heading_controller.compute_control(
                self.x, self.y, self.theta, t
            )
        elif self.mode == Mode.TELEOP:
            V = self.V_teleop
            om = self.om_teleop
        else:
            V = 0.0
            om = 0.0

        cmd_vel = Twist()
        cmd_vel.linear.x = V
        cmd_vel.angular.z = om
        self.nav_vel_pub.publish(cmd_vel)

    def get_current_plan_time(self):
        t = (rospy.get_rostime() - self.current_plan_start_time).to_sec()
        return max(0.0, t)  # clip negative time to 0

    def replan(self):
        """
        loads goal into pose controller
        runs planner based on current pose
        if plan long enough to track:
            smooths resulting traj, loads it into traj_controller
            sets self.current_plan_start_time
            sets mode to ALIGN
        else:
            sets mode to PARK
        """
        # Make sure we have a map
        if not self.occupancy:
            rospy.loginfo(
                "Navigator: replanning canceled, waiting for occupancy map."
            )
            self.switch_mode(Mode.IDLE)
            return

        # Attempt to plan a path
        state_min = self.snap_to_grid((0, 0))
        state_max = self.snap_to_grid((self.plan_horizon, self.plan_horizon))
        x_init = self.snap_to_grid((self.x, self.y))
        self.plan_start = x_init
        x_goal = self.snap_to_grid((self.x_g, self.y_g))
        problem = AStar(
            state_min,
            state_max,
            x_init,
            x_goal,
            self.occupancy,
            self.plan_resolution,
        )

        rospy.loginfo("Navigator: computing navigation plan")
        success = problem.solve()
        if not success:
            rospy.loginfo("Planning failed")
            #problem.plot_tree()
            return

        rospy.loginfo("Planning Succeeded")

        planned_path = problem.path

        # Check whether path is too short
        if len(planned_path) < 4:
            rospy.loginfo("Path too short to track")
            self.pose_controller.load_goal(self.x_g, self.y_g, self.theta_g)
            self.switch_mode(Mode.PARK)
            return

        # Smooth and generate a trajectory
        t_new, traj_new = compute_smoothed_traj(
            planned_path, self.v_des, self.spline_deg, self.spline_alpha, self.traj_dt
        )

        # If currently tracking a trajectory, check whether new trajectory will take more time to follow
        if self.mode == Mode.TRACK:
            t_remaining_curr = (
                self.current_plan_duration - self.get_current_plan_time()
            )

            # Estimate duration of new trajectory
            th_init_new = traj_new[0, 2]
            th_err = wrapToPi(th_init_new - self.theta)
            t_init_align = abs(th_err / self.om_max)
            t_remaining_new = t_init_align + t_new[-1]

            if t_remaining_new > t_remaining_curr:
                rospy.loginfo(
                    "New plan rejected (longer duration than current plan)"
                )
                self.publish_smoothed_path(
                    traj_new, self.nav_smoothed_path_rej_pub
                )
                return

        # Otherwise follow the new plan
        self.publish_planned_path(planned_path, self.nav_planned_path_pub)
        self.publish_smoothed_path(traj_new, self.nav_smoothed_path_pub)

        self.pose_controller.load_goal(self.x_g, self.y_g, self.theta_g)
        self.traj_controller.load_traj(t_new, traj_new)

        self.current_plan_start_time = rospy.get_rostime()
        self.current_plan_duration = t_new[-1]

        self.th_init = traj_new[0, 2]
        self.heading_controller.load_goal(self.th_init)

        if not self.aligned():
            rospy.loginfo("Not aligned with start direction")
            self.switch_mode(Mode.ALIGN)
            return

        rospy.loginfo("Ready to track")
        self.switch_mode(Mode.TRACK)

    def run(self):
        rate = rospy.Rate(10)  # 10 Hz
        while not rospy.is_shutdown():
            # try to get state information to update self.x, self.y, self.theta
            try:
                (translation, rotation) = self.trans_listener.lookupTransform(
                    "/map", "/base_footprint", rospy.Time(0)
                )
                euler = tf.transformations.euler_from_quaternion(rotation)
                if self.x_start is None:
                    self.x_start = translation[0]
                    self.y_start = translation[1]
                    self.theta_start = euler[2]
                self.x = translation[0]
                self.y = translation[1]
                self.theta = euler[2]
            except (
                tf.LookupException,
                tf.ConnectivityException,
                tf.ExtrapolationException,
            ) as e:
                self.current_plan = []
                rospy.loginfo("Navigator: waiting for state info")
                if self.mode != Mode.TELEOP:
                    self.switch_mode(Mode.IDLE)
                print(e)
                pass

            # STATE MACHINE LOGIC
            # some transitions handled by callbacks
            if self.mode == Mode.IDLE:
                pass
            elif self.mode == Mode.ALIGN:
                if self.aligned():
                    self.current_plan_start_time = rospy.get_rostime()
                    self.switch_mode(Mode.TRACK)
            elif self.mode == Mode.TRACK:
                if self.near_goal():
                    self.switch_mode(Mode.PARK)
                elif not self.close_to_plan_start():
                    rospy.loginfo("replanning because far from start")
                    self.replan()
                elif (
                    rospy.get_rostime() - self.current_plan_start_time
                ).to_sec() > self.current_plan_duration:
                    rospy.loginfo("replanning because out of time")
                    self.replan()  # we aren't near the goal but we thought we should have been, so replan
            elif self.mode == Mode.PARK:
                if self.at_goal():
                    # forget about goal:
                    self.x_g = None
                    self.y_g = None
                    self.theta_g = None
                    if not self.rescue_phase:
                        msg = String()
                        msg.data = "arrived"
                        self.arrived_pub.publish(msg)
                    if self.rescue_phase:
                        rospy.loginfo("STARTING PICKUP")
                        self.pickup_start = rospy.get_rostime()
                        self.switch_mode(Mode.PICKUP)
                    else:
                        # forget about goal:
                        self.switch_mode(Mode.IDLE)
            elif self.mode == Mode.PICKUP:
                # pick up animal
                #if not self.at_goal():
                    #self.replan()
                    #rospy.loginfo("Stuck trying to find path")
                if self.has_picked_up(): # robot has stopped for 3-5 sec
                    rospy.loginfo("PICKUP COMPLETE")
                    self.sorted_rescue_targets.remove(self.rescue_target)
                    if self.sorted_rescue_targets: # still more animals to rescue
                        self.rescue_target = self.sorted_rescue_targets[0]
                        rospy.loginfo("NEXT PICKUP TARGET: %s", self.rescue_target)
                        self.x_g, self.y_g, self.theta_g = self.rescue_targets[self.rescue_target]
                    else:
                        # return to start
                        rospy.loginfo("DONE RESCUING")
                        self.x_g = self.x_start
                        self.y_g = self.y_start
                        self.theta_g = self.theta_start
                    self.replan()
            elif self.mode == Mode.STOP:
                # stop at stop sign
                if self.has_stopped():
                    self.replan()

            self.publish_control()
            rate.sleep()

if __name__ == "__main__":
    nav = Navigator()
    rospy.on_shutdown(nav.shutdown_callback)
    nav.run()


"""
This module handles settings of system.

Platform-dependent, currently works on Unix only (uses '/proc' filesystem).

Part of 'Adaptor' framework.

Author: Michael Pankov, 2012-2013.

Please do not redistribute.
"""

import os
import recordtype as rt

from data_types import *

SettingsBase = rt.recordtype('SettingsBase',
    'program_name framework_root_dir benchmark_root_dir benchmark_bin_dir '
    'build_settings run_settings')

class Settings(PrintableStructure, SettingsBase):
    def __init__(self,
                 program_name,
                 benchmark_root_dir='data/sources/polybench-c-3.2/',
                 benchmark_bin_dir='data/bin/'):
        """ Initialize framework settings and return Settings object.

        program_name is name of directory to be found in benchmark_root_dir.
        benchmark_root_dir is directory, containing all sources of given
            benchmark.
        benchmark_bin_dir is directory, where binaries will be put.
        """
        self.program_name = program_name
        self.framework_root_dir = os.path.join(os.path.dirname(
            os.path.abspath(__file__)), '..')
        self.benchmark_root_dir = os.path.realpath(
            os.path.join(self.framework_root_dir, benchmark_root_dir))
        self.benchmark_bin_dir = os.path.realpath(
            os.path.join(self.framework_root_dir, benchmark_bin_dir))
        self.define_build_settings('src', '')
        self.define_run_settings()

    def define_build_settings(self, sources_path, other_flags):
        self.build_settings = BuildSettings(
            benchmark_source_dir=os.path.join(
                self.benchmark_root_dir, '', sources_path),
            program_source="{0}.c".format(self.program_name),
            compiler=None,
            base_opt=None,
            optimization_flags=None,
            other_flags=other_flags,
            linker_options='-lm')

    def define_run_settings(self):
        self.run_settings = RunSettings()