# -*- coding: utf-8 -*-
"""
Example: Read and visualize Level-1 GEO-IPP observations from NetCDF.

This example prints the dataset metadata, reads every preset variable,
selects one receiver-satellite line of sight, and plots its STEC and ROTI
time series in a 2-by-1 figure.
"""

import matplotlib.pyplot as plt
import numpy as np
from netCDF4 import Dataset


# -------------------------------------------------------------
# 1. File name and target line-of-sight ID
# -------------------------------------------------------------
filename = "TEC_BDSGEO_20242650000_01D_30S.nc"

# LOS_id uniquely identifies one receiver-to-GEO-satellite line of sight.
# Its six-digit format is 1XXXYY:
#   XXX = three-digit receiver station ID
#   YY  = two-digit BDS GEO satellite PRN
# For example, 113803 means receiver station 138 observing satellite C03.
target_los_id = 113803

EXPECTED_VARIABLES = ("time", "STEC", "elev", "ROTI", "lat", "lon", "LOS_id")


def show_var_info(ds, var_name):
    """Print dimensions, data type, and all NetCDF attributes of a variable."""
    var = ds.variables[var_name]
    print(var_name)
    print(f"  Size       : {var.shape}")
    print(f"  Dimensions : {var.dimensions}")
    print(f"  Datatype   : {var.dtype}")
    print("  Attributes :")
    for attr in var.ncattrs():
        print(f"    {attr} = {getattr(var, attr)}")


def as_nan_array(var):
    """Read a NetCDF variable and convert masked/fill values to NaN."""
    return np.ma.asarray(var[:], dtype=np.float64).filled(np.nan)


# -------------------------------------------------------------
# 2. Open the NetCDF file, inspect it, and read all preset variables
# -------------------------------------------------------------
with Dataset(filename, mode="r") as ds:
    print("=== Dataset summary ===")
    print(ds)

    print("\n=== Variables ===")
    print(list(ds.variables.keys()))

    missing_variables = [name for name in EXPECTED_VARIABLES if name not in ds.variables]
    if missing_variables:
        raise KeyError(f"Missing expected variables: {missing_variables}")

    print("\n====== Variable info ======")
    for variable_name in ds.variables:
        show_var_info(ds, variable_name)

    data = {name: as_nan_array(ds.variables[name]) for name in EXPECTED_VARIABLES}


# Descriptive local names matching the definitions in save_GEOIPP_nc.py
time_all = data["time"]       # seconds since 00:00:00 UTC of the file date
stec_all = data["STEC"]       # slant TEC along the receiver-to-GEO LOS, TECU
elev_all = data["elev"]       # receiver elevation angle, degrees
roti_all = data["ROTI"]       # rate of TEC index along the LOS, TECU/min
lat_all = data["lat"]         # GEO-IPP geographical latitude, degrees north
lon_all = data["lon"]         # GEO-IPP geographical longitude, degrees east
los_id_all = data["LOS_id"]   # six-digit LOS identifier: 1XXXYY


# -------------------------------------------------------------
# 3. Select and sort all samples belonging to the target LOS
# -------------------------------------------------------------
mask = los_id_all == target_los_id
if not np.any(mask):
    raise ValueError(f"LOS_id {target_los_id} is not present in {filename}")

order = np.argsort(time_all[mask])
los_time = time_all[mask][order]
los_stec = stec_all[mask][order]
los_roti = roti_all[mask][order]

# The remaining variables are also fully read and available for analysis.
los_elev = elev_all[mask][order]
los_lat = lat_all[mask][order]
los_lon = lon_all[mask][order]


# -------------------------------------------------------------
# 4. Plot STEC and the corresponding ROTI time series (2 rows x 1 column)
# -------------------------------------------------------------
fig, axes = plt.subplots(2, 1, figsize=(9, 7), sharex=True)

axes[0].plot(los_time, los_stec, ".", markersize=3)
axes[0].set_ylabel("STEC (TECU)", fontsize=13)
axes[0].set_title(f"Level-1 GEO-IPP time series for LOS_id {target_los_id}", fontsize=15)

axes[1].plot(los_time, los_roti, ".", markersize=3)
axes[1].set_xlabel("UT (hour)", fontsize=13)
axes[1].set_ylabel("ROTI (TECU/min)", fontsize=13)

xticks_sec = np.arange(0, 86401, 14400)
for ax in axes:
    ax.set_xlim(0, 86400)
    ax.set_xticks(xticks_sec)
    ax.set_xticklabels(["0", "4", "8", "12", "16", "20", "24"])
    ax.grid(True, linestyle="--", alpha=0.5)
    ax.tick_params(
        axis="both", which="major", direction="out",
        length=6, width=1, labelsize=11
    )

fig.tight_layout()


# -------------------------------------------------------------
# 5. Save and show the figure
# -------------------------------------------------------------
savename = "GEO_IPP_STEC_ROTI.png"
fig.savefig(savename, dpi=300, bbox_inches="tight")
plt.show()

print(f"Figure saved as: {savename}")