#!/usr/bin/env python
# -*- coding: utf-8 -*-

# ==================================================================== #
# Author: Sonia Labetoulle                                             #
# Contact: sonia.labetoulle _at_ ipsl.jussieu.fr                       #
# Created: 2016                                                        #
# History:                                                             #
# Modification:                                                        #
# ==================================================================== #

# ==================================================================== #
# ssh readonly@prodiguer-test-db.ipsl.upmc.fr                          #
# psql -U prodiguer_db_user prodiguer                                  #
#                                                                      #
# ssh readonly@prodiguer-test-db.ipsl.upmc.fr -L 5432:localhost:5432   #
# psql -h prodiguer-test-db.ipsl.upmc.fr -U prodiguer_db_user prodiguer
# ==================================================================== #

# this must come first
from __future__ import print_function, unicode_literals, division
# from __future__ import print_function, division

# standard library imports
from argparse import ArgumentParser
import datetime as dt
import numpy as np

# Application library imports
from libconso import *
import libconso_db as cdb
import db_data


########################################################################
class PlotData(object):
  # --------------------------------------------------------------------
  def __init__(self, date_min, date_max):
    self.date_min = date_min
    self.date_max = date_max

    self.data = DataDict()
    self.data.init_range(self.date_min, self.date_max)


########################################################################
class DataDict(dict):
  # --------------------------------------------------------------------
  def __init__(self):
    self = {}

  # --------------------------------------------------------------------
  def init_range(self, date_beg, date_end, inc=1):
    """
    """
    delta = date_end - date_beg

    (deb, fin) = (0, delta.days+1)

    dates = (date_beg + dt.timedelta(days=i)
             for i in xrange(deb, fin, inc))

    for date in dates:
      self.add_item(date)

  # --------------------------------------------------------------------
  def fill_data(self, projet):
    """
    """

    id_condition = (
      "({})".format(
        " OR ".join(
          [
            "allocation_id={}".format(item.alloc_id)
            for item in projet.alloc_items
          ]
        )
      )
    )

    if args.subproject:
      subp_condition = "sub_project = \'{}\'".format(args.subproject)
    else:
      subp_condition = "sub_project IS NULL"

    table_name = "conso.tbl_consumption"
    request = (
      "SELECT date, total_hrs, allocation_id "
      "FROM " + table_name + " "
      "WHERE login IS NULL "
      # "  AND sub_project IS NULL "
      "  AND " + subp_condition + " "
      "  AND " + id_condition + " "
      "ORDER BY date"
      ";"
    )

    if args.verbose:
      print("Access table \"{}\"".format(table_name))
    cdb.select_db(cursor, request)

    for date, conso, alloc_id in cursor:
      if date.date() in self:
        real_use = conso / projet.alloc
        self.add_item(
          date=date.date(),
          conso=conso,
          real_use=real_use,
        )
        self[date.date()].fill()

  # --------------------------------------------------------------------
  def add_item(self, date, conso=np.nan,
               real_use=np.nan, theo_use=np.nan,
               run_mean=np.nan, pen_mean=np.nan,
               run_std=np.nan, pen_std=np.nan):
    """
    """
    # self[date.date()] = Conso(
    self[date] = Conso(
      date, conso,
      real_use, theo_use,
      run_mean, pen_mean,
      run_std, pen_std
    )

  # --------------------------------------------------------------------
  def get_items_in_range(self, date_beg, date_end, inc=1):
    """
    """
    items = (item for item in self.itervalues()
                   if item.date >= date_beg and
                      item.date <= date_end)
    items = sorted(items, key=lambda item: item.date)

    return items[::inc]

  # --------------------------------------------------------------------
  def get_items_in_full_range(self, inc=1):
    """
    """

    items = (item for item in self.itervalues())
    items = sorted(items, key=lambda item: item.date)

    return items[::inc]

  # --------------------------------------------------------------------
  def get_items(self, inc=1):
    """
    """
    items = (item for item in self.itervalues()
                   if item.isfilled())
    items = sorted(items, key=lambda item: item.date)

    return items[::inc]


