trip planner

find.py

import json
from datetime import date, datetime, timedelta
from collections import defaultdict, namedtuple
from multiprocessing import Pool, cpu_count
import csv
from operator import itemgetter
from heapq import merge
from itertools import islice
from dateutil.parser import parse

MIN_START = date(2018, 9, 1)
MAX_START = date(2018, 10, 1)

MIN_STAY = 3
MAX_STAY = 5

MAX_TRIP = 28
MIN_VISITED = 4

MAX_FLIGHT_PRICE = 500
MAX_HOME_FLIGHT_PRICE = 700
MAX_TRIP_PRICE = 2500

RESULT_SIZE = 100_000

HOME_CITY = 'Amsterdam'

CITIES = [
    HOME_CITY,
    'Rio de Janeiro',
    'Montevideo',
    'Asunción',
    'Lima',
    'Mexico City',
    'Santiago',
    'La Paz',
    'Quito',
    'Bogota',
    'Buenos Aires',
]

Flight = namedtuple('Flight', ('from_id', 'to_id', 'day_number', 'price'))
Trip = namedtuple('Trip', ('price', 'flights'))

id2city = dict(enumerate(CITIES))
city2id = {city: id for id, city in id2city.items()}
city_ids = set(id2city.keys())

home_city_id = city2id[HOME_CITY]


def read_flights(path):
    with open(path) as f:
        data = json.load(f)

    for flight_data in data:
        if MIN_START <= parse(flight_data['date']).date() < MAX_START:
            yield Flight(from_id=city2id[flight_data['from']],
                         to_id=city2id[flight_data['to']],
                         day_number=(parse(flight_data['date']).date() - MIN_START).days,
                         price=flight_data['price'])

flights_with_stop = [flight for flight in read_flights('flights.json')
                     if flight.price < MAX_FLIGHT_PRICE
                     and (flight.to_id == home_city_id or flight.from_id == home_city_id)
                     and flight.price < MAX_HOME_FLIGHT_PRICE]
flights_nonstop = [flight for flight in read_flights('flights_nonstop.json')
                   if flight.price < MAX_FLIGHT_PRICE
                   or ((flight.to_id == home_city_id or flight.from_id == home_city_id)
                       and flight.price < MAX_HOME_FLIGHT_PRICE)]

flights = flights_with_stop + flights_nonstop

print(f"Read {len(flights)} flights with tolerable price")

from_id2day_number2to_id2flight = defaultdict(
    lambda: defaultdict(
        lambda: {}))
for flight in flights:
    from_id2day_number2to_id2flight[flight.from_id] \
        [flight.day_number][flight.to_id] = flight


def _generate_trips(can_visit, can_travel, can_spent, current_id,
                    current_day, trip_flights):
    if trip_flights[-1].to_id == home_city_id:
        yield Trip(
            price=sum(flight.price for flight in trip_flights),
            flights=trip_flights)
        return

    if not can_visit or can_travel < MIN_STAY or can_spent == 0:
        return

    if len(trip_flights) >= MIN_VISITED and home_city_id not in can_visit:
        can_visit.add(home_city_id)

    for to_id in can_visit:
        can_visit_next = can_visit.difference({to_id})
        for stay in range(MIN_STAY, min(MAX_STAY, can_travel) + 1):
            current_day_next = current_day + stay
            flight_next = from_id2day_number2to_id2flight \
                .get(current_id, {}).get(current_day_next, {}).get(to_id)
            if not flight_next:
                continue

            can_spent_next = can_spent - flight_next.price
            if can_spent_next < 0:
                continue

            yield from _generate_trips(
                can_visit_next, can_travel - stay, can_spent_next, to_id,
                                current_day + stay, trip_flights + [flight_next])


def _generator_stage(params):
    return sorted(_generate_trips(*params), key=itemgetter(0))


def generate_trips():
    generators_params = [(
        city_ids.difference({start_id, home_city_id}),
        MAX_TRIP,
        MAX_TRIP_PRICE - from_id2day_number2to_id2flight[home_city_id][start_day][start_id].price,
        start_id,
        start_day,
        [from_id2day_number2to_id2flight[home_city_id][start_day][start_id]])
        for start_day in range((MAX_START - MIN_START).days)
        for start_id in from_id2day_number2to_id2flight[home_city_id][start_day].keys()]

    start = datetime.now()

    with Pool(cpu_count() * 2) as pool:
        for n, stage_result in enumerate(pool.imap_unordered(_generator_stage, generators_params)):
            print(f'Generated {n + 1} of {len(generators_params)} in {datetime.now() - start}')
            yield stage_result


