"""Grid the daily/24H data onto a grid for IEMRE

This is tricky as some variables we can compute sooner than others.

    RUN_MIDNIGHT.sh
    RUN_NOON.sh
    RUN_0Z.sh
    RUN_10_AFTER.sh
"""

import subprocess
import warnings
from datetime import date, datetime, timedelta, timezone
from typing import cast

import click
import numpy as np
import pandas as pd
import xarray as xr
from metpy.calc import relative_humidity_from_dewpoint
from metpy.interpolate import inverse_distance_to_grid
from metpy.units import units
from pyiem.database import get_sqlalchemy_conn, sql_helper
from pyiem.grid.nav import get_nav
from pyiem.grid.util import grid_smear
from pyiem.iemre import (
    DOMAINS,
    get_grids,
    get_hourly_ncname,
    hourly_offset,
    set_grids,
)
from pyiem.util import convert_value, logger, ncopen, utc
from scipy.stats import zscore

LOG = logger()
# How much data is enough, this is a hack to figure out if we have too much
# missing data over land, ie the ASOS interpolator got used instead of model
# ERA5Land is the limiting constraint here
QUORUM_THRESHOLDS = {
    "conus": 90,
    "china": 79,
    "europe": 77,
    "sa": 51,
}


def generic_gridder(df, idx, domain: str = "conus"):
    """
    Generic gridding algorithm for easy variables
    """
    gridnav = get_nav("iemre", domain)
    if not idx.startswith("precip"):
        window = 2.0
        f1 = df[df[idx].notnull()]
        for lat in np.arange(gridnav.bottom, gridnav.top, window):
            for lon in np.arange(gridnav.left, gridnav.right, window):
                (west, east, south, north) = (
                    lon,
                    lon + window,
                    lat,
                    lat + window,
                )
                box = f1[
                    (f1["lat"] >= south)
                    & (f1["lat"] < north)
                    & (f1["lon"] >= west)
                    & (f1["lon"] < east)
                ]
                # can't QC data that is all equal
                if len(box.index) < 4 or box[idx].min() == box[idx].max():
                    continue
                z = np.abs(cast("np.ndarray", zscore(box[idx])))
                # Compute where the standard dev is +/- 2std
                bad = box[z > 1.5]
                df.loc[bad.index, idx] = np.nan

    df2 = df[df[idx].notnull()]
    if len(df2.index) < 4:
        LOG.info("Not enough data %s", idx)
        return None
    xi, yi = np.meshgrid(gridnav.x_points, gridnav.y_points)
    res = np.ones(xi.shape) * np.nan
    # do our gridding
    grid = inverse_distance_to_grid(
        df2["lon"].values, df2["lat"].values, df2[idx].values, xi, yi, 1.5
    )
    # replace nan values in res with whatever now is in grid
    res = np.where(np.isnan(res), grid, res)
    # Do we still have missing values?
    if np.isnan(res).any():
        # very aggressive with search radius
        grid = inverse_distance_to_grid(
            df2["lon"].values, df2["lat"].values, df2[idx].values, xi, yi, 5.5
        )
        # replace nan values in res with whatever now is in grid
        res = np.where(np.isnan(res), grid, res)
    # replace sentinel back to np.nan
    res = np.where(res == -9999, np.nan, res)
    return np.ma.array(res, mask=np.isnan(res))


def compute_daily_pairs(
    dt: date, domain: str, is_12z: bool
) -> list[tuple[datetime, datetime]]:
    """Figure out the pairs.

    The 12z variant is a bit different and will include the end timestamp
    """
    # Compute a timestamp at noon within the local calendar date
    noon_local = datetime(
        dt.year, dt.month, dt.day, 12, tzinfo=DOMAINS[domain]["tzinfo"]
    )
    if is_12z:
        ets = noon_local.replace(hour=7)
        sts = ets - timedelta(hours=23)
    else:
        sts = noon_local.replace(hour=1 if noon_local.dst() else 0)
        ets = sts + timedelta(hours=23)
    # Logic below needs UTC
    sts = sts.astimezone(timezone.utc)
    ets = ets.astimezone(timezone.utc)
    LOG.info("Using %s to %s for %s[is_12z: %s]", sts, ets, domain, is_12z)
    pairs = [(sts, ets)]
    if sts.year != ets.year:
        # These are inclusive
        pairs = [
            (sts, sts.replace(hour=23)),
            (ets.replace(hour=0), ets),
        ]
    return pairs


