Finding the cheapest flights for a multi-leg trip with Amadeus API and Python

min_start = date(2019, 7, 10)
max_start = date(2019, 7, 20)
max_days = 21

places_df = pd.DataFrame([('Amsterdam', 'NL', 0, (max_start - min_start).days, True),  # for enabling tentative start date
                          ('Moscow', 'RU', 3, 5, True),
                          ('Irkutsk', 'RU', 7, 10, True),
                          ('Beijing', 'CN', 3, 5, True),
                          ('Shanghai', 'CN', 3, 5, True),
                          ('Tokyo', 'JP', 3, 5, False),
                          ('Amsterdam', 'NL', 0, 0, True)],  # the final destination
                         columns=['city', 'cc', 'min_days', 'max_days', 'only_direct'])

places_df['min_day_of_dep'] = places_df.min_days.rolling(min_periods=1, window=len(places_df)).sum()
places_df['max_day_of_dep'] = places_df.max_days.rolling(min_periods=1, window=len(places_df)).sum()

places_df

def get_iata(city, cc):
    response = call_api('/reference-data/locations',  # full code in the notebook
                        keyword=city,
                        countryCode=cc,
                        subType='AIRPORT')

    return [result['iataCode'] for result in response['data']]

get_iata('Moscow', 'RU')

['DME', 'SVO', 'VKO']

places_df['iata'] = places_df.apply(lambda place: get_iata(place['city'], place['cc']), axis=1)

routes_df = places_df.assign(dest_iata=places_df.iloc[1:].reset_index().iata)
routes_df['routes'] = routes_df.apply(
    lambda row: [*product(row['iata'], row['dest_iata'])] if isinstance(row['dest_iata'], list) else [],
    axis=1)

routes_df = routes_df.routes \
    .apply(pd.Series) \
    .merge(routes_df, right_index=True, left_index=True) \
    .drop(['routes', 'min_days', 'max_days', 'iata', 'dest_iata'], axis=1) \
    .melt(id_vars=['city', 'cc', 'min_day_of_dep', 'max_day_of_dep', 'only_direct'], value_name="route") \
    .drop('variable', axis=1) \
    .dropna()

routes_df['origin'] = routes_df.route.apply(lambda route: route[0])
routes_df['destination'] = routes_df.route.apply(lambda route: route[1])
routes_df = routes_df \
    .drop('route', axis=1) \
    .rename(columns={'city': 'origin_city',
                     'cc': 'origin_cc'})

routes_df.head(10)

route_dates_df = routes_df.assign(
    dates=routes_df.apply(lambda row: [min_start + timedelta(days=days)
                                       for days in range(int(row.min_day_of_dep), int(row.max_day_of_dep) + 1)],
                          axis=1))

route_dates_df = route_dates_df.dates \
    .apply(pd.Series) \
    .merge(route_dates_df, right_index=True, left_index=True) \
    .drop(['dates', 'min_day_of_dep', 'max_day_of_dep'], axis=1) \
    .melt(id_vars=['origin_city', 'origin_cc', 'origin', 'destination', 'only_direct'], value_name="date") \
    .drop('variable', axis=1) \
    .dropna()

valid_routes_df = route_dates_df[route_dates_df.date <= max_start + timedelta(days=max_days)]

valid_routes_df.head(10)

def get_prices(origin, destination, date, only_direct):
    response = call_api('/shopping/flight-offers',
                         origin=origin,
                         destination=destination,
                         nonStop='true' if only_direct else 'false',
                         departureDate=date.strftime("%Y-%m-%d"))

    if 'data' not in response:
        print(response)
        return []

    return [(origin, destination, date,
             Decimal(offer_item['price']['total']),
             parse_date(offer_item['services'][0]['segments'][0]['flightSegment']['departure']['at']),
             parse_date(offer_item['services'][0]['segments'][-1]['flightSegment']['arrival']['at']),
             len(offer_item['services'][0]['segments']))
            for flight in response['data']
            for offer_item in flight['offerItems']]

get_prices('IKT', 'PEK', date(2019, 7, 20), True)[:5]

[('IKT',
  'PEK',
  datetime.date(2019, 7, 20),
  Decimal('209.11'),
  datetime.datetime(2019, 7, 20, 1, 50, tzinfo=tzoffset(None, 28800)),
  datetime.datetime(2019, 7, 20, 4, 40, tzinfo=tzoffset(None, 28800)),
  1),
 ('IKT',
  'PEK',
  datetime.date(2019, 7, 20),
  Decimal('262.98'),
  datetime.datetime(2019, 7, 20, 15, 15, tzinfo=tzoffset(None, 28800)),
  datetime.datetime(2019, 7, 20, 18, 5, tzinfo=tzoffset(None, 28800)),
  1)]

