# pylint: disable=g-bad-file-header
# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Utilities for plotting metrics."""

import pathlib
from typing import Any, List, Sequence

from absl import logging
from dm_c19_modelling.evaluation import base_indexing
from dm_c19_modelling.evaluation import constants
import matplotlib
matplotlib.use("agg")
import matplotlib.pyplot as plt  # pylint: disable=g-import-not-at-top
import numpy as np
import pandas as pd


def plot_metrics(metrics_df: pd.DataFrame, target_name: str,
                 last_observation_date: str, eval_dataset_creation_date: str,
                 forecast_horizon: int,
                 forecast_index_entries: Sequence[base_indexing.IndexEntryType],
                 num_dates: int, num_sites: int, cadence: int,
                 dropped_sites: np.ndarray) -> plt.Figure:
  """Plots metrics dataframe as a series of bar charts.

  Args:
    metrics_df: Dataframe of metrics, with columns [forecast_id, metric_name,
      metric_value, target_name].
    target_name: the target being predicted.
    last_observation_date: the last date in the training data.
    eval_dataset_creation_date: the creation date of the dataset used for
      evaluation.
    forecast_horizon: the number of days into the future that the forecasts
      extend to.
    forecast_index_entries: the entries in the forecast index for each of the
      forecasts that are included in the metrics dataframe.
    num_dates: the number of dates included in this evaluation.
    num_sites: the number of sites included in this evaluation.
    cadence: the cadence of the forecasts i.e. a cadence of 1 corresponds to
      daily forecasts, a cadence of 7 corresponds to weekly forecasts.
    dropped_sites: optional list of sites that were dropped during evaluation
      from at least one forecast to ensure that all forecasts are for the same
      sites.

  Returns:
      A series of bar plots, one for each metric calculated in the dataframe,
      evaluating different forecasts against each other.
  """
  fig = plt.figure(figsize=(4, 3))
  plot_width = 2
  offset = 0
  column_width = 0.8

  axes = []
  metric_names = metrics_df.metric_name.unique()
  for _ in metric_names:
    ax = fig.add_axes([offset, 0.1, plot_width, 1.])
    ax.grid(axis="y", alpha=0.3, which="both", zorder=0)
    axes.append(ax)
    offset += plot_width * 1.2

  colour_map = plt.get_cmap("tab20c")(
      np.linspace(0, 1.0, len(forecast_index_entries)))
  x_centers = np.arange(len(forecast_index_entries))

  for ax_idx, metric_name in enumerate(metric_names):
    x_offset = ax_idx * column_width - plot_width / 2 + column_width / 2
    x_values = x_centers + x_offset
    ax = axes[ax_idx]
    for bar_idx, forecast_entry in enumerate(forecast_index_entries):
      forecast_id = forecast_entry["forecast_id"]
      row = metrics_df.query(
          f"forecast_id=='{forecast_id}' and metric_name=='{metric_name}'")
      assert len(row) == 1, (
          "Duplicate entries found in metrics dataframe. "
          f"Found {len(row)} entries for {forecast_id} and {metric_name}")
      row = row.iloc[0]
      metric_value = row.metric_value

      ax.bar(
          x_values[bar_idx],