geol.py

"""Library for obtaining data from BGS WMS servers.

Unfortunately the BGS do not provide a simple text based API to access geology information for specific coordinates
or areas within a bounding box.

Available classes:
- GeoBGS: When called instantiates a GeoBGS object.
"""

from io import BytesIO
from PIL import Image, ImageDraw, ImageFont
from bs4 import BeautifulSoup
import requests
import numpy
import datetime

class GeoBGS(object):
    """Represents an area from which a BGS geology image tile may be formed or records of geologies shown in image maybe
    queried.

    Available functions:
    - get_map: Requests a map tile based on either a bounding box set by the user, or a bounding box generated from an
    OSGB coordinate and bounding box size.
    - get_map_geology: Analyses a user requested image or automatically calls the 'get_map' method, the image is then
    analysed for distinct colours and forms requests for geology relating to each colour.
    """

    def __init__(self, bounding_box={'minx': None, 'maxx': None, 'miny': None, 'maxy': None}, coordinate=None, buffer=None):
        """Initialisation of the GeoBGS class.

        Attributes, if not set on initialisation, maybe set through 'class.attribute = x'

        Args:
            bounding_box: Optional Dictionary containing keys [minx, maxx, miny, maxy]
            coordinate: Optional String of an osgb easting, northing coordinate e.g. '400000,410000'
            buffer: Optional Integer of a buffer in metres to form a bounding box around a given coordinate. Value must
            not exceed 1000.
        """
        self.coordinate = coordinate # make something to validate this
        self.buffer = buffer
        self.bounding_box = bounding_box
        self._baseurl = 'https://map.bgs.ac.uk/arcgis/services/BGS_Detailed_Geology/MapServer/WmsServer'
        self._map_params = {'SERVICE': 'WMS',
                                'VERSION': '1.3.0',
                                'REQUEST': '',
                                'BBOX': '',
                                'CRS': 'EPSG:27700',
                                'WIDTH': '',
                                'HEIGHT': '',
                                'LAYERS': '',
                                'STYLES': '',
                                'FORMAT': ''}
        self._query_params = {'QUERY_LAYERS': '',
                              'INFO_FORMAT': 'text/html',
                              'I': '',
                              'J': '',
                              'FEATURE_COUNT': ''}
        self._layers = ['BGS.50k.Linear.features',
                        'BGS.50k.Mass.movement',
                        'BGS.50k.Artificial.ground',
                        'BGS.50k.Superficial.deposits',
                        'BGS.50k.Bedrock']
        self._image_formats = ['image/png8', 'image/png24', 'image/png32', 'image/jpeg', 'image/bmp', 'image/gif',
                               'image/svg']

    def _bbmin(self, coord, buffer):
        """Simple function to return min value for bounding box.

        Args:
            coord: Integer or Float representing a single easting or northing coordinate
            buffer: Integer of a buffer in metres
        """
        return coord - buffer

    def _bbmax(self, coord, buffer):
        """Simple function to return max value for bounding box.

        Args:
            coord: Integer or Float representing a single easting or northing coordinate
            buffer: Integer of a buffer in metres
        """
        return coord + buffer

    def _create_bbox(self):
        """Help text
        """
        if self.coordinate is None:
            return 'Coordinate value not set - Cannot generate bounding box'
        elif self.buffer is None:
            return 'Buffer value not set - Cannot generate bounding box'
        else:
            x, y = [int(i.strip()) for i in self.coordinate.split(',')]
            return {'minx': self._bbmin(x, self.buffer),
                    'maxx': self._bbmax(x, self.buffer),
                    'miny': self._bbmin(y, self.buffer),
                    'maxy': self._bbmax(y, self.buffer)}

    def _memoize(self, f):
        """Help text"""
        memo = {}

        def helper(x):
            if x not in memo:
                memo[x] = f(x)
            return memo[x]

        return helper

    def _merge_dict(self, dictA, dictB):
        """python 3.5 only"""
        dict = {**dictA, **dictB}
        return dict

    def _string_bbox(self):
        """some text"""
        bbox = [self.bounding_box['minx'],
                self.bounding_box['miny'],
                self.bounding_box['maxx'],
                self.bounding_box['maxy']]
        return ','.join(map(str, bbox))

    def _make_request(self, params):
        """bit of text"""
        return requests.get(self._baseurl, params=params)

    def _prepare_request(self):
        """text"""
        self._map_params['BBOX'] = self._string_bbox()
        self._map_params['REQUEST'] = 'GetMap'
        self._map_params['WIDTH'] = 1000  # self.bounding_box['maxx'] - self.bounding_box['minx'] #TODO Add some sort of scaling if over 1000
        self._map_params['HEIGHT'] = 1000  # self.bounding_box['maxy'] - self.bounding_box['miny']
        self._map_params['LAYERS'] = ','.join([self._layers[4]])  # multiple layers can be used
        self._map_params['FORMAT'] = self._image_formats[1]
        return

    def _save_image(self, data):
        """TODO put the watermarking into its own method and call at end of get_map_geology
        as its been messing with geology id in some cases.
        """
        with Image.open(BytesIO(data.content), mode='r') as img:
            img = img.convert(mode='RGBA')
            width, height = img.size
            padding = 16
            font_size = 16
            rect_height = height - (font_size + padding * 2)
            rect_coords = (0, rect_height, width, height)
            message = 'Contains British Geological Survey materials © NERC {}'.format(datetime.date.today().year)

            # make a blank image for the text
            background = Image.new('RGBA', img.size, (180, 180, 180, 0))

            # get a font
            fnt = ImageFont.truetype('arialbd.ttf', font_size)
            # get a drawing context
            draw = ImageDraw.Draw(background)
            draw.rectangle(rect_coords, fill=(180, 180, 180, 100), outline=None)
            draw.text((padding, height - padding - font_size), message, font=fnt, fill=(0, 0, 0, 255))

            out = Image.alpha_composite(img, background)

            out.show()

            out.save(fp='geology.{}'.format('png'))
        return

    def _make_query(self):
        """text"""
        self._query_params['REQUEST'] = 'GetFeatureInfo'
        self._query_params['QUERY_LAYERS'] = ','.join([self._layers[4]])  # multiple layers can be used
        self._query_params['I'] = 1000  #TODO link to _map_params width setting and buffer
        self._query_params['J'] = 1000
        self._query_params['FEATURE_COUNT'] = 10  # Sensible default
        self._prepare_request()
        return

    def _extract_data(self):
        """right now not coping with multiple geologies in same data table"""
        if self.bounding_box is None:
            self._create_bbox()
        else:
            #self._make_query() call this from get_map_geology and overwrite I & J
            merged = self._merge_dict(self._map_params, self._query_params)
            r = self._make_request(merged)

            soup = BeautifulSoup(r.text, 'lxml')
            table_data = [[cell.text for cell in row(['th', 'td'])] for row in soup("tr")]
            geology_table = {}

            for i, header in enumerate(soup('th')):
                geology_table[header.text] = soup('td')[i].text

            return geology_table['LEX_RCS_D']

    def _open_image(self):
        """text"""
        img = Image.open('geology.{}'.format('png')).convert(mode='RGB')
        return img

    def _get_colours(self):
        """text"""
        img = self._open_image()
        data = numpy.asarray(img, dtype="int32")
        colours = [x[1] for x in img.getcolors() if (181, 181, 181) != x[1]]  # 181 ignore is for borderlines
        return data, colours

    def get_map(self):
        """"Make something up"""
        if self.bounding_box is None:
            self._create_bbox()
        else:
            self._prepare_request()
            r = self._make_request(self._map_params)
            self._save_image(r)

    def get_map_geology(self):
        """text"""
        try:
            data, colours = self._get_colours()
        except:
            self.get_map()

        geology_types = {}

        for colour in colours:

            indices = numpy.where(numpy.all(data == colour, axis=-1))

            self._make_query()
            self._query_params['I'] = indices[0][0]
            self._query_params['J'] = indices[1][0]

            try:
                geology_types[colour] = self._extract_data()
            except KeyError:
                print('Colour {} is invalid'.format(colour))
                pass

        print(geology_types)

    def __str__(self):
        """Help text"""
        return '<GeoBGS Object: Bounding Box = {}>'.format(self.bounding_box)