def copy_iemre_hourly(ts: date, ds: xr.Dataset, domain: str):
    """Lots of work to do here..."""
    pairs = compute_daily_pairs(ts, domain, is_12z=False)
    pairs12z = compute_daily_pairs(ts, domain, is_12z=True)

    # One Off
    for vname in ["min_rh", "max_rh"]:
        aggfunc = np.ma.max if vname == "max_rh" else np.ma.min
        res = None
        for pair in pairs:
            ncfn = get_hourly_ncname(pair[0].year, domain)
            tidx1 = hourly_offset(pair[0])
            tidx2 = hourly_offset(pair[1])
            with ncopen(ncfn, timeout=600) as nc:
                tmpk = nc.variables["tmpk"]
                dwpk = nc.variables["dwpk"]
                for offset in range(tidx1, tidx2 + 1):
                    with warnings.catch_warnings(
                        category=RuntimeWarning, action="ignore"
                    ):
                        rh = (
                            relative_humidity_from_dewpoint(
                                units("degK") * tmpk[offset],
                                units("degK") * dwpk[offset],
                            ).m
                            * 100.0
                        )
                    if res is None:
                        res = rh
                    else:
                        res = aggfunc([res, rh], axis=0)
        ds[vname].values = res

    # One off
    hours = 0
    runningsum = None
    for pair in pairs:
        ncfn = get_hourly_ncname(pair[0].year, domain)
        tidx1 = hourly_offset(pair[0])
        tidx2 = hourly_offset(pair[1])
        with ncopen(ncfn, timeout=600) as nc:
            uwnd = nc.variables["uwnd"]
            vwnd = nc.variables["vwnd"]
            for offset in range(tidx1, tidx2 + 1):
                val = np.hypot(uwnd[offset], vwnd[offset])
                quorum_pct = 100.0 - np.ma.count_masked(val) / val.size * 100.0
                if quorum_pct < QUORUM_THRESHOLDS[domain]:
                    dt = utc(pair[0].year) + timedelta(hours=offset)
                    LOG.info(
                        "Skipping wind at offset %s[%s] %.2f<%.2f required",
                        offset,
                        f"{dt:%Y-%m-%d %H}",
                        quorum_pct,
                        QUORUM_THRESHOLDS[domain],
                    )
                    continue
                if runningsum is None:
                    runningsum = val
                else:
                    runningsum += val
                hours += 1
    if hours > 0:
        ds["wind_speed"].values = runningsum / hours
    # -----------------------------------------------------------------
    aggfuncs = {
        "high_tmpk": np.ma.max,
        "low_tmpk": np.ma.min,
        "p01d": np.ma.sum,
        "high_soil4t": np.ma.max,
        "avg_dwpk": np.ma.mean,
        "rsds": np.ma.sum,
        "low_soil4t": np.ma.min,
        "high_tmpk_12z": np.ma.max,
        "low_tmpk_12z": np.ma.min,
        "p01d_12z": np.ma.sum,
    }
    for vname in (
        "high_tmpk low_tmpk p01d high_soil4t avg_dwpk rsds "
        "low_soil4t high_tmpk_12z low_tmpk_12z p01d_12z"
    ).split():
        if vname == "rsds" and domain == "conus":
            # Done via other means
            continue
        res = None
        aggfunc = aggfuncs[vname]
        ncvarname = (
            vname.replace("high_", "")
            .replace("low_", "")
            .replace("_12z", "")
            .replace("avg_", "")
        )
        ncvarname = "p01m" if ncvarname == "p01d" else ncvarname
        for pair in pairs12z if vname.endswith("12z") else pairs:
            ncfn = get_hourly_ncname(pair[0].year, domain)
            tidx1 = hourly_offset(pair[0])
            tidx2 = hourly_offset(pair[1])
            tslice = slice(tidx1, tidx2 + 1)
            if vname.startswith("p01d"):
                # Caution: precip is stored arrears, so we have ugliness
                tslice = slice(0 if tidx1 == 0 else tidx1 + 1, tidx2 + 2)
            LOG.info("Using %s[%s] for %s", ncfn, tslice, vname)
            with ncopen(ncfn, timeout=600) as nc:
                ncvar = nc.variables[ncvarname]
                ncval = aggfunc(ncvar[tslice], axis=0)
                if res is None:
                    res = ncval
                else:
                    res = aggfunc([res, ncval], axis=0)
        # NB: these are crude bias offsets involved when using hourly data
        if vname in ["high_tmpk", "high_tmpk_12z"]:
            res += 0.8
        if vname in ["low_tmpk", "low_tmpk_12z"]:
            res -= 0.8
        if vname == "rsds":
            # Convert from Wm-2
            res = res / 24.0
        ds[vname].values = res


