Problème d'algorithme simple

Bonjour,

Je précise que je n'ai jamais fait de Python, mais là mon avenir direct en dépend.

Je dois réaliser un mini script, voilà l'énoncé du problème :

Un socle est ainsi constitué d'étages, chaque étage ayant une hauteur égale à une unité et une base carrée. Le côté des carrés diminue de une unité à chaque étage.

Votre programme doit lire deux entiers, représentant respectivement la largeur du socle au niveau du sol et la largeur du socle au niveau de la face supérieure du socle. Il doit ensuite calculer et afficher le volume du socle.

Si on passe au script la valeur 7 puis ensuite la valeur 3, on doit obtenir en sortie 135.

Voila les explications pour arriver à ce résultat :

Le volume du premier étage est 7 × 7 = 49, le volume du second est 6 × 6 = 36, etc. et le volume du dernier étage est 3 × 3 = 9. Le volume total est donc : 7 × 7 + 6 × 6 + 5 × 5 + 4 × 4 + 3 × 3 = 135.

Mon algorithme (pourtant parfait ;) fonctionne avec certaines valeurs d'entrée mais pas avec tous, je ne comprends vraiment pas pourquoi :

largeur_Sol = int(input())
largeur_Sup = int(input())

count = largeur_Sol - largeur_Sup
result = largeur_Sol * largeur_Sol

for loop in range(count):
   etage = (largeur_Sol - count) * (largeur_Sol - count)
   result = result + etage

print(result)

Merci d'avance si vous pouvez m'expliquer ce qui cloche.