%matplotlib inline

from progpy.models.pneumatic_valve import PneumaticValve

def run_example(): 
    # Create a model object
    valv = PneumaticValve(process_noise= 0)

    # Define future loading function
    cycle_time = 20
    def future_loading(t, x=None):
            t = t % cycle_time
            if t < cycle_time/2:
                return valv.InputContainer({
                    'pL': 3.5e5,
                    'pR': 2.0e5,
                    # Open Valve
                    'uTop': False,
                    'uBot': True
                })
            return valv.InputContainer({
                'pL': 3.5e5,
                'pR': 2.0e5,
                # Close Valve
                'uTop': True,
                'uBot': False
            })

    # Simulate to threshold
    print('\n\n------------------------------------------------')
    print('Simulating to threshold\n\n')
    # Configure options
    config = {
        'dt': 0.01,
        'horizon': 800,
        'save_freq': 60,
        'print': True,
        'progress': True,
    }
    # Set wear parameter for spring to 1
    valv.parameters['x0']['wk'] = 1

    # Define first measured output. This is needed by the simulat_to_threshold method to initialize state
    first_output = valv.output(valv.initialize(future_loading(0)))
    # Simulate
    simulated_results = valv.simulate_to_threshold(future_loading, first_output, **config)

    # Simulate to threshold again but with a different wear mode
    print('\n\n------------------------------------------------')
    print('Simulating to threshold\n\n')
    # Configure options
    config = {
        'dt': 0.01,
        'horizon': 800,
        'save_freq': 60,
        'print': True,
        'progress': True
    }
    # Reset wear parameter for spring to 0, set wear parameter for friction to 1
    valv.parameters['x0']['wk'] = 0
    valv.parameters['x0']['wr'] = 1

    # Define first measured output. This is needed by the simulat_to_threshold method to initialize state
    first_output = valv.output(valv.initialize(future_loading(0)))
    # Simulate
    simulated_results = valv.simulate_to_threshold(future_loading, first_output, **config)

# This allows the module to be executed directly
if __name__ == '__main__':
    run_example()