#!/usr/bin/env python
from CBDMultipleOutput.Source.CBD import *

DELTA = 0.1
MASS_PERSON = 73
MASS_TRAIN = 6000
K = 300
C = 150



class TrainControlCBD(CBD):
    def __init__(self, block_name, kp, ki, kd):
        CBD.__init__(self, block_name)
        self.addBlock(FixedTimeCBD(block_name="time"))
        self.addBlock(LookUpBlockFixed(block_name="computer"))
        self.addBlock(AdderBlock(block_name="sum"))
        self.addBlock(NegatorBlock(block_name="neg"))
        self.addBlock(TrainPIDController("controller", kp, ki, kd))
        self.addBlock(TrainPerson(block_name="plant"))
        self.addBlock(CostFunctionCBD(block_name="cost"))
        
        # the connections follow the topological sort.
        self.addConnection("time", "computer", 
                            input_port_name="IN1", output_port_name="OutTime")
        self.addConnection("time", "plant", 
                            input_port_name="InDelta", output_port_name="OutDelta")
        self.addConnection("time", "controller", 
                            input_port_name="InDelta", output_port_name="OutDelta")
        self.addConnection("computer", "sum", 
                            input_port_name="IN1", output_port_name="OUT1")
        self.addConnection("plant", "neg", 
                            input_port_name="IN1", output_port_name="OutVTrain")
        self.addConnection("neg", "sum", 
                            input_port_name="IN2", output_port_name="OUT1")
        self.addConnection("sum", "controller", 
                            input_port_name="InError", output_port_name="OUT1")
        self.addConnection("controller", "plant", 
                            input_port_name="InTraction", output_port_name="OutTraction")
        self.addConnection("sum", "cost", 
                            input_port_name="InError")
        self.addConnection("time", "cost", 
                            input_port_name="InDelta", output_port_name="OutDelta")
        self.addConnection("plant", "cost", 
                            input_port_name="InXPerson", output_port_name="OutXPerson")
        
        

class CostFunctionCBD(CBD):        
    def __init__(self, block_name):
        CBD.__init__(self, block_name, ["InError","InDelta","InXPerson"], ["OutCost"])
        self.addBlock(ABSBlock(block_name="abs"))
        self.addBlock(ABSBlock(block_name="absError"))
        self.addBlock(IntegratorBlock(block_name="integral"))
        self.addBlock(ConstantBlock(block_name="zero", value=0.0))
        self.addBlock(AboveThresholdBlock(block_name="above04", threshold=0.4))
        self.addBlock(StopSimulationBlock(block_name="stop"))
        
        self.addConnection("InError", "absError")
        self.addConnection("absError", "integral", input_port_name="IN1")
        self.addConnection("InDelta", "integral", input_port_name="delta_t")
        self.addConnection("zero", "integral", input_port_name="IC")
        self.addConnection("integral", "OutCost")
        self.addConnection("InXPerson", "abs")
        self.addConnection("abs", "above04")
        self.addConnection("above04", "stop")
        

class ABSBlock(BaseBlock):
    def __init__(self, block_name):
        BaseBlock.__init__(self, block_name, ["IN1"], ["OUT1"])

    def compute(self, curIteration):
        self.appendToSignal(abs(self.getInputSignal(curIteration).value))

class LookUpBlockFixed(BaseBlock):
    def __init__(self, block_name):
        BaseBlock.__init__(self, block_name, ["IN1"], ["OUT1"])
        
    def compute(self, curIteration):
        currentTime = self.getInputSignal(curIteration, "IN1").value
        if currentTime < 0:
            result = 0
        elif currentTime < 160:
            result = 14
        elif currentTime < 200:
            result = 2
        elif currentTime < 260:
            result = 15
        else:
            result = 0
        
        self.appendToSignal(result)



class AboveThresholdBlock(BaseBlock):
    def __init__(self, block_name, threshold):
        BaseBlock.__init__(self, block_name, ["IN1"], ["OUT1"])
        self.threshold = threshold

    def compute(self, curIteration):
        self.appendToSignal(1.0 if self.getInputSignal(curIteration).value > self.threshold else -1.0)



class StopSimulationBlock(BaseBlock):
    def __init__(self, block_name):
        BaseBlock.__init__(self, block_name, ["IN1"], [])

    def compute(self, curIteration):
        inSignalValue = self.getInputSignal(curIteration).value
        if inSignalValue > 0.0:
            raise StopSimulationException()
        

