'use strict';

/*
 * yield_geo_near.js (extends yield.js)
 *
 * Intersperse geo $near queries with updates and deletes of documents they may match.
 */
load('jstests/concurrency/fsm_libs/extend_workload.js');  // for extendWorkload
load('jstests/concurrency/fsm_workloads/yield.js');       // for $config

var $config = extendWorkload($config, function($config, $super) {

    /*
     * Use geo $near query to find points near the origin. Note this should be done using the
     * geoNear command, rather than a $near query, as the $near query doesn't work in a sharded
     * environment. Unfortunately this means we cannot batch the request.
     */
    $config.states.query = function geoNear(db, collName) {
        // This distance gets about 80 docs around the origin. There is one doc inserted
        // every 1m^2 and the area scanned by a 5m radius is PI*(5m)^2 ~ 79.
        var maxDistance = 5;

        var res = db.runCommand({geoNear: collName, near: [0, 0], maxDistance: maxDistance});
        assertWhenOwnColl.commandWorked(res);  // Could fail if more than 1 2d index.
        assertWhenOwnColl(function verifyResults() {
            var results = res.results;
            var prevDoc = {dis: 0};  // distance should never be less than 0
            for (var i = 0; i < results.length; i++) {
                var doc = results[i];
                assertAlways.lte(NumberInt(doc.dis), maxDistance);  // satisfies query
                assertAlways.lte(prevDoc.dis, doc.dis);             // returned in the correct order
                prevDoc = doc;
            }
        });
    };

    $config.data.genUpdateDoc = function genUpdateDoc() {
        var P = Math.floor(Math.sqrt(this.nDocs));

        // Move the point to another location within the PxP grid.
        var newX = Random.randInt(P) - P / 2;
        var newY = Random.randInt(P) - P / 2;
        return {$set: {geo: [newX, newY]}};
    };

    $config.data.getIndexSpec = function getIndexSpec() {
        return {geo: '2d'};
    };

    $config.data.getReplaceSpec = function getReplaceSpec(i, coords) {
        return {_id: i, geo: coords};
    };

    /*
     * Insert some docs in geo form and make a 2d index.
     */
    $config.setup = function setup(db, collName, cluster) {
        $super.setup.apply(this, arguments);

        var P = Math.floor(Math.sqrt(this.nDocs));
        var i = 0;
        // Set up some points to query (in a PxP grid around 0,0).
        var bulk = db[collName].initializeUnorderedBulkOp();
        for (var x = 0; x < P; x++) {
            for (var y = 0; y < P; y++) {
                var coords = [x - P / 2, y - P / 2];
                bulk.find({_id: i}).upsert().replaceOne(this.getReplaceSpec(i, coords));
                i++;
            }
        }
        assertAlways.writeOK(bulk.execute());
        assertAlways.commandWorked(db[collName].ensureIndex(this.getIndexSpec()));
    };

    return $config;
});