import breve
import math

####################
# CONSTANTS
LOGGING = None

maxvel = 4
dT = 0.050000
NU = 0.08

MAXTHRUST = .1
MAXBOX = 27
SELECTION_LOCK = 10 #iterations
MAX_TRANSFER = 150        # guess?
POWER_LOSS_FACTOR = 0.9
# btw there are about 30 iterations per 1 "second" of getTime
##################
#           +Z
#   -XXXXXXZXXXXXXXXXX+
#           Z
#          -Z

# move (X, Y, Z)

## TODO:
### - add resource loss for attitude control (when sharing)
### - add real attitude control
### - get a more intelligent satellite control

#####################################
## Global Functions
##
def randomPosInBox():
    return breve.randomExpression( breve.vector( MAXBOX, MAXBOX, MAXBOX ) ) - breve.randomExpression( breve.vector( .5*MAXBOX, .5*MAXBOX, .5*MAXBOX ) )

def outOfBox(x, y, z):
    if not (-MAXBOX < x < MAXBOX) : return True 
    if not (-MAXBOX < y < MAXBOX) : return True 
    if not (-MAXBOX < z < MAXBOX) : return True 
    return False
    
    
breve.counted_blocking = 0

##################################
## Control class from breve
class myControl( breve.Control ):
    # variables
    # agents list of agents
    def __init__(self, time = 500):
        breve.Control.__init__(self)
        self.agentShape = None
        self.agents = []
        myControl.init(self)
        self.maxTime = time
        
    def getAgentShape(self):
        return self.agentShape

    def init(self):
        print "Initializing the simulation... "

        self.setIntegrationStep( dT )

        # init the satellites
        self.agentShape = breve.createInstances( breve.Cube, 1 ).initWith( breve.vector( 2,2,2 ) )
        self.agents = breve.createInstances( SimpleStayInBoxSatellite, 5 )
        self.agents.append ( SimplePowerSatellite () )
        self.agents.last
        self.agents.append ( SimplePowerSatellite () )
        self.pointCamera( breve.vector( 0, 0, 0 ), breve.vector( 0, -150, 0 ) )
        self.initSatellitePositions()

        # init the selection process        
        self.selectionLock = 0
        #self.selector = RandomSelector()
        #self.selector = ClosestSelector()
        #self.selector = LowestSelector()
        #self.selector = UtilitarianWelfareSelector()
        #self.selector = EgalitarianWelfareSelector()
        self.selector = NashProductSelector()

        
        # abbiamo finito init
        print "Start!"
    
    def initSatellitePositions(self):
        # organize the satellites in a pre-defined matter!
        print "  Initialize the satellites ...",
        
        # TODO loop over all agents
        #for sat in self.agents:
        self.agents[0].move( randomPosInBox() ) #breve.vector ( -1, 0, -1) )  #set location 
        self.agents[1].move( randomPosInBox() ) #breve.vector ( -1, 0, -1) )  #set location 
        self.agents[2].move( randomPosInBox() ) #breve.vector ( -1, 0, -1) )  #set location 
        self.agents[3].move( randomPosInBox() ) #breve.vector ( -1, 0, -1) )  #set location 
        self.agents[4].move( randomPosInBox() ) #breve.vector ( -1, 0, -1) )  #set location 
        
        self.agents[5].move( breve.vector ( -2, 0, -2) )  #set location  
        self.agents[6].move( breve.vector (  2, 0,  2) )  #set location  
        print "done!"

    def iterate(self):
        breve.Control.iterate(self)
        self.updateNeighbors()
        
        # select one satellite from the each of the power satellites neighborhoods!
        if breve.Control.getTime(self) > 0.01 :
            if self.selectionLock == SELECTION_LOCK :