def format_trip(trip):
    days = trip.flights[-1].day_number - trip.flights[0].day_number
    cities = len(trip.flights) - 1
    start_city = id2city[trip.flights[0].to_id]
    start_date = MIN_START + timedelta(days=trip.flights[0].day_number)
    end_city = id2city[trip.flights[-1].from_id]
    end_date = MIN_START + timedelta(days=trip.flights[-1].day_number)
    details = ' & '.join(id2city[flight.from_id] + ' -> ' + id2city[flight.to_id]
                         + ' ' + str(MIN_START + timedelta(days=flight.day_number))
                         + ' ' + str(flight.price)
                         for flight in trip.flights)
    return (trip.price,
            days, cities,
            start_city, str(start_date),
            end_city, str(end_date),
            details)


trips = [*merge(*generate_trips(), key=itemgetter(0))]
print('Generated all possible trips')

with open('trips.json', 'w') as f:
    json.dump(trips, f)

with open('trips.csv', 'w', newline='') as f:
    writer = csv.writer(f)
    writer.writerow(('price', 'days', 'cities', 'start city', 'start date', 'end city', 'end date', 'details'))
    for trip in trips:
        writer.writerow(format_trip(trip))

parse.js

const HOME_CITY = 'Amsterdam';

const CITIES = [
  'Rio de Janeiro',
  'Montevideo',
  'Asunción',
  'Lima',
  'Mexico City',
  'Santiago',
  'La Paz',
  'Quito',
  'Bogota',
  'Buenos Aires',
];

const FROM = 0;
const TO = 1;

const getElements = (selector, intervalSize = 300) => new Promise((resolve) => {
  const interval = setInterval(() => {
    console.log(`Getting ${selector}`);

    const elements = [...document.querySelectorAll(selector)];

    if (elements.length) {
      resolve(elements);
      clearInterval(interval);
    }
  }, intervalSize);
});

const delay = (delaySize) => new Promise((resolve) => setTimeout(() => resolve(), delaySize));

const changeInput = (input, value) => {
  input.value = value;

  const keypressEvent = document.createEvent("HTMLEvents");
  keypressEvent.initEvent('keypress');
  input.dispatchEvent(keypressEvent);
};

const setDestination = async (type, value) => {
  const holder = (await getElements('.gws-flights-form__location-text'))[type];
  holder.click();

  const [input] = await getElements('#flt-modaldialog input');
  changeInput(input, value);

  const [destination] = await getElements('#flt-modaldialog .fsapp-option-content');
  await delay(500);
  destination.click();

  await delay(3000);
};

const getPrices = async () => {
  const [holder] = await getElements('.gws-flights-form__calendar-input>div');
  holder.click();
  await delay(5000);

  await getElements('.gws-travel-calendar__active [data-price]');
  const dates = await getElements('[data-day] [data-price]');

  const prices = dates
    .filter(({textContent}) => textContent.length)
    .map(({textContent, parentElement}) => [
      parentElement.parentElement.getAttribute('data-day'),
      parseInt(textContent.replace(/[^\d]*/g, ''), 10)
    ]);

  const [underlay] = await getElements('.gws-flights__modal-underlay');
  underlay.click();

  await delay(1000);

  return prices;
};

const getFlightsData = async ([from, to]) => {
  await setDestination(FROM, from);
  await setDestination(TO, to);

  const prices = await getPrices();

  return prices.map(([date, price]) => ({
    date, price, from, to,
  }));
};

function* getAllPossibleFlights() {
  for (let from of CITIES) {
    for (let to of CITIES) {
      if (from !== to) {
        yield [from, to];
      }
    }

    yield [HOME_CITY, from];
    yield [from, HOME_CITY];
  }
}

const collectData = async () => {
  let result = [];
  for (let flight of getAllPossibleFlights()) {
    console.log(`Fetching data for ${flight}`);
    const flightsData = await getFlightsData(flight);
    console.log(flightsData);
    result = result.concat(flightsData);
  }
  return result;
};

const win = window.open('');

collectData().then(
  (data) => win.document.write(JSON.stringify(data)),
  (error) => console.error("Can't get flights", error),
);