from DEVS import *
from infinity import INFINITY

import random, math, sys, getopt

''' The simulation does not end. '''
def termination_condition(time, model, transitioned):
    time = time[0]
    return False
    
class Canvas:
    def __init__(self, resolution):
        self.width = resolution[0]
        self.height = resolution[1]
        
class ParticleState:
    def __init__(self, canvas, particle_id):
        self.r = random.uniform(5.0, 30.0)
        x = random.uniform(self.r, canvas.width - self.r)
        y = random.uniform(self.r, canvas.height - self.r)
        self.pos = (x, y)
        self.vel = {'x': random.uniform(-75.0, 75.0), 'y': random.uniform(-75.0, 75.0)}
        self.phase = "initializing" # phases: ('initializing', 'initialized', 'bouncing', 'spawning', 'selected', 'deleting', 'clicked')
        self.color = "red"
        self.prev_color = "red"
        self.collision_detected = False
        self.clicked = False
        self.frames_passed = 0
        self.frames_remaining = 0
        self.remaining = 0
        self.particle_id = particle_id

class Particle(AtomicDEVS):
    def __init__(self, particle_id, canvas, framerate, spawn_chance, die_chance):
        AtomicDEVS.__init__(self, "Particle[%i]" % particle_id)
        
        self.particle_id = particle_id
        self.canvas = canvas
        self.framerate = framerate
        self.spawn_chance = spawn_chance
        self.die_chance = die_chance
        
        self.COLLISION_DETECT = self.addInPort("COLLISION_DETECT")
        self.POS_OUT = self.addOutPort("POS_OUT")
        self.SPAWNER_COMM = self.addOutPort("SPAWNER_COMM")
        self.COLOR_OUT = self.addOutPort("COLOR_OUT")
        
        self.state = ParticleState(canvas, particle_id)
    
    def timeAdvance(self):
        return self.state.remaining
    
    def outputFnc(self):
        output = {}
        if self.state.phase in ("bouncing", "spawning", "selected") or self.state.clicked:
            output[self.POS_OUT] = (self.particle_id, self.state.pos, self.state.r, self.state.vel)
        if self.state.collision_detected and self.state.phase != "spawning" and random.random() < self.spawn_chance:
            output[self.SPAWNER_COMM] = ["spawn_particle"]
        if self.state.phase == "selected" and self.state.prev_color == self.state.color:
            output[self.POS_OUT] = ("delete", self.particle_id)
        if self.state.color != self.state.prev_color:
            output[self.COLOR_OUT] = (self.particle_id, self.state.color)
        if self.state.phase == "initializing":
            output[self.POS_OUT] = ("created", self.particle_id, self.state.pos, self.state.r, self.state.vel)
        return output
    
    def intTransition(self):
        if self.state.phase == "initializing":
            self.state.phase = "initialized"
            self.state.remaining = 0
        elif self.state.phase == "initialized":
            self.state.phase = "bouncing"
            self.state.remaining = 1
        elif self.state.phase == "selected":
            if self.state.prev_color == self.state.color:
                self.state.phase = "deleting"
                self.state.remaining = INFINITY
            else:
                self.state.prev_color = self.state.color
                self.state.remaining = int(2.5 * self.framerate)
        elif self.state.prev_color != self.state.color:
            self.state.prev_color = self.state.color
            self.state.remaining = 0
        elif self.state.phase == "clicked":
            self.state.remaining = INFINITY
        elif random.random() < self.die_chance:
            self.state.vel = {'x': 0, 'y': 0}
            self.state.color = "yellow"
            self.state.phase = "selected"
            self.state.remaining = 0
        else:
            if self.state.frames_passed < self.framerate and self.state.phase in ("bouncing", "spawning") and self.state.color != "black":
                self.state.frames_passed += 1
            if self.state.frames_passed >= self.framerate and self.state.phase == "bouncing":
                self.state.color = "black"
                self.state.phase = "bouncing"
                self.state.frames_passed = 0
            if self.state.frames_passed >= 1 * self.framerate and self.state.phase == "spawning":
                self.state.color = "black"
                self.state.phase = "bouncing"
                self.state.frames_passed = 0
            x = self.state.pos[0]
            y = self.state.pos[1]
            if (x - self.state.r) <= 0 or x + self.state.r >= self.canvas.width:
                self.state.vel['x'] = -self.state.vel['x']
            if (y - self.state.r) <= 0 or y + self.state.r >= self.canvas.height:
                self.state.vel['y'] = -self.state.vel['y']
            self.state.pos = (self.state.pos[0] + (self.state.vel['x'] / self.framerate), self.state.pos[1] + (self.state.vel['y'] / self.framerate))
            if self.state.collision_detected:
                self.state.phase = "spawning"
                self.state.color = "blue"
            elif self.state.clicked:
                self.state.vel = {'x': 0, 'y': 0}
                self.state.phase = "clicked"
                self.state.color = "orange"
                self.state.clicked = False
            self.state.remaining = 1
        self.state.collision_detected = False
        return self.state
    
    def extTransition(self, inputs):
        if self.COLLISION_DETECT in inputs:
            if inputs[self.COLLISION_DETECT] == "delete_if_selected":
                if self.state.phase == "clicked":
                    self.state.phase = "selected"
                    self.state.remaining = 0
            elif inputs[self.COLLISION_DETECT] == "clicked":
                if self.state.phase == "bouncing":
                    self.state.clicked = True
                    self.state.remaining = 0
            elif self.state.phase in ("bouncing", "clicked", "deleting"):
                self.state.phase = "bouncing"
                self.state.collision_detected = True
                new_speed = inputs[self.COLLISION_DETECT][1]
                self.state.vel = {'x': new_speed[0], 'y': new_speed[1]}
                self.state.remaining = 0
            #else:
                #self.state.remaining -= self.elapsed
        return self.state
    
    def modelTransition(self, passed_values):
        if self.state.phase == "deleting":
            passed_values['to_delete'] = self
            return True
        elif self.state.phase == "initialized":
            passed_values['connect_particle'] = (self.particle_id, self.COLLISION_DETECT)
            return True
        return False
        
