from CBDMultipleOutput.Source.CBD import ConstantBlock, NegatorBlock, \
    IntegratorBlock, CBD, ProductBlock, AdderBlock
from fsa_cbd_weaved_classes import InputConstantBlock, CBDState, WhenCrossesZeroTrigger
from fsaclasses import FSAModel, Transition, Event, State


class BouncingBallFSA(FSAModel):
    def __init__(self):
        initial = CBDState("Initial", BouncingBallInitial("BouncingBallInitial"))
        freeFall = CBDState("FreeFall", BouncingBallFreefall("BouncingBallFreefall"))
        collision = CBDState("Collision", BouncingBallCollision("BouncingBallCollision"))
        kicked = CBDState("Kicked", BouncingBallKicked("BouncingBallKicked"))
        end = State("End", True)
        
        initial2Freefall = Transition("Initial2Freefall", initial, freeFall)
        freefall2Collision = Transition("Freefall2Collision", freeFall, collision)
        freefall2Collision.trigger = WhenCrossesZeroTrigger("x", up_direction=False);
        collision2Freefall = Transition("Collision2Freefall", collision, freeFall)
        freefall2Kicked = Transition("Freefall2Kicked", freeFall, kicked)
        freefall2Kicked.trigger = Event("kick");
        kicked2Freefall = Transition("kicked2Freefall", kicked, freeFall)
        freefall2End = Transition("Freefall2End", freeFall, end)
        freefall2End.trigger = Event("enough");
        
        self.states = [
                       initial,
                       freeFall,
                       collision,
                       kicked,
                       end
                       ]
        
        self.initialState = initial
        
        self.transitions = [
                            initial2Freefall,
                            freefall2Collision,
                            collision2Freefall,
                            freefall2Kicked,
                            kicked2Freefall,
                            freefall2End
                            ]
        
       
    
        
        

class BouncingBallFreefall(CBD):
    def __init__(self, blockName, delta_t = 0.01):
        CBD.__init__(self,blockName, input_ports=[], output_ports=["out_v", "out_x"])
        
        self.addBlock(InputConstantBlock(block_name="in_x"))
        self.addBlock(InputConstantBlock(block_name="in_v"))
        
        self.addBlock(ConstantBlock(block_name="gravity", value=-9.81))
        self.addBlock(ConstantBlock(block_name="delta", value=delta_t))
        
        self.addBlock(IntegratorBlock(block_name="intgv"))
        self.addBlock(IntegratorBlock(block_name="intvx"))
        
        self.addConnection("delta", "intgv",  input_port_name="delta_t")
        self.addConnection("delta", "intvx",  input_port_name="delta_t")
        self.addConnection("in_v", "intgv",  input_port_name="IC")
        self.addConnection("in_x", "intvx",  input_port_name="IC")
        
        self.addConnection("gravity", "intgv")
        self.addConnection("intgv", "intvx")
        self.addConnection("intgv", "out_v")
        self.addConnection("intvx", "out_x")
    
    def setDeltaT(self, deltaT):
        CBD.setDeltaT(self, deltaT)
        self.getBlockByName("delta").setValue(deltaT);
    
    

class BouncingBallInitial(CBD):
    def __init__(self, blockName, delta_t=0.01):
        CBD.__init__(self,blockName, input_ports=[], output_ports=["out_v", "out_x"])
        
        self.addBlock(ConstantBlock(block_name="init_velocity", value=15.0))
        self.addBlock(ConstantBlock(block_name="init_position", value=10.0))
        
        self.addConnection("init_velocity", "out_v")
        self.addConnection("init_position", "out_x")
    


class BouncingBallCollision(CBD):
    def __init__(self, blockName, delta_t=0.01):
        CBD.__init__(self,blockName, input_ports=[], output_ports=["out_v"])
        
        self.addBlock(InputConstantBlock(block_name="in_v"))
        
        self.addBlock(ConstantBlock(block_name="coeficient", value=-0.8))
        self.addBlock(ProductBlock(block_name="mult"))
        
        self.addConnection("in_v", "mult")
        self.addConnection("coeficient", "mult")
        self.addConnection("mult", "out_v")


class BouncingBallKicked(CBD):
    def __init__(self, blockName, delta_t=0.01):
        CBD.__init__(self,blockName, input_ports=[], output_ports=["out_v"])
        
        self.addBlock(InputConstantBlock(block_name="in_v"))
        
        self.addBlock(ConstantBlock(block_name="kick_impact", value=20))
        self.addBlock(AdderBlock(block_name="sum"))
        
        self.addConnection("in_v", "sum")
        self.addConnection("kick_impact", "sum")
        self.addConnection("sum", "out_v")


