For checkouts or to view logs direct your SVN client to svn://svn.saintamh.org/code/azimuth/py/azimuth.py

#!/usr/bin/env python

"""
$Id: azimuth.py 2460 2013-09-14 08:54:19Z saintamh $
"""

#----------------------------------------------------------------------------------------------------------------------------------
# includes

# standards
import cPickle
import collections
import datetime
import os
import shutil
import sys

# saintamh
from saintamh.http import SimpleHTTPClient
from saintamh.maps import LatLng, LatLngRectangle
from saintamh.maps.ortho import MapEngine, map_engine_by_id
from saintamh.struct import *
from saintamh.util.cartesian import FloatPoint, IntDimension, IntPoint, IntRectangle, Line, iter_points_on_line
from saintamh.util.dates import DatetimeRange, days, english_timedelta, minutes
from saintamh.util.iterators import index, pairwise, uniq
from saintamh.util.paths import ensure_parent_dir_exists, here, makedirs
from saintamh.util.scrapers import parse_json

# other modules in this dir
import airports
import flightradar24

from PIL import Image,ImageDraw

#----------------------------------------------------------------------------------------------------------------------------------
# data structs

Contrail = collections.namedtuple (
    'Contrail', (
        'plane',
        'tracks',
        )
    )

#----------------------------------------------------------------------------------------------------------------------------------
# utils

http = SimpleHTTPClient (
    cache_path = here (__file__, '..', 'cache'),
    cache_requests = True,
    courtesy_delay = 5,
    )

class cached_iter:
    def __init__ (self, iterable):
        self.iterable = iterable
    def __iter__ (self):
        if not hasattr (self.iterable, '__len__'):
            self.iterable = tuple (self.iterable)
        return iter (self.iterable)


#----------------------------------------------------------------------------------------------------------------------------------
# utils

def compute_llrect (center_ll, map_engine, map_zoom, box_size_in_px):
    half_box_size = box_size_in_px/2
    airport_px = map_engine.latlng_to_pixel (center_ll, map_zoom)
    return LatLngRectangle (
        top_left  = map_engine.pixel_to_latlng (airport_px - half_box_size, map_zoom),
        bot_right = map_engine.pixel_to_latlng (airport_px + half_box_size, map_zoom),
        )

def dt_str (dt):
    return str(dt).replace(' ', 'T').replace('T00:00:00','')


#----------------------------------------------------------------------------------------------------------------------------------
# load and parse data

def _contrails_cache_file_path (llrect, dtrange, min_altitude):
    return here (
        __file__,
        '..',
        'cache',
        'contrails',
        '%s_%s_%s_%s_%s.pickle' % (
            dt_str (dtrange.start),
            dt_str (dtrange.stop),
            llrect.top_left,
            llrect.bot_right,
            min_altitude,
            )
        )


def _fetch_tracks_at_dt_indexed_by_plane (dt, selected_llrects):
    index = {}
    print "    %s" % dt,
    for _,obs in flightradar24.fetch_observations (dt, selected_llrects):
        #assert index.get(obs.plane) in (None,obs.track), "Two tracks for %s" % repr(obs.plane)
        index.setdefault (obs.plane, obs.track)
    return index


