# AUTOGENERATED! DO NOT EDIT! File to edit: ../../notebooks/05_datasets_ookla.ipynb.

# %% auto 0
__all__ = ['list_ookla_files', 'download_ookla_file', 'parallel_download', 'download_ookla_parallel', 'download_ookla_year_data',
           'lookup_ookla_file', 'compute_datakey', 'write_ookla_metajson', 'OoklaDataManager']

# %% ../../notebooks/05_datasets_ookla.ipynb 4
import hashlib
import json
import os
import shutil
import xml.etree.ElementTree as ET
from collections import namedtuple
from functools import lru_cache
from pathlib import Path
from typing import Any, Dict, List, Union
from urllib.parse import urlparse
from urllib.request import HTTPError

import geopandas as gpd
import numpy as np
import pandas as pd
import requests
from fastcore.all import L, parallel, patch
from loguru import logger

import geowrangler.area_zonal_stats as azs
from geowrangler import grids
from geowrangler.datasets.utils import make_report_hook, urlretrieve

# %% ../../notebooks/05_datasets_ookla.ipynb 5
OoklaFile = namedtuple("OoklaQuarter", ["type", "year", "quarter"])
DEFAULT_CACHE_DIR = "~/.cache/geowrangler"
OOKLA_QUADKEY_LEVEL = 16

# %% ../../notebooks/05_datasets_ookla.ipynb 6
@lru_cache(maxsize=None)
def list_ookla_files() -> dict:
    """Get list of ookla data"""
    # Query parquet files as they are easier to deal with then shapefiles
    resp = requests.get(
        "https://ookla-open-data.s3.us-west-2.amazonaws.com/?list-type=2&prefix=parquet"
    )
    resp.raise_for_status()
    root = ET.fromstring(resp.text)
    keys = {}
    # Get keys. This would require pagination once there are more than 1000 keys under the parquet folder
    # but that would only happen after ~125 years
    for child in root.findall("{http://s3.amazonaws.com/doc/2006-03-01/}Contents"):
        key = child.find("{http://s3.amazonaws.com/doc/2006-03-01/}Key").text
        path_key = Path(key)
        type_ = path_key.parts[2].rsplit("=")[-1]
        year = path_key.parts[3].rsplit("=")[-1]
        quarter = path_key.parts[4].rsplit("=")[-1]
        file = path_key.parts[5]
        keys.update({OoklaFile(type_, year, quarter): file})
    return keys

# %% ../../notebooks/05_datasets_ookla.ipynb 8
def download_ookla_file(
    type_: str,  # Internet connection type: 'fixed' or 'mobile'
    year: str,  # Year (e.g. '2020')
    quarter: str,  # Quarter (valid values: '1','2','3','4')
    directory: str = "data/",  # Download directory
    overwrite: bool = False,  # Overwrite if existing
    show_progress=True,  # show progres bar
    chunksize=8192,  # Download chunksize
    reporthook=None,  # Use custom progress bar
) -> Union[Path, None]:
    """Download ookla file to path"""
    if not os.path.isdir(directory):
        os.makedirs(directory)
    ookla_info = list_ookla_files()
    key = OoklaFile(type_, str(year), str(quarter))
    if key not in ookla_info:
        raise ValueError(
            f"{key} not found in ookla. Run list_ookla_data() to learn more about available files"
        )
    fname = ookla_info[key]
    url = f"https://ookla-open-data.s3.us-west-2.amazonaws.com/parquet/performance/type={type_}/year={year}/quarter={quarter}/{fname}"
    parsed_url = urlparse(url)
    filename = Path(os.path.basename(parsed_url.path))
    filepath = directory / filename
    if not filepath.exists() or overwrite:
        if reporthook is None:
            reporthook = make_report_hook(show_progress)

        try:
            filepath, _, _ = urlretrieve(
                url, filepath, reporthook=reporthook, chunksize=chunksize
            )
        except HTTPError as err:
            if err.code == 404:
                logger.warning(
                    f"No url found for type {type_} year {year} and {quarter} : {url} "
                )
                return None
            else:
                raise err

    return filepath

