"""
Utilities to locate and download GOES-16 ABI L1b radiance data
from the public NOAA S3 bucket, and generate GOES-16 Ash RGB imagery.
"""

import os
import re
from datetime import datetime, timedelta
from typing import Optional, Dict, Tuple, Iterable, List
from concurrent.futures import ThreadPoolExecutor, as_completed
import s3fs
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
from pyresample import geometry, kd_tree

# -----------------------------------------------------------------------------
# Configuration
# -----------------------------------------------------------------------------

DEFAULT_BUCKET = "noaa-goes16"
DEFAULT_PRODUCT = "ABI-L1b-RadF"
DEFAULT_ASH_CHANNELS = ("11", "13", "14", "15")
DEFAULT_MAX_TIME_DIFF = timedelta(minutes=30)

DEFAULT_OUT_ROOT = "images/goes16_l1b"
DEFAULT_ASH_RGB_ROOT = "images/goes16_ash_rgb"

DEFAULT_ASH_LAT_BOUNDS = (-40.0, 70.0)
DEFAULT_ASH_LON_BOUNDS = (-120.0, -30.0)
DEFAULT_ASH_RES_DEG = 0.02  # ~2 km

DEBUG_GOES16 = True
SCAN_RE = re.compile(r"_s(\d{4})(\d{3})(\d{2})(\d{2})(\d{2})")


def _log(msg: str):
    if DEBUG_GOES16:
        print(f"[GOES16] {msg}")


# -----------------------------------------------------------------------------
# Time helpers
# -----------------------------------------------------------------------------

def target_datetime(year: int, month: int, day: int, hhmm: str) -> datetime:
    hour = int(hhmm[:2])
    minute = int(hhmm[2:])
    dt = datetime(year, month, day, hour, minute)
    _log(f"Target datetime: {dt.isoformat()}Z")
    return dt


def day_of_year(dt: datetime) -> int:
    return (dt - datetime(dt.year, 1, 1)).days + 1


# -----------------------------------------------------------------------------
# Filename parsing
# -----------------------------------------------------------------------------

def extract_scan_time_from_name(name: str) -> Optional[datetime]:
    m = SCAN_RE.search(name)
    if not m:
        return None
    year, jjj, hh, mm, ss = map(int, m.groups())
    base = datetime(year, 1, 1) + timedelta(days=jjj - 1)
    return datetime(base.year, base.month, base.day, hh, mm, ss)


# -----------------------------------------------------------------------------
# S3 helpers
# -----------------------------------------------------------------------------

def _build_s3_filesystem():
    return s3fs.S3FileSystem(
        anon=True,
        config_kwargs={"max_pool_connections": 64},
    )


def _candidate_hour_prefixes(target, max_time_diff, bucket, product):
    if max_time_diff is None:
        doy = day_of_year(target)
        return [f"{bucket}/{product}/{target.year:04d}/{doy:03d}/"]

    tmin = target - max_time_diff
    tmax = target + max_time_diff

    prefixes = []
    current = datetime(tmin.year, tmin.month, tmin.day, tmin.hour)
    end = datetime(tmax.year, tmax.month, tmax.day, tmax.hour)

    seen = set()
    while current <= end:
        doy = day_of_year(current)
        prefix = f"{bucket}/{product}/{current.year:04d}/{doy:03d}/{current.hour:02d}/"
        if prefix not in seen:
            seen.add(prefix)
            prefixes.append(prefix)
        current += timedelta(hours=1)
    return prefixes


# -----------------------------------------------------------------------------
# Find matching L1b files
# -----------------------------------------------------------------------------

def find_ash_keys_for_datetime(
    year, month, day, hhmm,
    *,
    channels=DEFAULT_ASH_CHANNELS,
    bucket=DEFAULT_BUCKET,
    product=DEFAULT_PRODUCT,
    max_time_diff=DEFAULT_MAX_TIME_DIFF,
):
    target = target_datetime(year, month, day, hhmm)
    fs = _build_s3_filesystem()
    prefixes = _candidate_hour_prefixes(target, max_time_diff, bucket, product)

    ch_to_candidates = {ch: [] for ch in channels}

    def _glob_one(prefix, ch):
        patterns = [
            f"{prefix}OR_{product}-M6C{ch}_*.nc",
            f"{prefix}*{product}-M6C{ch}_*.nc",
        ]
        out = []
        for pat in patterns:
            try:
                out.extend(fs.glob(pat))
            except FileNotFoundError:
                pass
        return out

    with ThreadPoolExecutor(max_workers=32) as ex:
        futures = {
            ex.submit(_glob_one, p if p.endswith("/") else p + "/", ch): ch
            for p in prefixes for ch in channels
        }

        for fut in as_completed(futures):
            ch = futures[fut]
            try:
                ch_to_candidates[ch].extend(fut.result())
            except Exception:
                pass

    keys = {}
    for ch, candidates in ch_to_candidates.items():
        if not candidates:
            raise FileNotFoundError(f"No files found for band {ch}")

        best_file = None
        best_diff = None
        best_time = None

        for fpath in candidates:
            t = extract_scan_time_from_name(os.path.basename(fpath))
            if t is None:
                continue
            diff = abs(t - target)
            if best_diff is None or diff < best_diff:
                best_diff = diff
                best_file = fpath
                best_time = t

        if best_file is None:
            raise FileNotFoundError(f"No valid file for band {ch}")

        keys[ch] = best_file
        _log(f"Band {ch}: using file {os.path.basename(best_file)}")

    return keys, target


