#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# -----------------------------------------------------------------------------
# THIS FILE IS PART OF THE CYLC WORKFLOW ENGINE.
# Copyright (C) NIWA & British Crown (Met Office) & Contributors.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
# -----------------------------------------------------------------------------
# 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:
    DOMAIN: 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 h5py
from pathlib import Path
import urllib
import requests   # noqa: F401 - required implicitly by urllib.

from mercator import get_scale, pos_to_coord
import util

S3URL = (
    'https://met-office-radar-obs-data.s3-eu-west-2.amazonaws.com/radar/'
    '{Y}/{m}/{d}/{YYYYmmddHHMM}_ODIM_ng_radar_rainrate_composite_1km_UK.h5'
)
DEBUG=os.environ['CYLC_DEBUG']=='true'
CYLC_TASK_LOG_DIR=os.environ['CYLC_TASK_LOG_DIR']


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.

    """
    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 (tuple): The value of the reading.

        """
        itt_x, itt_y = util.get_grid_coordinates(lng, lat, self.domain,
                                                 self.resolution)
                    
        self.data[itt_y][itt_x].append(self.value_map(value))

    @staticmethod
    def value_map(v):
        """Convert rainfall rate values into colour space values.

        Checks if rainfall value above each threshold in turn.

        TODO:
            - Unit test this
        """
        thresholds = {32, 16, 8, 4, 2, 1, .5, .2}
        for i, threshold in enumerate(sorted(thresholds, reverse=True)):
            if v > threshold:
                return 8 - i
        return 0

    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_archived_radar_image(filename, time):
    """Retrieve a png image from the archived data in the workflow directory.

    Args:
        filename (str): The path to write the image file to.
        time (str): The datetime of the image to retrieve in ISO8601 format.

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


def get_amazon_radar_data(filename, time):
    time = datetime.strptime(time, '%Y%m%dT%H%MZ')
    url = S3URL.format(
        Y=time.strftime('%Y'),
        m=time.strftime('%m'),
        d=time.strftime('%d'),
        YYYYmmddHHMM=time.strftime('%Y%m%d%H%M')
    )
    print(f'[INFO] Getting data from {url=}')
    data = urllib.request.urlopen(url).read()
    with open(filename, 'wb') as fh:
        fh.write(data)


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

    Args:
        filename (str): Path to the png image to process.
        resolution (float): The length/weight 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.

    """
    print(f'[INFO] Analysing data from {filename}')
    data = h5py.File(filename)['dataset1']['data1']['data']
    rainfall = Rainfall(domain, resolution)

    _height, width = data.shape

    scale = get_scale(domain, width)
    offset = (-1100.8461538461539, 1400.6953225710452)

    if DEBUG:
        print(f'[INFO] {scale=}, {offset=}')
        from matplotlib import pyplot as plt
        Path(CYLC_TASK_LOG_DIR).mkdir(parents=True, exist_ok=True)
        plt.imshow(data)
        plt.savefig(f'{CYLC_TASK_LOG_DIR}/raw.data.png')

    for itt_y, row in enumerate(data):
        for itt_x, col in enumerate(row):
            lng, lat = pos_to_coord(
                itt_x,
                itt_y * (2. / 3.),  # Counter aspect ratio.
                offset,
                scale
            )
            val = float(col)
            # Original data uses -1 to indicate radar mask
            val = 0 if val == -1 else val
            rainfall.add(lng, lat, val)
    data = rainfall.compute_bins()

    if DEBUG:
        plt.imshow(data)
        plt.savefig(f'{CYLC_TASK_LOG_DIR}/processed.data.png')

    return data

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

    # Acts as a switch - if a file-name is given, use that file-name
    canned_data = os.environ.get('CANNED_DATA', 'fetch')

    if canned_data == 'fetch':
        print('[INFO] Attempting to get rainfall data from S3 bucket')
        get_amazon_radar_data('radardata.h5', time)
        canned_data = 'radardata.h5'
    else:
        raise Exception('Not currently provided functionality')
        # Not currently used.
        #print(f'[INFO] Canned data provided: {canned_data}')
        #get_archived_radar_data(canned_data, time)

    data = process_rainfall_data(canned_data, resolution, domain)

    util.write_csv('rainfall.csv', data)


if __name__ == '__main__':
    util.sleep(2)  # make the tutorial run a little slower
    main()
