#!/usr/bin/env python

import json
import os
import overpy
import numpy as np
import pandas as pd

from pathlib import Path

# brandname : overpass query filters
BRANDS: dict[str, str] = {
    "greggs": "[\"brand:wikidata\"=\"Q3403981\"]",
    "tesco": "[\"brand:wikidata\"~\"^(Q487494|Q98456772|Q25172225|Q65954217)$\"]",  # Includes Tesco Express, Tesco Extra, and One Stop
}

CACHE_FOLDER = Path(".cache")

LOCS_COUNT = 3
DISTS_COUNT = 100

FORMAT_FACTOR = 1e6  # μm

EncodedLocation = list[tuple[float, list[float]]]


def fetch_data(brand: str, cache: bool = True) -> list[tuple[float, float]]:
    """Fetch a list of locations from OSM."""
    cache_loc = (CACHE_FOLDER / f"{brand}.json")

    # Try load from cache
    if cache and cache_loc.exists():
        with open(cache_loc, "r") as f:
            data = json.load(f)

        return data

    api = overpy.Overpass()

    filters = BRANDS[brand]
    query = api.query(f"nwr{filters}; out center;")

    result = []

    for way in query.ways:
        result.append((float(way.center_lat), float(way.center_lon)))

    for node in query.nodes:
        result.append((float(node.lat), float(node.lon)))

    for (lat, lon) in result:
        if (lat is None) or (lon is None):
            raise ValueError("Item missing coords!")

    # Save to cache
    if cache:
        if not CACHE_FOLDER.exists():
            os.makedirs(CACHE_FOLDER)

        with open(cache_loc, "w") as f:
            json.dump(result, f)

    print(f"Got {len(result)} {brand}s")

    return result


def spherical_dist(pos1, pos2, r=6378137):
    """Calculate sperical distances between two arrays of coordinates."""
    pos1 = pos1 * np.pi / 180
    pos2 = pos2 * np.pi / 180

    cos_lat1 = np.cos(pos1[..., 0])
    cos_lat2 = np.cos(pos2[..., 0])
    cos_lat_d = np.cos(pos1[..., 0] - pos2[..., 0])
    cos_lon_d = np.cos(pos1[..., 1] - pos2[..., 1])
    return r * np.arccos(cos_lat_d - cos_lat1 * cos_lat2 * (1 - cos_lon_d))


def encode(location: tuple[float, float]) -> EncodedLocation:
    """Encode a location."""

    greggs = np.array(fetch_data("greggs"))

    repeat_rows = np.tile(greggs, (len(greggs), 1, 1))
    repeat_cols = np.transpose(repeat_rows, (1, 0, 2))
    dist_matrix = spherical_dist(repeat_rows, repeat_cols)

    repeated = np.tile(location, (len(greggs), 1))
    distances = spherical_dist(repeated, greggs)
    distances = pd.Series(distances)
    distances = distances.sort_values()

    closest = distances.head(LOCS_COUNT)
    closest_dist = list(closest.values)
    closest_ind  = list(closest.index)

    result: EncodedLocation = []
    for v, i in zip(closest.values, closest.index):
        greggs_distances = np.sort(dist_matrix[i])[1:DISTS_COUNT+1]

        result.append((v, list(map(float, greggs_distances))))

    # Stub
    return result


def decode(location: EncodedLocation) -> tuple[float, float]:
    """Decode into a location."""

    # Stub
    return (0.091659, 52.210796)


def format_dist(dist: float) -> str:
    return f"{int(round(dist * FORMAT_FACTOR))}"


def parse_dist(dist: str) -> float:
    return float(dist) / FORMAT_FACTOR


def format_location(location: EncodedLocation) -> str:
    """Format an encoded location as a string."""
    return ";\n".join([f"{format_dist(a)}:{','.join(map(format_dist, b))}" for (a, b) in location])


def parse_location(location: str) -> EncodedLocation:
    """Parse a location string into an EncodedLocation."""

    # Stub
    return [
        (5., [1., 2., 3.]),
        (6., [4., 5., 6.]),
    ]


def main():
    """Testing."""
    #print("Running query...")
    #greggs = fetch_data("greggs")
    #print(f"Query done - got {len(greggs)} Greggs!")

    # print(format_location(encode((52.210796, 0.091659))))


if __name__ == "__main__":
    main()