class Conso(object):
  # --------------------------------------------------------------------
  def __init__(self, date, conso=np.nan,
               real_use=np.nan, theo_use=np.nan,
               run_mean=np.nan, pen_mean=np.nan,
               run_std=np.nan, pen_std=np.nan):
    self.date       = date
    self.conso      = conso
    self.real_use   = real_use
    self.theo_use   = theo_use
    self.poly_theo  = np.poly1d([])
    self.poly_delay = np.poly1d([])
    self.run_mean   = run_mean
    self.pen_mean   = pen_mean
    self.run_std    = run_std
    self.pen_std    = pen_std
    self.filled     = False

  # --------------------------------------------------------------------
  def __repr__(self):
    return "{:.2f} ({:.2%})".format(self.conso, self.real_use)

  # --------------------------------------------------------------------
  def isfilled(self):
    return self.filled

  # --------------------------------------------------------------------
  def fill(self):
    self.filled = True


########################################################################
def plot_init():
  paper_size  = np.array([29.7, 21.0])
  fig, ax_conso = plt.subplots(figsize=(paper_size * cm2in))
  # fig, ax_conso = plt.subplots(figsize=(paper_size/2.54))
  ax_theo = ax_conso.twinx()

  return fig, ax_conso, ax_theo


########################################################################
def plot_data(
  ax_conso, ax_theo, xcoord, dates,
  consos, real_uses, theo_uses, delay_uses,
  run_mean, pen_mean, run_std, pen_std
):
  """
  """
  line_style = "-"
  if args.full:
    line_width = 0.05
  else:
    line_width = 0.1

  ax_conso.bar(
    xcoord, consos, width=1, align="center", color="linen",
    linewidth=line_width, label="conso (heures)"
  )

  ax_theo.plot(
    xcoord, real_uses, line_style,
    color="forestgreen", linewidth=1, markersize=8,
    solid_capstyle="round", solid_joinstyle="round",
    label="conso\nréelle (%)"
  )

  ax_theo.plot(
    xcoord, theo_uses, "--",
    color="firebrick", linewidth=0.5,
    solid_capstyle="round", solid_joinstyle="round",
    label="conso\nthéorique (%)"
  )

  ax_theo.plot(
    xcoord, delay_uses, ":",
    color="firebrick", linewidth=0.5,
    solid_capstyle="round", solid_joinstyle="round",
    label="retard de\ndeux mois (%)"
  )


