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

__all__ = ['load_table', 'FWIe_calc']

# Cell
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from pathlib import Path
import scipy.io as sio
from tqdm import tqdm
import requests

# Cell
def load_table(table_path):
    path = Path(table_path)
    table_file = path/'FWIeCalc_table.mat'
    if not table_file.is_file():
        print('Downloading FWIe table...')
        path.mkdir(exist_ok=True)
        url = 'https://github.com/mnpinto/FWIe/raw/master/data/FWIeCalc_table.mat'
        file = requests.get(url)
        open(str(table_file), 'wb').write(file.content)
    return sio.loadmat(str(table_file))

def FWIe_calc(FWI:np.ndarray, CHI:np.ndarray, table_path='data'):
    table = load_table(table_path)
    assert FWI.shape == CHI.shape
    rs, cs = FWI.shape
    FWIe = np.zeros_like(FWI)*np.nan
    for i in tqdm(range(rs)):
        for j in range(cs):
            fwi_argmin = np.argmin(np.abs(table['FWI'].reshape(-1) - FWI[i,j]))
            chi_argmin = np.argmin(np.abs(table['CHI'].reshape(-1) - CHI[i,j]))
            FWIe[i,j] = table['FWIe'][fwi_argmin, chi_argmin]
    return np.round(FWIe,1)