from pathlib import Path
import json
import warnings

import numpy as np
import xarray as xr

try:
    import cfgrib
except ModuleNotFoundError:  # pragma: no cover - dependency guard for runtime environments
    cfgrib = None

PROJECT_ROOT = Path(__file__).resolve().parent.parent
INPUT_DIR = PROJECT_ROOT / "data" / "grib"
OUTPUT_DIR = PROJECT_ROOT / "data" / "grid"

OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

REGION = {
    "lon_min": 120,
    "lon_max": 150,
    "lat_min": 20,
    "lat_max": 50,
}

VARIABLE_CANDIDATES = {
    "u10": ("u10", "u"),
    "v10": ("v10", "v"),
    "tp": ("tp", "prate", "unknown"),
    "msl": ("prmsl", "msl", "pres"),
    "temp": ("t2m", "t"),
    "wave_h": ("htsgw", "swh", "wvhgt"),
    "wave_dir": ("dirpw", "mwd", "wvdir"),
    "wave_period": ("perpw", "mwp", "wvper"),
}


def compute_wind(u_component, v_component):
    speed = np.sqrt(u_component ** 2 + v_component ** 2)
    # Meteorological direction: where the wind comes from, in degrees clockwise from north.
    direction = (270 - np.degrees(np.arctan2(v_component, u_component))) % 360
    return speed, direction


def load_datasets(path):
    if cfgrib is None:
        raise RuntimeError("cfgrib is required to build grids from GRIB2 files")
    with xr.set_options(use_new_combine_kwarg_defaults=True):
        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore",
                message="In a future version of xarray the default value for compat will change",
                category=FutureWarning,
            )
            return cfgrib.xarray_store.open_datasets(str(path))


def find_variable(datasets, candidates):
    for candidate in candidates:
        for dataset in datasets:
            if candidate in dataset:
                return dataset[candidate]
    return None


def get_lat_lon(datasets):
    for dataset in datasets:
        if "latitude" in dataset and "longitude" in dataset:
            return dataset["latitude"].values, dataset["longitude"].values
    raise ValueError("no latitude/longitude coordinates found in GRIB datasets")


def to_2d_values(data_array, lat_size, lon_size):
    if data_array is None:
        return np.zeros((lat_size, lon_size), dtype=float)

    values = np.asarray(data_array.squeeze().values)
    if values.ndim != 2:
        raise ValueError(f"expected 2D field, got shape {values.shape} for {data_array.name}")
    if values.shape != (lat_size, lon_size):
        raise ValueError(
            f"field {data_array.name} shape {values.shape} does not match coordinates {(lat_size, lon_size)}"
        )
    return values


def normalize_longitudes(lon_values, fields):
    lon_values = np.asarray(lon_values, dtype=float)
    normalized_lon = ((lon_values + 180) % 360) - 180
    sort_idx = np.argsort(normalized_lon)
    normalized_lon = normalized_lon[sort_idx]
    normalized_fields = [field[:, sort_idx] for field in fields]
    return normalized_lon, normalized_fields


def select_region(lat_values, lon_values, fields):
    lat_mask = (lat_values >= REGION["lat_min"]) & (lat_values <= REGION["lat_max"])
    lon_mask = (lon_values >= REGION["lon_min"]) & (lon_values <= REGION["lon_max"])

    lat_idx = np.where(lat_mask)[0]
    lon_idx = np.where(lon_mask)[0]
    if lat_idx.size == 0 or lon_idx.size == 0:
        raise ValueError("selected region is empty after coordinate filtering")

    return (
        lat_values[lat_idx],
        lon_values[lon_idx],
        [field[np.ix_(lat_idx, lon_idx)] for field in fields],
    )


def process_file(path):
    datasets = load_datasets(path)
    try:
        wave_path = path.with_name(f"{path.stem}_wave.grib2")
        wave_datasets = load_datasets(wave_path) if wave_path.exists() else []
        all_datasets = [*datasets, *wave_datasets]

        lat_values, lon_values = get_lat_lon(datasets)
        lat_size = len(lat_values)
        lon_size = len(lon_values)

        u10 = to_2d_values(find_variable(all_datasets, VARIABLE_CANDIDATES["u10"]), lat_size, lon_size)
        v10 = to_2d_values(find_variable(all_datasets, VARIABLE_CANDIDATES["v10"]), lat_size, lon_size)
        rain = to_2d_values(find_variable(all_datasets, VARIABLE_CANDIDATES["tp"]), lat_size, lon_size)
        pressure = to_2d_values(find_variable(all_datasets, VARIABLE_CANDIDATES["msl"]), lat_size, lon_size)
        temp = to_2d_values(find_variable(all_datasets, VARIABLE_CANDIDATES["temp"]), lat_size, lon_size)
        wave_h = to_2d_values(find_variable(all_datasets, VARIABLE_CANDIDATES["wave_h"]), lat_size, lon_size)
        wave_dir = to_2d_values(find_variable(all_datasets, VARIABLE_CANDIDATES["wave_dir"]), lat_size, lon_size)
        wave_period = to_2d_values(find_variable(all_datasets, VARIABLE_CANDIDATES["wave_period"]), lat_size, lon_size)

        wind_speed, wind_dir = compute_wind(u10, v10)

        lon_values, normalized_fields = normalize_longitudes(
            lon_values,
            [wind_speed, wind_dir, rain, temp, pressure, wave_h, wave_dir, wave_period],
        )
        wind_speed, wind_dir, rain, temp, pressure, wave_h, wave_dir, wave_period = normalized_fields

        lat_region, lon_region, region_fields = select_region(
            lat_values,
            lon_values,
            [wind_speed, wind_dir, rain, temp, pressure, wave_h, wave_dir, wave_period],
        )
        wind_speed, wind_dir, rain, temp, pressure, wave_h, wave_dir, wave_period = region_fields

        return {
            "lat": lat_region.tolist(),
            "lon": lon_region.tolist(),
            "wind_speed": wind_speed.tolist(),
            "wind_dir": wind_dir.tolist(),
            "rain": rain.tolist(),
            "temp": temp.tolist(),
            "pressure": pressure.tolist(),
            "wave_h": wave_h.tolist(),
            "wave_dir": wave_dir.tolist(),
            "wave_period": wave_period.tolist(),
        }
    finally:
        for dataset in datasets:
            dataset.close()
        if 'wave_datasets' in locals():
            for dataset in wave_datasets:
                dataset.close()


def main():
    for path in sorted(INPUT_DIR.glob("*.grib2")):
        if path.stem.endswith("_wave"):
            continue
        print("processing", path.name)
        try:
            grid = process_file(path)
        except Exception as exc:
            print("  error processing", path.name, exc)
            continue

        output_path = OUTPUT_DIR / f"grid_{path.stem}.json"
        payload = {
            "time": path.stem,
            "grid": grid,
        }
        with output_path.open("w", encoding="utf-8") as file_handle:
            json.dump(payload, file_handle)
        print("saved", output_path)


if __name__ == "__main__":
    main()
