# -*- coding: utf-8 -*-
#
# Copyright 2010-2014 Will Barton. 
# All rights reserved.
# 
# Redistribution and use in source and binary forms, with or without 
# modification, are permitted provided that the following conditions
# are met:
# 
#   1. Redistributions of source code must retain the above copyright 
#      notice, this list of conditions and the following disclaimer.
#   2. Redistributions in binary form must reproduce the above copyright 
#      notice, this list of conditions and the following disclaimer in the 
#      documentation and/or other materials provided with the distribution.
#   3. The name of the author may not be used to endorse or promote products
#      derived from this software without specific prior written permission.
# 
# THIS SOFTWARE IS PROVIDED ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES,
# INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY
# AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL
# THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
# 

from mpl_toolkits.basemap import Basemap

try:
    import Image
except ImportError:
    from PIL import Image

from observation.conditions.orderedset import OrderedSet
from observation.conditions.utils import current_update

import datetime
from dateutil import tz
from dateutil.relativedelta import *


import urlparse
import os, os.path
import shutil

import logging
logger = logging.getLogger(__name__)

# The map projection that everything here is based on.
__map__ = Basemap(projection='npstere', resolution='l',
                  boundinglat=27.0, lon_0=-111)

TYPES = {
        'seeing':'seeing',
        'transp':'transparency',
        'uv':'wind',
        'tt':'temperature',
        'hr':'humidity',
        'nt':'clouds',
    }

SEEING_COLOR_TABLE = {
        0: (254, 254, 254), # None
        1: (198, 198, 198), # Bad
        2: (148, 212, 212), # Poor
        3: (98, 162, 226), # Average
        4: (43,107, 171), # Good
        5: (0, 62, 126), # Excellent
    }

TRANSPARENCY_COLOR_TABLE = {
        0: (254, 254, 254), # None
        1: (198, 198, 198), # Bad
        2: (148, 212, 212), # Poor
        3: (98, 162, 226), # Average
        4: (43,107, 171), # Good
        5: (0, 62, 126), # Excellent
    }

WIND_COLOR_TABLE = {
        5: (254, 254, 254), # Very Strong
        4: (198, 198, 198), # Strong
        3: (148, 212, 212), # Moderate 
        2: (98, 162, 226), # Light to Moderate 
        1: (43,107, 171), # Light
        0: (0, 62, 126), # Calm
    }

TEMPERATURE_COLOR_TABLE = {
        50: (198, 198, 198),
        45: (125, 0, 0),
        40: (168, 0, 0),
        35: (225, 0, 0),
        30: (253, 29, 0),
        25: (253, 89, 0),
        20: (253, 157, 0),
        15: (253, 221, 0),
        10: (161, 253, 89),
        5: (93, 253, 157),
        0: (250, 250, 250),
        -5: (29, 253, 221),
        -10: (0, 221, 253),
        -15: (0, 136, 253),
        -20: (0, 51, 253),
        -25: (0, 0, 235),
        -30: (0, 0, 178),
        -35: (0, 0, 132),
        -40: (252, 0, 252),
    }

HUMIDITY_COLOR_TABLE = {
        100: (225, 0, 0),
        95: (183, 0, 0),
        90: (234, 0, 0),
        85: (254, 52, 1),
        80: (252, 134, 2),
        75: (254, 198, 0),
        70: (234, 251, 22),
        65: (148, 254, 106),
        60: (85, 250, 173),
        55: (8, 254, 237),
        50: (128, 192, 192),
        45: (113, 177, 241),
        40: (78, 142, 206),
        35: (48, 112, 176),
        30: (13, 77, 141),
        25: (8, 3, 93),
    }