# %% ../../notebooks/05_datasets_ookla.ipynb 9
def parallel_download(item):
    (
        quarter,
        type_,
        year,
        directory,
        overwrite,
        show_progress,
        chunksize,
        reporthook,
    ) = item  # unpack tuple
    logger.info(
        f"Ookla Data: Downloading Ookla parquet file for quarter {quarter}... type: {type_} year: {year} in {directory}"
    )
    return download_ookla_file(
        type_=type_,
        year=year,
        quarter=quarter,
        directory=directory,
        overwrite=overwrite,
        show_progress=show_progress,
        chunksize=chunksize,
        reporthook=reporthook,
    )

# %% ../../notebooks/05_datasets_ookla.ipynb 10
def download_ookla_parallel(
    num_expected_ookla_files,
    type_,
    year,
    directory,
    overwrite,
    show_progress,
    chunksize,
    reporthook,
):
    items = [
        (
            str(i),
            type_,
            year,
            directory,
            overwrite,
            show_progress,
            chunksize,
            reporthook,
        )
        for i in range(1, num_expected_ookla_files + 1, 1)
    ]
    parallel(parallel_download, items, threadpool=True, progress=True)

# %% ../../notebooks/05_datasets_ookla.ipynb 11
def download_ookla_year_data(
    type_,
    year,
    cache_dir,
    use_cache=True,
    show_progress=True,
    chunksize=8192,
    reporthook=None,
):

    "Download ookla data for a specifed type (fixed or mobile) and year. Data for all 4 quarters will be downloaded."

    # Determine number of expected data for type_ and year, specified by OoklaFile(type, year, quarter)
    available_ookla_files = list_ookla_files()
    expected_ookla_files = {}
    for quarter in [1, 2, 3, 4]:
        quarter_ookla_file = OoklaFile(str(type_), str(year), str(quarter))
        if quarter_ookla_file in available_ookla_files.keys():
            expected_ookla_files[quarter_ookla_file] = available_ookla_files[
                quarter_ookla_file
            ]
    num_expected_ookla_files = len(expected_ookla_files)

    if num_expected_ookla_files == 0:
        logger.warning(f"Ookla data: No data available for {type_} and {year}")
    else:
        logger.info(
            f"Ookla Data: Number of available files for {type_} and {year}: {num_expected_ookla_files}"
        )

    ookla_cache_dir = os.path.join(cache_dir, "ookla/")
    type_year_cache_dir = os.path.join(ookla_cache_dir, type_, str(year))

    # Check if the cached data is valid. Otherwise, we have to re-download.
    # For Ookla, we need to check if we've downloaded all expected files for that year.
    cached_data_available = (
        os.path.exists(type_year_cache_dir)
        and len(os.listdir(type_year_cache_dir)) == num_expected_ookla_files
    )

    logger.info(
        f"Ookla Data: Cached data available for {type_} and {year} at {type_year_cache_dir}? {cached_data_available}"
    )

    # Download if cache is invalid or user specified use_cache = False
    if not cached_data_available or not use_cache:
        logger.info(
            f"Ookla Data: Re-initializing Ookla type/year cache dir at {type_year_cache_dir}..."
        )
        # Re-create the country cache dir and start over to fix any corrupted states
        shutil.rmtree(type_year_cache_dir, ignore_errors=True)
        Path(type_year_cache_dir).mkdir(parents=True, exist_ok=True)

        # This downloads a parquet file to the type_year_dir for each quarter
        overwrite = not use_cache
        download_ookla_parallel(
            num_expected_ookla_files,
            type_,
            year,
            type_year_cache_dir,
            overwrite,
            show_progress,
            chunksize,
            reporthook,
        )

        logger.info(
            f"Ookla Data: Successfully downloaded and cached Ookla data for {type_} and {year} at {type_year_cache_dir}!"
        )

    return type_year_cache_dir