#                print "Selecting ... "
                s1, s2 = self.selector.selectTransferAgents( self.agents[5], self.agents[5].getNeighbors(), self.agents[6], self.agents[6].getNeighbors() )
                self.selectionTime = breve.Control.getTime(self)
                
                self.selection1 = s1
                self.selection2 = s2
                
                # do exchange
                if not self.selection1 is None: self.exchange( [self.agents[5], self.selection1] )
                if not self.selection2 is None: self.exchange( [self.agents[6], self.selection2] )
                
                # start lock
                self.selectionLock -= 1
                
                
                
            else :
                if self.selectionLock == 0 :
                    # free lock
                    self.selectionLock = SELECTION_LOCK
                else: 
                    self.selectionLock -= 1
                    # do exchange over again
                    if not self.selection1 is None: self.exchange( [self.agents[5], self.selection1] )
                    if not self.selection2 is None: self.exchange( [self.agents[6], self.selection2] )
                    
            
        if LOGGING :
            # print status            
            print "T:", breve.Control.getTime(self), "; ",
            for i in range(len(self.agents)) :
                print "S", i, ": ", self.agents[i].printStatus(), "; ", 
            print "==";
        
        # ending
        if breve.Control.getTime(self) > self.maxTime: 
            print "Stop simulation"
            print "Counted blocked's : ", breve.counted_blocking
            breve.Control.saveSnapshot (self, "BREVE_screencap" + str(breve.Control.getRealTime(self)) + ".png")
            # throws error ;) is there a better way?
            breve.Control.stop(self)
            
            
    def exchange(self, ex):
        if len(ex) == 2 :
            # print "In Exchange func"
            # simple exchange function
            distance = ex[0].getDistance(ex[1])
            
            if distance > MAXBOX: return    # not too far for power transfer
            
            if ex[0].res.level > ex[0].res.usagePerIteration + MAX_TRANSFER and ex[1].res.level < ex[1].res.maxLevel :
               ex[0].res.level -= MAX_TRANSFER
               ex[1].res.level += math.pow(POWER_LOSS_FACTOR, distance) * MAX_TRANSFER
               if ex[1].res.level > ex[1].res.maxLevel : 
                   ex[1].res.level = ex[1].res.maxLevel
            
            
            
#################################
## Selector class
class Selector :
    # variables
    # agents list of agents
    def __init__(self):
        pass
        
    def selectTransferAgents(self, s1, n1, s2, n2):
        pass

class SimpleSelector :
    '''Selects the first neighbor of a power satellite to be the one satellite to transfer the power to!'''
    # variables
    # agents list of agents
    def __init__(self):
        print "  -- SIMPLE SELECTOR selected ;) --"
    
    # simple selection
    def selectTransferAgents(self, s1, n1, s2, n2):
        r1 = None
        r2 = None
        if n1: r1 = n1[0]
        if n2: r2 = n2[0]
        return [ r1, r2 ]

class RandomSelector :
    '''Selects a random neighbor of a power satellite to be the one satellite to transfer the power to!'''
    # variables
    # agents list of agents
    def __init__(self):
        print "  -- RANDOM SELECTOR selected ;) --"
    
    # simple selection
    def selectTransferAgents(self, s1, n1, s2, n2):
        r1 = None
        r2 = None
        if n1: r1 = n1[breve.randomExpression(len(n1) - 1)]
        if n2: r2 = n2[breve.randomExpression(len(n2) - 1)]
        return [ r1, r2 ]
            
class ClosestSelector :
    '''Selects the closest neighbor of a power satellite to be the one satellite to transfer the power to!'''
    # variables
    # agents list of agents
    def __init__(self):
        print "  -- CLOSEST SELECTOR selected ;) --"
    
    # simple selection
    def selectTransferAgents(self, s1, n1, s2, n2):
        r1 = None
        r2 = None
        if n1: 
            dist = MAXBOX + 1
            r1 = None
            for s in n1:
                if s1.getDistance(s) < dist :
                    dist = s1.getDistance(s)
                    r1 = s
        if n2:
            dist = MAXBOX + 1
            r2 = None
            for s in n2:
                if s2.getDistance(s) < dist :
                    dist = s2.getDistance(s)
                    r2 = s

        return [ r1, r2 ]
            

class UtilitarianWelfareSelector :
    '''Selects the neighbor based on utilitarian welfare theory to transfer the power to!'''
    # variables
    # agents list of agents
    def __init__(self):
        print "  -- UTILITARIAN WELFARE SELECTOR selected ;) --"
    
    # simple selection
    def selectTransferAgents(self, s1, n1, s2, n2):
        r1 = None
        r2 = None
        if n1: 
            r1 = None
            highest_estimate = 0
            for s in n1:
                h = (s.usageCycle - s.cycle)*1.0/SELECTION_LOCK * s.usagePowerMultiplier
                if h > highest_estimate :
                    highest_estimate = h
                    r1 = s
        if n2:
            r2 = None
            highest_estimate = 0
            for s in n2:
                h = (s.usageCycle - s.cycle)*1.0/SELECTION_LOCK * s.usagePowerMultiplier
                if h > highest_estimate :
                    highest_estimate = h
                    r2 = s

        return [ r1, r2 ]


