#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Created on 30 May 2024, atelier DEPHY Cap Ferret
Based on MPACE REF (E. Vignon)

@author: Lea Raillard

Modifications:
  2024/05/31, C. Agosta, add xarray and automatic renaming

LEA original case definition
From RALI ThinkIce campaign and ICO-LMDZ-LAM simulation
"""
import sys
import numpy as np

sys.path = ['../../utils/', ] + sys.path

# import netCDF4 as nc
import xarray as xr
# import numpy as np

from Case import Case

################################################
# 0. General configuration of the present script
################################################

lplot = True  # plot all the variables
lverbose = False  # print information about variables and case

################################################
# 1. General information about the case
################################################
# This is an idealized case so date is arbritrary
# lat/lon are fixed to a lat representative of Arctic conditions
# 8h duration with a constant surface cooling

lon = 2.32
lat = 82.5
startDate = '2022-08-13 05:00:00'
endDate = '2022-08-13 18:00:00'

case = Case('LEA/REF',
            lat=lat,
            lon=lon,
            startDate=startDate,
            endDate=endDate,
            surfaceType="ocean",
            zorog=0.)

case.set_title("Forcing and initial conditions for LEA case - Original defileinition")
case.set_reference("Raillard et al. (future)")
case.set_author("L. Raillard and E. Vignon")
case.set_script("driver_DEF.py")
case.set_comment("Use of forcing file from ICO LMDZ LAM LEA JZ")

################################################
# 2. Input netCDF file
################################################
filename = '/data/lraillard/These/ORLAM-ARCTIC-LUDO-nudg/ORLAM-ARCTIC-LUDO-nudg_20220801_20220831_HF_histhf.nc'
filein = xr.open_dataset(filename)
# for LMDZ, rename time axis
filein = filein.rename({'time_counter': 'time'})

# varnames in LMDZ
varname = {'ps': 'slp',  # Surface pressure (hPa) # WARNING should be psol
           'pressure': 'pres',  # Pressure (Pa)
           'temp': 'temp',  # Temperature (K)
           'qv': 'ovap',  # Water vapor mixing ratio (kg kg-1)
           'qcond': 'ocond',  # Cloud condensate mixing ratio (kg kg-1)
           'u': 'vitu',  # Zonal wind (m s-1)
           'v': 'vitv',  # Meridional wind (m s-1)
           'omega': 'vitw',  # Large scale vertical pressure velocity (Pa s-1)
           'tadv': 'dtdyn',  # Large-scale advection of temperature (K s-1)
           'qvadv': 'dqdyn',  # Large-scale advection of specific humidity (kg kg-1 s-1)
           'hl': 'flat',  # Surface latent fluxes (W m-2)
           'hs': 'sens'  # Surface sensible fluxes (W m-2)
           }

################################################
# 3. Initial state
################################################

# Surface pressure (
ps = filein[varname['ps']].sel(lon=lon, lat=lat, time=startDate, method='nearest').values
case.add_init_ps(ps)

# Pressure
pressure = filein[varname['pressure']].sel(lon=lon, lat=lat, time=startDate, method='nearest').values
# pressure vertical axis plev
plev = pressure
# pressure[0] = ps # replace surface value by the surface pressure
case.add_init_pressure(pressure, lev=plev, levtype='pressure', levid='lev')

# Temperature in K and moisture (vapor and total mixing ratio)
# Note that in the ref paper, the ice-liquid potential temperature is initially set 
# but the initial state corresponds to a purely-liquid cloud
temp = filein[varname['temp']].sel(lon=lon, lat=lat, time=startDate, method='nearest').values
case.add_init_temp(temp, lev=plev, levtype='pressure', levid='lev')

# Water vapor  in kg/kg
# ----------------------
# in LMDZ, we directly have the mixing ratio
qv = filein[varname['qv']].sel(lon=lon, lat=lat, time=startDate, method='nearest').values
case.add_init_qv(qv, lev=plev, levtype='pressure', levid='lev')

# Total water mixing ratio in kg/kg
# ---------------------------------
qcond = filein[varname['qcond']].sel(lon=lon, lat=lat, time=startDate, method='nearest').values
qt = qv + qcond

case.add_init_qt(qt, lev=plev, levtype='pressure', levid='lev')

# Zonal and meridional wind
# -------------------------
u = filein[varname['u']].sel(lon=lon, lat=lat, time=startDate, method='nearest').values
v = filein[varname['v']].sel(lon=lon, lat=lat, time=startDate, method='nearest').values

case.add_init_wind(u=u, v=v, lev=plev, levtype='pressure', levid='lev')

################################################
# 3. Forcing
################################################
time_axis = filein.time.sel(time=slice(startDate, endDate)).values
# pressure levels
# ---------------
# surface pressure
ps = filein[varname['ps']].sel(lon=lon, lat=lat, method='nearest').sel(time=slice(startDate, endDate)).values
case.add_surface_pressure_forcing(ps, timeid='time', time=time_axis)

# pressure
pressure = filein[varname['pressure']].sel(lon=lon, lat=lat, method='nearest').sel(time=slice(startDate, endDate)).values
case.add_pressure_forcing(pressure, timeid='time', time=time_axis, lev=plev, levtype='pressure', levid='lev')

# nudging of wind
u = filein[varname['u']].sel(lon=lon, lat=lat, method='nearest').sel(time=slice(startDate, endDate)).values
v = filein[varname['v']].sel(lon=lon, lat=lat, method='nearest').sel(time=slice(startDate, endDate)).values
case.add_wind_nudging(unudg=u, vnudg=v, timescale=3600., p_nudging=110000.,
                      timeid='time', time=time_axis, lev=plev, levtype='pressure', levid='lev')

# Large scale vertical pressure velocity
omega = filein[varname['omega']].sel(lon=lon, lat=lat, method='nearest').sel(time=slice(startDate, endDate)).values
case.add_vertical_velocity(omega=omega, timeid='time', time=time_axis, lev=plev, levtype='pressure', levid='lev')

# Large-scale advection of temperature in K s-1
tadv = filein[varname['tadv']].sel(lon=lon, lat=lat, method='nearest').sel(time=slice(startDate, endDate)).values
case.add_temp_advection(tadv, timeid='time', time=time_axis, lev=plev, levtype='pressure', levid='lev')

# Large-scale advection of specific humidity in kg kg-1 s-1
qvadv = filein[varname['qvadv']].sel(lon=lon, lat=lat, method='nearest').sel(time=slice(startDate, endDate)).values
case.add_qv_advection(qvadv, timeid='time', time=time_axis, lev=plev, levtype='pressure', levid='lev')

# Surface Forcing
# constant surface latent and sensible fluxes [W m-2] (be careful sign convention)
hs = filein[varname['hs']].sel(lon=lon, lat=lat, method='nearest').sel(time=slice(startDate, endDate)).values
hl = filein[varname['hl']].sel(lon=lon, lat=lat, method='nearest').sel(time=slice(startDate, endDate)).values
case.add_surface_fluxes(sens=hs, lat=hl, timeid='time', time=time_axis, forc_wind='z0', z0=0.01)

# Constant SST [K]
ts = 274.01 + np.zeros_like(ps)
case.add_surface_temp(ts, timeid='time', time=time_axis)

################################################
# 4. Writing file
################################################

case.write('LEA_REF_DEF_driver.nc')

if lverbose:
    case.info()

################################################
# 5. Ploting, if asked
################################################

if lplot:
    case.plot(rep_images='./images/driver_DEF/', timeunits='hours')