# %% ../../notebooks/05_datasets_ookla.ipynb 12
def lookup_ookla_file(filename):
    """Get OoklaFile for the given filename"""
    return next((k for k, v in list_ookla_files().items() if v == filename), None)

# %% ../../notebooks/05_datasets_ookla.ipynb 14
def compute_datakey(aoi_bounds, type_, year, return_geometry):
    data_tuple = (
        np.array2string(aoi_bounds, precision=6),
        str(type_),
        str(year),
        str(return_geometry),
    )
    m = hashlib.md5()
    for item in data_tuple:
        m.update(item.encode())
    data_key = m.hexdigest()
    return data_key

# %% ../../notebooks/05_datasets_ookla.ipynb 15
def write_ookla_metajson(
    cache_dir, data_key, total_bounds, type_, year, return_geometry
):
    extension = "geojson" if return_geometry else "csv"
    cached_metajson_file_path = cache_dir / f"{data_key}.{extension}.metadata.json"
    with open(cached_metajson_file_path, "w") as f:
        f.write(
            json.dumps(
                {
                    "bounds": np.array2string(total_bounds, precision=6),
                    "type_": type_,
                    "year": year,
                    "with_geom": return_geometry,
                }
            )
        )

# %% ../../notebooks/05_datasets_ookla.ipynb 16
class OoklaDataManager:
    """An instance of this class provides convenience functoins for loading and caching Ookla data"""

    def __init__(self, cache_dir=DEFAULT_CACHE_DIR):
        self.aggregated_cache = {}
        self.cache_dir = os.path.expanduser(cache_dir)
        processed_cache_dir = os.path.join(self.cache_dir, "ookla/processed/")
        aggregated_cache_dir = os.path.join(self.cache_dir, "ookla/aggregated/")
        self.processed_cache_dir = Path(processed_cache_dir)
        self.processed_cache_dir.mkdir(parents=True, exist_ok=True)
        self.aggregated_cache_dir = Path(aggregated_cache_dir)
        self.aggregated_cache_dir.mkdir(parents=True, exist_ok=True)

# %% ../../notebooks/05_datasets_ookla.ipynb 17
@patch
def reinitialize_processed_cache(self: OoklaDataManager):
    "Reinitialize processed_cache_dir to start over from scratch."
    shutil.rmtree(self.processed_cache_dir, ignore_errors=True)
    self.processed_cache_dir.mkdir(parents=True, exist_ok=True)
    logger.info(
        f"{self.processed_cache_dir} reintialized. All cached processed data in this folder has been deleted."
    )

# %% ../../notebooks/05_datasets_ookla.ipynb 18
@patch
def reinitialize_aggregated_cache(self: OoklaDataManager):
    "Reinitialize aggregated_cache_dir to start over from scratch."
    shutil.rmtree(self.aggregated_cache_dir, ignore_errors=True)
    self.aggregated_cache_dir.mkdir(parents=True, exist_ok=True)
    logger.info(
        f"{self.aggregated_cache_dir} reintialized. All cached aggregated data in this folder has been deleted."
    )

