# Set Hyperparameters
n = 100  # number of areas to randomly sample
DATA_DIR = "data/tutorial_data/"  # Set the directory for the data
d = 0.05  # degree measurement size

NPARTITIONS = int(n / 50)
# Set dataset filepaths
LAND_MAP = "map.json"  # Set the directory for the Polygons of the the land masses

# Directory to store the .csv file of randomly sampled grids
GRIDS = f"{DATA_DIR}grids.csv"
# Store the query .json file path in .csv
GRIDS_TAGS = f"{DATA_DIR}grids_tags.csv"
# Store collected tags and metadata
GRIDS_TAGS_IMAGES = f"{DATA_DIR}grids_tags_images.csv"

IMAGES_TAGS = f"{DATA_DIR}images_tags.csv"  # Store the collected images and tag lists
IMAGES_CAPTIONS = f"{DATA_DIR}images_captions.csv"  # Store the generated captions
IMAGES_DIR = f"{DATA_DIR}images/"  # Directory for cropped images
MODEL_DIR = f"{DATA_DIR}model/"  # Directory for the model files

print(f"Master data directory: {DATA_DIR}")
print(f"Generate {n} grids with size {d} x {d}")

os.makedirs(DATA_DIR, exist_ok=True)
os.makedirs(IMAGES_DIR, exist_ok=True)
os.makedirs(MODEL_DIR, exist_ok=True)

print(f"Directories {DATA_DIR}, {IMAGES_DIR}, and {MODEL_DIR} created")

# Load GeoJSON of land polygons
with open(LAND_MAP) as file:
    land = json.load(file)

# Create a MultiPolygon from the land geometries by extracting the coordinates
ob = MultiPolygon(
    [Polygon((feature["coordinates"])[0]) for feature in land["geometries"]]
)

# Plot Multipolygon
fig, ax = plt.subplots()
ax.set_aspect("equal")
for polygon in ob.geoms:
    x, y = polygon.exterior.xy
    ax.plot(x, y)
plt.show()

def generate_random_coordinate() -> tuple[float, float]:
    """
    This function generates a random coordinate on the Earth's
    surface using a uniform normal distribution on the unit sphere.

    Returns:
        (longitude, latitude): A tuple of the longitude and latitude
        coordinates of the generated coordinate, in this order.

    """
    gx = random.gauss(0.0, 1.0)
    gy = random.gauss(0.0, 1.0)
    gz = random.gauss(0.0, 1.0)

    norm2 = gx**2 + gy**2 + gz**2
    norm1 = 1.0 / math.sqrt(norm2)

    x = gx * norm1
    y = gy * norm1
    z = gz * norm1

    radLat = math.asin(z)
    radLon = math.atan2(y, x)

    lat = math.degrees(radLat)
    lon = math.degrees(radLon)

    return lon, lat


print("Function generate_random_coordinate defined!")

def not_crossing_antimeridian(lon: float, lat: float, distance: float) -> bool:
    """
    This function ensures that the generated coordinate would not cross the
    antimeridian when a square of a given distance is centered at the (lon, lat) coordinate.

    Args:
        lon: float representing the decimal longitude of the coordinate you wish to check
        lat: float representing the decimal latitude of the coordinate you wish to check

    Returns:
        bool: True if generated coordinate would not cross the antimeridian, False otherwise.

    """
    distance = (
        distance / 2
    )  # distance is the length of the side of the square, so we need to divide by 2
    distance = distance * LONGITUDE  # Grid dimensions in miles

    # convert latitude to radians
    lat_rad = math.radians(lat)
    long_degree = math.cos(lat_rad) * LONGITUDE
    longitude = distance / long_degree
    latitude = distance / LATITUDE

    minX = lon - longitude
    minY = lat - latitude
    maxX = lon + longitude
    maxY = lat + latitude

    if (
        (minX < -180) or (minY < -90) or (maxX > 180) or (maxY > 90)
    ):  # Goes out of bounds
        return False
    return True


print("Function not_crossing_antimeridian defined!")

