Trying to make it so function main() recognizes specific integers when calling upon function problem_solving

I am a total novice so bear with me here. Basically, I want the main() function to be able to pick up on a specific integer. As in, if my input for step_graded was ‘1’, I could write gradeforstep = problem_solving(‘1’) and be taken to stepgrade=input(‘Please enter a score 0 – 5 earned for understanding step 1:’). Any help would be greatly appreciated.

def problem_solving(step_graded):
    if step_graded == '1':
        stepgrade=input('Please enter a score 0 - 5 earned for understanding step 1:')
        finalstepgrade1 = f"Step: {step_graded} Grade: {stepgrade}."
        return finalstepgrade1
    if step_graded == '2':
        stepgrade=input('Please enter a score 0 - 5 earned for understanding step 2:')
        finalstepgrade = f"Step: {step_graded} Grade: {stepgrade}."
        return finalstepgrade
    if step_graded == '3':
        stepgrade=input('Please enter a score 0 - 5 earned for understanding step 3:')
        finalstepgrade = f"Step: {step_graded} Grade: {stepgrade}."
        return finalstepgrade
    if step_graded == '4':
        stepgrade=input('Please enter a score 0 - 5 earned for understanding step 4:')
        finalstepgrade = f"Step: {step_graded} Grade: {stepgrade}."
        return finalstepgrade

def main():
    gradeforstep = problem_solving(input)
    print (gradeforstep)
main()

>Solution :

You need to get the value of step_graded from the user before calling problem_solving.

You can simplify problem_solving because all the cases are the same except for the step number, which you can substitute using the variable.

def problem_solving(step_graded):
    stepgrade=input(f'Please enter a score 0 - 5 earned for understanding step {step_graded}:')
    finalstepgrade = f"Step: {step_graded} Grade: {stepgrade}."
    return finalstepgrade

def main():
    step = input("What step have you done:")
    gradeforstep = problem_solving(step)
    print (gradeforstep)

main()