import os
import types

import h5py
import numpy as np

# ======================================================================

def ztmd_loader(fn,names):
    '''
    Loader function for ZTMD files, return blank object with datasets attached
    '''
    
    if not os.path.isfile(fn):
        raise FileNotFoundError(fn)
    
    obj = types.SimpleNamespace()
    
    with h5py.File(fn,'r') as hf:
        
        ## 0D: Attach top-level attributes
        for k,v in hf.attrs.items():
            setattr(obj, k, v)
        
        ## tmp
        obj.mu_inf    = obj.mu_Suth_ref*(obj.T_inf/obj.T_Suth_ref)**(3/2) * ((obj.T_Suth_ref+obj.S_Suth)/(obj.T_inf+obj.S_Suth))
        obj.rho_inf   = obj.p_inf/(obj.R*obj.T_inf)
        obj.nu_inf    = obj.mu_inf/obj.rho_inf
        obj.a_inf     = np.sqrt(obj.kappa*obj.R*obj.T_inf)
        obj.U_inf     = obj.Ma*obj.a_inf
        obj.cp        = obj.R*obj.kappa/(obj.kappa-1.)
        obj.cv        = obj.cp/obj.kappa
        obj.recov_fac = obj.Pr**(1/3)
        obj.Taw       = obj.T_inf + obj.recov_fac*obj.U_inf**2/(2*obj.cp)
        obj.lchar     = obj.Re*obj.nu_inf/obj.U_inf
        
        ## 1D: Attach contents of group dims/
        if 'dims' in hf:
            for dsn in hf['dims']:
                ds = hf[f'dims/{dsn}']
                setattr(obj, dsn, ds[()])
        
        ## 1D: Attach contents of group data_1Dx/
        if 'data_1Dx' in hf:
            for dsn in hf['data_1Dx']:
                ds = hf[f'data_1Dx/{dsn}']
                setattr(obj, dsn, ds[()])
        
        ## 2D : Attach contents of group data/
        if 'data' in hf:
            for dsn in hf['data']:
                if dsn in names: ## only load if explicitly asked for
                    ds = hf[f'data/{dsn}']
                    setattr(obj, dsn, ds[()].T)
    
    return obj