# -----------------------------------------------------------------------------
# Download L1b files (always used)
# -----------------------------------------------------------------------------

def download_ash_bands_for_datetime(
    year, month, day, hhmm,
    *,
    out_root=DEFAULT_OUT_ROOT,
    channels=DEFAULT_ASH_CHANNELS,
    bucket=DEFAULT_BUCKET,
    product=DEFAULT_PRODUCT,
    max_time_diff=DEFAULT_MAX_TIME_DIFF,
):
    keys, target_dt = find_ash_keys_for_datetime(
        year, month, day, hhmm,
        channels=channels,
        bucket=bucket,
        product=product,
        max_time_diff=max_time_diff,
    )

    fs = _build_s3_filesystem()

    out_dir = os.path.join(out_root, f"{year:04d}", f"{month:02d}", f"{day:02d}", hhmm)
    os.makedirs(out_dir, exist_ok=True)

    local_paths = {}
    for ch, key in keys.items():
        fname = os.path.basename(key)
        local_path = os.path.join(out_dir, fname)

        if os.path.exists(local_path):
            _log(f"Using cached file for band {ch}: {local_path}")
        else:
            _log(f"Downloading band {ch} → {local_path}")
            fs.get(key, local_path)

        local_paths[ch] = local_path

    return local_paths, target_dt


# -----------------------------------------------------------------------------
# Load L1b radiances (LOCAL-ONLY)
# -----------------------------------------------------------------------------

def _load_bt(path: str):
    _log(f"Loading brightness temperature from {path}")

    with xr.open_dataset(path) as ds:
        rad = ds["Rad"]
        if "t" in rad.dims:
            rad = rad.isel(t=0)
        rad = rad.values

        fk1 = float(ds["planck_fk1"])
        fk2 = float(ds["planck_fk2"])
        bc1 = float(ds["planck_bc1"])
        bc2 = float(ds["planck_bc2"])

        bt = fk2 / np.log(fk1 / (rad + bc1) + 1.0) + bc2

        proj = ds["goes_imager_projection"]
        x = ds["x"]
        y = ds["y"]

    return bt.astype(np.float32), proj, x, y


# -----------------------------------------------------------------------------
# RGB construction
# -----------------------------------------------------------------------------

def _goes_latlon(proj, x, y):
    _log("Converting GOES fixed grid to lat/lon")
    Re = float(proj.semi_major_axis)
    Rp = float(proj.semi_minor_axis)
    H = float(proj.perspective_point_height) + Re
    lon0 = np.deg2rad(float(proj.longitude_of_projection_origin))

    xx, yy = np.meshgrid(x.values, y.values)

    sinx = np.sin(xx)
    cosx = np.cos(xx)
    siny = np.sin(yy)
    cosy = np.cos(yy)

    a = sinx**2 + cosx**2 * (cosy**2 + (Re**2 / Rp**2) * siny**2)
    b = -2 * H * cosx * cosy
    c = H**2 - Re**2
    disc = b**2 - 4 * a * c

    mask = disc <= 0
    disc = np.where(mask, np.nan, disc)

    rs = (-b - np.sqrt(disc)) / (2 * a)

    Sx = rs * cosx * cosy
    Sy = rs * sinx * cosy
    Sz = rs * siny

    lat = np.arctan((Re**2 / Rp**2) * (Sz / np.sqrt((H - Sx)**2 + Sy**2)))
    lon = lon0 - np.arctan2(Sy, H - Sx)

    latdeg = np.rad2deg(lat)
    londeg = np.rad2deg(lon)

    latdeg[mask] = np.nan
    londeg[mask] = np.nan
    return londeg, latdeg


def _clip_scale(arr, vmin, vmax):
    return np.clip((arr - vmin) / (vmax - vmin), 0, 1)


def _build_area_def(latmin, latmax, lonmin, lonmax, res):
    _log("Building target grid")
    lats = np.arange(latmin, latmax + res, res)
    lons = np.arange(lonmin, lonmax + res, res)
    width, height = len(lons), len(lats)

    return geometry.AreaDefinition(
        "ash_latlon",
        "GOES-16 Ash RGB",
        "latlon",
        {"proj": "longlat", "datum": "WGS84"},
        width,
        height,
        (lonmin, latmin, lonmax, latmax),
    )


def _resample_channel(channel, swath, area):
    return kd_tree.resample_nearest(
        swath, channel, area,
        radius_of_influence=8000,
        fill_value=np.nan
    )


