"""
This plot presents the trailing X number of days
temperature or precipitation departure from long term average. You can
express this departure either in Absolute Departure or as a Standard
Deviation.  The Standard Deviation option along with precipitation is
typically called the "Standardized Precipitation Index".

<p>The plot also contains an underlay with the weekly US Drought Monitor
that is valid for the station location.  If you plot a climate district
station, you get the US Drought Monitor valid for the district centroid.
If you plot a statewide average, you get no USDM included.
"""

from datetime import date, datetime, timedelta

import matplotlib.dates as mdates
import pandas as pd
import requests
from matplotlib.axes import Axes
from matplotlib.patches import Rectangle
from pyiem.database import get_sqlalchemy_conn, sql_helper
from pyiem.exceptions import NoDataFound
from pyiem.plot import figure_axes
from pyiem.util import LOG

UNITS = {"precip": "inch", "avgt": "°F", "high": "°F", "low": "°F"}
PDICT = {
    "precip": "Precipitation",
    "avgt": "Daily Average Temperature",
    "high": "Daily High Temperature",
    "low": "Daily Low Temperature",
}
PDICT2 = {"diff": "Absolute Departure", "sigma": "Standard Deviation"}
COLORS = ["#ffff00", "#fcd37f", "#ffaa00", "#e60000", "#730000"]


def get_description():
    """Return a dict describing how to call this plotter"""
    desc = {"description": __doc__, "data": True, "cache": 86400}
    today = date.today()
    sts = today - timedelta(days=720)
    desc["arguments"] = [
        dict(
            type="station",
            name="station",
            default="IATAME",
            label="Select Station:",
            network="IACLIMATE",
        ),
        dict(
            min=1,
            type="int",
            name="p1",
            default=31,
            label="First Period of Days",
        ),
        dict(
            min=1,
            type="int",
            name="p2",
            default=91,
            label="Second Period of Days",
        ),
        dict(
            min=1,
            type="int",
            name="p3",
            default=365,
            label="Third Period of Days",
        ),
        dict(
            type="date",
            name="sdate",
            default=sts.strftime("%Y/%m/%d"),
            min="1893/01/01",
            label="Start Date of Plot",
        ),
        dict(
            type="date",
            name="edate",
            default=today.strftime("%Y/%m/%d"),
            min="1893/01/01",
            label="End Date of Plot",
        ),
        dict(
            type="select",
            name="pvar",
            default="precip",
            options=PDICT,
            label="Which variable to plot?",
        ),
        dict(
            type="select",
            name="how",
            default="diff",
            options=PDICT2,
            label="How to Express Departure?",
        ),
    ]
    return desc


def underlay_usdm(axis: Axes, sts, ets, lon, lat):
    """Underlay the USDM as pretty bars, somehow"""
    if ets < date(2000, 1, 1):
        axis.text(
            0.0,
            1.03,
            "No Drought Information Prior to 2000",
            transform=axis.transAxes,
        )
        return
    rects = [Rectangle((0, 0), 1, 1, fc=color) for color in COLORS]
    axis.text(
        0.0, 1.03, "Drought Category Underlain", transform=axis.transAxes
    )
    axis.legend(
        rects,
        [f"D{cat}" for cat in range(5)],
        ncol=5,
        fontsize=11,
        loc=(0.3, 1.01),
    )
    uri = (
        "http://mesonet.agron.iastate.edu/"
        f"api/1/usdm_bypoint.json?sdate={sts:%Y-%m-%d}&edate={ets:%Y-%m-%d}"
        f"&lon={lon}&lat={lat}"
    )
    try:
        resp = requests.get(uri, timeout=30)
        resp.raise_for_status()
        data = resp.json()
    except Exception as exp:
        LOG.info("usdm fetch failed: %s", exp)
        return
    for row in data["data"]:
        if row["category"] is None:
            continue
        ts = datetime.strptime(row["valid"], "%Y-%m-%d")
        dt = date(ts.year, ts.month, ts.day)
        axis.axvspan(
            dt,
            dt + timedelta(days=7),
            color=COLORS[row["category"]],
            zorder=-3,
        )


