#!/usr/bin/env python3
"""Plot seasonal latitude–lag regressions of MSE and AR occurrence on BAM."""

from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from netCDF4 import Dataset
import numpy as np


BASE = Path(__file__).resolve().parent


def add_ar_contours(axis, lag, latitude, field):
    negative = [x for x in (-2, -1.5, -1, -0.5) if np.nanmin(field) < x < np.nanmax(field)]
    positive = [x for x in (0.5, 1, 1.5, 2) if np.nanmin(field) < x < np.nanmax(field)]
    if negative:
        axis.contour(lag, latitude, field, levels=negative, colors="black",
                     linewidths=0.8, linestyles="dashed", zorder=5)
    if positive:
        axis.contour(lag, latitude, field, levels=positive, colors="black",
                     linewidths=0.8, linestyles="solid", zorder=5)


def main():
    with Dataset(BASE / "data_figure2.nc") as dataset:
        lag = np.asarray(dataset.variables["lag"][:], dtype=float)
        latitude = np.asarray(dataset.variables["lat"][:], dtype=float)
        ar_lag = np.asarray(dataset.variables["ar_lag"][:], dtype=float)
        ar_latitude = np.asarray(dataset.variables["ar_lat"][:], dtype=float)
        mse = np.asarray(dataset.variables["mse_regression"][:], dtype=float)
        significant = np.asarray(dataset.variables["significant_95"][:], dtype=bool)
        ar = np.asarray(dataset.variables["ar_regression"][:], dtype=float)
    figure, axes = plt.subplots(2, 2, figsize=(13.0, 7.0), sharex=True,
                                layout="constrained")
    levels = np.linspace(-0.6, 0.6, 13)
    titles = ("(a) NH DJF", "(b) NH JJA", "(c) SH JJA", "(d) SH DJF")
    image = None
    for index, axis in enumerate(axes.flat):
        image = axis.contourf(lag, latitude, mse[index], levels=levels,
                              cmap="RdBu_r", extend="both")
        with matplotlib.rc_context({"hatch.color": "0.5", "hatch.linewidth": 0.35}):
            axis.contourf(lag, latitude, significant[index].astype(np.int8),
                          levels=[0.5, 1.5], colors="none", hatches=[".."])
        add_ar_contours(axis, ar_lag, ar_latitude, ar[index])
        axis.axvline(0, color="black", linewidth=0.9, linestyle="--")
        axis.grid(axis="x", color="0.85", linewidth=0.45)
        axis.set_xticks(np.arange(-10, 11, 5))
        axis.set_xlim(-12, 12)
        if index < 2:
            axis.set_ylim(10, 80)
            axis.set_yticks(np.arange(15, 76, 15))
        else:
            axis.set_ylim(-80, -10)
            axis.set_yticks(np.arange(-75, -14, 15))
        axis.tick_params(labelsize=12)
        axis.set_title(titles[index], loc="left", fontsize=17, pad=5)
        if index % 2 == 0:
            axis.set_ylabel("Latitude (°)", fontsize=14)
        if index >= 2:
            axis.set_xlabel("Lag (days)", fontsize=14)
    colorbar = figure.colorbar(image, ax=axes, orientation="horizontal", pad=0.04,
                              fraction=0.045, shrink=0.7, aspect=50)
    colorbar.set_ticks(levels)
    colorbar.set_label(
        "MSE transport anomaly regression (PW per BAM standard deviation)",
        fontsize=14,
    )
    colorbar.ax.tick_params(labelsize=12)
    figure.suptitle("MSE & AR regression onto BAM", fontsize=18)
    figure.canvas.draw()
    box = colorbar.ax.get_position()
    figure.legend(
        handles=[Line2D([0], [0], color="black", linewidth=1.0,
                        label="AR occurrence")],
        loc="center left", bbox_to_anchor=(box.x1 + 0.035,
                                           box.y0 + 0.5 * box.height),
        frameon=False, fontsize=9,
    )
    for suffix in ("png", "pdf"):
        path = BASE / f"Figure_2.{suffix}"
        figure.savefig(path, dpi=300, bbox_inches="tight")
        print(f"Saved {path.name}")
    plt.close(figure)


if __name__ == "__main__":
    main()