CLOUDS_COLOR_TABLE = {
        100.0: (255, 255, 255),
        97.5: (248, 248, 248),
        95.0: (238, 238, 238),
        92.5: (233, 233, 233),
        90.0: (223, 223, 223),
        87.5: (218, 218, 218),
        85.0: (203, 203, 203),
        82.5: (193, 193, 193),
        80.0: (188, 252, 252),
        77.5: (183, 247, 247),
        75.0: (178, 242, 242),
        72.5: (173, 237, 237),
        70.0: (168, 232, 232),
        67.5: (163, 227, 227),
        65.0: (158, 222, 222),
        62.5: (153, 217, 217),
        60.0: (148, 212, 212),
        57.5: (128, 192, 192),
        55.0: (123, 187, 251),
        52.5: (118, 182, 246),
        50.0: (113, 177, 241),
        47.5: (108, 172, 236),
        45.0: (103, 167, 231),
        42.5: (98, 162, 226),
        40.0: (93, 157, 221), 
        37.5: (88, 152, 216),
        35.0: (83, 147, 211),
        32.5: (78, 142, 206),
        30.0: (68, 132, 196),
        27.5: (48, 112, 176),
        25.0: (43, 107, 171),
        22.5: (38, 102, 166),
        20.0: (33, 97, 161),
        17.5: (28, 92, 156),
        15.0: (23, 87, 151),
        12.5: (18, 82, 146),
        10.0: (13, 77, 141),
        7.5: (8, 72, 136),
        5.0: (3, 67, 131),
        2.5: (0, 62, 126), 
    }

class point(object):
    def __init__(self, x, y):
        self.x = x
        self.y = y

    def __repr__(self):
        return "%s, %s" % (str(self.x), str(self.y))


class bounds(object):
    def __init__(self, x, y, width, height):
        self.x = x
        self.y = y
        self.width = width
        self.height = height

    def __contains__(self, point):
        if point.x > self.x and point.y > self.y and \
                point.x < self.x + self.width and \
                point.y < self.y + self.height:
            return True
        return False

    def __repr__(self):
        return str(self.y, self.y, self.width, self.height)


class Forecast(object):
    """
    A forecast is the value of a specific longitude and latitude at a
    specific date/time as extracted from a ForecastMapSeries.
    """
    latitude = None
    longitude = None
    date = None

    def __init__(self, latitude, longitude, date, **values):
        self.latitude = latitude
        self.longitude = longitude
        self.date = date

        self.__values = values

        for v in values:
            # t = TYPES[v]
            setattr(self, v, values[v])

    def __repr__(self):
        return u'<Forecast at %f, %f: %s %s>' % (self.latitude,
                self.longitude, self.date, str(self.__values))