def random_points_in_bounds(n: int) -> list[tuple[float, float]]:
    """
    This function generates n random coordinates located in the land boundaries
    of Earth's surface.

    Args:
        n: The number of random coordinates to generate.

    Returns:
        A list of tuples of the longitude and latitude coordinates of the
        generated coordinates.

    """
    # Load GeoJSON of land polygons
    with open(LAND_MAP) as file:
        land = json.load(file)

    # Create a MultiPolygon from the land geometries by extracting the coordinates
    land_areas = MultiPolygon(
        [Polygon((feature["coordinates"])[0]) for feature in land["geometries"]]
    )
    points = []

    # Continuously generate a random coordinate and check if it is in the MultiPolygon
    while len(points) < n:
        point = generate_random_coordinate()

        valid = False
        valid = shapely.contains_xy(
            land_areas, point[0], point[1]
        ) and not_crossing_antimeridian(point[0], point[1], d)
        if valid:
            points.append(point)

    return points


print("Function random_points_in_bounds defined!")

def create_bbox(lon: float, lat: float, distance: float) -> list[float]:
    """
    This function generates a square bounding box of degree d x d (normalized
    to latitude and longitude ratios) with the provided point located in the center.

    Args:
        lon: The longitude of the coordinate.
        lat: The latitude of the coordinate.
        distance: The length of the side of the square bounding box in degrees.

    Returns:
        A bounding box as a list of the form [minX, minY, maxX, maxY]

    """
    distance = (
        distance / 2
    )  # distance is the length of the side of the square, so we need to divide by 2
    distance = distance * LONGITUDE  # Grid dimensions in miles

    # convert latitude to radians
    lat_rad = math.radians(lat)
    long_degree = math.cos(lat_rad) * LONGITUDE
    longitude = distance / long_degree
    latitude = distance / LATITUDE

    minX = lon - longitude
    minY = lat - latitude
    maxX = lon + longitude
    maxY = lat + latitude

    # Point cannot go out of bounds because it is already checked in the
    # random_points_in_bounds function
    return [minX, minY, maxX, maxY]


print("Function create_bbox defined!")

# Generate n random square grids
grids = []
for point in random_points_in_bounds(n):
    minX, minY, maxX, maxY = create_bbox(point[0], point[1], d)
    grids.append([minX, minY, maxX, maxY])

# Create .csv file for sampled coordinates to be saved
df = pandas.DataFrame(grids, columns=["minX", "minY", "maxX", "maxY"])
df["uid"] = range(1, n + 1)
df.to_csv(GRIDS, index=False)
print(f"Dataframe saved to {GRIDS}")

# Open the csv file in reader mode
with open(GRIDS, "r") as csvfile:
    reader = csv.reader(csvfile)
    grids = list(reader)[1:]

# Create a map centered at a specific location
m = folium.Map(location=[0, 0], zoom_start=2)

# Iterate over the grids and create a rectangle for each bounding box
for bbox in grids:
    min_lon, min_lat, max_lon, max_lat = bbox[0], bbox[1], bbox[2], bbox[3]
    rectangle = folium.Rectangle(
        bounds=[(min_lat, min_lon), (max_lat, max_lon)],
        color="blue",
        fill=True,
        fill_color="blue",
        fill_opacity=0.3,
    )
    rectangle.add_to(m)

# Display the map
m

# Read the CSV file into a DataFrame
grids_df = pandas.read_csv(GRIDS)

# Display the DataFrame
grids_df

# Rename certain tag names
def rename_tags(tag_list: list[list[str, str]]) -> list[list[str, str]]:
    """
    This function renames keys and values in the tag_list.

    Args:
        tag_list: A list of [key, value] pairs that represent object tags.

    Returns:
        (longitude, latitude): A tuple of the longitude and latitude
        coordinates of the generated coordinate, in this order.

    """
    for tag in tag_list:
        if tag[0] == "highway":
            if tag[1] != "primary" and tag[1] != "motorway" and tag[1] != "truck road":
                tag[0] = "road"
        elif tag[0] == "aeroway":
            tag[0] = "airport"
        elif tag[0] == "lit":
            tag[0] = "light"
        elif tag[0] == "leisure":
            tag[0] = "leisure land"
        elif tag[0] == "addr:city":
            tag[0] = "city"
            tag[1] = "yes"
    return tag_list


print("Function rename_tags defined!")

MILE = 1609  # 1 mile is 1609 meters

