from netCDF4 import Dataset
import numpy as np
import time
import collections
from os import path, remove


def export_netCDF(md, filename):
    if path.exists(filename):
        print('File {} allready exist'.format(filename))
        newname = input('Give a new name or "delete" to replace: ')
        if newname == 'delete':
            remove(filename)
        else:
            print(('New file name is {}'.format(newname)))
            filename = newname
    NCData = Dataset(filename, 'w', format='NETCDF4')
    NCData.description = 'Results for run' + md.miscellaneous.name
    NCData.history = 'Created ' + time.ctime(time.time())
    # define netCDF dimensions
    try:
        StepNum = np.shape(dict.values(md.results.__dict__))[1]
    except IndexError:
        StepNum = 1
    Dimension1 = NCData.createDimension('DimNum1', StepNum)  # time is first
    DimDict = {len(Dimension1): 'DimNum1'}
    dimindex = 1
    dimlist = [2, md.mesh.numberofelements, md.mesh.numberofvertices, np.shape(md.mesh.elements)[1]]
    print('===Creating dimensions ===')
    for i in range(0, 4):
        if dimlist[i] not in list(DimDict.keys()):
            dimindex += 1
            NewDim = NCData.createDimension('DimNum' + str(dimindex), dimlist[i])
            DimDict[len(NewDim)] = 'DimNum' + str(dimindex)
    typelist = [bool, str, str, int, float, complex,
                collections.OrderedDict,
                np.int64, np.ndarray, np.float64]
    groups = dict.keys(md.__dict__)
    # get all model classes and create respective groups
    print('===Creating and populating groups===')
    for group in groups:
        NCgroup = NCData.createGroup(str(group))
    # In each group gather the fields of the class
        fields = dict.keys(md.__dict__[group].__dict__)
    # looping on fields
        for field in fields:
            # Special treatment for list fields
            if type(md.__dict__[group].__dict__[field]) == list:
                StdList = False
                if len(md.__dict__[group].__dict__[field]) == 0:
                    StdList = True
                else:
                    StdList = type(md.__dict__[group].__dict__[field][0]) in typelist
                NCgroup.__setattr__('classtype', md.__dict__[group].__class__.__name__)
                if StdList:  # this is a standard or empty list just proceed
                    Var = md.__dict__[group].__dict__[field]
                    DimDict = CreateVar(NCData, Var, field, NCgroup, DimDict)
                else:  # this is a list of fields, specific treatment needed
                    Listsize = len(md.__dict__[group].__dict__[field])
                    Subgroup = NCgroup.createGroup(str(field))
                    Subgroup.__setattr__('classtype', md.__dict__[group].__dict__[field].__class__.__name__)
                    for listindex in range(0, Listsize):
                        try:
                            Listgroup = Subgroup.createGroup(str(md.__dict__[group].__dict__[field].__getitem__(listindex).__dict__['name']))
                        except KeyError:
                            for naming in ['step']:
                                Listgroup = Subgroup.createGroup(str(md.__dict__[group].__dict__[field].__getitem__(listindex).__dict__[naming]))
                        except AttributeError:
                            Listgroup = Subgroup.createGroup(str(md.__dict__[group].__dict__[field].__class__.__name__) + str(listindex))
                        Listgroup.__setattr__('classtype', md.__dict__[group].__dict__[field].__getitem__(listindex).__class__.__name__)
                        try:
                            subfields = dict.keys(md.__dict__[group].__dict__[field].__getitem__(listindex).__dict__)
                        except AttributeError:
                            subfields = dict.keys(md.__dict__[group].__dict__[field].__getitem__(listindex))
                        for subfield in subfields:
                            if subfield != 'outlog':
                                try:
                                    Var = md.__dict__[group].__dict__[field].__getitem__(listindex).__dict__[subfield]
                                except AttributeError:
                                    Var = md.__dict__[group].__dict__[field].__getitem__(listindex)[subfield]
                                DimDict = CreateVar(NCData, Var, subfield, Listgroup, DimDict, md.__dict__[group], field, listindex)
            # No subgroup, we directly treat the variable
            elif type(md.__dict__[group].__dict__[field]) in typelist or field == 'bamg':
                NCgroup.__setattr__('classtype', md.__dict__[group].__class__.__name__)
                Var = md.__dict__[group].__dict__[field]
                DimDict = CreateVar(NCData, Var, field, NCgroup, DimDict)
            elif md.__dict__[group].__dict__[field] is None:
                print('field md.{}.{} is None'.format(group, field))
            # do nothing
            # if it is a masked array
            elif type(md.__dict__[group].__dict__[field]) is np.ma.core.MaskedArray:
                NCgroup.__setattr__('classtype', md.__dict__[group].__class__.__name__)
                Var = md.__dict__[group].__dict__[field].data
                DimDict = CreateVar(NCData, Var, field, NCgroup, DimDict)
            else:
                NCgroup.__setattr__('classtype', str(group))
                Subgroup = NCgroup.createGroup(str(field))
                Subgroup.__setattr__('classtype', md.__dict__[group].__class__.__name__)
                subfields = dict.keys(md.__dict__[group].__dict__[field].__dict__)
                for subfield in subfields:
                    if str(subfield) != 'outlog':
                        Var = md.__dict__[group].__dict__[field].__dict__[subfield]
                        DimDict = CreateVar(NCData, Var, subfield, Subgroup, DimDict)
    NCData.close()

    #============================================================================
    # Define the variables