class EgalitarianWelfareSelector :
    '''Selects the neighbor based on egalitarian welfare theory to transfer the power to!'''
    # variables
    # agents list of agents
    def __init__(self):
        print "  -- EGALITARIAN WELFARE SELECTOR selected ;) --"
    
    # simple selection
    def selectTransferAgents(self, s1, n1, s2, n2):
        r1 = None
        r2 = None
        if n1: 
            r1 = None
            lowest = 99999
            # find lowest that is above the producer
            for s in n1:
                if s.res.level + MAX_TRANSFER < s1.res.level - MAX_TRANSFER and s.res.level < lowest:
                    lowest = s.res.level
                    r1 = s
        if n2:
            r2 = None
            lowest = 99999
            # find lowest that is above th
            for s in n2:
                if s.res.level + MAX_TRANSFER < s2.res.level - MAX_TRANSFER and s.res.level < lowest:
                    lowest = s.res.level
                    r2 = s

        return [ r1, r2 ]

class NashProductSelector :
    '''Selects the neighbor based on resource value, first trying to allocate it to the agent that makes the producer not the weakest!'''
    # variables
    # agents list of agents
    def __init__(self):
        print "  -- NASHPRODUCT WELFARE SELECTOR selected ;) --"
    
    def selectTransferAgents(self, s1, n1, s2, n2):
        r1 = None
        r2 = None
        if n1: 
            r1 = None
            highest_estimate = 0
            for s in n1:
                if s.res.level == 0.0 :
                    r1 = s
                    break
                    
                h = (s.usageCycle - s.cycle)*1.0/SELECTION_LOCK * s.usagePowerMultiplier
                if h > highest_estimate :
                    highest_estimate = h
                    r1 = s
        if n2:
            r2 = None
            highest_estimate = 0
            for s in n2:
                if s.res.level == 0.0 :
                    r2 = s
                    break
                    
                h = (s.usageCycle - s.cycle)*1.0/SELECTION_LOCK * s.usagePowerMultiplier
                if h > highest_estimate :
                    highest_estimate = h
                    r2 = s

        return [ r1, r2 ]


class LowestSelector :
    '''Selects the neighbor that is the lowest neighbor!'''
    # variables
    # agents list of agents
    def __init__(self):
        print "  -- LOWEST SELECTOR selected ;) --"
    
    def selectTransferAgents(self, s1, n1, s2, n2):
        r1 = None
        r2 = None
        if n1: 
            r1 = None
            lowest = 99999
            for s in n1:
                if s.res.level < lowest :        
                    lowest = s.res.level
                    r1 = s
        if n2:
            r2 = None
            lowest = 99999
            for s in n2:
                if s.res.level < lowest :
                    lowest = s.res.level
                    r2 = s

        return [ r1, r2 ]

            

##################################
## Satellite control classes
##
## Random movement satelite
class RandomWalker( breve.Mobile ):
    def __init__(self):
        breve.Mobile.__init__(self)
        RandomWalker.init(self)
    
    def init(self):
        self.setShape( self.controller.getAgentShape() )
        self.setColor( breve.vector( 1,1,1 ))
        
    def iterate(self):
        self.setVelocity( ( breve.randomExpression( breve.vector( maxvel, maxvel, maxvel ) ) - breve.vector( 0.5 * maxvel, 0.5 * maxvel, 0.5 * maxvel) ) )
  
#########      
## HCW controlled satellite
class HCW( breve.Mobile ):
    '''Represents a satellite in the Hill-Cohessey-Wiltshire environment.'''
    xdot = 0.0
    zdot = 0.0
    ydot = 0.0
    
    mass = 0.0 #kg
    
    def __init__(self):
        breve.Mobile.__init__(self)
        HCW.init(self)
    
    def init(self):
        self.setShape( self.controller.getAgentShape() )
                
    def iterate(self):
        # use HCW equations for motion here
        # calculate xdot, ydot and zdot
#    y and z are switched ... maybe we can change that by chaning the viewpoint?!              
        x, y, z = self.getLocation()
        self.xdot += dT * (2 * NU * self.zdot + 3 * NU * NU * x) 
        self.zdot += dT * (-2 * NU * self.xdot)
        self.ydot += dT * (-1 * NU * NU * y)
        self.setVelocity( breve.vector( self.xdot, self.ydot, self.zdot) )
        
    def thrust(self):
        # maybe work on that? better thrusting strategy
        x, y, z = self.getLocation()
        if x < -MAXBOX: self.xdot += MAXTHRUST 
        if x >  MAXBOX: self.xdot -= MAXTHRUST
        if y < -MAXBOX: self.ydot += MAXTHRUST 
        if y >  MAXBOX: self.ydot -= MAXTHRUST
        if z < -MAXBOX: self.zdot += MAXTHRUST 
        if z >  MAXBOX: self.zdot -= MAXTHRUST      
        