for i in range(len(grids_df)):  # Query each grid saved in grids.csv for OSM objects

    if len(grids_df) > 10000:
        print(
            "Too many queries! The API for OpenStreetMap, OverPass, has a limit"
            "of 10,000 queries per day. Please do not overload the API."
        )
        break

    tags = set()

    # Need to switch our lon, lat convention for OSM query
    # This query gathers all nodes that fall within the grid
    minX, minY, maxX, maxY = (
        grids_df["minX"][i],
        grids_df["minY"][i],
        grids_df["maxX"][i],
        grids_df["maxY"][i],
    )
    query = f"(node({minY},{minX},{maxY},{maxX});<;);"

    # Get the result of the query
    result_json = overpass_api.get(query, responseformat="json")

    # parse the JSON file into a readable format using OverPy
    result = overpy.Result.from_json(result_json)
    ways = result.get_elements(filter_cls=overpy.Way)

    for way in ways:
        locations = way._node_ids
        locations = [result.get_nodes(x) for x in locations]
        locations = [[node[0].lat, node[0].lon] for node in locations if node != []]
        # Check if way is closed, then calculate the area ==>
        # remove tags from objects that are too small to see given gsd
        latitude_ratio = math.cos(math.radians((maxY + minY) / 2))
        minimum_area = (gsd * gsd) / (latitude_ratio * MILE * LATITUDE) ** 2
        if not (
            (len(locations) > 3)
            and (LineString(locations).is_closed)
            and (Polygon(locations).area > minimum_area)
        ):
            for tag in way.tags:
                if tag in accepted_keys:
                    pair = (tag, way.tags[tag])
                    tags.add(pair)

    tag_list = rename_tags([list(tag) for tag in list(tags)])
    tag_lists.append(tag_list)

# Add the tag lists to the dataframe
grids_df["tags"] = tag_lists
grids_df

# Save the dataframe to a CSV file
grids_df.to_csv(GRIDS_TAGS, index=False)
print(f"Dataframe saved to {GRIDS_TAGS}")

# Select a grid to visualize with the OSM data overlayed
i = 1
minY, minX, maxY, maxX = (
    grids_df["minY"][i],
    grids_df["minX"][i],
    grids_df["maxY"][i],
    grids_df["maxX"][i],
)
# Create a map centered at a specific location
m = folium.Map(
    location=[
        (maxY + minY) / 2,
        (maxX + minX) / 2,
    ],  # Center the map at the center of the bounding box
    zoom_start=14,
    min_lat=minY,
    max_lat=maxY,
    min_lon=minX,
    max_lon=maxX,
    control_scale=True,
)

# Create a rectangle for the bounding box and display it on the map
rectangle = folium.Rectangle(
    bounds=[(minY, minX), (maxY, maxX)],
    color="blue",
    fill=True,
    fill_color="blue",
    fill_opacity=0.3,
)

rectangle.add_to(m)

# Load the cached query results for the bounding box
query = f"(node({minY},{minX},{maxY},{maxX});<;);"
print(query)
data = overpass_api.get(query, responseformat="json")
result = overpy.Result.from_json(data)

# Iterate over the ways in the query results and create a polygon for each way
ways = result.get_elements(filter_cls=overpy.Way)
for way in ways:
    locations = way._node_ids
    locations = [result.get_nodes(x) for x in locations]
    locations = [[node[0].lat, node[0].lon] for node in locations if node != []]

    for location in locations:
        folium.Circle(
            location=location,
            radius=20,
        ).add_to(m)
    if len(locations) != 1:
        linestring = LineString(locations)
        if linestring.is_closed:
            folium.Polygon(
                locations=locations,
                color="red",
                weight=2,
                fill=True,
                fill_color="red",
                fill_opacity=0.3,
            ).add_to(m)
        else:
            folium.PolyLine(
                locations=locations,
                color="blue",
                weight=2,
            ).add_to(m)
    else:
        folium.PolyLine(
            locations=locations,
            color="blue",
            weight=2,
        ).add_to(m)


# Display the map
m

# Read the CSV file into a DataFrame
df = pandas.read_csv(GRIDS_TAGS)

# Display the DataFrame
df

# Use dask_geopandas to handle large dataframes of geospatial data
ddf = dask_geopandas.from_geopandas(gdf, npartitions=1)
hd = ddf.hilbert_distance().compute()
gdf["hd"] = hd
gdf = gdf.sort_values(
    "hd"
)  # Sort the data by the Hilbert distance (spatial locality -- similar hd values are close)

