import argparse
import numpy as np
from pathlib import Path
from google_map import GoogleMap
from simulator import Simulator
from yandex_map import YandexMap
from constants import CHUNK_WIDTH

LAT_MIN, LAT_MAX = 49.134520, 49.235065
LON_MIN, LON_MAX = 55.767660, 55.825204



def create_new_asset(yandex_map, google_map, dname, tile_y=None, tile_x=None):
    folder = Path(f'dataset_{dname}')

    id = 0
    print(id)
    while (folder / f"{id:0{4}}_google.png").exists():
        id += 1
    
    google_file = folder / f"{id:0{4}}_google.png"
    yandex_file = folder / f"{id:0{4}}_yandex.png"

    lat = np.random.rand() * (LAT_MAX - LAT_MIN) + LAT_MIN
    lon = np.random.rand() * (LON_MAX - LON_MIN) + LON_MIN
    if tile_y is not None and tile_x is not None:
        lat = tile_y * (LAT_MAX - LAT_MIN) + LAT_MIN
        lon = tile_x * (LON_MAX - LON_MIN) + LON_MIN

    yandex_map.open(lat, lon, 18)
    google_map.open(lat, lon, 18)

    simulator = Simulator()
    im_ya = simulator._apply_perspective_transform(yandex_map.make_screenshot())
    im_go = simulator._apply_perspective_transform(google_map.make_screenshot())
    im_ya.resize((CHUNK_WIDTH // 2, CHUNK_WIDTH // 2)).save(yandex_file)
    im_go.resize((CHUNK_WIDTH // 2, CHUNK_WIDTH // 2)).save(google_file)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--dname', type=str, default='ya_go_maps', help='Dataset name')
    args = parser.parse_args()

    folder = Path(f'dataset_{args.dname}')
    if not folder.exists():
        folder.mkdir()

    yandex_map = YandexMap(initial_zoom=18)
    google_map = GoogleMap(initial_zoom=18)

    for tile_y in np.linspace(0, 1, 20):
        for tile_x in np.linspace(0, 1, 20):
            create_new_asset(yandex_map, google_map, args.dname, tile_y, tile_x)

if __name__ == '__main__':
    main()