class ForecastMap(object):
    """
    A weather map is a single map tied to an image. It allows for the
    lookup of pixel values on that image based on GPS coordinates.
    """

    pixel_dimensions = point(1436, 1436)

    def __init__(self, image, region = None):
        global __map__

        self.region = region

        self.image = image

        self.map = __map__

        self.pixel_dimensions = None
        self.image_bounds = None
        self.legend_bounds = []

        # Based on the region, we map the particular projection used by
        # the Canadian Weather Office to a basemap projection.

        if not self.region:
            # Resolution of map at 300 dpi
            self.pixel_dimensions = point(1436, 1436)

            # Upper left corner, bottom right corner, relative to
            # pixel_dimensions
            self.image_bounds = bounds(479, 835, 718, 600)
            
            # Bounds of any legends within the image_bounds
            self.legend_bounds = [bounds(4, 4, 395, 42),]
            
        elif region == 'northeast':
            # Eastern CA
            # Resolution at 510 dpi  
            self.pixel_dimensions = point(2441, 2441)

            # Upper left corner, bottom right corner, relative to
            # pixel_dimensions
            self.image_bounds = bounds(1299, 1498, 718, 600)
            
            # Bounds of any legends within the image_bounds
            self.legend_bounds = [bounds(4, 4, 395, 42),
                    bounds(660, 196, 56, 401)]

        elif region == 'northwest':
            # Eastern CA
            # Resolution at 530 dpi  
            self.pixel_dimensions = point(2537, 2537)

            # Upper left corner, bottom right corner, relative to
            # pixel_dimensions
            self.image_bounds = bounds(840, 1548, 718, 600)
            
            # Bounds of any legends within the image_bounds
            self.legend_bounds = [bounds(4, 4, 395, 42),
                    bounds(660, 196, 56, 401)]

        elif region == 'southeast':
            # Eastern US
            # Resolution at 555 dpi  
            self.pixel_dimensions = point(2656, 2656)

            # Upper left corner, bottom right corner, relative to
            # pixel_dimensions
            self.image_bounds = bounds(1430, 2023, 718, 600)
            
            # Bounds of any legends within the image_bounds
            self.legend_bounds = [bounds(4, 4, 395, 42),
                    bounds(660, 196, 56, 401)]

        elif region == 'southwest':
            # Western US
            # Resolution at 530 dpi  
            self.pixel_dimensions = point(2999, 3132)

            # Upper left corner, bottom right corner, relative to
            # pixel_dimensions
            self.image_bounds = bounds(1149, 2412, 718, 600)
            
            # Bounds of any legends within the image_bounds
            self.legend_bounds = [bounds(4, 4, 395, 42),
                    bounds(660, 196, 56, 401)]
        else:
            raise ValueError("Bad value for region")

        # The map's dimensions in meters
        self.dimensions = point(self.map.urcrnrx - self.map.llcrnrx, 
                           self.map.urcrnry - self.map.llcrnry)

        # The ratio of map meters to pixels
        self.dimensions_ratio = point(self.pixel_dimensions.x/self.dimensions.x, 
                                      self.pixel_dimensions.y/self.dimensions.y)
    
    def coordinatePixel(self, latitude, longitude):
        """
        Get the pixel coordinate, given our class variable pixel
        dimensions, of the given GPS coordinate.
        """

        # X, Y, in meters
        x_m, y_m = self.map(longitude, latitude)

        # Origin is lower left. Convert it to upper left.
        y_m_c = self.dimensions.y - y_m

        # X, Y, in pixels
        # NOTE: These are the coordinates on the generated map
        # projection, NOT on the given weather map image.
        x_p, y_p = (self.dimensions_ratio.x * x_m, 
                    self.dimensions_ratio.y * y_m_c)

        # Make sure the point is within the given weather map image
        # bounds, and is not within any legend bounds in the image.
        if point(x_p, y_p) not in self.image_bounds:
            raise ValueError("Location outside coverage area")
        else:
            for legend in self.legend_bounds:
                if point(x_p, y_p) in legend:
                    raise ValueError("Location outside coverage area")

        # Convert these pixel coordinates into weather map image
        # coordinates
        x = x_p - self.image_bounds.x
        y = y_p - self.image_bounds.y

        return point(x, y)

    def coordinateValue(self, latitude, longitude):
        """
        Get the color value of the pixel on the map's image at
        the given coordinate.
        """
        
        # Get the x, y coordinates for the image.
        p = self.coordinatePixel(latitude, longitude)
        
        # Use an intermediate RGB image?
        i = Image.open(self.image)
        try: 
            rgbi = i.convert('RGB')
            pixels = rgbi.load()
            color = pixels[int(p.x), int(p.y)]
            del i, rgbi
        except Exception, e: # XXX: Catch specific exception
            return None

        return color


class ForecastMapSet(object):
    """
    A set of maps that view different regions but share the same
    forecast information.
    """

    def __init__(self, images, regions):
        self.images = images
        self.regions = regions
        
        self.mapset = OrderedSet()

        for region in regions:
            image = self.images[region]
            map = ForecastMap(image, region=region)
            self.mapset.add(map)
    
    def coordinateValue(self, latitude, longitude):
        """
        Get the color value of the pixel on the map's image at
        the given coordinate.
        """
        # Find the first map in our map set to contain this value.

        value = None
        for map in self.mapset:
            try:
                value = map.coordinateValue(latitude, longitude)
            except ValueError:
                continue

        if value is None:
            raise ValueError("Coordinates not in bounds")

        return value


