#!/usr/bin/env python

from engine import encode_greggs, decode_greggs, encode, decode, format_location, parse_location, spherical_dist, fetch_data, format_dist, parse_dist
from tqdm import tqdm
import numpy as np
import time

DIST_THRESH = 0.1


def gen_locations() -> list[tuple[float, float]]:
    """Generate a list of locations for testing."""
    # return [(52.210796, 0.091659)]
    return fetch_data("tesco")


def test_encode_decode():
    locations = gen_locations()
    passed = 0
    fails = 0
    dists = []

    with open("fails.txt", "w") as f:
        f.write("lat, lon, guesslat, guesslon\n")

    with open("success.txt", "w") as f:
        f.write("lat, lon, guesslat, guesslon\n")

    try:
        for loc in tqdm(locations):
            encoded = format_location(encode(loc))
            decoded = decode(parse_location(encoded))
            dist = spherical_dist(np.array(loc), np.array(decoded))

            if dist > DIST_THRESH:
                tqdm.write(f"FAIL:\n\tloc: {loc}\n\tdecoded: {decoded}\n\tdist: {dist}")
                fails += 1
                dists.append(dist)

                with open("fails.txt", "a") as f:
                    f.write(f"{loc[0]}, {loc[1]}, {decoded[0]}, {decoded[1]}\n")

            else:
                passed += 1

                with open("success.txt", "a") as f:
                    f.write(f"{loc[0]}, {loc[1]}, {decoded[0]}, {decoded[1]}\n")

    except KeyboardInterrupt:
        print("KeyboardInterrupt - halting...")

    avg_dist = sum(dists) / len(dists)
    print(f"\nDONE. Passed: {passed}. Failed: {fails}. Average error: {avg_dist:.3f}m")


def test_encode_decode_greggs():
    num_greggs = len(fetch_data("greggs"))
    fails = 0

    for i in tqdm(range(num_greggs)):
        # t = time.monotonic()
        encoded = encode_greggs(i)
        # tqdm.write(f"Encoding took {time.monotonic() - t:.3f}s")

        stringified = ",".join(map(format_dist, encoded))
        parsed = [parse_dist(s.strip()) for s in stringified.split(",")]

        # t = time.monotonic()
        decoded = decode_greggs(parsed)
        # tqdm.write(f"Decoding took {time.monotonic() - t:.3f}s")

        if i != decoded:
            tqdm.write(f"FAIL: {i} != {decoded}")
            fails += 1

    print(f"Encode/Decode Greggs: {fails} failures.")


if __name__ == "__main__":
    test_encode_decode()