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':1},
    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):

    ts_dict = {}

    xs = []
    ys = []
    lats = []
    lons = []
    buoy_ids  = []
    for lat0, lon0, bid in buoys:

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

        # U-point
        u_col = ds["xcur"].isel(y=j, x_u=i)
        
        # V-point
        v_col = ds["ycur"].isel(y_v=j, x=i)

        # =====================================================
        # U COMPONENT (xcur)
        # =====================================================

        z_u = gop.get_z(
            croco_model,
            ds=ds,
            z_sfc=ds.z_sfc,
            xgrid=grid,
            hgrid="u"
        ).isel(y=j, x_u=i)

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

        # =====================================================
        # V COMPONENT (ycur)
        # =====================================================

        z_v = gop.get_z(
            croco_model,
            ds=ds,
            z_sfc=ds.z_sfc,
            xgrid=grid,
            hgrid="v"
        ).isel(y_v=j, x=i)

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

        # =====================================================
        # ROTATION TO GEOGRAPHIC FRAME
        # =====================================================
        # =====================================================
        u = gop.to_grid_point(u_3m, grid, hcoord="r", vcoord="r")
        v = gop.to_grid_point(v_3m, grid, hcoord="r", vcoord="r")

        angle = float(gop.to_grid_point(ds["angle"], grid, hcoord="r", vcoord="r").isel(y=j, x=i))

        u_rot = u * np.cos(angle) - v * np.sin(angle)
        v_rot = u * np.sin(angle) + v * np.cos(angle)
        # -----------------------------
        # Store
        # -----------------------------
        ts_dict[f"u_{bid}"] = u
        ts_dict[f"v_{bid}"] = v
        ts_dict[f"u_rot_{bid}"] = u_rot
        ts_dict[f"v_rot_{bid}"] = v_rot

        xs.append(i)
        ys.append(j)
        lats.append(float(ds.lat[j, i]))
        lons.append(float(ds.lon[j, i]))
        buoy_ids.append(bid)

    ts_ds = xr.Dataset(ts_dict)

    ts_ds = ts_ds.assign_coords(
    buoy=("buoy", buoy_ids),x=("buoy", xs),y=("buoy", ys),lat=("buoy", lats),lon=("buoy", lons))
    return ts_ds

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

ts_ds = extract_buoys(
    ds,
    depth,
    grid,
    croco,
    buoys
)

# ==========================
# SAVE
# ==========================
output_file = os.path.join(
    fig_dir,
    f"buoys_uv_rot_3m_S01_2019_first.nc"
)

io.store_netcdf(ts_ds, filename=output_file)

print("Saved:", output_file)