class ParticleSpawnerState:
    def __init__(self):        
        self.id_ctr = 0
        self.new_particle = None
        self.remaining = 0

class ParticleSpawner(AtomicDEVS):
    def __init__(self, canvas, framerate, spawn_chance, die_chance):
        AtomicDEVS.__init__(self, "ParticleSpawner")
        
        self.canvas = canvas
        self.framerate = framerate
        self.spawn_chance = spawn_chance
        self.die_chance = die_chance
        
        self.REQUEST = self.addInPort("REQUEST")
        
        self.state = ParticleSpawnerState()
    
    def timeAdvance(self):
        return self.state.remaining
    
    def intTransition(self):
        self.state.new_particle = Particle(self.state.id_ctr, self.canvas, self.framerate, self.spawn_chance, self.die_chance)
        self.state.id_ctr += 1
        self.state.remaining = self.framerate
        return self.state
    
    def extTransition(self, inputs):
        self.state.remaining -= self.elapsed
        if self.REQUEST in inputs:
            self.state.new_particle = Particle(self.state.id_ctr, self.canvas, self.framerate, self.spawn_chance, self.die_chance)
            self.state.id_ctr += 1
        return self.state
    
    def modelTransition(self, passed_values):
        if self.state.new_particle:
            passed_values['new_particle'] = self.state.new_particle
            return True
        return False

class PositionManagerState:
    def __init__(self):        
        self.positions = {}
        self.collisions = set([])
        self.phase = "detecting"
        self.clicked = None
        self.delete_selected = False
        self.frames = 0
        self.remaining = 0
    
    def __str__(self):
        return str((self.phase, self.frames, self.remaining))

