"""
Gridcell sample the NLDAS NetCDF files to save srad to climodat database.

Run from RUN_0Z.sh for six UTC days ago.
"""

from datetime import datetime, timedelta

import click
import geopandas as gpd
import numpy as np
import pandas as pd
from affine import Affine
from pyiem.database import get_dbconnc, get_sqlalchemy_conn
from pyiem.grid.zs import CachingZonalStats
from pyiem.iemre import hourly_offset
from pyiem.util import convert_value, logger, ncopen, utc

LOG = logger()


def compute_regions(data, varname, df):
    """Do the spatial averaging work."""
    with get_sqlalchemy_conn("coop") as conn:
        gdf = gpd.read_postgis(
            """
            SELECT t.id, c.geom from stations t JOIN climodat_regions c on
            (t.iemid = c.iemid) ORDER by t.id ASC
            """,
            conn,
            index_col="id",
            geom_col="geom",
        )
    affine = Affine(0.125, 0, -125.0, 0, -0.125, 53.0)
    czs = CachingZonalStats(affine)
    data = czs.gen_stats(np.flipud(data), gdf["geom"])
    for i, sid in enumerate(gdf.index.values):
        df.at[sid, varname] = data[i]


def build_stations(dt) -> pd.DataFrame:
    """Figure out what we need data for."""
    with get_sqlalchemy_conn("coop") as conn:
        # NLDAS does not cover PR
        df = pd.read_sql(
            """
            SELECT station, st_x(geom) as lon, st_y(geom) as lat, temp_hour
            from alldata a JOIN stations t on (a.station = t.id) WHERE
            t.network ~* 'CLIMATE' and a.day = %s and
            st_x(geom) between -127 and -68
            ORDER by station ASC
            """,
            conn,
            params=(dt,),
            index_col="station",
        )
    for col in [
        "nldas_soilt4_min",
        "nldas_soilt4_avg",
        "nldas_soilt4_max",
        "nldas_soilm4_avg",
        "nldas_soilm1m_avg",
    ]:
        df[col] = np.nan
    df["i"] = np.nan
    df["j"] = np.nan
    LOG.info("Found %s database entries", len(df.index))
    return df


def compute(df, sids, dt, do_regions=False):
    """Do the magic."""
    # Life choice is to run 6z to 6z
    sts = utc(dt.year, dt.month, dt.day, 6)
    ets = sts + timedelta(hours=24)

    ncfn = f"/mesonet/data/nldas/{sts.year}_hourly.nc"
    idx0 = hourly_offset(sts)
    idx1 = hourly_offset(ets)
    with ncopen(ncfn) as nc:
        lons = nc.variables["lon"][:]
        lats = nc.variables["lat"][:]
        tsel = slice(idx0, idx1)
        if f"{dt:%m%d}" == "1231":
            # Close enough
            tsel = slice(idx0, None)
        soilm = np.mean(nc.variables["soilm"][tsel, 0], 0)
        soilm1m = (
            np.mean(nc.variables["soilm"][tsel, 0], 0) * 10.0
            + np.mean(nc.variables["soilm"][tsel, 1], 0) * 30.0
            + np.mean(nc.variables["soilm"][tsel, 2], 0) * 60.0
        ) / 100.0
        soiltmin = np.min(nc.variables["soilt"][tsel, 0], 0)
        soiltmax = np.max(nc.variables["soilt"][tsel, 0], 0)
        soiltavg = np.mean(nc.variables["soilt"][tsel, 0], 0)

    df["i"] = np.digitize(df["lon"].values, lons)
    df["j"] = np.digitize(df["lat"].values, lats)
    soilm = soilm.filled(np.nan)
    soilm1m = soilm1m.filled(np.nan)
    soiltmin = soiltmin.filled(np.nan)
    soiltmax = soiltmax.filled(np.nan)
    soiltavg = soiltavg.filled(np.nan)

    for sid, row in df.loc[sids].iterrows():
        df.at[sid, "nldas_soilt4_avg"] = soiltavg[int(row["j"]), int(row["i"])]
        df.at[sid, "nldas_soilt4_min"] = soiltmin[int(row["j"]), int(row["i"])]
        df.at[sid, "nldas_soilt4_max"] = soiltmax[int(row["j"]), int(row["i"])]
        df.at[sid, "nldas_soilm4_avg"] = soilm[int(row["j"]), int(row["i"])]
        df.at[sid, "nldas_soilm1m_avg"] = soilm1m[int(row["j"]), int(row["i"])]

    if do_regions:
        compute_regions(soiltavg, "nldas_soilt4_avg", df)
        compute_regions(soiltmin, "nldas_soilt4_min", df)
        compute_regions(soiltmax, "nldas_soilt4_max", df)
        compute_regions(soilm, "nldas_soilm4_avg", df)
        compute_regions(soilm1m, "nldas_soilm1m_avg", df)

    if "IA0200" in df.index:
        LOG.info("IA0200 %s", df.loc["IA0200"])


def do(dt):
    """Process for a given date."""
    LOG.info("do(%s)", dt)
    df = build_stations(dt)
    if df.empty:
        LOG.warning("Aborting as no database obs to update for %s", dt)
        return
    df["day"] = dt
    # We currently do two options
    # 1. For morning sites 1-11 AM, they get yesterday's values
    sids = df[(df["temp_hour"] > 0) & (df["temp_hour"] < 12)].index.values
    compute(df, sids, dt - timedelta(days=1), True)
    # 2. All other sites get today
    sids = df[df["nldas_soilt4_avg"].isna()].index.values
    compute(df, sids, dt)

    df["station"] = df.index.values
    df["nldas_soilt4_avg"] = convert_value(
        df["nldas_soilt4_avg"].values, "degK", "degF"
    )
    df["nldas_soilt4_min"] = convert_value(
        df["nldas_soilt4_min"].values, "degK", "degF"
    )
    df["nldas_soilt4_max"] = convert_value(
        df["nldas_soilt4_max"].values, "degK", "degF"
    )

    # prevent NaN from being inserted
    df = df.replace({np.nan: None})
    pgconn, cursor = get_dbconnc("coop")
    cursor.executemany(
        """
        UPDATE alldata set
        nldas_soilt4_avg = %(nldas_soilt4_avg)s,
        nldas_soilt4_min = %(nldas_soilt4_min)s,
        nldas_soilt4_max = %(nldas_soilt4_max)s,
        nldas_soilm4_avg = %(nldas_soilm4_avg)s,
        nldas_soilm1m_avg = %(nldas_soilm1m_avg)s
        where station = %(station)s and day = %(day)s
        """,
        df.to_dict("records"),
    )
    cursor.close()
    pgconn.commit()


@click.command()
@click.option("--valid", type=click.DateTime(), required=True)
def main(valid: datetime):
    """Go Main Go"""
    do(valid.date())


if __name__ == "__main__":
    main()
