import os
import warnings
import glob
import dask
import xarray as xr
import inout as io
import gridop as gop
import diags as dg
import plot as cplt
from model import Model
from tools import wait_cluster_ready
from config_tools import load_ini
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

warnings.filterwarnings(action="ignore")

# ==========================
# CONFIG
# ==========================

config_file = "xcroco.ini"
config = load_ini(config_file)

scratch = config["Xcroco_Outputs"]["scratch_dir"]

files_dir = config["Croco_Files"]["croco_files_dir"]
croco_output_files = config["Croco_Files"]["croco_output_files"]
croco_grid_file = config["Croco_Files"]["croco_grid_file"]
model = config["Croco_Files"]["croco_output_type"]

filenames = glob.glob(os.path.join(files_dir, croco_output_files))
gridname = os.path.join(files_dir, croco_grid_file)

croco = Model(model)

fig_dir = config["Xcroco_Outputs"]["fig_dir"]


drop_variables = []

ds = io.open_files(
    croco,
    gridname,
    filenames,
    grid_metrics=2,
    drop_variables=drop_variables,
    chunks={'t':20},
    remove_ghost_pts=True,
    xperiodic=False,
    yperiodic=False,
    verbose=False,
)


grid = gop.fast_xgcm_grid(ds, croco)

# ---  ---
def find_nearest_point(ds, lat0, lon0):
    lat = np.deg2rad(ds.lat)
    lon = np.deg2rad(ds.lon)
    lat0 = np.deg2rad(lat0)
    lon0 = np.deg2rad(lon0)
    dlat = lat - lat0
    dlon = lon - lon0
    a = np.sin(dlat/2)**2 + np.cos(lat)*np.cos(lat0)*np.sin(dlon/2)**2
    j, i = np.unravel_index(a.argmin(), a.shape)
    return j, i

# --- ---
buoys = [
    [36.4897,-6.9647,6200085],
    [42.1177,-9.4312,6200084],
    [43.4946,-9.2114,6200083],
    [44.1200,-7.6843,6200082],
    [43.7510,-6.1755,6200025],
    [43.6338,-3.0499,6200024]
]


# ==========================
# MAIN FUNCTION
# ==========================
def extract_buoys(ds, depth, grid, croco_model, buoys):

    out = {}

    for lat0, lon0, bid in buoys:

        # --------------------------
        # nearest point
        # --------------------------
        j, i = find_nearest_point(ds, lat0, lon0)

        ds_pt = ds.isel(x=i, y=j)

        # =====================================================
        # U COMPONENT (xcur)
        # =====================================================
        u_col = ds_pt["xcur"]

        z_u = gop.get_z(
            croco_model,
            ds=ds_pt,
            z_sfc=ds_pt.z_sfc,
            xgrid=grid,
            hgrid="u"
        )

        u_3m = gop.isoslice(
            u_col,
            target=depth,
            xgrid=grid,
            target_data=z_u,
            axis="z"
        ).reset_coords(drop=True)

        # =====================================================
        # V COMPONENT (ycur)
        # =====================================================
        v_col = ds_pt["ycur"]

        z_v = gop.get_z(
            croco_model,
            ds=ds_pt,
            z_sfc=ds_pt.z_sfc,
            xgrid=grid,
            hgrid="v"
        )

        v_3m = gop.isoslice(
            v_col,
            target=depth,
            xgrid=grid,
            target_data=z_v,
            axis="z"
        ).reset_coords(drop=True)

        # =====================================================
        # ROTATION TO GEOGRAPHIC FRAME
        # =====================================================
        # =====================================================
        angle = float(ds["angle"].isel(x=i, y=j))
        u = u_3m.values
        v = v_3m.values

        u_rot = u * np.cos(angle) - v * np.sin(angle)
        v_rot = u * np.sin(angle) + v * np.cos(angle)

        # =====================================================
        # STORE (1D time series only)
        # =====================================================
        out[f"u_{bid}"] = ("t", u)
        out[f"v_{bid}"] = ("t", v)
        out[f"u_rot_{bid}"] = ("t", u_rot)
        out[f"v_rot_{bid}"] = ("t", v_rot)

    return xr.Dataset(out)

# ==========================
# RUN
# ==========================
depth = -3

ts_ds = extract_buoys(
    ds,
    depth,
    grid,
    croco,
    buoys
)
#def clean_attrs(ds):
 #   for v in ds.data_vars:
   #     ds[v].attrs = {
  #          k: val
   #         for k, val in ds[v].attrs.items()
       #     if val is not None
   #     }
    #return ds

#ts_ds = clean_attrs(ts_ds)
# ==========================
# SAVE
# ==========================
output_file = os.path.join(
    fig_dir,
    f"buoys_uv_rot_3m_S01_2019.nc"
)

io.store_netcdf(ts_ds, filename=output_file)

print("Saved:", output_file)