#!/usr/bin/env python2
# -*- coding: utf-8 -*-
# -----------------------------------------------------------------------------
# Copyright (C) 2012-2019 British Crown (Met Office) & Contributors - GNU V3+.
# This is illustrative code developed for tutorial purposes, it is not
# intended for scientific use and is not guarantied to be accurate or correct.
# -----------------------------------------------------------------------------
"""
Usage:
    get-rainfall

Environment Variables:
    API_KEY: The DataPoint API key, required for getting live weather data.
        If un-specified then get-observations will fall back to archive data
        from the suite directory.
    DOMIAN: The area in which to generate forecasts for in the format
        (lng1, lat1, lng2, lat2).
    RESOLUTION: The length/width of each grid cell in degrees.

"""

from datetime import datetime
import math
import os
import shutil
import urllib2

from mecator import get_offset, get_scale, coord_to_pos, pos_to_coord
from png import Reader
import util


URL = ('http://datapoint.metoffice.gov.uk/public/data/layer/wxobs/'
       'RADAR_UK_Composite_Highres/png?TIME={time}&key={api_key}')


class Rainfall(object):
    """Class for holding rainfall data.

    Args:
        domain (dict): Domain as returned by util.parse_domain()
        resolution (float): The length of each grid cell in degrees.

    """
    VALUE_MAP = {
        (0, 0, 254, 255): 1,
        (50, 101, 254, 255): 2,
        (127, 127, 0, 255): 3,
        (254, 203, 0, 255): 4,
        (254, 152, 0, 255): 5,
        (254, 0, 0, 255): 6,
        (254, 0, 254, 255): 7
    }

    def __init__(self, domain, resolution):
        self.resolution = resolution
        self.domain = domain

        rows = int(math.ceil(abs(domain['lat1'] - domain['lat2']) / resolution))
        cols = int(math.ceil(abs(domain['lng1'] - domain['lng2']) / resolution))

        self.data = []
        for itt_y in range(rows):
            self.data.append([])
            for _ in range(cols):
                self.data[itt_y].append([])

    def add(self, lng, lat, value):
        """Add a data point to this data set.

        Args:
            lng (float): The longitude for this reading.
            lat (float): The latitude fo this reading.
            value (float): The value of the reading.

        """
        itt_x, itt_y = util.get_grid_coordinates(lng, lat, self.domain,
                                                 self.resolution)
        try:
            self.data[itt_y][itt_x].append(self.VALUE_MAP[value])
        except KeyError:
            pass

    def compute_bins(self):
        """Return this dataset as a 2D matrix."""
        for row in self.data:
            for itt, col in enumerate(row):
                if col:
                    row[itt] = sum(col) / float(len(col))
                else:
                    row[itt] = 0
        return self.data


def get_datapoint_radar_image(filename, time, api_key):
    """Retrieve a png image of rainfall from the DataPoint service.

    Args:
        filename (str): The path to write the image file to.
        time (str): The date-time of the image to reteieve in ISO8601 format.
        api_key (str): Datapoint API key.

    """
    time = datetime.strptime(time, '%Y%m%dT%H%MZ'
                            ).strftime('%Y-%m-%dT%H:%M:%SZ')
    url = URL.format(time=time, api_key=api_key)
    print 'Opening URL: %s' % url
    with open(filename, 'w+') as png_file:
        png_file.write(urllib2.urlopen(url).read())


def get_archived_radar_image(filename, time):
    """Retrieve a png image from the archived data in the suite directory.

    Args:
        filename (str): The path to write the image file to.
        time (str): The date-time of the image to reteieve in ISO8601 format.

    """
    shutil.copyfile(
        os.path.join(os.environ['CYLC_SUITE_DEF_PATH'], 'data', time, filename),
        filename)


def process_rainfall_data(filename, resolution, domain):
    """Generate a 2D matric of data from the rainfall data in the image.

    Args:
        filename (str): Path to the png image to process.
        resolution (float): The length/wieght of each grid cell in degrees.
        domain (dict): The bounds of the domain as returned by
            util.parse_domain.

    Return:
        list - A 2D matrix of rainfall data.

    """
    reader = Reader(file(filename, 'r'))
    data = reader.read()

    scale = get_scale(domain, data[0])
    offset = get_offset(domain, scale)

    rainfall = Rainfall(domain, resolution)

    for itt_x in range(0, data[0]):
        for itt_y in range(0, data[1]):
            lng, lat = pos_to_coord(itt_x,
                                    itt_y * (2. / 3.),  # Counter aspect ratio.
                                    offset, scale)
            datum = data[2][itt_y][itt_x * 4:(itt_x * 4) + 4]
            rainfall.add(lng, lat, tuple(datum))
    return rainfall.compute_bins()


def main():
    time = os.environ['CYLC_TASK_CYCLE_POINT']
    resolution = float(os.environ['RESOLUTION'])
    domain = util.parse_domain(os.environ['DOMAIN'])
    api_key = os.environ.get('API_KEY')

    if api_key:
        print 'Attempting to get weather data from the DataPoint service.'
        get_datapoint_radar_image('rainfall-radar.png', time, api_key)
    else:
        print 'No API key provided, falling back to arhived data'
        get_archived_radar_image('rainfall-radar.png', time)

    data = process_rainfall_data('rainfall-radar.png', resolution, domain)
    util.write_csv('rainfall.csv', data)


if __name__ == '__main__':
    main()
