#!/usr/bin/env python

# vim: set et ts=4 sw=4 fileencoding=utf-8 :

u"""

Éleconso : sources et capacité de production électrique de la France

--------------------------------------------------------------------

*Ou: dois-je lancer mon sèche-linge à 19h...*

Génère un graphe de la consommation électrique (au quart d'heure près) en

surimpression de la capacité de production électrique française, de façon à

connaître d'un coup d'oeil la nature de l'électricité consommée :

* peu émettrice de CO2

* très émettrice de CO2

* très émettrice de CO2 et importée depuis nos voisins

Les données sources sont publiées sur le site de RTE, le réseau de transport

d'électricité français :

* http://clients.rte-france.com/lang/fr/visiteurs/vie/courbes.jsp

* http://clients.rte-france.com/lang/fr/visiteurs/vie/prod/realisation_production.jsp

.. :Authors:

       Aurélien Bompard <aurelien@bompard.org> <http://aurelien.bompard.org>

.. :License:

       GNU GPL v3 or later

"""

import os

import re

import datetime

import urllib2

import optparse

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt

from matplotlib.dates import  HourLocator, DateFormatter

from matplotlib.font_manager import FontProperties

import numpy

class Production(object):

    """

    La production électrique française, par source. Un point toutes les heures

    mais les données peuvent dater d'un jour ou deux.

    """

    def __init__(self):

        self.page_url = "http://clients.rte-france.com/lang/fr/visiteurs/vie/prod/realisation_production.jsp"

        self.data_line = re.compile('<param\s+name="data"\s*value="([^"]+)"\s*/>')

        self.data = {}

        self.date = None

        self.considered_clean = [u"Nucléaire", u"Hydraulique"]

    def read(self, debug=False):

        if debug:

            page = open(os.path.basename(self.page_url))

        else:

            page = urllib2.urlopen(self.page_url)

        for line in page:

            if '<param name="data" ' not in line:

                continue

            line_match = self.data_line.match(line.strip())

            data = line_match.group(1)

            break

        for source in data.split("|"):

            if not source:

                continue

            source_data = source.split(";")

            source_name = source_data.pop(0).decode("iso-8859-1")

            self.data[source_name] = {}

            self.date = source_data.pop(0)

            self.date = datetime.date(int(self.date[6:10]),

                                      int(self.date[3:5]),

                                      int(self.date[0:2]))

            for hour, prod in enumerate(source_data):

                prod_time = datetime.time(hour, 0, 0)

                if prod.endswith(".0"):

                    prod = prod[:-2]

                try:

                    prod = int(prod)

                except ValueError:

                    prod = None

                self.data[source_name][prod_time] = prod

        page.close()

    def get_averages(self):

        avg = {}

        for source, data in self.data.iteritems():

            prod_values = [v for v in data.values() if v is not None]

            if not prod_values:

                avg[source] = None

            else:

                avg[source] = sum(prod_values) / len(prod_values)

        return avg

    def get_limit(self, sourcetype, time=None):

        if sourcetype == "clean":

            return self._get_limit(self.considered_clean, time)

        elif sourcetype == "cheap":

            return self._get_limit(["Total"], time)

        raise KeyError()

    def _get_limit(self, sources, time=None):

        if time is None:

            values = self.get_averages()

        else:

            dt = datetime.time(time.hour, 0, 0)

            values = {}

            for source in sources:

                values[source] = self.data[source][dt]

        return sum([values[source] for source in sources])

    def get_aggregate(self, sourcetype):

        if sourcetype == "clean":

            return self._get_aggregate(self.considered_clean)

        elif sourcetype == "cheap":

            return self._get_aggregate(["Total"])

        raise KeyError()

    def _get_aggregate(self, sources):

        data = {}

        for source in sources:

            for date, value in self.data[source].iteritems():

                if date not in data:

                    data[date] = 0

                data[date] += value

        return data