def download(url, filepath):
    import urllib2

    request = urllib2.urlopen(urllib2.Request(url))
    try:
        with open(filepath, 'wb') as f:
            shutil.copyfileobj(request, f)
    finally:
        request.close()


class ForecastMapSeries(object):
    """
    A Forecast Map Series is a collection of forecast maps containing 
    the same information for different time periods. 
    """

    base_url = 'http://www.weatheroffice.gc.ca/data/prog/regional/'

    def __init__(self, period, color_table, type,
            image_path, regions = None):
        """
        Create a forecast of the specified period (the time between
        images in the imageset) in hours from the specified date, and 
        optionally for the given regions. 
        """
        self.period = period
        self.regions = regions
        self.color_table = color_table
        self.type = type
        
        self.image_path = image_path
        self.__cachedpath__ = None

        # The refresh function will set these.
        self.mapset = OrderedSet()
        self.date = None

        # Get the image set. Unless we're specifically asked to, load
        # the most recent set.
        self.load()


    def __repr__(self):
        return """<ForecastMapSeries %s, %d maps>""" % \
                (self.type, len(str(self.mapset)))

    @property
    def range(self):
        """
        Provide the ending date and time covered by this map series.
        """
        hour_range = len(self.mapset) * self.period
        start_dt = self.date + relativedelta(hours=+3)
        end_dt = self.date + relativedelta(hours=+hour_range)
        end_dt = end_dt.replace(tzinfo=tz.tzlocal()).astimezone(tz.tzutc())
        return (start_dt, end_dt)

    def forTime(self, time):
        """ 
        Get the set of forecast maps for the specified time. Time 
        should be an absolute datetime object. 
        """

        # If there's no tzinfo, assume it's localtime.
        if time.tzinfo == None:
            time = time.replace(tzinfo=tz.tzlocal()).astimezone(tz.tzutc())

        start, end = self.range
        if (time < start) or (time > end):
            raise ValueError("Time is outside of forecast range", (time,
                start, end))

        delta = relativedelta(time, self.date).hours
        index = delta / self.period
        return list(self.mapset)[index]

    def coordinateValueAtTime(self, lat, lon, time, nudged=False):
        """
        Get the color value of the pixel on the map's image at
        the given coordinate at the given datetime.
        """
        time = time.astimezone(tz.tzutc())

        if self.mapset is None or len(self.mapset) == 0:
            raise ValueError("Forecast data unavailable at this time")

        # Lookup the map for the given time
        map = self.forTime(time)

        # Get the coordinate value
        color = map.coordinateValue(lat, lon)

        try:
            value_index = self.color_table.values().index(color)
            value = self.color_table.keys()[value_index]
        except ValueError:
            # XXX: Border colors are black and red. If we're one of
            # those colors, we're on a border. Nudge the coordinates
            # ever so slightly and see what we get.
            if (color == (0, 0, 0) or color == (254, 0, 0)) and not nudged:
                return self.coordinateValueAtTime(
                        lat, lon-0.04, time, nudged=True)
            value = -1

        return value

    def load(self, date=None):
        """
        Discover what image file names we have, if any, and the
        appropriate dates for those files. 

        We want to try to update at 18:00 UTC and 06:00 UTC (the
        forecasts are updated at "approximately" 17:30 and 5:30).
        """

        if date is None and os.path.isdir(self.image_path):
            # If we're not given a date, look for the last directory in 
            # the image_path. Because we name the directory based on the
            # date and hour of the forecast, this should give us the 
            # most recent. 
            dirs = [d for d in os.listdir(self.image_path) if
                    os.path.isdir(os.path.join(self.image_path, d))]
            dirname = dirs[-1]
            date = datetime.datetime.strptime(dirname, '%Y%m%d%H').replace(
                    tzinfo=tz.tzutc())

        else:
            # Find the immediately previous update 
            date = datetime.datetime.now(tz=tz.tzlocal()).astimezone(
                    tz.tzutc()).replace(
                            minute=0, second=0, microsecond=0)
            if date.hour > 18:
                date = date + datetime.timedelta(hours=-(date.hour-18))
            elif date.hour < 6:
                date = date + datetime.timedelta(hours=(-6 - date.hour))
            elif date.hour > 6 and date.hour < 18:
                date = date + datetime.timedelta(hours=-(date.hour-6))
            else:
                date = date.replace(hour=6)

        # Get the imageset for the date.
        imageset = self.__imageset(date)
        self.date = date

        self.mapset = OrderedSet()
        if not self.regions:
            for image in imageset:
                map = ForecastMap(image)
                self.mapset.add(map)
        else:
            for images in list(imageset):
                map = ForecastMapSet(images, self.regions)
                self.mapset.add(map)

        return date

    def refresh(self):
        """ Refresh the map series from the server. """

        date = current_update()

        # The foldername for the update
        dirname = date.strftime('%Y%m%d%H')

        # The download path, create it if it doesn't exist.
        download_path = os.path.abspath(
                os.path.join(self.image_path, dirname))
        if not os.path.exists(download_path):
            os.makedirs(download_path)

        def download_func(image_url, image_path):
            # The download path, create it if it doesn't exist.
            if not os.path.exists(os.path.dirname(image_path)):
                os.makedirs(os.path.dirname(image_path))
                
            # If the image already exists, don't worry about it.
            if not os.path.exists(image_path):
                logger.debug("Downloading %s" % image_url)
                download(image_url, image_path)

        imageset = self.__imageset(date, download_func)
        self.date = date

        self.mapset = OrderedSet()
        if not self.regions:
            for image in imageset:
                map = ForecastMap(image)
                self.mapset.add(map)
        else:
            for images in list(imageset):
                map = ForecastMapSet(images, self.regions)
                self.mapset.add(map)

        return date

    def __imageset(self, date, download_func=None):
        """ Get the imageset for a given date. If download_func is
            provided, it is called as:

                download_func(image_url, image_path)

            and expected to download the image at the given URL to the
            given path. """

        dirname = date.strftime('%Y%m%d%H')

        images_path = os.path.abspath(
                os.path.join(self.image_path, dirname))

        # The image name
        if not self.regions:
            image_format = '%(dirname)s_054_R1_north@america@astro_I_ASTRO_%(type)s_%(period)s.png'
        else:
            image_format = '%(dirname)s_054_R1_north@america@%(region)s_I_ASTRO_%(type)s_%(period)s.png'

        # The URL path
        url = urlparse.urljoin(self.base_url, dirname + '/')

        # Each period within a 48 hour range has an appropriate image
        imageset = []

        type = TYPES.keys()[TYPES.values().index(self.type)]

        for p in range(3, 49, self.period):
            period = str(p).zfill(3)

            if not self.regions:
                image_name = image_format % {'dirname': dirname,
                                             'type': type,
                                             'period': period,}
                image_path = os.path.join(images_path, image_name)
                image_url = urlparse.urljoin(url, image_name)

                if callable(download_func):
                    download_func(image_url, image_path)

                # Add the image path to the set.
                imageset.append(image_path)
            else:
                rdict = {}
                for region in self.regions:
                    image_name = image_format % {'dirname': dirname,
                                                 'type': type,
                                                 'period': period,
                                                 'region': region,}
                    image_path = os.path.join(images_path, image_name)
                    image_url = urlparse.urljoin(url, image_name)

                    if callable(download_func):
                        download_func(image_url, image_path)

                    rdict[region] = image_path

                imageset.append(rdict)

        return imageset
       
