import json
import numpy as np
import os
import re
import sys

from tqdm import tqdm
# from PIL import Image


weights = (9, 4, 5, 4, 3, 2)
VALID_CHECKSUMS = 'AZYXUTSRPMLKJHGEDCB'
denominion = len(VALID_CHECKSUMS)


# def main():
#     data = json.load(open(sys.argv[1], 'rb'))
#     id2path = {im['id']:im['path'] for im in data['images']}
#     id2cat = {c['id']:c['name'] for c in data['categories']}

#     if os.path.exists('keypoints.json'):
#         with open('keypoints.json', 'rb') as f:
#             kpt_dict = json.load(f)
#     else:
#         kpt_dict = {}

#     for a in tqdm(data['annotations'], leave=False):
#         im_id = a['image_id']
#         path = id2path[im_id]

#         try:
#             cat = id2cat[a['category_id']]
#             pts = np.array(a['keypoints'], np.float32).reshape([4, 3])[:, :2]
#             box = a['bbox']
#         except:
#             continue

#         if 'single' not in cat and 'double' not in cat:
#             continue

#         im = Image.open(path)
#         try:
#             number = a['metadata']['name']
#             assert is_valid(number)
#             l, t, w, h = box
#             r = l + w
#             b = t + h
#             image_rect = im.crop([l, t, r, b])
#             pts[:, 0] -= l
#             pts[:, 1] -= t
#             image_name = f'{number}_{im_id}.jpg'
#             kpt_dict[image_name] = pts.flatten().tolist()
#             image_rect.save(os.path.join(sys.argv[2], image_name))
#         except:
#             continue
#     with open('keypoints.json', 'w') as f:
#         json.dump(kpt_dict, f, indent=2)


def is_valid(license_number):
    if not isinstance(license_number, str):
        return False
    elif len(license_number) in [0, 1, 2] or len(license_number) > 8:
        return False
    elif license_number.islower():
        return False
    elif 'I' in license_number or 'O' in license_number:
        return False
    # partition the number
    try:
        alphas, numbers, checksum = partition(license_number)
        return get_checksum(alphas, numbers) == checksum
    except ValueError as ve:
        return False


regex1 = re.compile(r'^[A-Z]{1,3}')
regex2 = re.compile(r'^[1-9][0-9]{0,3}')
regex3 = re.compile(r'^[AZYXUTSRPMLKJHGEDCB]')
def partition(license_number):
    try:
        alphas = regex1.match(license_number).group(0)
        license_number = license_number[len(alphas):]
        numbers = regex2.match(license_number).group(0)
        license_number = license_number[len(numbers):]
        checksum = regex3.match(license_number).group(0)
        return alphas, numbers, checksum
    except AttributeError:
        raise ValueError('illegal license number: %s' % license_number)


def get_checksum(alphas, numbers):
    l_alpha = len(alphas)
    if l_alpha == 1:
        alphas = alphas * 2
    elif l_alpha == 2:
        pass
    elif l_alpha == 3:
        alphas = alphas[1:]
    else:
        raise ValueError('too many numbers : %s' % numbers)

    numbers = '%04d' % int(numbers)
    code = list(ord(x) - 64 for x in alphas)
    code += list(map(int, numbers))
    check = sum(x * y for x,y in zip(code, weights))
    return VALID_CHECKSUMS[check % denominion]


# if __name__ == '__main__':
#     main()