######################################
## Resource classes
##
## Simple Resource
class Resource :
    ''' Represents a resource to be used and/or produced by the satellites.'''
    # variables
    level = 0.0
    maxLevel = 0.0
    increasePerIteration = 0.0
    usagePerIteration = 0.0
    
    # methods
    def __init__ (self, usage, inc, maxlevel):
        #print "Initialization of Resource"
        self.usagePerIteration = usage
        self.increasePerIteration = inc
        self.maxLevel = maxlevel
        
    def iterate(self):
        if self.level < self.maxLevel :
            self.level = self.level + self.increasePerIteration - self.usagePerIteration

    # use the resource n times (e.g. we wanna take a picture and need power for the next n iterations)
    def use(self, n = 1):
        if self.level >= self.usagePerIteration * n :
            self.level -= self.usagePerIteration * n
            return True
        else:
            return False

    

######################################
## Satellite Definition classes
##
## Simple Satellite
class SimpleSatellite( HCW ):
    # variables
    stdColor = None  # the standard color
    res = None       # the resource
    usagePowerMultiplier = 0   # how much energy do we need
    usageRepetition = 0        # we weann use more resources every usageCycles iterations 
    
    cycle = 0        # iNTERNAL counting the cycles between resource use
    
    def __init__(self):
        HCW.__init__(self)
        SimpleSatellite.init(self)
        
    def init(self):
        self.mass = 75    #kg
        self.stdColor = breve.vector( 1,1,1 )
        self.setColor( self.stdColor )
        self.res = Resource(90, 100, 500)    # usage, generation, maxLevel
        self.usageCycle = int(5 + breve.randomExpression ( 5 ))   # between 5 and 10
        self.usagePowerMultiplier = int(1 + breve.randomExpression ( 2 ))    # between 1 and 3
        self.usageBlocked = False        # used to indicate (in the log) that we could not use the rousource (not enough here)
        
    def iterate(self):
        super(SimpleSatellite, self).iterate()
        self.res.iterate()
        if self.cycle >= self.usageCycle:
            self.cycle = 0
            if not self.res.use(self.usagePowerMultiplier):
                self.usageBlocked = True
                breve.counted_blocking = breve.counted_blocking + 1                
                #print "could not do it!"
            else :
                self.usageBlocked = False
        else:
            self.cycle += 1
        
    def setResource(self, resource):
        self.res = resource
        
    def printStatus(self):
        s = "SimpSat, L: " + str( self.res.level ) + " usCyc: " + str(self.usageCycle) + " usMult: " + str(self.usagePowerMultiplier)
        if self.usageBlocked : 
            s += " BL!"        # blocked
        return s
#        print "SimplSat: usCyc.", self.usageCycle, "usMult.", self.usagePowerMultiplier
        

## Simple Satellite that (tries to) stay in the maxbox
class SimpleStayInBoxSatellite( SimpleSatellite ):
    def __init__(self):
        SimpleSatellite.__init__(self)
        SimpleStayInBoxSatellite.init(self)
        
    def init(self):
        self.thrust = breve.vector ( 0, 0, 0 )
        
    def iterate(self):
        super(SimpleStayInBoxSatellite, self).iterate()
        x, y, z = self.getLocation()
        # thrust backwards if we run out of the box
        if ( outOfBox(x,y,z) ): 
            HCW.thrust(self)
            self.setColor( breve.vector ( 1, 0, 0) )
        else:
            self.setColor( self.stdColor )
        
## Simple Power Satellite
#class SimplePowerSatellite( SimpleSatellite ):
#    def __init__(self):
#        SimpleSatellite.__init__(self)
#        SimplePowerSatellite.init(self)
#        
#    def init(self):
#        self.setColor( breve.vector( 0.3,0.3,0.3 ))
#        self.setNeighborhoodSize( MAXBOX )        
#        
#    def iterate(self):
#        super(SimplePowerSatellite, self).iterate()
##        print 'Iteration: ', self.getVelocity()
#        self.showNeighborLines()

class SimplePowerSatellite( SimpleStayInBoxSatellite ):
    def __init__(self):
        SimpleStayInBoxSatellite.__init__(self)
        SimplePowerSatellite.init(self)
        
    def init(self):
        self.mass = 225    #kg        
        self.stdColor = breve.vector( 0.3,0.3,0.3 )        
        self.setColor( self.stdColor )
        self.res = Resource(400, 600, 3000)    # usage, generation, maxLevel        
        self.usagePowerMultiplier = 0
        self.usageCycle = 0
        self.setNeighborhoodSize( MAXBOX )        
        
    def iterate(self):
        super(SimplePowerSatellite, self).iterate()
        self.showNeighborLines()

     
    
##################################
# start the sim
myControl(500)