def CreateVar(NCData, var, field, Group, DimDict, * step_args):
    # grab type
    try:
        val_type = str(var.dtype)
    except AttributeError:
        val_type = type(var)
    # grab dimension
    if val_type in [collections.OrderedDict, dict]:
        val_shape = len(var)
        val_dim = 2
    else:
        val_shape = np.shape(var)
        val_dim = np.shape(val_shape)[0]
    TypeDict = {float: 'f8',
                'float64': 'f8',
                np.float64: 'f8',
                int: 'i8',
                'int64': 'i8',
                np.int64: 'i8',
                str: str,
                dict: str}
    # Now define and fill up variable
    # treating scalar string or bool as atribute
    if val_type in [str, bool]:
        Group.__setattr__(str(field).swapcase(), str(var))
    # treating list as string table
    elif val_type == list:
        dimensions, DimDict = GetDim(NCData, val_shape, val_type, DimDict, val_dim)
    # try to get the type from the first element
        try:
            nctype = TypeDict[type(var[0])]
        except IndexError:
            nctype = str  # most probably an empty list take str for that
            ncvar = Group.createVariable(str(field), nctype, dimensions, zlib=True)
            if val_shape == 0:
                ncvar = []
            else:
                for elt in range(0, val_shape[0]):
                    ncvar[elt] = var[elt]
    # treating bool tables as string tables
    elif val_type == 'bool':
        dimensions, DimDict = GetDim(NCData, val_shape, val_type, DimDict, val_dim)
        ncvar = Group.createVariable(str(field), str, dimensions, zlib=True)
        for elt in range(0, val_shape[0]):
            ncvar[elt] = str(var[elt])
    # treating dictionaries as tables of strings
    elif val_type in [collections.OrderedDict, dict]:
        dimensions, DimDict = GetDim(NCData, val_shape, val_type, DimDict, val_dim)
        ncvar = Group.createVariable(str(field), str, dimensions, zlib=True)
        for elt, key in enumerate(dict.keys(var)):
            ncvar[elt, 0] = key
            ncvar[elt, 1] = str(var[key])  # converting to str to avoid potential problems
    # Now dealing with numeric variables
    elif val_type in [float, 'float64', np.float64, int, 'int64']:
        dimensions, DimDict = GetDim(NCData, val_shape, val_type, DimDict, val_dim)
        ncvar = Group.createVariable(str(field), TypeDict[val_type], dimensions, zlib=True)
        try:
            nan_val = np.isnan(var)
            if nan_val.all():
                ncvar[:] = 'NaN'
            else:
                ncvar[:] = var
        except TypeError:  # type does not accept nan, get vallue of the variable
            ncvar[:] = var
    else:
        print(('WARNING type "{}" is unknown for "{}.{}"'.format(val_type, Group.name, field)))
    return DimDict


# ============================================================================
# retriev the dimension tuple from a dictionnary
def GetDim(NCData, val_shape, val_type, DimDict, val_dim):
    output = []
    if val_type in [collections.OrderedDict, dict]:  # dealling with a dictionnary
        try:
            output = [str(DimDict[val_shape])]  # first try to get the coresponding dimension if ti exists
            output = output + [DimDict[2]]  # dimension5 is 2 to treat with dict
        except KeyError:
            index = len(DimDict) + 1  # if the dimension does not exist, increment naming
            NewDim = NCData.createDimension('DimNum' + str(index), val_shape)  # create dimension
            DimDict[len(NewDim)] = 'DimNum' + str(index)  # and update the dimension dictionary
            output = [str(DimDict[val_shape])] + [DimDict[2]]  # now proceed with the shape of the value
    else:
        # loop on dimensions
        for dim in range(0, val_dim):  # loop on the dimensions
            try:
                output = output + [str(DimDict[val_shape[dim]])]  # test if the dimension allready exist
            except KeyError:  # if not create it
                if (val_shape[dim]) > 0:
                    index = len(DimDict) + 1
                    NewDim = NCData.createDimension('DimNum' + str(index), (val_shape[dim]))
                    DimDict[len(NewDim)] = 'DimNum' + str(index)
                    output = output + [str(DimDict[val_shape[dim]])]
    return tuple(output), DimDict
