%%capture
!pip install pandas_path loguru nb_black typing

# load extension for automatic formatting
%load_ext lab_black

import os

import matplotlib.pyplot as plt
import pandas as pd
from pandas_path import path  # noqa
from pathlib import Path
import requests
from tqdm import tqdm
from typing import Optional, List

DATA_DIR = Path("/driven-data/cloud-cover")
TRAIN_FEATURES = DATA_DIR / "train_features"
TRAIN_LABELS = DATA_DIR / "train_labels"

assert TRAIN_FEATURES.exists()

BANDS = ["B02", "B03", "B04", "B08"]

import rioxarray
import xrspatial.multispectral as ms


def true_color_img(chip_id: str, data_dir: os.PathLike = TRAIN_FEATURES):
    """Given the path to the directory of Sentinel-2 chip feature images
    and a chip id, plots the true color image of the chip"""
    chip_dir = data_dir / chip_id
    red = rioxarray.open_rasterio(chip_dir / "B04.tif").squeeze()
    green = rioxarray.open_rasterio(chip_dir / "B03.tif").squeeze()
    blue = rioxarray.open_rasterio(chip_dir / "B02.tif").squeeze()

    return ms.true_color(r=red, g=green, b=blue)

import rasterio.warp


def lat_lon_bounds(filepath: os.PathLike):
    """Given the path to a GeoTIFF, returns the image bounds in latitude and
    longitude coordinates.

    Returns points as a tuple of (left, bottom, right, top)
    """
    with rasterio.open(filepath) as im:
        bounds = im.bounds
        meta = im.meta
    # create a converter starting with the current projection
    return rasterio.warp.transform_bounds(
        meta["crs"],
        4326,  # code for the lat-lon coordinate system
        *bounds,
    )

def add_paths(
    df: pd.DataFrame,
    feature_dir: os.PathLike,
    label_dir: Optional[os.PathLike] = None,
    bands: List[str] = BANDS,
):
    """
    Given dataframe with a column for chip_id, returns a dataframe with a column for
    each of the bands provided as "{band}_path", eg "B02_path". Each band column is
    the path to that band saved as a TIF image. If the path to the labels directory
    is provided, a column is also added to the dataframe with paths to the label TIF.
    """
    for band in bands:
        df[f"{band}_path"] = feature_dir / df["chip_id"] / f"{band}.tif"
        # make sure a random sample of paths exist
        assert df.sample(n=40, random_state=5)[f"{band}_path"].path.exists().all()
    if label_dir is not None:
        df["label_path"] = label_dir / (df["chip_id"] + ".tif")
        # make sure a random sample of paths exist
        assert df.sample(n=40, random_state=5)["label_path"].path.exists().all()

    return df

from datetime import timedelta

import planetary_computer as pc
from pystac_client import Client

# Establish a connection to the STAC API
catalog = Client.open("https://planetarycomputer.microsoft.com/api/stac/v1")
collection = catalog.get_collection("sentinel-2-l2a")

from bs4 import BeautifulSoup

from datetime import datetime, timedelta
from typing import Dict, Optional, Sequence, Tuple, Union

from loguru import logger
import numpy as np
import pandas as pd

from PIL import Image
import planetary_computer as pc
from pystac_client import Client
import shapely.geometry
import rasterio
import requests
import rioxarray

DATETIME_FORMAT = "%Y-%m-%dT%H:%M:%SZ"


catalog = Client.open("https://planetarycomputer.microsoft.com/api/stac/v1")
collection = catalog.get_collection("sentinel-2-l2a")


def get_closest_item(
    items: Sequence[dict],
    polygon: shapely.geometry.polygon.Polygon,
    timestamp: datetime,
) -> Optional[dict]:
    """
    Returns the item with maximum overlap and closest timestamp.

    Args:
        items (Sequence[dict]): items returned by a PySTAC catalog search
        polygon (shapely.geometry.polygon.Poylgon): polygon of the latitude/longitude
            coordinates for the original chip to be matched
        timestamp (datetime.timestamp): timestamp for the original chip to be matched

    Returns:
        pystac.item.Item: PySTAC item with the closest timestamp and/or maximum
            geographic overlap to the original chip to be matched.
    """
    if len(items) == 1:
        return items[0]

    # Check for exactly matching timestamp (to the minute) with sensing_time
    sensing_times = []
    for item in items:
        granule_metadata_href = pc.sign(item.assets["granule-metadata"].href)
        granule_metadata = requests.get(granule_metadata_href)
        soup = BeautifulSoup(granule_metadata.text)
        sensing_times.append(
            pd.to_datetime(soup.find("sensing_time").text).round(freq="T")
        )
    exact_timestamp_matches = [
        item
        for item, sensing_time in zip(items, sensing_times)
        if sensing_time == pd.to_datetime(timestamp).round(freq="T")
    ]

    # If only one item has a perfect match, return it
    if len(exact_timestamp_matches) == 1:
        return exact_timestamp_matches[0]

    # Compute overlap between each query result and the geotiff polygon
    overlaps = [
        shapely.geometry.shape(item.geometry).intersection(polygon).area / polygon.area
        for item in items
    ]
    max_overlap = max(overlaps)
    items_overlaps = [
        (item, overlap)
        for item, overlap in zip(items, overlaps)
        if overlap == max_overlap
    ]

    # If one item has higher overlap than the rest, return it
    if len(items_overlaps) == 1:
        return items_overlaps[0][0]

    # If multiple have equally high overlap, return item with closest sensing_time
    max_overlap_sensing_times = [
        sensing_time
        for sensing_time, overlap in zip(sensing_times, overlaps)
        if overlap == max_overlap
    ]
    time_deltas = [
        abs(sensing_time - timestamp) for sensing_time in max_overlap_sensing_times
    ]
    return items_overlaps[np.argmin(time_deltas)][0]