dgdf = dask_geopandas.from_geopandas(gdf, npartitions=NPARTITIONS, sort=False)
dgdf.compute()  # Compute the dask dataframe

for i in range(NPARTITIONS):
    print(f"Calculating partition {i}...")
    chunk = dgdf.to_delayed()[i].compute()
    query = catalog.search(
        collections=["sentinel-2-l2a"],
        intersects=chunk.unary_union,
        datetime="2024-01-01/2024-06-01",
    )
    items = query.item_collection()
    if items != []:
        geodf = geopandas.GeoDataFrame.from_features(items.to_dict()["features"])
        geodf["assets"] = [item.assets for item in items]
        print(f"Saving images to {IMAGES_DIR}...")
        for j in range(len(chunk)):
            uid, minY, minX, maxY, maxX = (
                chunk["uid"].iloc[j],
                chunk["minY"].iloc[j],
                chunk["minX"].iloc[j],
                chunk["maxY"].iloc[j],
                chunk["maxX"].iloc[j],
            )
            grid_bbox = box(minX, minY, maxX, maxY)
            filtered_items = geodf[geodf.contains(grid_bbox)]
            if not filtered_items.empty:
                selected_item_index = filtered_items["eo:cloud_cover"].idxmin(axis=0)
                selected_item = filtered_items.loc[selected_item_index]

                water = selected_item["s2:water_percentage"]
                snow = selected_item["s2:snow_ice_percentage"]
                vegetated = selected_item["s2:vegetation_percentage"]
                cloud = selected_item["eo:cloud_cover"]
                metadata_df = pandas.concat(
                    [
                        metadata_df,
                        pandas.DataFrame([[uid, water, snow, vegetated, cloud]]),
                    ],
                    ignore_index=True,
                )

                if not os.path.exists(f"{IMAGES_DIR}/{uid}.tif"):
                    with rasterio.open(selected_item.assets["visual"].href) as src:
                        # Create a Window and calculate the transform from the dataset
                        project = pyproj.Transformer.from_crs(
                            pyproj.CRS("EPSG:4326"), src.crs, always_xy=True
                        ).transform
                        grid_bbox = transform(project, grid_bbox)
                        out_image, out_transform = rasterio.mask.mask(
                            src, [grid_bbox], crop=True
                        )
                        out_meta = src.meta

                        # Create a new cropped raster to write to
                        out_meta.update(
                            {
                                "driver": "GTiff",
                                "height": out_image.shape[1],
                                "width": out_image.shape[2],
                                "transform": out_transform,
                            }
                        )

                        # Save the cropped image into a new .tif file in IMAGES_DIR
                        with rasterio.open(
                            f"{IMAGES_DIR}/{uid}.tif", "w", **out_meta
                        ) as dest:
                            dest.write(out_image)

print(f"Images saved to {IMAGES_DIR}")

metadata_df.columns = [
    "uid",
    "water_percentage",
    "snow_ice_percentage",
    "vegetated_percentage",
    "cloud_cover",
]
metadata_df

df = pandas.read_csv(GRIDS_TAGS)
images = []
for i in range(n):
    uid = df["uid"].iloc[i]
    images.append(f"{uid}.tif")
df = df.assign(image=images)
df

# Merge the metadata DataFrame with the original DataFrame (GRIDS_QUERIES_TAGS)
df = pandas.merge(df, metadata_df, on=["uid"])
df

# Save the dataframe with uid, minX, minY, maxX, maxY, tags, image,
# water_percentage, snow_ice_percentage, vegetated_percentage,
# cloud_cover to a .csv file
df.to_csv(GRIDS_TAGS_IMAGES, index=False)
print(f"Dataframe saved to {GRIDS_TAGS_IMAGES}")

# Read the CSV file into a DataFrame
grids_images_df = pandas.read_csv(GRIDS_TAGS_IMAGES)

# Display the DataFrame
grids_images_df