def ensure_contrail_data_preparsed_and_cached (selected_llrects, dtrange, min_altitude):

    all_cache_files_by_llrect = dict (
        (llrect, _contrails_cache_file_path (llrect, dtrange, min_altitude))
        for llrect in selected_llrects
        )

    if not all (map (os.path.isfile, all_cache_files_by_llrect.itervalues())):

        lookback = datetime.datetime.now() - dtrange.start
        if lookback > days(20):
            print "WARNING - You're requesting data %s in the past, you'll prob. get empty data" % english_timedelta (lookback)

        cache_fh_by_llrect = dict (
            (llrect, open (ensure_parent_dir_exists('%s.part' % cache_file), 'wb'))
            for llrect,cache_file in all_cache_files_by_llrect.iteritems()
            )

        def save_contrail (llrect, plane, tracks):
            if min(t.altitude_feet for t in tracks) <= min_altitude:
                # We could've used struct.py's built-in serialization modules, but they're quite inefficent, making it painfully slow
                # to tweak the output graphics. This should be faster.
                fh = cache_fh_by_llrect[llrect]
                cPickle.dump (plane, fh)
                cPickle.dump (tracks, fh)

        partial_trails_by_llrect = {}

        # for every pair of consecutive datetimes
        for (_,all_tracks_by_plane_at_prev_dt),(dt,all_tracks_by_plane_at_dt) in pairwise (
            (dt, _fetch_tracks_at_dt_indexed_by_plane(dt,selected_llrects))
            for dt in dtrange
            ):

            active_partial_trails_by_llrect = {}

            # for every plane that appears in both datetimes
            for plane,track in all_tracks_by_plane_at_dt.iteritems():
                prev_track = all_tracks_by_plane_at_prev_dt.get (plane)
                if prev_track is not None:

                    # for every selected llrect that this plane is passing through (including segments that cross over the llrect
                    # boundary, i.e. where only one of the two obs is in the llrect)
                    for llrect in selected_llrects:
                        if track.ll in llrect or prev_track.ll in llrect:

                            # record the trail
                            trail = partial_trails_by_llrect.get(llrect,{}).pop (plane,None) or [prev_track]
                            trail.append (track)
                            active_partial_trails_by_llrect.setdefault(llrect,{})[plane] = trail

                            # NB don't `break' here -- a single obs can appear in several llrects, when we monitor nearby airports

            # to save memory, dump to disk right away contrails that weren't seen at this dt
            for llrect,partial_trails in partial_trails_by_llrect.iteritems():
                plane_is_active = active_partial_trails_by_llrect.get(llrect,{}).__contains__
                for plane,tracks in partial_trails.iteritems():
                    if not plane_is_active(plane):
                        save_contrail (llrect, plane, tracks)

            print "      Tracking %d planes" % sum(len(trails) for trails in active_partial_trails_by_llrect.itervalues())
            partial_trails_by_llrect = active_partial_trails_by_llrect

        # flush what's left at the end
        for llrect,partial_trails in partial_trails_by_llrect.iteritems():
            for plane,tracks in partial_trails.iteritems():
                save_contrail (llrect, plane, tracks)

        for llrect,fh in cache_fh_by_llrect.iteritems():
            cache_file = all_cache_files_by_llrect[llrect]
            cPickle.dump (None, fh)
            fh.close()
            shutil.move ('%s.part' % cache_file, cache_file)


def load_preparsed_contrails (llrect, dtrange, min_altitude):

    cache_file = _contrails_cache_file_path (llrect, dtrange, min_altitude)
    print "Loading contrails from %s" % cache_file

    with open (cache_file, 'rb') as fh:
        while True:
            plane = cPickle.load (fh)
            if plane is None:
                break
            else:
                yield Contrail (
                    plane = plane,
                    tracks = cPickle.load(fh),
                    )


#----------------------------------------------------------------------------------------------------------------------------------
# draw the contrails and the map

