# AUTOGENERATED! DO NOT EDIT! File to edit: 00_core.ipynb (unless otherwise specified).

__all__ = ['chi', 'time_from_wrf', 'coords_from_wrf', 'chi_from_wrf', 'wrf2msg']

# Cell
import numpy as np
from pathlib import Path
import matplotlib.pyplot as plt
import pandas as pd
from netCDF4 import Dataset
import scipy.io as sio
from tqdm import tqdm
from wrf import getvar, interplevel

# Cell
def chi(t_850:np.ndarray, td_850:np.ndarray, t_700:np.ndarray):
    """
    Computes Continuous Haines Index (CHI)
    t_850 is the temperature at 850 hPa
    t_700 is the temperature at 700 hPa
    td_850 is the dew point temperature at 850 hPa
    """
    A = t_850 - t_700
    B = t_850 - td_850
    B[B>30]=30
    CA = 0.5 * A - 2
    CB = 0.3333 * B - 1
    CB[CB>5] = 5 + (CB[CB>5] - 5) / 2
    return CA + CB

def time_from_wrf(file, time_index):
    ncfile = Dataset(file, mode='r')
    return pd.Timestamp(str(getvar(ncfile, "times", timeidx=time_index).data))

def coords_from_wrf(file):
    ncfile = Dataset(file, mode='r')
    return getvar(ncfile, "lon").data, getvar(ncfile, "lat").data

def chi_from_wrf(file, time_index):
    ncfile = Dataset(file, mode='r')
    td = getvar(ncfile, "td", timeidx=time_index)
    tc  = getvar(ncfile, "tc", timeidx=time_index)
    p = getvar(ncfile, "pressure", timeidx=time_index)
    t850, t700 = [interplevel(tc, p, level).data for level in [850, 700]]
    td850 = interplevel(td, p, 850).data
    chaines = chi(t850, td850, t700)
    return chaines

def wrf2msg(lon_wrf, lat_wrf, data_wrf, lon_msg, lat_msg):
    rs, cs = lat_msg.shape
    data_msg = np.zeros_like(lat_msg)*np.nan
    for i in range(1,rs-1):
        for j in range(1,cs-1):
            lat_o = [lat_msg[i-1,j-1], lat_msg[i-1,j+1], lat_msg[i+1,j-1], lat_msg[i+1, j+1]]
            lon_o = [lon_msg[i-1,j-1], lon_msg[i-1,j+1], lon_msg[i+1,j-1], lon_msg[i+1, j+1]]
            I = ((lat_wrf >= min(lat_o)) & (lat_wrf < max(lat_o)) &
                 (lon_wrf >= min(lon_o)) & (lon_wrf < max(lon_o)))
            data_msg[i,j] = np.nanmean(data_wrf[I])
    return data_msg