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},
)


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

# --- ---
buoysPositions = [
    [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]
]

#  ---
def extract_buoy_timeseries(ds, varname, depth, grid, croco_model, buoysPositions):

    ts_dict = {}

    xs = []
    ys = []
    lats = []
    lons = []
    buoy_ids = []

    zvar = gop.get_z(croco_model, ds=ds, z_sfc=ds.z_sfc, xgrid=grid, hgrid="u") #inicial hgrid add

    for lat, lon, buoy_id in buoysPositions:

        j, i = find_nearest_point(ds, lat, lon)
        #ds_point = ds.isel(x=i, y=j) #added do inicial
        var_col = ds[varname].isel(x=i, y=j) #inicial 
        z_col = zvar.isel(x=i, y=j) #inicial
        #var_col = ds_point[varname] #added inicial
        #z_col = gop.get_z(
           # croco_model,
           # ds=ds_point,
           # z_sfc=ds_point.z_sfc,
           # xgrid=grid,
           # hgrid="u"
        #) #added inicial

        var_at_depth = gop.isoslice(
            var_col,
            target=depth,
            xgrid=grid,
            target_data=z_col,
            axis="z"
        )

        
        var_at_depth = var_at_depth.reset_coords(drop=True)

        ts_dict[f"buoy_{buoy_id}"] = var_at_depth

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

    ts_ds = xr.Dataset(ts_dict)

    # buoys coordinates
    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

#_____________
depth = -3  # depth [m]
varname = "xcur"
ts_ds = extract_buoy_timeseries(ds, varname, depth, grid, croco, buoysPositions)

#save netcdf ---
output_file = os.path.join(fig_dir, f"buoys_{varname}_timeseries_{depth}_S01.nc")
io.store_netcdf(
    ts_ds,
    filename=output_file,
)
print(f"Temporal Series: {output_file}")