frostfs-locode-db/pkg/locode/db/continents/geojson/calls.go

package continentsdb

import (
	"fmt"
	"os"

	locodedb "git.frostfs.info/TrueCloudLab/frostfs-locode-db/pkg/locode/db"
	"github.com/paulmach/orb"
	"github.com/paulmach/orb/geojson"
	"github.com/paulmach/orb/planar"
	"github.com/paulmach/orb/quadtree"
)

const continentProperty = "Continent"

// PointContinent goes through all polygons and returns the continent
// in which the point is located.
//
// Returns locodedb.ContinentUnknown if no entry matches.
//
// All GeoJSON feature are parsed from file once and stored in memory.
func (db *DB) PointContinent(point *locodedb.Point) (*locodedb.Continent, error) {
	var err error

	db.once.Do(func() {
		err = db.init()
	})

	if err != nil {
		return nil, err
	}

	planarPoint := orb.Point{point.Longitude(), point.Latitude()}

	var (
		continent string
		minDst    float64
	)

	pointer := db.tree.Matching(planarPoint, func(p orb.Pointer) bool {
		return planar.PolygonContains(
			p.(*geojson.Feature).Geometry.(orb.Polygon),
			planarPoint,
		)
	})

	if pointer != nil {
		continent = pointer.(*geojson.Feature).Properties.MustString(continentProperty)
	}

	if continent == "" {
		for _, feature := range db.features {
			distance := planar.DistanceFrom(feature.Geometry, planarPoint)
			if minDst == 0 || minDst > distance {
				minDst = distance
				continent = feature.Properties.MustString(continentProperty)
			}
		}
	}

	c := continentFromString(continent)

	return &c, nil
}

func (db *DB) init() error {
	data, err := os.ReadFile(db.path)
	if err != nil {
		return fmt.Errorf("could not read data file: %w", err)
	}

	features, err := geojson.UnmarshalFeatureCollection(data)
	if err != nil {
		return fmt.Errorf("could not unmarshal GeoJSON feature collection: %w", err)
	}

	db.features = features.Features

	err = db.buildQuadtree()
	if err != nil {
		return fmt.Errorf("could not build quadtree: %w", err)
	}

	return nil
}

func (db *DB) buildQuadtree() error {
	db.tree = quadtree.New(orb.Bound{
		Min: orb.Point{-180, -180},
		Max: orb.Point{180, 180},
	})

	for _, feature := range db.features {
		var multiPolygon orb.MultiPolygon

		if polygon, ok := feature.Geometry.(orb.Polygon); ok {
			multiPolygon = append(multiPolygon, polygon)
		} else {
			multiPolygon = feature.Geometry.(orb.MultiPolygon)
		}

		for _, polygon := range multiPolygon {
			newFeature := geojson.NewFeature(polygon)
			newFeature.Properties = feature.Properties.Clone()
			err := db.tree.Add(newFeature)
			if err != nil {
				return err
			}
		}
	}

	return nil
}

func continentFromString(c string) locodedb.Continent {
	switch c {
	default:
		return locodedb.ContinentUnknown
	case "Africa":
		return locodedb.ContinentAfrica
	case "Asia":
		return locodedb.ContinentAsia
	case "Europe":
		return locodedb.ContinentEurope
	case "North America":
		return locodedb.ContinentNorthAmerica
	case "South America":
		return locodedb.ContinentSouthAmerica
	case "Antarctica":
		return locodedb.ContinentAntarctica
	case "Australia", "Oceania":
		return locodedb.ContinentOceania
	}
}