def query_bands(
    geotiff: rasterio.io.DatasetReader,
    timestamp: Union[datetime, pd.Timestamp, str],
    asset_keys: Sequence[str],
    query_range_minutes: int = 120,
    output_shape: Optional[Tuple[int, int]] = None,
    verbose: Optional[bool] = True,
) -> Dict[str, np.ndarray]:
    """
    Queries the Planetary Computer STAC API for additional imagery that
    corresponds to the same spatial extent as a provided GeoTIFF.

    Args:
        geotiff (rasterio.io.DatasetReader): A rasterio GeoTIFF
        timestamp (datetime or str): Timestamp for GeoTIFF acquisition used
            in the STAC API query to find the corresponding scene
        asset_keys (sequence of str): A sequence (list, tuple, set, etc.) of
            keys specifying the desired STAC assets to return
        query_range_minutes (int): Duration of the time range for the STAC API
            query. You can increase this if the query does not return any results.
        output_shape (tuple of ints, optional): If provided, reshape the output
            to this (height, width)
        verbose (bool, Optional): Whether to print logging messages. Defaults to True

    Returns:
        dict {str: np.ndarray}: A dictionary where each key is an asset name, and each value
            is the array of values for that asset from the PySTAC item that most closely
            matches the original chip's location and time
    """
    # Convert bounds to regular lat/long coordinates
    left, bottom, right, top = rasterio.warp.transform_bounds(
        geotiff.meta["crs"],
        4326,  # code for the lat-lon coordinate system
        *geotiff.bounds,
    )

    # Get a polygon for the area to search
    area_of_interest = shapely.geometry.shape(
        {
            "type": "Polygon",
            "coordinates": [
                [
                    [left, bottom],
                    [left, top],
                    [right, top],
                    [right, bottom],
                    [left, bottom],
                ]
            ],
        }
    )

    # Get the timestamp range to search
    if isinstance(timestamp, str):
        timestamp = pd.to_datetime(timestamp)
    if isinstance(timestamp, pd.Timestamp):
        timestamp = timestamp.to_pydatetime()
    range_start = timestamp - timedelta(minutes=query_range_minutes // 2)
    range_end = timestamp + timedelta(minutes=query_range_minutes // 2)
    time_range = (
        f"{range_start.strftime(DATETIME_FORMAT)}/{range_end.strftime(DATETIME_FORMAT)}"
    )

    # Search the Sentinel-2 catalog
    search = catalog.search(
        collections=["sentinel-2-l2a"],
        intersects=area_of_interest,
        datetime=time_range,
    )

    # Filter to the best-matching item
    items = list(search.get_items())
    if len(items) == 0:
        raise ValueError(
            "Query returned no results. Check that the bounding box is correct "
            f"or try increasing the query time range. Chip: {geotiff.name}"
        )
    item = get_closest_item(items, area_of_interest, timestamp)

    # Ensure that original chip and PySTAC item have the same coordinate projection
    if geotiff.meta["crs"] == item.properties["proj:epsg"]:
        if verbose:
            logger.debug(
                f"""GeoTIFF and STAC item have same CRS {geotiff.meta["crs"]}"""
            )
        bounds = geotiff.bounds
    else:
        if verbose:
            logger.debug(
                f"""Transforming from {geotiff.meta["crs"]} """
                f"""to {item.properties["proj:epsg"]}"""
            )
        bounds = rasterio.warp.transform_bounds(
            geotiff.meta["crs"],
            item.properties["proj:epsg"],
            geotiff.bounds.left,
            geotiff.bounds.bottom,
            geotiff.bounds.right,
            geotiff.bounds.top,
        )

    # Load the matching PySTAC asset
    assets = {}
    for asset_key in asset_keys:
        asset = np.array(
            rioxarray.open_rasterio(pc.sign(item.assets[asset_key].href))
            .rio.clip_box(*bounds)
            .load()
            .transpose("y", "x", "band")
        )

        # Reshape to singe-band image and resize if needed
        asset = Image.fromarray(asset.squeeze())
        if output_shape:
            asset = asset.resize(output_shape)
        assets[asset_key] = np.array(asset)

    return assets