def plotter(ctx: dict):
    """Go"""
    station = ctx["station"]
    p1 = ctx["p1"]
    p2 = ctx["p2"]
    p3 = ctx["p3"]
    pvar = ctx["pvar"]
    sts = ctx["sdate"]
    ets = ctx["edate"]
    how = ctx["how"]
    maxdays = max([p1, p2, p3])

    with get_sqlalchemy_conn("coop") as conn:
        df = pd.read_sql(
            sql_helper("""
        -- Get all period averages
        with avgs as (
            SELECT day, sday,
            count(high) OVER (ORDER by day ASC ROWS :md1 PRECEDING) as counts,
            avg(high) OVER (ORDER by day ASC ROWS :d1 PRECEDING) as p1_high,
            avg(high) OVER (ORDER by day ASC ROWS :d2 PRECEDING) as p2_high,
            avg(high) OVER (ORDER by day ASC ROWS :d3 PRECEDING) as p3_high,
            avg(low) OVER (ORDER by day ASC ROWS :d1 PRECEDING) as p1_low,
            avg(low) OVER (ORDER by day ASC ROWS :d2 PRECEDING) as p2_low,
            avg(low) OVER (ORDER by day ASC ROWS :d3 PRECEDING) as p3_low,
            avg((high+low)/2.)
                OVER (ORDER by day ASC ROWS :d1 PRECEDING) as p1_avgt,
            avg((high+low)/2.)
                OVER (ORDER by day ASC ROWS :d2 PRECEDING) as p2_avgt,
            avg((high+low)/2.)
                OVER (ORDER by day ASC ROWS :d3 PRECEDING) as p3_avgt,
    sum(precip) OVER (ORDER by day ASC ROWS :d1 PRECEDING) as p1_precip,
    sum(precip) OVER (ORDER by day ASC ROWS :d2 PRECEDING) as p2_precip,
    sum(precip) OVER (ORDER by day ASC ROWS :d3 PRECEDING) as p3_precip
    from alldata WHERE station = :station
        ),
        -- Get sday composites
        sdays as (
            SELECT sday,
            avg(p1_high) as p1_high_avg, stddev(p1_high) as p1_high_stddev,
            avg(p2_high) as p2_high_avg, stddev(p2_high) as p2_high_stddev,
            avg(p3_high) as p3_high_avg, stddev(p3_high) as p3_high_stddev,
            avg(p1_low) as p1_low_avg, stddev(p1_low) as p1_low_stddev,
            avg(p2_low) as p2_low_avg, stddev(p2_low) as p2_low_stddev,
            avg(p3_low) as p3_low_avg, stddev(p3_low) as p3_low_stddev,
            avg(p1_avgt) as p1_avgt_avg, stddev(p1_avgt) as p1_avgt_stddev,
            avg(p2_avgt) as p2_avgt_avg, stddev(p2_avgt) as p2_avgt_stddev,
            avg(p3_avgt) as p3_avgt_avg, stddev(p3_avgt) as p3_avgt_stddev,
            avg(p1_precip) as p1_precip_avg,
            stddev(p1_precip) as p1_precip_stddev,
            avg(p2_precip) as p2_precip_avg,
            stddev(p2_precip) as p2_precip_stddev,
            avg(p3_precip) as p3_precip_avg,
            stddev(p3_precip) as p3_precip_stddev
            from avgs WHERE counts = :maxdays GROUP by sday
        )
        -- Now merge to get obs
            SELECT day, s.sday,
            p1_high - p1_high_avg as p1_high_diff,
            p2_high - p2_high_avg as p2_high_diff,
            p3_high - p3_high_avg as p3_high_diff,
            p1_low - p1_low_avg as p1_low_diff,
            p2_low - p2_low_avg as p2_low_diff,
            p3_low - p3_low_avg as p3_low_diff,
            p1_avgt - p1_avgt_avg as p1_avgt_diff,
            p2_avgt - p2_avgt_avg as p2_avgt_diff,
            p3_avgt - p3_avgt_avg as p3_avgt_diff,
            p1_precip - p1_precip_avg as p1_precip_diff,
            p2_precip - p2_precip_avg as p2_precip_diff,
            p3_precip - p3_precip_avg as p3_precip_diff,
            (p1_high - p1_high_avg) / p1_high_stddev as p1_high_sigma,
            (p2_high - p2_high_avg) / p2_high_stddev as p2_high_sigma,
            (p3_high - p3_high_avg) / p3_high_stddev as p3_high_sigma,
            (p1_low - p1_low_avg) / p1_low_stddev as p1_low_sigma,
            (p2_low - p2_low_avg) / p2_low_stddev as p2_low_sigma,
            (p3_low - p3_low_avg) / p3_low_stddev as p3_low_sigma,
            (p1_avgt - p1_avgt_avg) / p1_avgt_stddev as p1_avgt_sigma,
            (p2_avgt - p2_avgt_avg) / p2_avgt_stddev as p2_avgt_sigma,
            (p3_avgt - p3_avgt_avg) / p3_avgt_stddev as p3_avgt_sigma,
            (p1_precip - p1_precip_avg) / p1_precip_stddev as p1_precip_sigma,
            (p2_precip - p2_precip_avg) / p2_precip_stddev as p2_precip_sigma,
            (p3_precip - p3_precip_avg) / p3_precip_stddev as p3_precip_sigma
            from avgs a JOIN sdays s on (a.sday = s.sday) WHERE
            day >= :sts and day <= :ets ORDER by day ASC
        """),
            conn,
            params={
                "md1": maxdays - 1,
                "d1": p1 - 1,
                "d2": p2 - 1,
                "d3": p3 - 1,
                "station": station,
                "maxdays": maxdays,
                "sts": sts,
                "ets": ets,
            },
            index_col="day",
        )
    if df.empty:
        raise NoDataFound("No Data Found.")

    (fig, ax) = figure_axes(apctx=ctx)
    ax.set_position([0.1, 0.14, 0.85, 0.71])

    (l1,) = ax.plot(
        df.index.values,
        df[f"p1_{pvar}_{how}"],
        lw=2,
        label=f"{p1} Day",
        zorder=5,
    )
    (l2,) = ax.plot(
        df.index.values,
        df["p2_" + pvar + "_" + how],
        lw=2,
        label=f"{p2} Day",
        zorder=5,
    )
    (l3,) = ax.plot(
        df.index.values,
        df["p3_" + pvar + "_" + how],
        lw=2,
        label=f"{p3} Day",
        zorder=5,
    )
    fig.text(
        0.5,
        0.93,
        f"{ctx['_sname']}\n"
        f"Trailing {p1}, {p2}, {p3} Day Departures & "
        "US Drought Monitor",
        ha="center",
        fontsize=14,
    )
    ax.xaxis.set_major_formatter(mdates.DateFormatter("%b\n%Y"))
    ax.set_ylabel(
        ("%s [%s] %s")
        % (
            PDICT.get(pvar),
            UNITS[pvar] if how == "diff" else r"$\sigma$",
            PDICT2[how],
        )
    )
    ax.grid(True)
    leg = ax.legend(
        handles=[l1, l2, l3], ncol=3, fontsize=12, loc="lower left"
    )
    ax.add_artist(leg)
    ax.text(
        1,
        -0.14,
        f"{sts:%-d %b %Y} to {ets:%-d %b %Y}",
        va="bottom",
        ha="right",
        fontsize=12,
        transform=ax.transAxes,
    )
    if station[2:] != "0000":
        underlay_usdm(
            ax,
            sts,
            ets,
            ctx["_nt"].sts[station]["lon"],
            ctx["_nt"].sts[station]["lat"],
        )
    offset = timedelta(days=2)
    ax.set_xlim(df.index.min() - offset, df.index.max() + offset)

    return fig, df
