# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/api/nc_template.ipynb.

# %% auto 0
__all__ = ['NCTemplate', 'derive']

# %% ../nbs/api/nc_template.ipynb 2
from netCDF4 import Dataset
import pandas as pd
from pathlib import Path
from fastcore.basics import patch, store_attr
from fastcore.test import *
from typing import Dict

from .utils import read_toml
from .configs import BASE_PATH

# %% ../nbs/api/nc_template.ipynb 3
class NCTemplate:
    def __init__(self, 
                 tpl_fname:str, # CDL file name
                 vars_fname:str, # File name and path of MARIS nuclide look up table
                 dest_dir:str, # Destination directory for generated NetCDF template files
                 cdl:Dict, # Pseudo CDL (`.toml`)
                ):
        store_attr()
        self.dim = self.cdl['dim']

# %% ../nbs/api/nc_template.ipynb 6
@patch
def get_analytes(self:NCTemplate,
                 col_varnames:str='nc_name', # Column name containing the NetCDF variable names
                 col_stdnames:str='nusymbol', # Column name containing the NetCDF standard names
                ):
    "Return the name of the variables analysed"
    df = pd.read_excel(self.vars_fname, index_col=0)
    df = df[df.nuclide != 'NOT AVAILABLE']
    var_names = df[col_varnames].tolist()
    std_names = df[col_stdnames].tolist()
    long_names = df[['nuclide', 'massnb']].apply(lambda row: ' '.join(row.values.astype(str)), 
                                                 axis=1).tolist()
    long_names = [name.capitalize() for name in long_names]

    return [{'name': n, 
             'long_name': ln,
             'standard_name': sn
            } for n, ln, sn in zip(*(var_names, long_names, std_names))]

# %% ../nbs/api/nc_template.ipynb 8
def derive(analyte, suffix):
    derived = analyte.copy()
    for k, v in suffix.items():
        derived[k] += v
    return derived

# %% ../nbs/api/nc_template.ipynb 10
@patch
def create_variable(self:NCTemplate, 
               nc, # NetCDF file
               var:Dict, # Variable
               dtype:str='f4', # Type of the variable
           ):
    name = var['name']
    attrs = {k:v for k, v in var.items() if k != name}
    nc_var = nc.createVariable(name, dtype, self.dim['name'])
    nc_var.setncatts(attrs)    
    return nc

# %% ../nbs/api/nc_template.ipynb 12
@patch
def generate(self:NCTemplate,
             common_vars:list=['lon', 'lat', 'depth', 'time'], # Common variables
            ):
    "Generate CDL"
    fname = Path(self.dest_dir)/self.tpl_fname
    
    common_vars = self.cdl['vars']['defaults'].keys()
    
    with Dataset(fname, 'w', format='NETCDF4') as nc:
        # Create dataset attributes
        nc.setncatts(self.cdl['global_attrs']) 
        
        # Create shared `sample` dimension
        nc.createDimension(self.dim['name'], None)
        
        # Create grps
        grp_names = [v['name'] for k, v in self.cdl['grps'].items()]
        for grp_name in grp_names:
            grp = nc.createGroup(grp_name)

            # Create 'dim' variable
            self.create_variable(grp, self.dim, 'i4')
            
            # Create default variables
            for var in self.cdl['vars']['defaults'].values(): 
                self.create_variable(grp, var)

            # Create analyte variables
            for analyte in self.get_analytes():
                analyte['units'] = self.cdl['placeholder']
                self.create_variable(grp, analyte)
            
                # Derived uncertainty and detection limit variables
                for k, v in self.cdl['vars']['suffixes'].items():
                    self.create_variable(grp, derive(analyte, v))
                #for related_var in ['uncertainty', 'detection_limit']:
                #    cfg = self.cfgs[related_var]
                #    attrs['long_name'] += cfg['long_name']
                #    attrs['standard_name'] += cfg['standard_name']
                #    self.create_variable(grp, analyte['name'] + cfg['var_suffix'], attrs)
