import argparse
import sys

try:
    from ai_edge_litert.interpreter import Interpreter
except ImportError as e:
    print("\n\n")
    print("Install the ""ai_edge_litert"" package with the command: 'pip install ai_edge_litert'")
    print("\n\n")
    sys.exit()

try:
    import numpy as np
except ImportError as e:
    print("\n\n")
    print("Install the ""numpy"" package with the command 'pip install numpy'")
    print("\n\n")
    sys.exit()    

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'



parser = argparse.ArgumentParser(prog='Get Top N birds')
parser.add_argument('-n', help="Number of Top N to list", default = 10, type=int)
parser.add_argument('--lat', help='Latitude', default= 42.2, type=float)
parser.add_argument('--long', help='Longitude', default= -83.0, type=float)
parser.add_argument('--week_number', help='Week number', default=1, type=int)

args, _ = parser.parse_known_args()


top_n_print = args.n
lat = args.lat
lng = args.long
week_number = args.week_number


print('\n\n')
print(f'Querying with:')
print(f'\tLatitude:  {lat}')
print(f'\tLongitude: {lng}')
print(f'\tWeek Number: {week_number}')
print('\n')






interpreter = Interpreter(model_path="geo_v46.tflite")
interpreter.allocate_tensors()

with open('labels_species.txt','r') as ff:
    output_label = ff.read().split('\n')


interpreter.set_tensor(0, [np.float32( lng ) ])
interpreter.set_tensor(1, [np.float32( week_number ) ])   #First week of the year
interpreter.set_tensor(2, [np.float32(  lat )] )
interpreter.invoke()

output_probabilities_ish = interpreter.get_tensor(88).squeeze()
sorted_idx = output_probabilities_ish.argsort()[::-1]

print('\n')
print(f'Showing {top_n_print} results')
print('Rank\t"Probability"\tSpecies')
for id_rank, rank in zip(sorted_idx, range(top_n_print)):
    print(rank,'\t', '{0:2.2f}'.format(output_probabilities_ish[id_rank]) , '\t\t',output_label[id_rank] )
print('\n')