class Consommation(object):

    """La consommation électrique française, au quart d'heure près"""

    def __init__(self):

        self.page_url = "http://clients.rte-france.com/lang/fr/visiteurs/vie/courbes.jsp"

        self.page_params = {"Chaine2": "Consommation",

                            "Chaine4": "Prev J-1",

                            "Chaine6": "Prev J",

                           }

        self.data_line_re = '<param\s+name\s*=\s*"%s"\s*value\s*=\s*"([^"]+);"\s*/>'

        self.date_line_re = re.compile('<param\s+name\s*=\s*"Jour"\s*value\s*=\s*"(\d\d)/(\d\d)/(\d\d\d\d)"\s*/>')

        self.data = {}

        self.date = None

    def read(self, debug=False):

        if debug:

            page = open(os.path.basename(self.page_url))

        else:

            page = urllib2.urlopen(self.page_url)

        data = {}

        for line in page:

            if '<param ' not in line:

                continue # légère optimisation inutile

            for param_page, param_name in self.page_params.iteritems():

                line_match = re.match(self.data_line_re % param_page,

                                      line.strip())

                if line_match is None:

                    continue

                data[param_name] = line_match.group(1)

            line_match = self.date_line_re.match(line.strip())

            if line_match is not None:

                self.date = datetime.date(int(line_match.group(3)),

                                          int(line_match.group(2)),

                                          int(line_match.group(1)))

        for data_name, data_values in data.iteritems():

            if data_values:

                self.data[data_name] = {}

            for index, value in enumerate(data_values.split(",")):

                data_time = index * 15 # une valeur tous les quarts d'heure

                data_time = datetime.time(data_time / 60, data_time % 60, 0)

                try:

                    conso_value = int(value)

                except ValueError:

                    conso_value = None

                self.data[data_name][data_time] = conso_value

        page.close()

    def get_last_time(self):

        now = datetime.datetime.now()

        now = datetime.time(hour=now.hour, minute=now.minute/15*15) # arrondir à 15 min

        for i in range(10): # 4 essais pour trouver une valeur

            if now in self.data["Consommation"]:

                return now

            now = datetime.datetime.combine(self.date, now) - datetime.timedelta(minutes=15)

            if now.date() < self.date:

                break

            now = now.time()

        return None