########################################################################
def plot_config(
  fig, ax_conso, ax_theo,
  xcoord, dates, max_real_use,
  title, faitle, projet, bilan_plot
):
  """
  """
  from matplotlib.ticker import AutoMinorLocator

  # ... Config axes ...
  # -------------------
  # 1) Range
  conso_max = np.nanmax(consos)
  if args.max:
    ymax = conso_max  # + conso_max*.1
  else:
    ymax = 3. * projet.max_daily_conso

  if conso_max > ymax:
    ax_conso.annotate(
      "{:.2e} heures".format(conso_max),
      ha="left",
      va="top",
      fontsize="xx-small",
      bbox=dict(boxstyle="round", fc="w", ec="0.5", color="gray",),
      xy=(np.nanargmax(consos)+1.2, ymax),
      textcoords="axes fraction",
      xytext=(0.01, 0.9),
      arrowprops=dict(
        arrowstyle="->",
        shrinkA=0,
        shrinkB=0,
        color="gray",
      ),
    )

  xmin, xmax = xcoord[0]-1, xcoord[-1]+1
  ax_conso.set_xlim(xmin, xmax)
  ax_conso.set_ylim(0., ymax)
  ax_theo.set_ylim(0., 100)

  # 2) Plot ideal daily consumption in hours
  #    Plot last real use value
  x_list = []
  y_list = []
  conso_yticks = list(ax_conso.get_yticks())
  for item in projet.alloc_items:
    x_list.extend([item.xi, item.xf])
    y_list.extend([item.daily_conso, item.daily_conso])
    conso_yticks.append(item.daily_conso)
  line_alpha = 0.5
  line_label = "conso journalière\nidéale (heures)"
  ax_conso.plot(
    x_list, y_list,
    linewidth=0.5,
    color="blue",
    alpha=line_alpha,
    label=line_label,
  )

  theo_yticks = list(ax_theo.get_yticks())
  theo_yticks.append(max_real_use)
  ax_theo.axhline(
    y=max_real_use,
    linestyle=":", linewidth=0.5,
    color="green", alpha=line_alpha,
  )

  # 3) Ticks labels
  (date_beg, date_end) = (dates[0], dates[-1])
  date_fmt = "{:%d-%m}"

  if date_end - date_beg > dt.timedelta(weeks=9):
    maj_xticks = [x for x, d in zip(xcoord, dates)
                     if d.weekday() == 0]
    maj_xlabs  = [date_fmt.format(d) for d in dates
                     if d.weekday() == 0]
  else:
    maj_xticks = [x for x, d in zip(xcoord, dates)]
    maj_xlabs  = [date_fmt.format(d) for d in dates]

  ax_conso.ticklabel_format(axis="y", style="sci", scilimits=(0, 0))

  ax_conso.set_xticks(xcoord, minor=True)
  ax_conso.set_xticks(maj_xticks, minor=False)
  ax_conso.set_xticklabels(
    maj_xlabs, rotation="vertical", size="x-small"
  )

  minor_locator = AutoMinorLocator()
  ax_conso.yaxis.set_minor_locator(minor_locator)
  minor_locator = AutoMinorLocator()
  ax_theo.yaxis.set_minor_locator(minor_locator)

  ax_conso.set_yticks(conso_yticks)
  ax_theo.set_yticks(theo_yticks)

  ax_theo.spines["right"].set_color("firebrick")
  ax_theo.tick_params(colors="firebrick")
  ax_theo.yaxis.label.set_color("firebrick")

  for x, d in zip(xcoord, dates):
    if d.weekday() == 0:
      ax_conso.axvline(
        x=x, color="gray", alpha=0.5,
        linewidth=0.5, linestyle=":"
      )

  # 4) Define axes title
  for ax, label in (
    (ax_conso, "heures"),
    (ax_theo, "%"),
  ):
    ax.set_ylabel(label, fontweight="bold")
    ax.tick_params(axis="y", labelsize="small")

  # 5) Define plot size
  fig.subplots_adjust(
    left=0.08,
    bottom=0.09,
    right=0.93,
    top=0.92,
  )

  # ... Main title and legend ...
  # -----------------------------
  fig.suptitle(title, fontweight="bold", size="large")
  for ax, loc in (
    (ax_conso, "upper left"),
    (ax_theo, "upper right"),
  ):
    ax.legend(loc=loc, fontsize="x-small", frameon=False)

  alloc_label = (
    "Allocation(s):\n" +
    "\n".join([
      "{:%d/%m/%Y}-{:%d/%m/%Y} : {:10,.0f}h".format(
        item.start_date, item.end_date, item.alloc
      ) for item in projet.alloc_items
    ])
  )
  plt.figtext(
    x=0.92, y=0.93, s=alloc_label.replace(",", " "),
    backgroundcolor="linen",
    ha="right", va="bottom", fontsize="x-small"
  )

  range_label = "{:%d/%m/%Y} - {:%d/%m/%Y}".format(
    bilan_plot.date_min,
    bilan_plot.date_max,
  )
  plt.figtext(
    x=0.5, y=0.93, s=range_label,
    ha="center", va="bottom", fontsize="large"
  )

  faitle_label = "fait le {:%d/%m/%Y}".format(bilan_plot.faitle)
  plt.figtext(
    x=0.02, y=0.98, s=faitle_label,
    ha="left", va="top", fontsize="x-small"
  )


