#!/usr/bin/env python

"""
test writing + reading distances
"""

import os
import shutil
import tempfile
import unittest
from common import datafile
from globalneighbors.distance import calculate_neighbors
from globalneighbors.distance import write_neighbors
from globalneighbors.locations import locations
from globalneighbors.neighbors import read_neighbors_file
from globalneighbors.read import read_cities


class TestDistanceReadWrite(unittest.TestCase):

    def test_10000(self):
        """test 10000 cities"""

        # read locations
        citiesfile = datafile('10000cities.tsv')
        assert os.path.exists(citiesfile)
        with open(citiesfile) as f:
            city_locations = locations(read_cities(f))

        # calculate neighbors
        neighbors = calculate_neighbors(city_locations,
                                        k=50,
                                        lat_tol=2.,
                                        lon_tol=2.)

        # make a staging area
        tmpdir = tempfile.mkdtemp()
        try:
            # write the neighbors
            outfile = os.path.join(tmpdir, 'neighbors.dat')
            assert not os.path.exists(outfile)
            with open(outfile, 'w') as f:
                write_neighbors(f, neighbors)
            assert os.path.exists(outfile)

            # read the neighbors
            with open(outfile) as f:
                new_neighbors = read_neighbors_file(f)
        finally:
            shutil.rmtree(tmpdir, ignore_errors=True)

        # they should be equal
        assert len(neighbors) == len(new_neighbors)
        assert sorted(neighbors.keys()) == sorted(new_neighbors.keys())
        for key in neighbors.keys():
            valueA = neighbors[key]
            valueB = new_neighbors[key]
            assert valueA == valueB


if __name__ == '__main__':
    unittest.main()