# %% ../../notebooks/05_datasets_ookla.ipynb 19
@patch
def load_type_year_data(
    self: OoklaDataManager,
    aoi: gpd.GeoDataFrame,  # area of interest
    type_: str,  # ookla data type: fixed or mobile
    year: str,  # year
    use_cache=True,  # use cache dir
    return_geometry=False,  # include geometry in returned values
    show_progress=True,  # display progress bar
    chunksize=8192,  # download buffer size
    reporthook=None,  # custom progress bar
):
    "Load Ookla data across all quarters for a specified aoi, type (fixed, mobile) and year"

    # Generate hash from aoi, type_, and year, which will act as a hash key for the cache
    data_key = compute_datakey(aoi.total_bounds, type_, year, return_geometry)

    # Get cached data from filesystem if saved
    extension = "geojson" if return_geometry else "csv"
    cached_file_path = self.processed_cache_dir / f"{data_key}.{extension}"

    if cached_file_path.exists() and use_cache:
        logger.debug(
            f"Processed Ookla data for aoi, {type_} {year} found in {cached_file_path}. Loading in cache."
        )
        if not return_geometry:
            df = pd.read_csv(cached_file_path)
        else:
            df = gpd.read_file(cached_file_path, driver="GeoJSON")
        return df
    if not cached_file_path.exists():
        logger.debug("No cached data found. Processing Ookla data from scratch.")
    else:
        logger.debug("Overwriting cache contents. Processing Ookla data from scratch")

    # Otherwise, load from raw file and add to RAM cache
    type_year_cache_dir = download_ookla_year_data(
        type_,
        year,
        cache_dir=self.cache_dir,
        use_cache=use_cache,
        show_progress=show_progress,
        chunksize=chunksize,
        reporthook=reporthook,
    )

    # Generate the bing tile quadkeys that intersect with the input aoi
    logger.debug(
        f"Generating bing tile grids for aoi total bounds {np.array2string(aoi.total_bounds, precision=6)}"
    )
    bing_tile_grid_generator_no_geom = grids.BingTileGridGenerator(
        OOKLA_QUADKEY_LEVEL, return_geometry=False
    )
    aoi_quadkeys = bing_tile_grid_generator_no_geom.generate_grid(aoi)[
        "quadkey"
    ].tolist()

    # Combine quarterly data for the specified year, filtered to the aoi using quadkey
    # Quarter is inferred from the Ookla filename
    quarter_df_list = []
    for ookla_filename in sorted(os.listdir(type_year_cache_dir)):
        ooklaFile = lookup_ookla_file(ookla_filename)
        if ooklaFile:
            ookla_quarter_filepath = os.path.join(type_year_cache_dir, ookla_filename)
            logger.debug(
                f"Ookla data for aoi, {type_} {year} {ooklaFile.quarter} being loaded from {ookla_quarter_filepath}"
            )
            quarter_df = pd.read_parquet(ookla_quarter_filepath)
            # TODO optimize filtering via merge
            quarter_df = quarter_df[quarter_df["quadkey"].isin(aoi_quadkeys)]
            quarter_df["quarter"] = int(ooklaFile.quarter)
            quarter_df_list.append(quarter_df)

            # Free memory after processing
            del quarter_df

    logger.debug(
        f"Concatenating quarterly Ookla data for {type_} and {year} into one dataframe"
    )
    df = pd.concat(quarter_df_list, ignore_index=True)

    # NOTE: Since there will be groupby operations in processing, we don't return
    #       a geodataframe by default since it does not work well with aggregations
    #       by quadkey.
    write_ookla_metajson(
        self.processed_cache_dir,
        data_key,
        aoi.total_bounds,
        type_,
        year,
        return_geometry,
    )

    if not return_geometry:
        logger.debug(f"Saving Ookla data into csv file {cached_file_path}")
        df.to_csv(cached_file_path, index=False)
    else:
        logger.debug(f"Converting Ookla data into geojson file {cached_file_path}")
        df = gpd.GeoDataFrame(df, geometry=gpd.GeoSeries.from_wkt(df["tile"]))
        df.to_file(cached_file_path, driver="GeoJSON")
    return df