########################################################################
def get_arguments():
  parser = ArgumentParser()
  parser.add_argument("project", action="store",
                      help="Project name")
  parser.add_argument("centre", action="store",
                      choices=["idris", "tgcc"],
                      help="Centre name")
  parser.add_argument(
    "node", action="store", default="thin/standard",
    choices=["thin/standard", "standard", "fat/large", "hybrid"],
    help="Node type"
  )

  parser.add_argument("--subproject", action="store",
                      type=str, default=None,
                      help="Sub-project")

  parser.add_argument("-v", "--verbose", action="store_true",
                      help="verbose mode")
  parser.add_argument("-f", "--full", action="store_true",
                      help="plot the whole period")
  parser.add_argument("--png", action="store_true",
                      help="pgn output in addition to pdf")
  parser.add_argument("-i", "--increment", action="store",
                      type=int, default=1, dest="inc",
                      help="sampling increment")
  parser.add_argument("-r", "--range", action="store", nargs=2,
                      type=string_to_date,
                      help="date range: ssaa-mm-jj ssaa-mm-jj")
  parser.add_argument("-m", "--max", action="store_true",
                      help="plot with y_max = allocation")
  parser.add_argument("-s", "--show", action="store_true",
                      help="interactive mode")
  parser.add_argument("-d", "--dods", action="store_true",
                      help="copy output on dods")

  return parser.parse_args()