def use_climodat_12z(ts, ds):
    """Look at what we have in climodat."""
    mybuf = 2
    gridnav = get_nav("iemre")
    with get_sqlalchemy_conn("coop") as conn:
        df = pd.read_sql(
            sql_helper("""
        WITH mystations as (
            SELECT id, ST_X(geom) as lon, ST_Y(geom) as lat, state, name
            from stations where ST_Contains(
            ST_MakeEnvelope(:west, :south, :east, :north, 4326), geom) and
            network ~* 'CLIMATE' and substr(id, 3, 1) not in ('C', 'D', 'T')
            and substr(id, 3, 4) != '0000'
        )
        SELECT m.lon, m.lat, m.state, m.id as station, m.name as name,
        case when (precip_estimated or precip_hour is null or
            precip_hour < 4 or precip_hour > 11) then null
            else precip end as precipdata,
        snow as snowdata, snowd as snowddata,
        case when temp_estimated then null else high end as highdata_all,
        case when temp_estimated then null else low end as lowdata_all
        from alldata a JOIN mystations m
        ON (a.station = m.id) WHERE a.day = :ts
        """),
            conn,
            params={
                "west": gridnav.left - mybuf,
                "south": gridnav.bottom - mybuf,
                "east": gridnav.right + mybuf,
                "north": gridnav.top + mybuf,
                "ts": ts,
            },
        )
    LOG.info("loaded %s rows from climodat database", len(df.index))
    if len(df.index) < 50:
        if ts != date.today():
            LOG.warning("Failed quorum")
        return
    if ts.year < 1951:
        res = generic_gridder(df, "highdata_all")
        ds["high_tmpk_12z"].values = convert_value(res, "degF", "degK")

        res = generic_gridder(df, "lowdata_all")
        ds["low_tmpk_12z"].values = convert_value(res, "degF", "degK")

    res = generic_gridder(df, "snowdata")
    if res is not None and ts < date(2008, 10, 1):  # NOHRSC covers
        ds["snow_12z"].values = convert_value(res, "inch", "millimeter")

    res = generic_gridder(df, "snowddata")
    if res is not None:
        ds["snowd_12z"].values = convert_value(res, "inch", "millimeter")


def use_asos_daily(ts, ds, domain):
    """Grid out available ASOS data."""
    mybuf = 2.0
    gridnav = get_nav("iemre", domain)
    with get_sqlalchemy_conn("iem") as conn:
        df = pd.read_sql(
            sql_helper(
                """
           SELECT ST_x(s.geom) as lon, ST_y(s.geom) as lat, s.state,
           s.name, s.id as station,
           (CASE WHEN pday >= 0 then pday else null end) as precipdata,
           (CASE WHEN max_tmpf > -50 and max_tmpf < 130
               then max_tmpf else null end) as highdata,
           (CASE WHEN min_tmpf > -50 and min_tmpf < 95
               then min_tmpf else null end) as lowdata,
           (CASE WHEN max_dwpf > -50 and max_dwpf < 130
               then max_dwpf else null end) as highdwpf,
           (CASE WHEN min_dwpf > -50 and min_dwpf < 95
               then min_dwpf else null end) as lowdwpf,
            (CASE WHEN min_rh > 0 and min_rh < 101
             then min_rh else null end) as minrh,
            (CASE WHEN max_rh > 0 and max_rh < 101
             then max_rh else null end) as maxrh
           from {table} c, stations s WHERE day = :ts and
           ST_Contains(ST_MakeEnvelope(:west, :south, :east, :north, 4326),
           geom) and s.network ~* 'ASOS'
           and c.iemid = s.iemid
            """,
                table=f"summary_{ts:%Y}",
            ),
            conn,
            params={
                "west": gridnav.left - mybuf,
                "south": gridnav.bottom - mybuf,
                "east": gridnav.right + mybuf,
                "north": gridnav.top + mybuf,
                "ts": ts,
            },
        )
    if len(df.index) < 4:
        LOG.warning("Failed data quorum")
        return
    if len(df.index) > 300:
        LOG.info("Using ASOS for high/low")
        res = generic_gridder(df, "highdata", domain)
        ds["high_tmpk"].values = convert_value(res, "degF", "degK")
        res = generic_gridder(df, "lowdata", domain)
        ds["low_tmpk"].values = convert_value(res, "degF", "degK")

    # We have alternative
    hres = generic_gridder(df, "highdwpf", domain)
    lres = generic_gridder(df, "lowdwpf", domain)
    if hres is not None and lres is not None:
        ds["avg_dwpk"].values = convert_value(
            (hres + lres) / 2.0, "degF", "degK"
        )
    res = generic_gridder(df, "minrh", domain)
    if res is not None:
        ds["min_rh"].values = res
    res = generic_gridder(df, "maxrh", domain)
    if res is not None:
        ds["max_rh"].values = res