def build_ash_rgb_from_paths(
    paths, *,
    lat_bounds=DEFAULT_ASH_LAT_BOUNDS,
    lon_bounds=DEFAULT_ASH_LON_BOUNDS,
    res_deg=DEFAULT_ASH_RES_DEG,
):
    bt11, proj, x, y = _load_bt(paths["11"])
    bt13, _, _, _ = _load_bt(paths["13"])
    bt14, _, _, _ = _load_bt(paths["14"])
    bt15, _, _, _ = _load_bt(paths["15"])

    lon, lat = _goes_latlon(proj, x, y)

    red = _clip_scale(bt15 - bt13, -4, 2)
    green = _clip_scale(bt14 - bt11, -4, 5)
    blue = _clip_scale(bt13, 243, 303)

    mask = np.isnan(lat) | np.isnan(lon)
    red[mask] = np.nan
    green[mask] = np.nan
    blue[mask] = np.nan

    latmin, latmax = lat_bounds
    lonmin, lonmax = lon_bounds

    area = _build_area_def(latmin, latmax, lonmin, lonmax, res_deg)
    swath = geometry.SwathDefinition(lons=lon, lats=lat)

    Rr = _resample_channel(red, swath, area)
    Gr = _resample_channel(green, swath, area)
    Br = _resample_channel(blue, swath, area)

    rgb = np.dstack([
        np.nan_to_num(Rr),
        np.nan_to_num(Gr),
        np.nan_to_num(Br)
    ])
    return np.clip(rgb, 0, 1)


# -----------------------------------------------------------------------------
# NEW: Geolocations of each pixel on the Ash RGB grid
# -----------------------------------------------------------------------------
def get_ash_rgb_pixel_geolocations() -> Tuple[np.ndarray, np.ndarray]:
    """
    Return the latitude/longitude of the *center* of each Ash RGB pixel.

    Assumes:
      - the Ash RGB grid is defined over
        DEFAULT_ASH_LAT_BOUNDS, DEFAULT_ASH_LON_BOUNDS
      - the grid size is the same as used in _build_area_def
      - images are plotted with extent = [lon_min, lon_max, lat_min, lat_max]

    Returns
    -------
    lat_grid : (H, W) array of float
    lon_grid : (H, W) array of float
        These shapes should match rgb.shape[:2].
    """
    lat_min, lat_max = DEFAULT_ASH_LAT_BOUNDS
    lon_min, lon_max = DEFAULT_ASH_LON_BOUNDS
    res = DEFAULT_ASH_RES_DEG

    # Reproduce the target grid size used in _build_area_def
    # (this is what ultimately sets rgb.shape)
    lats_edge = np.arange(lat_min, lat_max + res, res)
    lons_edge = np.arange(lon_min, lon_max + res, res)
    height = len(lats_edge)     # number of rows (y)
    width = len(lons_edge)      # number of cols (x)

    # Given Matplotlib's imshow + extent semantics and pyresample's AreaDefinition,
    # the pixel centers are offset half a pixel inside the outer extent.
    dlat = (lat_max - lat_min) / height
    dlon = (lon_max - lon_min) / width

    lat_centers = lat_min + (np.arange(height) + 0.5) * dlat
    lon_centers = lon_min + (np.arange(width) + 0.5) * dlon

    lon_grid, lat_grid = np.meshgrid(lon_centers, lat_centers)

    return lat_grid, lon_grid



# -----------------------------------------------------------------------------
# High-level function (ALWAYS downloads)
# -----------------------------------------------------------------------------

def generate_ash_rgb_for_datetime(
    year, month, day, hhmm,
    *,
    l1b_root=DEFAULT_OUT_ROOT,
    rgb_root=DEFAULT_ASH_RGB_ROOT,
    lat_bounds=DEFAULT_ASH_LAT_BOUNDS,
    lon_bounds=DEFAULT_ASH_LON_BOUNDS,
    res_deg=DEFAULT_ASH_RES_DEG,
    save_png=True,
):

    _log(f"Generating Ash RGB for {year}-{month}-{day} {hhmm}Z")

    local_paths, target_dt = download_ash_bands_for_datetime(
        year, month, day, hhmm,
        out_root=l1b_root,
        channels=DEFAULT_ASH_CHANNELS,
        bucket=DEFAULT_BUCKET,
        product=DEFAULT_PRODUCT,
        max_time_diff=DEFAULT_MAX_TIME_DIFF,
    )

    rgb = build_ash_rgb_from_paths(
        local_paths,
        lat_bounds=lat_bounds,
        lon_bounds=lon_bounds,
        res_deg=res_deg,
    )

    res_str = str(res_deg).replace(".", "p")
    out_dir = os.path.join(rgb_root, f"{year:04d}", f"{month:02d}", f"{day:02d}", hhmm)
    os.makedirs(out_dir, exist_ok=True)

    png_path = os.path.join(out_dir, f"ash_rgb_{res_str}deg.png")

    if save_png:
        plt.imsave(png_path, rgb)

    return png_path, rgb, target_dt


__all__ = [
    "target_datetime",
    "day_of_year",
    "extract_scan_time_from_name",
    "find_ash_keys_for_datetime",
    "download_ash_bands_for_datetime",
    "build_ash_rgb_from_paths",
    "generate_ash_rgb_for_datetime",
    "get_ash_rgb_pixel_geolocations",
]