class StopSimulationException(Exception):
    pass
        
class FixedTimeCBD(CBD):        
    def __init__(self, block_name):
        CBD.__init__(self, block_name, [], ["OutTime", "OutDelta"])
        self.addBlock(ConstantBlock(block_name="delta", value=DELTA))
        self.addBlock(ConstantBlock(block_name="zero", value=0.0))
        self.addBlock(AdderBlock(block_name="sum"))
        self.addBlock(DelayBlock(block_name="delay"))
        
        self.addConnection("delta", "OutDelta")
        self.addConnection("delay", "OutTime")
        self.addConnection("delay", "sum")
        self.addConnection("delta", "sum")
        self.addConnection("sum", "delay")
        self.addConnection("zero", "delay",
                            input_port_name="IC")
    

class TrainPIDController(CBD):
    def __init__(self, block_name, kp, ki, kd):
        CBD.__init__(self, block_name, ["InError","InDelta"],["OutTraction"])
        self.addBlock(IntegratorBlock(block_name="integral"))
        self.addBlock(DerivatorBlock(block_name="derivative"))
        self.addBlock(ConstantBlock(block_name="Kp", value=kp))
        self.addBlock(ProductBlock(block_name="multKp"))
        self.addBlock(ConstantBlock(block_name="Ki", value=ki))
        self.addBlock(ProductBlock(block_name="multKi"))
        self.addBlock(ConstantBlock(block_name="Kd", value=kd))
        self.addBlock(ProductBlock(block_name="multKd"))
        self.addBlock(ConstantBlock(block_name="zero", value=0.0))
        self.addBlock(AdderBlock(block_name="sum1"))
        self.addBlock(AdderBlock(block_name="sum2"))
        
        self.addConnection("InError", "multKp")
        self.addConnection("Kp", "multKp")
        self.addConnection("InError", "integral")
        self.addConnection("InError", "derivative")
        self.addConnection("InDelta", "derivative", 
                            input_port_name="delta_t")
        self.addConnection("InDelta", "integral", 
                            input_port_name="delta_t")
        self.addConnection("zero", "derivative", 
                            input_port_name="IC")
        self.addConnection("zero", "integral", 
                            input_port_name="IC")
        self.addConnection("integral", "multKi")
        self.addConnection("Ki", "multKi")
        self.addConnection("derivative", "multKd")
        self.addConnection("Kd", "multKd")
        self.addConnection("multKd", "sum1")
        self.addConnection("multKi", "sum1")
        self.addConnection("sum1", "sum2")
        self.addConnection("multKp", "sum2")
        self.addConnection("sum2", "OutTraction")
        
        
class TrainPerson(CBD):
    def __init__(self, block_name):
        CBD.__init__(self, block_name, ["InTraction","InDelta"],["OutVTrain", "OutXPerson", "OutATrain"])
        self.addBlock(ConstantBlock(block_name="zero", value=0.0))
        self.addBlock(AccelerationPerson(block_name="accelerationPerson"))
        self.addBlock(AccelerationTrain(block_name="accelerationTrain"))
        self.addBlock(IntegratorBlock(block_name="integral_person_a_v"))
        self.addBlock(IntegratorBlock(block_name="integral_person_v_x"))
        self.addBlock(IntegratorBlock(block_name="integral_train_a_v"))
        
        self.addConnection("zero", "integral_person_a_v", 
                            input_port_name="IC")
        self.addConnection("zero", "integral_person_v_x", 
                            input_port_name="IC")
        self.addConnection("zero", "integral_train_a_v", 
                            input_port_name="IC")
        
        self.addConnection("InDelta", "integral_person_a_v", 
                            input_port_name="delta_t")
        self.addConnection("InDelta", "integral_person_v_x", 
                            input_port_name="delta_t")
        self.addConnection("InDelta", "integral_train_a_v", 
                            input_port_name="delta_t")
        
        self.addConnection("InTraction", "accelerationPerson", 
                            input_port_name="InTraction")
        self.addConnection("InTraction", "accelerationTrain", 
                            input_port_name="InTraction")
        
        self.addConnection("accelerationPerson", "integral_person_a_v", 
                            input_port_name="IN1", output_port_name="OutAPerson")
        self.addConnection("integral_person_a_v", "accelerationPerson", 
                            input_port_name="InVPerson", output_port_name="OUT1")
        self.addConnection("integral_person_a_v", "integral_person_v_x", 
                            input_port_name="IN1", output_port_name="OUT1")
        self.addConnection("integral_person_v_x", "accelerationPerson", 
                            input_port_name="InXPerson", output_port_name="OUT1")
        
        self.addConnection("accelerationTrain", "integral_train_a_v", 
                            input_port_name="IN1", output_port_name="OutATrain")
        
        self.addConnection("integral_train_a_v", "OutVTrain")
        
        self.addConnection("integral_person_v_x", "OutXPerson")
        
        self.addConnection("accelerationTrain", "OutATrain", input_port_name=None, output_port_name="OutATrain")
        
        
        