def draw_contrail_img (job, iter_contrails_in_extended_llrect):

    img_file = here (
        __file__, '..', 'cache', 'contrails', '%s.png' % '-'.join (
            map (str, (
                    job.top_left_pixel.x,
                    job.top_left_pixel.y,
                    job.map_img_px_rect.bot_right.x,
                    job.map_img_px_rect.bot_right.y,
                    job.map_engine.projection.__name__,
                    job.color_policy,
                    job.contrail_alpha,
                    ))
            )
        )
    if os.path.isfile (img_file):
        print "Loading contrail images from %s" % img_file
        return Image.open (img_file)

    get_segment_color = COLOR_POLICIES[job.color_policy]

    print "Drawing %s" % img_file
    contrail_img = Image.new ('RGBA', job.output_img_size_in_pixels.as_tuple, (0,0,0,0))
    contrail_px = contrail_img.load()

    num_outboard = 0
    for contrail_i,contrail in enumerate(iter_contrails_in_extended_llrect):
        is_outboard = False
        for prev_track,track in pairwise (contrail.tracks):
            red_inc,green_inc,blue_inc = get_segment_color (prev_track, track)
            alpha_inc = job.contrail_alpha

            # print (prev_track,track)
            # print tuple (
            #     (job.map_engine.latlng_to_pixel(t.ll,job.map_zoom)-job.top_left_pixel)
            #     for t in (prev_track,track)
            #     )

            for pt in iter_points_on_line (
                *tuple (
                    (job.map_engine.latlng_to_pixel(t.ll,job.map_zoom)-job.top_left_pixel)
                    for t in (prev_track,track)
                    )
                 ):
                if pt in job.map_img_px_rect:
                    p = contrail_px[pt.as_tuple]
                    contrail_px[pt.as_tuple] = (
                        p[0] + red_inc,
                        p[1] + green_inc,
                        p[2] + blue_inc,
                        p[3] + alpha_inc,
                        )
                else:
                    is_outboard = True
        num_outboard += int(is_outboard)
    print "%s contrails, %d outboard" % (contrail_i+1, num_outboard)

    contrail_img.save (ensure_parent_dir_exists (img_file))
    return contrail_img


def draw_map_img (job):

    img_file = here (
        __file__, '..', 'cache', 'maps', '%s.png' % '_'.join (
            str(e) for e in (
                job.llrect.top_left,
                job.llrect.bot_right,
                job.top_left_pixel.x,
                job.top_left_pixel.y,
                job.map_img_px_rect.bot_right.x,
                job.map_img_px_rect.bot_right.y,
                job.map_engine.id,
                )
            )
        )
    if os.path.isfile (img_file):
        print "Loading map image from %s" % img_file
        return Image.open (img_file)

    # fetch the satellite picture that covers this area
    print "Assembling %s" % img_file
    map_img = job.map_engine.draw_map (http, job.llrect, zoom=job.map_zoom, ignore_missing_tiles=True)

    if hasattr (job, 'airport_info'):
        map_img_draw = ImageDraw.Draw (map_img)
        for rwy in job.airport_info.runways:
            rwy_px_coords = tuple (
                job.map_engine.latlng_to_pixel(ll,job.map_zoom,force_ints=False)-job.top_left_pixel
                for ll in rwy.latlng_at_both_ends
                )
            rwy_line = Line.through (*tuple(FloatPoint(*c.as_tuple) for c in rwy_px_coords))
            map_img_draw.line (
                (0, rwy_line.y(0), job.map_img_px_rect.bot_right.x, rwy_line.y(job.map_img_px_rect.bot_right.x)),
                fill='white'
                )

    map_img.save (ensure_parent_dir_exists (img_file))
    return map_img


#----------------------------------------------------------------------------------------------------------------------------------
# Jobs config. Each "job" corresponds to one output image.

COLOR_POLICIES = {
    'all-red':            lambda prev_track,track: (16,0,0),
    'all-yellow':         lambda prev_track,track: (16,16,0),
    'green-asc-red-desc': lambda prev_track,track: (
        (64,0,0)
        if prev_track.altitude_feet > track.altitude_feet else
        (0,64,0)
        ),
    }

class Job (struct (
        id = str,
        center_latlng = LatLng,
        map_zoom = int,
        map_zoom_for_catchment_area = int,
        output_img_size_in_pixels = IntDimension,
        map_engine = MapEngine,
        datetime_range = DatetimeRange,
        min_altitude = int,
        color_policy = one_of (*COLOR_POLICIES.keys()),
        contrail_alpha = int,

        # We want to only load plane data around the airports we're interested in, but if we only load the data points that are
        # within the pictures we produce, we lose the plane segments that cross the picture boundary, since they all have, by
        # definition, one endpoint outside the picture frame.
        # 
        # FIXME - 2012-02-08 - Ideally we wouldn't need to specify the catchment_area -- all it needs to be is an area big enough
        # that we catch all contrails that have at least one endpoint in the map. Haven't gotten around to implementing that yet.
        #
        catchment_area = lambda self: compute_llrect (
            self.center_latlng,
            self.map_engine,
            self.map_zoom_for_catchment_area,
            self.output_img_size_in_pixels + IntDimension(200,200),
            ),
        llrect = lambda self: compute_llrect (
            self.center_latlng,
            self.map_engine,
            self.map_zoom,
            self.output_img_size_in_pixels,
            ),

        # Figure out the pixel size and position of the image
        top_left_pixel = lambda self: self.map_engine.latlng_to_pixel (self.llrect.top_left, self.map_zoom),
        map_img_px_rect = lambda self: IntRectangle (
            (0,0),
            self.output_img_size_in_pixels.as_tuple,
            ),
        )):
    pass


