import os
import requests
import time
from math import log, tan, cos, pi
from tqdm import tqdm
import shutil

# Define the bounding boxes and zoom levels   <------------- uncomment one of these, OR add your own bounding box
regions = {
#    "Victoria, central": (-38.82,143.74,-37.12,146.09)
#    "Victoria, SE": (-38.51,144.84,-36.69,147.25)
#     "Melbourne, East": (-38.03,144.85,-37.67,145.44)

#    "USA": (24.28,-125.57,49.62,-59.08)
#    "Alaska": (54.82,-168.99,71.53,-136.24)

#     "World": (-80.0,-170.0,80.0,170.0)

#     "Europe": (34.93,-11.36,62.46,39.98)

    "UK": (48.89,-10.77,59.86,1.93)
#    "Greater London": (50.62,-1.7,52.24,1.53)

#    "cbd": (-37.87, 144.91, -37.77, 145.03)
#    "Australia": (-44.46, 113.07, -10.21, 155.02)
#    "Argentina": (-55.96, -76.08, -21.09, -53.22)
#     "France": (41.78,-5.29,51.38,8.45),
#     "Germany": (47.24,5.6,54.95,15.45),
#     "Switzerland": (45.76,5.95,47.81,10.71)

#    "South-West USA": (28.63, -125.08, 42.26, -102.74)

#      "New Mexico": (31.27,-109.08,37.05,-102.91),
#      "Oklahoma and Nth Texas": (31.3,-102.91,37.09,-93.54),
#      "South Texas": (25.82,-100.35,31.3,-93.45),
#      "South-west Texas": (28.51,-106.29,31.3,-100.35)

#      "Florida-N": (29.03,-87.73,31.15,-80.8),
#      "Florisa-S": (24.23,-83.38,29.03,-79.76),

#      "Arkansas, Mississippi, Louisiana": (28.81,-94.63,36.7,-87.22),
#      "Georgia, Alabama, Sth Carolina, Tennessee": (30.64,-87.22,36.7,-78.91),
#      "Nth Carolina": (33.43,-78.91,36.64,-75.3)

#      "Missouri, Illinois": (36.5,-94.71,42.47,-87.49),
#      "Missouri, Iowa": (39.03,-96.69,43.57,-90.03),
#      "Kentucky, Indiana, Ohio": (36.44,-87.63,41.98,-80.38),
#      "West Virginia, Virginia, Maryland, Delaware": (36.44,-80.43,39.72,-74)

#      "New York, Pennsylvania": (39.61,-80.63,45.01,-73.27),
#      "Pennsylvania 2": (38.36,-75.78,39.77,-74.04),
#      "New England states": (40.35,-73.34,45.14,-69.49),
#      "Maine": (42.82,-71.71,47.55,-66.77)

#      "North-west USA (Oregon, Washington)": (41.92,-125.06,49.06,-116.6)

#      "Minnesota": (43.42,-97.19,49.08,-91.08),
#      "Wisconsin": (42.55,-91.17,48.3,-86.81),
#      "Michigan": (41.7,-87.97,46.64,-82.19)

#     "Idaho, Montana 1": (41.95,-117.15,49.04,-110.92),
#     "Montana 2, Wyoming": (40.92,-110.96,49.07,-101.9),
#     "Nth Dakota, Sth Dakota, Nebraska": (39.86,-102.14,49.05,-95.38),
#     "Colorado": (37.03,-109.12,41.01,-101.96),
#     "Kansas": (36.92,-102.05,40.08,-94.57)

}
zoom_levels = range(1, 5)  # <----- EDIT: start zoom, end zoom (second is exclusive)

# mapstyle = "cycle"
# mapstyle = "transport"
# mapstyle = "landscape"
mapstyle = "outdoors"
# mapstyle = "transport-dark"
# mapstyle = "spinal-map"
# mapstyle = "pioneer"
# mapstyle = "mobile-atlas"
# mapstyle = "neighbourhood"
# mapstyle = "atlas"

# API Key and output directory
api_key = ""   # <-------------- SET YOUR API KEY HERE

global_dir = os.path.join(os.path.expanduser("~"), "Documents", "tiles-greater-london")
output_dir = os.path.join(os.path.expanduser("~"), "Desktop", "tiles")
os.makedirs(output_dir, exist_ok=True)

def lon2tilex(lon, zoom):
    return int((lon + 180.0) / 360.0 * (1 << zoom))

def lat2tiley(lat, zoom):
    return int((1.0 - log(tan(lat * pi / 180.0) + 1.0 / cos(lat * pi / 180.0)) / pi) / 2.0 * (1 << zoom))

def download_tile(zoom, x, y):
    #url = f"https://tile.thunderforest.com/{mapstyle}/{zoom}/{x}/{y}.png?apikey={api_key}"
    #url = f"https://a.tile.openstreetmap.org/{zoom}/{x}/{y}.png"
    url = f"https://tiles.stadiamaps.com/tiles/stamen_toner/{zoom}/{x}/{y}.png"
    tile_dir = os.path.join(output_dir, str(zoom), str(x))
    tile_path = os.path.join(tile_dir, f"{y}.png")
    os.makedirs(tile_dir, exist_ok=True)

    if not os.path.exists(tile_path):
        g_dir = os.path.join(global_dir, str(zoom), str(x))
        g_path = os.path.join(g_dir, f"{y}.png")
        if os.path.exists(g_path):
            shutil.copy2(g_path, tile_path)
        else:
            headers = {'Connection': 'close', 'Accept-Encoding': None, 'User-Agent': None, 'Authorization': f"Stadia-Auth {api_key}" }
            response = requests.get(url, headers=headers)
            if response.status_code == 200:
                with open(tile_path, "wb") as file:
                    file.write(response.content)
            else:
                print(f"Failed to download tile {zoom}/{x}/{y}: {response.status_code} {response.reason}")

def main():
    total_tiles = 0

    for zoom in zoom_levels:
        for min_lat, min_lon, max_lat, max_lon in regions.values():
            start_x = lon2tilex(min_lon, zoom)
            end_x = lon2tilex(max_lon, zoom)
            start_y = lat2tiley(max_lat, zoom)
            end_y = lat2tiley(min_lat, zoom)

            total_tiles += (end_x - start_x + 1) * (end_y - start_y + 1)

    with tqdm(total=total_tiles, desc="Downloading tiles") as pbar:
        for zoom in zoom_levels:
            for min_lat, min_lon, max_lat, max_lon in regions.values():
                start_x = lon2tilex(min_lon, zoom)
                end_x = lon2tilex(max_lon, zoom)
                start_y = lat2tiley(max_lat, zoom)
                end_y = lat2tiley(min_lat, zoom)

                for x in range(start_x, end_x + 1):
                    for y in range(start_y, end_y + 1):
                        download_tile(zoom, x, y)
                        pbar.update(1)
                        # time.sleep(1)

if __name__ == "__main__":
    main()
