#!/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
# ==================================================

# Vertical shift applied to the flight density, in feet.
# Negative = shifted DOWN (descend) -- e.g. -2000 for "2000 ft below".
# Optional second CLI arg lets you reuse this script for other shifts
# (e.g. -1000, +2000) without editing the file.
SHIFT_FT = float(sys.argv[2]) if len(sys.argv) > 2 else -2000.0

INPUT = f"2024{month}_map.nc4"
OUTPUT = f"2024{month}_ONERA_NA_domain_shift{SHIFT_FT:+.0f}ft.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
# Apply the vertical shift BEFORE the conversion, so each flight's
# recorded traffic gets relabeled to the pressure level it would sit
# at after moving SHIFT_FT feet vertically. The underlying data values
# (Number_of_kms, Number_of_planes) are untouched -- only the pressure
# each native alt level maps to changes, which is what makes this file
# directly comparable, level-by-level, against the unshifted one.
# Keep all native levels between 500 and 100 hPa (of the SHIFTED grid)
# ==================================================

print(f"Converting altitude to pressure (shift = {SHIFT_FT:+.0f} ft)...")

pressure = ft_to_hpa(ds.alt.values + SHIFT_FT)

# 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)