class ConsoGraph(object):

    """

    Le graphe de la consommation sur la production

    """

    def __init__(self, prod, conso):

        self.prod = prod

        self.conso = conso

        self.date = self.conso.date

        self.title = "Consommation le %s" % self.date.strftime("%d/%m/%Y")

        self.size = (6.4, 4.8)

        self.fontsize = 10

        self.status_message = ["OK", u"électricité polluante", u"électricité polluante et importée"]

        self.status_color = ["green", "orange", "red"]

    def get_status(self):

        clean_limit = self.prod.get_limit("clean", datetime.datetime.now())

        cheap_limit = self.prod.get_limit("cheap", datetime.datetime.now())

        last_time = self.conso.get_last_time()

        if last_time is None:

            return None

        cur_conso = self.conso.data["Consommation"][last_time]

        status = 0

        if cur_conso > clean_limit:

            status += 1

        if cur_conso > cheap_limit:

            status += 1

        #print "(conso:%s prod:%s)" % (cur_conso, cheap_limit)

        return status

    def get_series(self):

        x = []

        series = {"clean": [], "cheap": [], "conso": [], "prev": []}

        clean_values = self.prod.get_aggregate("clean")

        cheap_values = self.prod.get_aggregate("cheap")

        for time in sorted(self.conso.data["Prev J-1"]):

            x.append(datetime.datetime.combine(self.conso.date, time))

            # clean

            if time in clean_values:

                series["clean"].append(clean_values[time])

            elif time.replace(minute=0) in clean_values:

                series["clean"].append(clean_values[time.replace(minute=0)])

            # cheap

            if time in clean_values:

                series["cheap"].append(cheap_values[time])

            elif time.replace(minute=0) in cheap_values:

                series["cheap"].append(cheap_values[time.replace(minute=0)])

            # conso

            if time in self.conso.data["Consommation"]:

                series["conso"].append(self.conso.data["Consommation"][time])

            series["prev"].append(self.conso.data["Prev J-1"][time])

        series["clean"] = numpy.array(series["clean"])

        series["cheap"] = numpy.array(series["cheap"])

        series["conso"] = numpy.array(series["conso"])

        series["prev"] = numpy.array(series["prev"])

        return x, series

    def graph(self, filename=None):

        x, series = self.get_series()

        matplotlib.rc('font', size=self.fontsize)

        fig = plt.figure(figsize=self.size)

        ax = fig.add_subplot(111)

        ax.xaxis.set_minor_locator(HourLocator())

        ax.xaxis.set_major_formatter(DateFormatter("%H:%M"))

        plt.suptitle(self.title, fontsize="large", weight="bold")

        plt.ylabel('MW', weight="bold")

        ax.fill_between(x, 0, series["clean"], facecolor="g", axes=ax, label="Production sans CO2")

        ax.fill_between(x, series["clean"], series["cheap"], facecolor="orange", axes=ax, label="Production avec CO2")

        ax.plot(x[:len(series["conso"])], series["conso"], axes=ax, lw=2, color="blue", label="Consommation")

        ax.fill_between(x[:len(series["conso"])],

                        series["cheap"][:len(series["conso"])],

                        series["conso"],

                        where=series["cheap"][:len(series["conso"])]<series["conso"],

                        facecolor="red", interpolate=True)

        ax.fill_between(x[:len(series["conso"])],

                        series["conso"],

                        series["cheap"][:len(series["conso"])],

                        where=series["conso"]<series["cheap"][:len(series["conso"])],

                        facecolor="white", alpha=0.6, interpolate=True)

        ax.plot(x, series["prev"], axes=ax, lw=1, color="blue", linestyle="dashed", label=u"Prévision")

        #plt.axis([0, None, 0, None])

        ax.xaxis_date()

        #ax.autoscale_view(scaley=False)

        plt.axis([None, None, 55000, None])

        plt.setp(plt.gca().get_xticklabels(), rotation=45, horizontalalignment='right')

        plt.legend(loc="lower left", fancybox=True, prop=FontProperties(size=11), handlelength=4)

        plt.grid(True)

        status = self.get_status()

        if status is not None:

            props = dict(boxstyle='round', facecolor=self.status_color[status], alpha=0.5)

            last_time = self.conso.get_last_time()

            ax.text(0.98, 0.97, u"À %s : %s" %

                    (last_time.strftime("%Hh%M"), self.status_message[status]),

                    transform=ax.transAxes, fontsize="large", ha="right",

                    verticalalignment='top', bbox=props)

            # On laisse de la place pour la boîte

            old_ylim = ax.get_ylim()

            ax.set_ylim(top=old_ylim[1]+5000)

        if filename is None:

            plt.show()

        else:

            fig.savefig(filename, dpi=96)

def parse_opts():

    """Command-line options"""

    parser = optparse.OptionParser()

    parser.add_option("-o", "--output", dest="output",

                      help=u"Fichier image à générer")

    parser.add_option("-d", "--debug", dest="debug", action="store_true",

                      help=u"Mode déboguage")

    options, args = parser.parse_args()

    if args:

        parser.error(u"Pas d'argument autorisé")

    if not options.output:

        parser.error("Il faut choisir un fichier de sortie")

    return options

def main():

    opts = parse_opts()

    prod = Production()

    prod.read(opts.debug)

    conso = Consommation()

    conso.read(opts.debug)

    graph = ConsoGraph(prod, conso)

    graph.graph(opts.output)

if __name__ == "__main__":

    main()