#!/usr/bin/env python3
import argparse
import json
from pathlib import Path
import sys

import matplotlib.pyplot as plt
import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from model.streamflow_lstm import metrics, write_json

parser = argparse.ArgumentParser(description="Evaluate streamflow forecasts and baselines")
parser.add_argument("--config", default="conf/config.yaml")
args = parser.parse_args()
with open(args.config, encoding="utf-8") as handle:
    config = json.load(handle)
with np.load(config["paths"]["predictions"]) as data:
    methods = {name: data[name] for name in ("prediction", "persistence", "glofas")}
    target, gauges, lead_hours = data["target"], data["gauges"].astype(str), data["lead_hours"]
lead_days = {"2_day": 7, "5_day": 19, "8_day": 31}
report = {"lead_definition": "zero-based indices 7/19/31 equal 48/120/192 hours", "gauges": {}}
for gauge_index, gauge in enumerate(gauges):
    report["gauges"][gauge] = {}
    for method, values in methods.items():
        report["gauges"][gauge][method] = {
            label: metrics(values[gauge_index, :, index], target[gauge_index, :, index])
            for label, index in lead_days.items()
        }
numbers = [value for gauge in report["gauges"].values() for method in gauge.values()
           for lead in method.values() for value in lead.values()]
if not np.all(np.isfinite(numbers)):
    raise FloatingPointError("evaluation metrics are not finite")
write_json(config["paths"]["evaluation_metrics"], report)

colors = {"prediction": "#146c94", "persistence": "#d17a22", "glofas": "#6a8e3a"}
figure, axes = plt.subplots(2, 1, figsize=(11, 8), constrained_layout=True)
for method, values in methods.items():
    rmse = [np.mean([metrics(values[g, :, lead], target[g, :, lead])["rmse_m3_s"] for g in range(len(gauges))])
            for lead in range(len(lead_hours))]
    axes[0].plot(lead_hours / 24, rmse, label=method, color=colors[method], linewidth=2)
axes[0].set(title="Mean gauge RMSE across forecast lead", xlabel="Lead (days)", ylabel="RMSE (m3 s-1)")
axes[0].grid(alpha=0.25)
axes[0].legend()
x = np.arange(len(gauges))
width = 0.25
for offset, (method, values) in enumerate(methods.items()):
    kge = [metrics(values[g, :, 19], target[g, :, 19])["kge"] for g in range(len(gauges))]
    axes[1].bar(x + (offset - 1) * width, kge, width, label=method, color=colors[method])
axes[1].axhline(1 - np.sqrt(2), color="black", linestyle="--", linewidth=1, label="KGE skill threshold")
axes[1].set(title="Gauge KGE at 5-day lead", xlabel="Gauge", ylabel="KGE", xticks=x, xticklabels=gauges)
axes[1].legend(ncol=4, fontsize=8)
comparison = Path(config["paths"]["comparison"])
comparison.parent.mkdir(parents=True, exist_ok=True)
figure.savefig(comparison, dpi=150)
plt.close(figure)
print(config["paths"]["evaluation_metrics"], comparison)