class AirportJob (struct (
        __extends__ = Job,
        airport_info = airports.AirportInfo,
        )):

    def __init__ (self, **kwargs):
        kwargs['center_latlng'] = kwargs['airport_info'].latlng
        super(AirportJob,self).__init__ (**kwargs)


ALL_JOBS = tuple (
    AirportJob (
        id = '%s_%s-%s_z%d_%s_%s_%s' % (
            airport_info.iata_id,
            dt_str (datetime_range.start),
            dt_str (datetime_range.stop),
            map_zoom,
            map_engine_id,
            color_policy,
            contrail_alpha,
            ),
        airport_info = airport_info,
        map_zoom = map_zoom,
        map_zoom_for_catchment_area = 10, # to avoid having to reparse the files for every zoom
        map_engine = map_engine_by_id (map_engine_id),
        datetime_range = datetime_range,
        output_img_size_in_pixels = IntDimension (2378,1682), # A4,A0,etc paper ratio
        min_altitude = 5000,
        color_policy = color_policy,
        contrail_alpha = contrail_alpha,
        )
    for color_policy,contrail_alpha in (
        ('green-asc-red-desc', 32),
        )
    for airport_info in (
        airports.fetch_airport_info (airport_iata_id, http.simple_get)
        for airport_iata_id in re.findall (
            r'\w+',
            # 'AMS LHR INN CDG DUB EDI VIE FRA LGW DME GVA BRU YUL LPA' # IST
            # SOF YTZ YVR YYZ ZQW SCN HHN HKG AGP TFN
            # 
            # SDU (Rio de Janeiro) apparently has an interesting approach, but as of 2012-02-06 FR24 coverage doesn't go all the
            # way to the ground
            'PEK',
            )
        )
    for datetime_range in {
        'EDI': (
            # For EDI, extend coverage, because there aren't as many flights
            DatetimeRange (datetime.datetime(2012,1,25), datetime.datetime(2012,2,10), minutes(2)),
            # Also try this week, during which it seemed to me from my window that the flights were mostly going west-east
            DatetimeRange (datetime.datetime(2012,01,27), datetime.datetime(2012,02,03), minutes(2)),
            ),
        'PEK': (
            # Extend coverage for PEK as well, because data is spotty
            DatetimeRange (datetime.datetime(2013,05,13), datetime.datetime(2013,06,03), minutes(2)),
            ),
        }.get (airport_info.iata_id, (
            # For all other airports, just this week here is enough:
            DatetimeRange (datetime.datetime(2013,05,13), datetime.datetime(2013,05,20), minutes(2)),
            ))
    for map_zoom in (10,11,12)
    for map_engine_id in (
        #'\x67\x6F\x6F\x67\x6C\x65-labelled-terrain',
        '\x67\x6F\x6F\x67\x6C\x65-satellite',
        '\x62\x69\x6e\x67-satellite',
        )
    #for CANCEL_THAT_HOUSTON in []

    # ) + tuple (
    #     Job (
    #         id = 'europa_%s_%s-%s_z%d-%s-%s-a%d' % (
    #             center_ll,
    #             dt_str(dtrange.start),
    #             dt_str(dtrange.stop),
    #             map_zoom,
    #             engine.id,
    #             color_policy,
    #             contrail_alpha,
    #             ),
    #         center_latlng = center_ll,
    #         map_zoom = map_zoom,
    #         map_zoom_for_catchment_area = map_zoom,
    #         output_img_size_in_pixels = IntDimension (4756,3364), # A1 sheet at 4 dots/mm
    #         map_engine = engine,
    #         datetime_range = dtrange,
    #         min_altitude = 10000000,
    #         contrail_alpha = contrail_alpha,
    #         color_policy = color_policy,
    #         )
    #     for map_zoom,center_ll,dtrange,all_contrail_alphas in (
    #         (
    #             7,
    #             LatLng ('47.842658,13.044433'),
    #             DatetimeRange (datetime.datetime(2012,02,14), datetime.datetime(2012,02,21), minutes(2)),
    #             (4,),
    #             ),
    #         (
    #             8,
    #             LatLng ('50.457504,8.254394'),
    #             DatetimeRange (datetime.datetime(2012,02,14), datetime.datetime(2012,02,21), minutes(2)),
    #             (4,),
    #             ),
    #         (
    #             9,
    #             LatLng ('49.912325,5.229492'),
    #             DatetimeRange (datetime.datetime(2012,02,14), datetime.datetime(2012,02,21), minutes(2)),
    #             (4,8),
    #             ),
    #         )
    #     for engine,color_policy in (
    #         (engine, 'all-red' if 'terrain' in engine.id else 'all-yellow')
    #         for engine in map (map_engine_by_id, (
    #                 "\x67\x6F\x6F\x67\x6C\x65-satellite",
    #                 #"\x67\x6F\x6F\x67\x6C\x65-labelled-terrain",
    #                 "\x67\x6F\x6F\x67\x6C\x65-terrain-few-labels",
    #                 "\x62\x69\x6e\x67-satellite",
    #                 ))
    #         )
    #     for contrail_alpha in all_contrail_alphas
        )