# Add tags for metadata columns
def add_metadata_tags(
    tag_list: list[list[str, str]], csv_filepath, row
) -> list[list[str, str]]:
    df = pandas.read_csv(csv_filepath)
    metadata_keys = [
        "water_percentage",
        "snow_ice_percentage",
        "vegetated_percentage",
        "cloud_cover",
    ]
    for key in metadata_keys:
        value = df[key][row]
        if value > 95:
            tag_list.append((key, "very high"))
        if value > 80:
            tag_list.append((key, "high"))
        elif value > 40:
            tag_list.append((key, "medium"))
        elif value > 10:
            tag_list.append((key, "low"))
        elif value > 1:
            tag_list.append((key, "very low"))
        elif value > 0.01:
            tag_list.append((key, "almost none"))
        else:
            tag_list.append((key, "none"))
    return tag_list


print("Function add_metadata_tags defined!")

# Create a new pandas dataframe
images_tags_df = pandas.DataFrame()
# Save the descriptions as a new column in the DataFrame
images_tags_df = images_tags_df.assign(caption=captions)
# Save the images as a new column in the DataFrame
images_tags_df = images_tags_df.assign(image=grids_images_df["image"])

# Final DataFrame with Images and their captions
images_tags_df.to_csv(IMAGES_TAGS, index=False)
print(f"Dataframe saved to {IMAGES_TAGS}")

df = pandas.read_csv(IMAGES_TAGS)
df

print("Generating captions...")
for i in range(len(df)):
    try:
        prompt = (
            "Give a concise natural language caption of an image"
            "using the following image tag descriptors:"
        ) + df["caption"].iloc[i]
        completion = get_completion(model, prompt, client)
        image_url = df["image"].iloc[i]
        images_captions_df = pandas.concat(
            [images_captions_df, pandas.DataFrame([[image_url, completion]])],
            ignore_index=True,
        )
    except Exception:
        print("Token Expired. Refreshing Token...")
        azure_ad_token = credential.get_token(
            "https://cognitiveservices.azure.com/.default"
        ).token
        client = AzureOpenAI(
            azure_ad_token=azure_ad_token,
            api_version="2024-02-01",
            azure_endpoint=endpoint,
        )
        prompt = (
            "Give a concise natural language caption of an image"
            "using the following image tag descriptors:"
        ) + df["caption"].iloc[i]
        completion = get_completion(model, prompt, client)
        image_url = df["image"].iloc[i]
        images_captions_df = pandas.concat(
            [images_captions_df, pandas.DataFrame([[image_url, completion]])],
            ignore_index=True,
        )
print("Generation complete!")

images_captions_df.rename(columns={0: "image", 1: "caption"}, inplace=True)
images_captions_df

# Test if torch is able to detect your compute GPU
torch.cuda.is_available()

# Rename captions to answers
images_captions_df.rename(columns={"caption": "answers"}, inplace=True)
# Wrap answers column in a list
images_captions_df["answers"] = images_captions_df["answers"].apply(lambda x: [x])

# Add question id column
images_captions_df["question_id"] = range(1, len(images_captions_df) + 1)

# Add question column
images_captions_df["question"] = "What does the image describe?"

# Add question_type column
images_captions_df["question_types"] = [["form"]] * len(images_captions_df)
images_captions_df

# Split the data into training and validation
train_df, val_df = train_test_split(images_captions_df, test_size=0.1, random_state=42)

# Transform train_df and val_df into Dataset()
train_dataset = Dataset.from_pandas(train_df)
val_dataset = Dataset.from_pandas(val_df)

# Wrap train_dataset and val_dataset in DatasetDict()
data = DatasetDict({"train": train_dataset, "validation": val_dataset})

# View the data structure
data

state_dict = torch.load(f"{MODEL_DIR}/final-florence-2-state-dict")
nl_trained_model = AutoModelForCausalLM.from_pretrained(
    "microsoft/Florence-2-base-ft", trust_remote_code=True, state_dict=state_dict
).eval()

nl_trained_processor = AutoProcessor.from_pretrained(
    f"{MODEL_DIR}/final-florence-2-processor", trust_remote_code=True
)

test_df = pandas.read_csv(f"{DATA_DIR}test.csv")
test_df

i = 0  # Select an example to look at from your test dataframe

image = Image.open(f"{DATA_DIR}images/" + test_df.iloc[i]["image"])

# Display the image
image

# Print out the ground truth caption (the caption from the test dataframe)
print(test_df.iloc[i]["caption"])

# Model fine-tuned on natural language captions generated by GPT 3.5
task_prompt = "<CAPTION>"
run_example(nl_trained_model, nl_trained_processor, task_prompt)