#!/usr/bin/env python3

import xarray as xr
import numpy as np
import pandas as pd
import sys

month=f"{sys.argv[1]}".zfill(2)
print(month)

# ==================================================
# Configuration
# ==================================================

INPUT = f"2024{month}_map.nc4"
OUTPUT = f"2024{month}_ONERA_NA_domain.nc4"

LAT_MIN = 20
LAT_MAX = 80

LON_MIN = -90
LON_MAX = 30

LEVEL_RANGE = [
    100, 125, 150, 175, 200, 225,
    250, 300, 350, 400, 450, 500
]


# ==================================================
# Conversion ft -> hPa
# ==================================================

def ft_to_hpa(alt_ft):
    alt_m = alt_ft * 0.3048
    return (
        1013.25
        * (1 - 2.2557695644629538e-5 * alt_m)
        ** 5.255932362359814
    )


# ==================================================
# Open lazily
# ==================================================

print("Opening file...")

ds = xr.open_dataset(
    INPUT,
    decode_times=False,
    chunks="auto"
)

####ds = ds.chunk(
####    {
####        "hours": 24,
####        "alt": 1,
####        "lat": 61,
####        "lon": 121
####    }
####)

# Remove first and last day (spin-up / overlap)
#ds = ds.isel(hours=slice(24, -24))


# ==================================================
# Keep only required variables
# ==================================================

print("Selecting variables...")

ds = ds[
    [
        "Number_of_kms",
        "Number_of_planes"
    ]
]


# ==================================================
# Spatial selection
# ==================================================

print("Selecting spatial domain...")

ds = ds.sel(
    lat=slice(LAT_MIN, LAT_MAX),
    lon=slice(LON_MIN, LON_MAX)
)

# ==================================================
# Convert altitude ft -> pressure hPa
# Keep all native levels between 500 and 100 hPa
# ==================================================

print("Converting altitude to pressure...")

pressure = ft_to_hpa(ds.alt.values)

# Add pressure as a coordinate
ds = ds.assign_coords(
    pressure=("alt", pressure)
)

# Select native levels between 500 and 100 hPa
print("Selecting pressure range 500-100 hPa...")

ds = ds.where(
    (ds.pressure <= 500) & (ds.pressure >= 100),
    drop=True
)

# Rename altitude dimension to level
ds = ds.rename(
    {"alt": "level"}
)

# Replace vertical coordinate with pressure
ds = ds.assign_coords(
    level=ds.pressure
)

ds = ds.drop_vars("pressure")

ds.level.attrs = {
    "standard_name": "air_pressure",
    "long_name": "pressure level",
    "units": "hPa"
}

# ==================================================
# Time axis
# ==================================================

print("Creating time coordinate...")

ntime = ds.sizes["hours"]

time = pd.date_range(
    start=f"2024-{month}-01 00:00:00",
    periods=ntime,
    freq="1h"
)

ds = ds.rename({"hours": "time"})
ds = ds.assign_coords(time=time)

# Remove last days
if month=="01" or month=="03":
    print(f"Month {month} has 31 days.")
    ds = ds.isel(time=slice(0, -24))
elif month=="02":
    print(f"Month {month} has 29 days.")
    ds = ds.isel(time=slice(0, -72))

# Convert to integer hours since 1975-01-01
reference = np.datetime64("1975-01-01T00:00:00")

hours_since = (
    (ds.time.values - reference)
    / np.timedelta64(1, "h")
).astype(np.int32)

ds = ds.assign_coords(time=("time", hours_since))

ds.time.attrs = {
    "standard_name": "time",
    "long_name": "time",
    "units": "hours since 1975-1-1 00:00:00",
    "calendar": "gregorian",
    "axis": "T",
}

# ==================================================
# Missing values
# ==================================================

print("Replacing missing values...")

fill_value = 9.96921e36

for var in ds.data_vars:
    ds[var] = ds[var].where(
        ds[var] != fill_value,
        0
    )


# ==================================================
# Save
# ==================================================

print("Writing output...")

encoding = {}

for var in ds.data_vars:
    encoding[var] = {
        "zlib": True,
        "complevel": 4,
        "_FillValue": 0
    }

encoding["time"] = {
    "dtype": "int32",
}

ds.time.attrs = {
    "standard_name": "time",
    "long_name": "time",
    "units": "hours since 1975-1-1 00:00:00",
    "calendar": "gregorian",
    "axis": "T",
}

ds.to_netcdf(
    OUTPUT,
    format="NETCDF4",
    engine="netcdf4",
    unlimited_dims=["time"],
    encoding=encoding,
)

print("Finished:")
print(OUTPUT)
print(ds)