#----------------------------------------------------------------------------------------------------------------------------------

def main (selected_job_ids):

    selected_jobs = tuple (
        job
        for job in ALL_JOBS
        if job.id in selected_job_ids or not selected_job_ids
        )
    if selected_job_ids and len(selected_jobs) < len(selected_job_ids):
        print "Bad job IDs. Available IDs:"
        for job in sorted (ALL_JOBS, key=lambda j: j.id):
            print "    %s" % job.id
        return False

    for dt_range,min_altitude in uniq ((job.datetime_range,job.min_altitude) for job in selected_jobs):
        ensure_contrail_data_preparsed_and_cached (
            tuple (uniq (job.catchment_area for job in selected_jobs if job.datetime_range == dt_range)),
            dt_range,
            min_altitude,
            )

    for job in sorted (selected_jobs, key=lambda j: j.id):
        print
        print job.id

        # Create an initially blank, transparent image, of our output's dimensions. Every plane then "scratches" away the
        # transparency over the track it flies, with the result that the lines are less visible where only one plane has flown, and
        # more visible where many have flown.
        contrail_img = draw_contrail_img (
            job,

            # For every map, we open a new iterator over the data that has been pickled to disk, rather than keeping it in memory.
            # This means we use more CPU, but rather little memory, allowing us to draw huge maps (e.g. all of Europe).
            load_preparsed_contrails (job.catchment_area, job.datetime_range, job.min_altitude)

            )

        # Paint the map background
        map_img = draw_map_img (job)

        # superimpose the airplane contrails
        assert job.map_img_px_rect.bot_right.as_tuple == map_img.size, repr ((job.map_img_px_rect, map_img.size))
        map_img.paste (contrail_img, (0,0), contrail_img)

        # save
        output_dir = makedirs (here (__file__, '..', 'out'))
        output_file = os.path.join (output_dir, '%s.jpg' % job.id)
        map_img.save (output_file)
        print "Saved %s" % output_file


#----------------------------------------------------------------------------------------------------------------------------------
# parse cmd line

if __name__ == '__main__':
    main (
        selected_job_ids = sys.argv[1:],
        )

#----------------------------------------------------------------------------------------------------------------------------------