# %% ../../notebooks/05_datasets_ookla.ipynb 20
@patch
def aggregate_ookla_features(
    self: OoklaDataManager,  # Ookla Data Manager Instance
    aoi: gpd.GeoDataFrame,  # Area of interest
    type_: str,  # Ookla speed type: 'fixed` or `mobile`
    year: str,  # Year to aggregate (over 4 quarters)
    use_cache=True,  # Use cached data in cache dir as specified in ookla_data_manager
    return_geometry=False,  # Save aggregated data as geojson
    output_crs="epsg:4326",  # crs to use in creating aggregated geodataframe
    aggregations: Dict[  # Aggregation functions on ookla data (see https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.agg.html)
        str, Any
    ] = dict(
        mean_avg_d_kbps=("avg_d_kbps", "mean"),
        mean_avg_u_kbps=("avg_u_kbps", "mean"),
        mean_avg_lat_ms=("avg_lat_ms", "mean"),
        mean_num_tests=("tests", "mean"),
        mean_num_devices=("devices", "mean"),
    ),
    show_progress=True,  # display progress bar
    chunksize=8192,  # download buffer size
    reporthook=None,  # custom progress bar
):
    """Generates yearly aggregate features for the AOI based on Ookla data for a given type (fixed, mobile) and year."""

    data_key = compute_datakey(aoi.total_bounds, type_, year, return_geometry)
    logger.debug(f"Contents of aggregated ache: {list(self.aggregated_cache.keys())}")
    if data_key in self.aggregated_cache:
        logger.debug(
            f"Ookla aggregated data for aoi, {type_} {year} (key: {data_key}) found in cache."
        )
        return self.aggregated_cache[data_key]

    # see if in file system
    ## Get cached data from filesystem if saved
    extension = "geojson" if return_geometry else "csv"
    cached_file_path = self.aggregated_cache_dir / f"{data_key}.{extension}"

    if cached_file_path.exists() and use_cache:
        logger.debug(
            f"Aggregated Ookla data for aoi, {type_} {year} (key: {data_key}) found in filesystem. Loading in cache."
        )
        if return_geometry:
            df = gpd.read_file(cached_file_path, driver="GeoJSON")
        else:
            df = pd.read_csv(cached_file_path)
        return df

    if cached_file_path.exists():
        logger.debug("Overwriting cached data. Processing Ookla data from scratch.")
    else:
        logger.debug("No aggregated data found. Processing Ookla data from scratch.")

    ookla = self.load_type_year_data(
        aoi,
        type_,
        year,
        use_cache=use_cache,
        return_geometry=return_geometry,
        show_progress=show_progress,
        chunksize=chunksize,
        reporthook=reporthook,
    )

    # Combine quarterly data from Ookla into yearly aggregate data
    # Geometries are stored separately and rejoined after aggregation by quadkey
    # TODO: incorporate parametrized aggregations, take inspiration from GeoWrangler agg spec

    logger.info(
        f"Aggregating ookla data for bounds {np.array2string(aoi.total_bounds, precision=6)} type {type_} year {year} "
    )
    ookla_geoms = ookla[["quadkey", "tile"]].drop_duplicates().reset_index(drop=True)
    ookla_yearly = ookla.groupby("quadkey").agg(**aggregations).reset_index()
    # Add type_year prefix to feature names
    ookla_yearly = ookla_yearly.rename(
        {
            col: f"{type_}_{year}_" + col
            for col in ookla_yearly.columns[~ookla_yearly.columns.isin(["quadkey"])]
        },
        axis=1,
    )
    ookla_yearly = ookla_yearly.merge(ookla_geoms, on="quadkey", how="left")

    # write ookla_yearly to cache/aggregated

    if return_geometry:
        logger.debug(
            f"Saving Ookla aggregated data into geojson file {cached_file_path}"
        )
        ookla_yearly = gpd.GeoDataFrame(
            ookla_yearly,
            geometry=gpd.GeoSeries.from_wkt(ookla_yearly["tile"], crs=output_crs),
        )
        ookla_yearly.to_file(cached_file_path, driver="GeoJSON")
    else:
        logger.debug(f"Saving Ookla aggregated data into csv file {cached_file_path}")
        ookla_yearly.to_csv(cached_file_path)
    write_ookla_metajson(
        self.aggregated_cache_dir,
        data_key,
        aoi.total_bounds,
        type_,
        year,
        return_geometry,
    )

    self.aggregated_cache[data_key] = ookla_yearly

    return ookla_yearly
