import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import TransferFunction, step, convolve

def step_response(numerator, denominator, t_end, step_amplitude):
    # Define the transfer function H(s) = numerator / denominator
    system = TransferFunction(numerator, denominator)

    # Generate time points
    t = np.linspace(0, t_end, 1000)
    
    # Calculate the step response
    t, response = step(system, T=t)
    
    # Scale the response by the step amplitude
    response *= step_amplitude
    
    # Plot the step response
    plt.plot(t, response)
    plt.title('Step Response of the System')
    plt.xlabel('Time (s)')
    plt.ylabel('Response')
    plt.grid(True)
    plt.show()
    
    return t, response

# Example usage
T = 1  # Time constant
numerator = [1]  # Numerator of the transfer function
den1 = [1, 1, 1]  # First part of the denominator
den2 = [T, 1]  # Second part of the denominator

# Convolution of the two denominator parts
denominator = convolve(den1, den2)

t_end = 20  # Duration of the simulation
step_amplitude = 0.1  # Step amplitude

t, response = step_response(numerator, denominator, t_end, step_amplitude)