class AccelerationPerson(CBD):
    def __init__(self, block_name):
        CBD.__init__(self, block_name, ["InTraction","InVPerson","InXPerson"],["OutAPerson"])
        self.addBlock(ConstantBlock(block_name="m_train", value=MASS_TRAIN))
        self.addBlock(ConstantBlock(block_name="m_person", value=MASS_PERSON))
        self.addBlock(ConstantBlock(block_name="c", value=C))
        self.addBlock(ConstantBlock(block_name="k", value=K))
        self.addBlock(AdderBlock(block_name="sum_masses"))
        self.addBlock(InverterBlock(block_name="inv_sum_masses"))
        self.addBlock(ProductBlock(block_name="mult_inv_sum_masses"))
        self.addBlock(ProductBlock(block_name="mult_mult_inv_sum_masses"))
        self.addBlock(InverterBlock(block_name="inv_mperson"))
        self.addBlock(NegatorBlock(block_name="neg_mult_mult_inv_sum_masses"))
        self.addBlock(NegatorBlock(block_name="neg_vperson"))
        self.addBlock(NegatorBlock(block_name="neg_xperson"))
        self.addBlock(ProductBlock(block_name="mult_neg_vperson"))
        self.addBlock(ProductBlock(block_name="mult_neg_xperson"))
        self.addBlock(AdderBlock(block_name="sum_mults_person"))
        self.addBlock(AdderBlock(block_name="sum_sum_mults_person"))
        self.addBlock(ProductBlock(block_name="mult_inv_mperson"))
        
        self.addConnection("m_train", "sum_masses")
        self.addConnection("m_person", "sum_masses")
        self.addConnection("sum_masses", "inv_sum_masses")
        self.addConnection("inv_sum_masses", "mult_inv_sum_masses")
        self.addConnection("InTraction", "mult_inv_sum_masses")
        self.addConnection("mult_inv_sum_masses", "mult_mult_inv_sum_masses")
        
        self.addConnection("m_person", "mult_mult_inv_sum_masses")
        
        self.addConnection("m_person", "inv_mperson")
        self.addConnection("inv_mperson", "mult_inv_mperson")
        
        self.addConnection("mult_mult_inv_sum_masses", "neg_mult_mult_inv_sum_masses")
        self.addConnection("neg_mult_mult_inv_sum_masses", "sum_sum_mults_person")
        
        self.addConnection("InVPerson", "neg_vperson")
        self.addConnection("neg_vperson", "mult_neg_vperson")
        self.addConnection("c", "mult_neg_vperson")
        
        self.addConnection("InXPerson", "neg_xperson")
        self.addConnection("neg_xperson", "mult_neg_xperson")
        self.addConnection("k", "mult_neg_xperson")
        
        self.addConnection("mult_neg_vperson", "sum_mults_person")
        self.addConnection("mult_neg_xperson", "sum_mults_person")
        
        self.addConnection("sum_mults_person", "sum_sum_mults_person")
        self.addConnection("sum_sum_mults_person", "mult_inv_mperson")
        self.addConnection("mult_inv_mperson", "OutAPerson")
        
        
        
class AccelerationTrain(CBD):
    def __init__(self, block_name):
        CBD.__init__(self, block_name, ["InTraction"],["OutATrain"])
        self.addBlock(ConstantBlock(block_name="m_train", value=MASS_TRAIN))
        self.addBlock(ConstantBlock(block_name="m_person", value=MASS_PERSON))
        self.addBlock(AdderBlock(block_name="sum"))
        self.addBlock(InverterBlock(block_name="inv"))
        self.addBlock(ProductBlock(block_name="mult"))
        
        self.addConnection("InTraction", "mult")
        self.addConnection("m_train", "sum")
        self.addConnection("m_person", "sum")
        self.addConnection("sum", "inv")
        self.addConnection("inv", "mult")
        self.addConnection("mult", "OutATrain")









