import PIL 
from PIL import Image 
import os
import shutil

# Global Vars
# Folders' name. Each folder serves as one kind of level for image upscale
UPSCALE_FOLDERS = ['X2', 'X3', 'X4']
X4_FOLDER = 'X4'
X3_FOLDER = 'X3'
X2_FOLDER = 'X2'
# Image's required format 
imgFormat = '.png'
# Min and Max limit for different upscale level
X2_Min = 2000**2
X2_Max = 3200**2
X3_Min = 1500**2
X3_Max = X2_Min
X4_Max = X3_Min
############################################################################

def getImgRes(imgName):
    curImg = PIL.Image.open(imgName)
    curImgWid, curImgHgt = curImg.size
    return curImgWid, curImgHgt

def getAllImgs():
    files = [f for f in os.listdir() if os.path.isfile(f)]
    imgs = []
    for eachFile in files:
        if len(eachFile) >= 4 and eachFile[-4:] == imgFormat:
            #print("Added: ", eachFile)
            imgs.append(eachFile)
            imgWid, imgHgt = getImgRes(eachFile)
            #print('Width: ' + str(imgWid) + ' Height: ' + str(imgHgt))
            res = imgWid*imgHgt
            if(res <= X4_Max):
                shutil.move(eachFile, X4_FOLDER)
            elif(res > X3_Min and res <= X3_Max):
                shutil.move(eachFile, X3_FOLDER)
            elif(res > X2_Min and res <= X2_Max):
                shutil.move(eachFile, X2_FOLDER)
    return imgs


#############################################
for folder in UPSCALE_FOLDERS:
    if not os.path.exists(folder):
        os.makedirs(folder)

print(getAllImgs())