def use_climodat_daily(ts: date, ds):
    """Do our gridding"""
    mybuf = 2.0
    gridnav = get_nav("iemre")
    with get_sqlalchemy_conn("coop") as conn:
        df = pd.read_sql(
            sql_helper("""
        WITH mystations as (
            SELECT id, ST_X(geom) as lon, ST_Y(geom) as lat, state, name
            from stations where ST_Contains(
            ST_MakeEnvelope(:west, :south, :east, :north, 4326), geom) and
            network ~* 'CLIMATE' and substr(id, 3, 1) not in ('C', 'D', 'T')
            and substr(id, 3, 4) != '0000'
        )
        SELECT m.lon, m.lat, m.state, m.id as station, m.name as name,
        case when not temp_estimated then high else null end as highdata_all,
        case when not temp_estimated then low else null end as lowdata_all,
        case when not precip_estimated then precip else null end as
            precipdata_all,
        case when temp_estimated or temp_hour is null or high < -50 or
            high > 130 or (temp_hour > 4 and temp_hour < 13) then null
            else high end as highdata,
        case when temp_estimated or temp_hour is null or low < -90 or low > 90
            or (temp_hour > 4 and temp_hour < 13) then null else low end
            as lowdata,
        case when precip_estimated or precip_hour is null or
        (precip_hour > 4 and precip_hour < 13) then null else precip end
            as precipdata
        from alldata a JOIN mystations m
        ON (a.station = m.id) WHERE a.day = :ts
        """),
            conn,
            params={
                "west": gridnav.left - mybuf,
                "south": gridnav.bottom - mybuf,
                "east": gridnav.right + mybuf,
                "north": gridnav.top + mybuf,
                "ts": ts,
            },
        )
    if len(df.index) < 4:
        if ts != date.today():
            LOG.warning("Failed quorum")
        return
    suffix = "_all" if ts.year < 1951 else ""
    res = generic_gridder(df, f"highdata{suffix}")
    if res is not None:
        ds["high_tmpk"].values = convert_value(res, "degF", "degK")
    res = generic_gridder(df, f"lowdata{suffix}")
    if res is not None:
        ds["low_tmpk"].values = convert_value(res, "degF", "degK")
    res = generic_gridder(df, f"precipdata{suffix}")
    if res is not None:
        ds["p01d"].values = convert_value(res, "inch", "mm")


def workflow(ts: date, domain: str, forecast: bool):
    """Do Work"""
    # load up our current data
    ds = get_grids(ts, domain=domain)
    LOG.info("loaded %s variables from IEMRE database", len(ds))

    # rsds -> grid_rsds.py
    # power_swdn -> ../dl/fetch_power.py

    # high_tmpk, low_tmpk, p01d, wind_speed, min_rh, max_rh, high_soil4t,
    # low_soil4t, high_tmpk_12z low_tmpk_12z p01d_12z
    if ts.year > 1949:
        copy_iemre_hourly(ts, ds, domain)
    else:
        # While data does exist, we have practical availability issues
        if ts.year > 1940:
            LOG.info("Using ASOS for daily summary variables")
            use_asos_daily(ts, ds, domain)
        if domain == "conus":
            use_climodat_daily(ts, ds)

    if domain == "conus" and not forecast:
        # snow_12z snowd_12z
        use_climodat_12z(ts, ds)

    for vname in list(ds.keys()):
        # Some grids are too tight to the land boundary, so we smear as a hack
        vals = ds[vname].to_masked_array()
        vals = grid_smear(vals, shift=4)
        ds[vname].values = vals
        # Compute the precentage of points within vals that are missing/NaN
        pct_missing = 100.0 * vals.mask.sum() / vals.size
        msg = (
            f"{vname:14s} {ds[vname].min():6.2f} {ds[vname].max():6.2f} "
            f"{pct_missing:6.2f}% missing"
        )
        LOG.info(msg)
    set_grids(ts, ds, domain=domain)
    subprocess.call(
        [
            "python",
            "db_to_netcdf.py",
            f"--date={ts:%Y-%m-%d}",
            f"--domain={domain}",
        ]
    )


@click.command()
@click.option(
    "--date", "valid", type=click.DateTime(), help="Date", required=True
)
@click.option("--domain", default="conus", help="Domain to process", type=str)
@click.option(
    "--forecast",
    is_flag=True,
    help="This is a forecast run, so we only use IEMRE hourly.",
)
def main(valid: datetime, domain: str, forecast: bool):
    """Go Main Go"""
    dt = valid.date()
    LOG.info("Run %s for domain: '%s' forecast: %s", dt, domain, forecast)
    workflow(dt, domain, forecast)
    LOG.info("Done.")


if __name__ == "__main__":
    main()