class PositionManager(AtomicDEVS):
    def __init__(self, framerate):
        AtomicDEVS.__init__(self, "PositionManager")
        
        self.framerate = framerate
        
        self.TIME_OUT = self.addOutPort("TIME_OUT")
        self.POS_IN = self.addInPort("POS_IN")
        self.INTERRUPT = self.addInPort("INTERRUPT")
        
        self.state = PositionManagerState()
    
    def timeAdvance(self):
        return self.state.remaining
    
    def outputFnc(self):
        output = {}
        if self.state.collisions:
            output.update({self.ports['PARTICLES_OUT[%i]' % particle_id]: ("collision_detected", deltap) for (particle_id, deltap) in self.state.collisions})
        if self.state.clicked is not None:
            portname = 'PARTICLES_OUT[%i]' % self.state.clicked
            if portname in self.ports:
                output[self.ports[portname]] = "clicked"
        if self.state.delete_selected:
            output.update({self.ports['PARTICLES_OUT[%i]' % particle_id]: "delete_if_selected" for particle_id in self.state.positions})
        output[self.TIME_OUT] = self.state.frames
        return output
    
    def intTransition(self):
        self.state.frames += self.timeAdvance()
        if self.state.clicked is not None:
            self.state.clicked = None
            self.state.remaining = 1
        if self.state.delete_selected:
            self.state.delete_selected = False
            self.state.remaining = 1
        elif self.state.phase == "detecting":
            for k1, v1 in self.state.positions.iteritems():
                for k2, v2 in self.state.positions.iteritems():
                    if k1 > k2:
                        dx = v2[0][0] - v1[0][0]
                        dy = v2[0][1] - v1[0][1]
                        distance = math.sqrt(dx**2 + dy**2)
                        if (distance < (v1[1] + v2[1])) and (k1 not in self.state.collisions or k2 not in self.state.collisions):
                            new_speed_1 = (v2[2]['x'], v2[2]['y'])
                            new_speed_2 = (v1[2]['x'], v1[2]['y'])
                            self.state.collisions.add((k1, new_speed_1))
                            self.state.collisions.add((k2, new_speed_2))
            if self.state.collisions:
                self.state.phase = "detected"
                self.state.remaining = 0
            else:
                self.state.phase = "detecting"
                self.state.remaining = 1
        elif self.state.phase == "detected":
            self.state.collisions = set([])
            self.state.phase = "detecting"
            self.state.remaining = 1
        return self.state
    
    def extTransition(self, inputs):
        self.state.remaining -= self.elapsed
        self.state.frames += self.elapsed
        if self.POS_IN in inputs:
            msg = inputs[self.POS_IN]
            if msg[0] == "created":
                new_particle_id = msg[1]
                self.addOutPort("PARTICLES_OUT[%s]" % new_particle_id)
                self.state.positions[new_particle_id] = (msg[2], msg[3], msg[4])
            elif msg[0] == "delete":
                self.removePort(self.ports["PARTICLES_OUT[%s]" % msg[1]])
                del self.state.positions[msg[1]]
            else:
                particle_id = msg[0]
                r = msg[2]
                self.state.positions[particle_id] = (msg[1], msg[2], msg[3])
        # TODO: Fix this, needs to be rounded to integer.
        if self.INTERRUPT in inputs:
            msg = inputs[self.INTERRUPT]
            if msg == "delete_selected":
                self.state.delete_selected = True
            else:
                self.state.clicked = int(msg)
        return self.state

class Root(CoupledDEVS):
    def __init__(self, canvas=Canvas((800, 600)), framerate=30.0, spawn_chance=0.2, die_chance=0.01):
        CoupledDEVS.__init__(self, "Field")
        
        self.INTERRUPT = self.addInPort("INTERRUPT")
        self.POS_OUT = self.addOutPort("POS_OUT")
        self.COLOR_OUT = self.addOutPort("COLOR_OUT")
        self.TIME_OUT = self.addOutPort("TIME_OUT")
        
        self.particle_spawner = self.addSubModel(ParticleSpawner(canvas, framerate, spawn_chance, die_chance))
        self.position_manager = self.addSubModel(PositionManager(framerate))
        self.connectPorts(self.INTERRUPT, self.position_manager.INTERRUPT)
        self.connectPorts(self.position_manager.TIME_OUT, self.TIME_OUT)
            
    def modelTransition(self, passed_values):
        if 'new_particle' in passed_values:
            new_particle = passed_values['new_particle']
            self.addSubModel(new_particle)
            self.connectPorts(new_particle.POS_OUT, self.position_manager.POS_IN)
            self.connectPorts(new_particle.POS_OUT, self.POS_OUT)
            self.connectPorts(new_particle.COLOR_OUT, self.COLOR_OUT)
            self.connectPorts(new_particle.SPAWNER_COMM, self.particle_spawner.REQUEST)
            del passed_values['new_particle']
        if 'to_delete' in passed_values:
            self.removeSubModel(passed_values['to_delete'])
            del passed_values['to_delete']
        if 'connect_particle' in passed_values:
            particle_id, COLLISION_DETECT = passed_values['connect_particle']
            self.connectPorts(self.position_manager.ports['PARTICLES_OUT[%s]' % particle_id], COLLISION_DETECT)
            del passed_values['connect_particle']
        return False
        
    def select(self, imm_children):
        pos_mgr, particle_spawner = [None] * 2
        for c in imm_children:
            if c.getModelName() == "ParticleSpawner":
                particle_spawner = c
        if particle_spawner:
            return particle_spawner
        else:
            return imm_children[0]