prices_df = pd.DataFrame([price
                          for route in valid_routes_df.to_dict('record')
                          for price in get_prices(route['origin'], route['destination'], route['date'], route['only_direct'])],
                         columns=['origin', 'destination', 'date', 'price', 'departure_at', 'arrival_at', 'segments'])

airport_to_city = dict(zip(routes_df.origin, routes_df.origin_city))

prices_with_city_df = prices_df \
    .assign(duration=prices_df.arrival_at - prices_df.departure_at,
            origin_city=prices_df.origin.apply(airport_to_city.__getitem__),
            destination_city=prices_df.destination.apply(airport_to_city.__getitem__))
prices_with_city_df['route'] = prices_with_city_df.origin_city + " ✈️ " + prices_with_city_df.destination_city

valid_prices_with_city_df = prices_with_city_df[prices_with_city_df.price <= 800]

valid_prices_with_city_df.head()

next_flight_origin_city = dict(zip(places_df.city.iloc[:-2], places_df.city.iloc[1:-1]))
place_min_days = dict(zip(places_df.city.iloc[:-1], places_df.min_days.iloc[:-1]))
place_max_days = dict(zip(places_df.city.iloc[:-1], places_df.max_days.iloc[:-1]))

def build_itinerary(place, date):
    if place is None:
        return

    next_place = next_flight_origin_city.get(place)

    for days in range(place_min_days[place], place_max_days[place] + 1):
        flight_date = date + timedelta(days=days)
        for rest_flights in build_itinerary(next_place, flight_date):
            yield [(place, flight_date), *rest_flights]

        if next_place is None:
            yield [(place, flight_date)]

itinerary = [*build_itinerary('Amsterdam', min_start)]
itinerary[:3]

[[('Amsterdam', datetime.date(2019, 7, 10)),
  ('Moscow', datetime.date(2019, 7, 13)),
  ('Irkutsk', datetime.date(2019, 7, 20)),
  ('Beijing', datetime.date(2019, 7, 23)),
  ('Shanghai', datetime.date(2019, 7, 26)),
  ('Tokyo', datetime.date(2019, 7, 29))],
 [('Amsterdam', datetime.date(2019, 7, 10)),
  ('Moscow', datetime.date(2019, 7, 13)),
  ('Irkutsk', datetime.date(2019, 7, 20)),
  ('Beijing', datetime.date(2019, 7, 23)),
  ('Shanghai', datetime.date(2019, 7, 26)),
  ('Tokyo', datetime.date(2019, 7, 30))],
 [('Amsterdam', datetime.date(2019, 7, 10)),
  ('Moscow', datetime.date(2019, 7, 13)),
  ('Irkutsk', datetime.date(2019, 7, 20)),
  ('Beijing', datetime.date(2019, 7, 23)),
  ('Shanghai', datetime.date(2019, 7, 26)),
  ('Tokyo', datetime.date(2019, 7, 31))]]

def find_flights(prices_with_city_df, itinerary_route, n_cheapest):
    result_df = None
    for place, date in itinerary_route:
        place_df = prices_with_city_df \
            [(prices_with_city_df.origin_city == place) & (prices_with_city_df.date == date)] \
            .sort_values('price', ascending=True) \
            .head(n_cheapest) \
            .add_prefix(f'{place}_')

        if result_df is None:
            result_df = place_df
        else:
            result_df = result_df \
                .assign(key=1) \
                .merge(place_df.assign(key=1), on="key") \
                .drop("key", axis=1)

            result_df['total_price'] = reduce(operator.add, (
                result_df[column] for column in result_df.columns
                if 'price' in column and column != 'total_price'
            ))

            result_df = result_df \
                .sort_values('total_price', ascending=True) \
                .head(n_cheapest)

    result_df['total_flights_duration'] = reduce(operator.add, (
        result_df[column] for column in result_df.columns
        if 'duration' in column
    ))

    return result_df[['total_price', 'total_flights_duration'] + [
        column for column in result_df.columns
        if 'total_' not in column]]

find_flights(prices_with_city_df, itinerary[0], 100).head(5)

itinerary_df = reduce(pd.DataFrame.append, (find_flights(prices_with_city_df, itinerary_route, 10)
                                           for itinerary_route in itinerary))

itinerary_df \
    .sort_values(['total_price', 'total_flights_duration']) \
    .head(10)