########################################################################
if __name__ == '__main__':

  # .. Initialization ..
  # ====================
  # ... Command line arguments ...
  # ------------------------------
  args = get_arguments()

  print(os.getcwd())
  print(os.path.abspath(__file__))

  # ... Turn interactive mode off ...
  # ---------------------------------
  if not args.show:
    import matplotlib
    matplotlib.use('Agg')

  import matplotlib.pyplot as plt
  # from matplotlib.backends.backend_pdf import PdfPages

  if not args.show:
    plt.ioff()

  # ... Variables and constants ...
  # -------------------------------
  delay = 60  # 2 months delay

  # ... Files and directories ...
  # -----------------------------
  config_file = os.path.join(
    "card",
    "config_{}_{}.ini".format(args.centre, args.project)
  )

  if not os.path.isfile(config_file):
    print("File {} missing ".format(config_file))
    exit(9)

  (DIR, DODS) = parse_config(config_file)

  # (file_param, file_utheo, file_data) = \
  #     get_input_files(DIR["SAVEDATA"],
  #                     [OUT["PARAM"], OUT["UTHEO"], OUT["BILAN"]])

  img_name = "{}_{}_{}_{}_{}".format(
    os.path.splitext(
      os.path.basename(__file__)
    )[0].replace("plot_", ""),
    args.centre,
    args.project,
    args.subproject if args.subproject else args.project,
    args.node.replace("/", ""),
  )

  conn, cursor = cdb.connect_db(
    db_data.db_host,
    db_data.db_name,
    db_data.db_user,
    db_data.db_pwd,
  )

  today = dt.date.today()

  if args.verbose:
    fmt_str = "{:10s} : {}"
    print(fmt_str.format("args", args))
    print(fmt_str.format("today", today))
    print(fmt_str.format("img_name", img_name))
    print("")

  # .. Range of dates ..
  # ====================
  if args.full:
    (date_min, date_max) = (
      dt.date(today.year, 1, 1),
      dt.date(today.year, 12, 31),
    )
  elif args.range:
    (date_min, date_max) = (args.range)
  else:
    (date_min, date_max) = (
      dt.date(today.year, 1, 1),
      dt.date(today.year, 12, 31),
    )
  bilan_plot = PlotData(date_min, date_max)
  bilan_plot.faitle = today

  # .. Get project info ..
  # ======================
  table_name = "conso.tbl_allocation"
  request = (
    "SELECT * "
    "FROM " + table_name + " "
    "WHERE project = '" + args.project + "' "
    "AND node_type = '" + args.node + "' "
    "AND is_active = TRUE "
    "ORDER BY start_date"
    ";"
  )

  cdb.select_db(cursor, request)
  # print(cursor.rowcount)
  # print(cursor.fetchall())
  # print(cursor.description)

  projet = Project(args.project, args.centre)
  for row in cursor:
    if (bilan_plot.date_min >= row["start_date"].date() and
        bilan_plot.date_min <= row["end_date"].date()) or \
       (bilan_plot.date_max >= row["start_date"].date() and
        bilan_plot.date_max <= row["end_date"].date()) or \
       (bilan_plot.date_min <= row["start_date"].date() and
        bilan_plot.date_max >= row["end_date"].date()):
      projet.add_alloc(
        row["id"],
        row["machine"],
        row["node_type"],
        row["start_date"],
        row["end_date"],
        row["total_hrs"],
      )

  # .. Fill in data ..
  # ==================
  # ... Initialization ...
  # ----------------------

  bilan = DataDict()
  bilan.init_range(projet.start_date.date(), projet.end_date.date())

  # ... Extract data from table ...
  # -------------------------------
  bilan.fill_data(projet)

  # .. Extract data depending on C.L. arguments ..
  # ==============================================
  if args.full:
    selected_items = bilan.get_items_in_full_range(args.inc)
  elif args.range:
    selected_items = bilan.get_items_in_range(
      args.range[0], args.range[1], args.inc
    )
  else:
    selected_items = bilan.get_items(args.inc)

  last_filled_date = max([
    item.date for item in selected_items
    if item.filled
  ])

  # .. Compute data to be plotted ..
  # ================================
  nb_items = len(selected_items)

  xcoord = np.linspace(1, nb_items, num=nb_items)
  dates = [item.date for item in selected_items]

  projet.get_theo_eq(dates)

  cumul = np.array(
    [item.conso for item in selected_items],
    dtype=float
  )
  consos = []
  consos.append(cumul[0])
  consos[1:nb_items] = cumul[1:nb_items] - cumul[0:nb_items-1]
  consos = np.array(consos, dtype=float)

  real_uses = np.array(
    [100.*item.real_use for item in selected_items],
    dtype=float
  )

  theo_uses = []
  real_uses2 = []
  for idx, date in enumerate(dates):
    for item in projet.alloc_items:
      if date >= item.start_date.date() and \
         date <= item.end_date.date():
        poly_theo  = item.theo_eq
        alloc_val = item.alloc
        break
    theo_uses.append(100. * poly_theo(dates.index(date)))
    real_uses2.append((100 * cumul[idx]) / alloc_val)

  if nb_items > delay:
    delay_uses = delay * [0., ] + theo_uses[:-delay]
  else:
    delay_uses = nb_items * [0., ]

  theo_uses = np.array(theo_uses, dtype=float)
  delay_uses = np.array(delay_uses, dtype=float)

  run_mean = np.array(
    [item.run_mean for item in selected_items],
    dtype=float
  )
  pen_mean = np.array(
    [item.pen_mean for item in selected_items],
    dtype=float
  )
  run_std  = np.array(
    [item.run_std for item in selected_items],
    dtype=float
  )
  pen_std  = np.array(
    [item.pen_std for item in selected_items],
    dtype=float
  )

  # .. Plot stuff ..
  # ================
  # ... Initialize figure ...
  # -------------------------
  (fig, ax_conso, ax_theo) = plot_init()

  # ... Plot data ...
  # -----------------
  plot_data(
    ax_conso, ax_theo, xcoord, dates,
    consos, real_uses, theo_uses, delay_uses,
    run_mean, pen_mean, run_std, pen_std
  )

  # ... Tweak figure ...
  # --------------------
  proj_title = args.subproject if args.subproject else projet.project
  title = "Consommation {}".format(
    # projet.project.upper(),
    proj_title.upper()
  )
  faitle = "(le {:%d/%m/%Y})".format(
    today,
  )

  plot_config(
    fig, ax_conso, ax_theo,
    xcoord, dates, real_uses[dates.index(last_filled_date)],
    title, faitle, projet, bilan_plot
  )

  # ... Save figure ...
  # -------------------
  if args.range:
    img_out = os.path.join(
      DIR["PLOT"], args.centre, args.project,
      "{}_{:%Y%m%d}_{:%Y%m%d}.pdf".format(img_name, *args.range)
    )
  else:
    img_out = os.path.join(
      DIR["PLOT"], args.centre, args.project,
      "{}_{:%Y%m%d}.pdf".format(img_name, today)
    )
  if args.verbose:
    print("Save image as {}".format(img_out))
    if args.png:
      print("Produce png")

  title = "Consommation {} ({}) au {:%Y%m%d}".format(
    args.project.upper(),
    args.centre.upper(),
    today
  )
  plot_save(img_out, title, png=args.png)

  # ... Publish figure on dods ...
  # ------------------------------
  if args.dods:
    if args.verbose:
      print("Publish figure on dods")
    dods_cp(img_out, img_name, DODS)

  # if args.show:
  #   plt.show()

  exit(0)