import shelve
import numpy as np
from netCDF4 import Dataset
from re import findall
from collections import OrderedDict
from model import *
#hack to keep python 2 compatibility
try:
    #py3 import
    from dbm import whichdb
except ImportError:
    #py2 import
    from whichdb import whichdb


def loadvars(*args):
    """
    LOADVARS - function to load variables to a file.

    This function loads one or more variables from a file.  The names of the variables
    must be supplied.  If more than one variable is specified, it may be done with
    a list of names or a dictionary of name as keys.  The output type will correspond
    to the input type.  All the variables in the file may be loaded by specifying only
    the file name.

    Usage:
        a = loadvars('shelve.dat', 'a')
        [a, b] = loadvars('shelve.dat', ['a', 'b'])
        nvdict = loadvars('shelve.dat', {'a':None, 'b':None})
        nvdict = loadvars('shelve.dat')

    """

    filename = ''
    nvdict = {}

    if len(args) >= 1 and isinstance(args[0], str):
        filename = args[0]
        if not filename:
            filename = '/tmp/shelve.dat'

    else:
        raise TypeError("Missing file name.")

    if len(args) >= 2 and isinstance(args[1], str):  # (filename, name)
        for name in args[1:]:
            nvdict[name] = None

    elif len(args) == 2 and isinstance(args[1], list):  # (filename, [names])
        for name in args[1]:
            nvdict[name] = None

    elif len(args) == 2 and isinstance(args[1], dict):  # (filename, {names:values})
        nvdict = args[1]

    elif len(args) == 1:  #  (filename)
        pass

    else:
        raise TypeError("Unrecognized input arguments.")

    if whichdb(filename):
        print("Loading variables from file {}.".format(filename))
        my_shelf = shelve.open(filename, 'r')  # 'r' for read - only
        if nvdict:
            for name in list(nvdict.keys()):
                try:
                    nvdict[name] = my_shelf[name]
                    print(("Variable '%s' loaded." % name))
                except KeyError:
                    value = None
                    print("Variable '{}' not found.".format(name))

        else:
            for name in list(my_shelf.keys()):
                nvdict[name] = my_shelf[name]
                print(("Variable '%s' loaded." % name))
        my_shelf.close()

    else:
        try:
            NCFile = Dataset(filename, mode='r')
            NCFile.close()
        except RuntimeError:
            raise IOError("File '{}' not found.".format(filename))

        classtype, classtree = netCDFread(filename)
        nvdict['md'] = model()
        NCFile = Dataset(filename, mode='r')
        for mod in dict.keys(classtype):
            #print(' - Now treating classtype {}'.format(mod))
            if np.size(classtree[mod]) > 1:
                curclass = NCFile.groups[classtree[mod][0]].groups[classtree[mod][1]]
                if classtype[mod][0] == 'list':
                    keylist = [key for key in curclass.groups]
                    try:
                        steplist = [int(key) for key in curclass.groups]
                    except ValueError:
                        steplist = [int(findall(r'\d + ', key)[0]) for key in keylist]
                    indexlist = [index * (len(curclass.groups) - 1) / max(1, max(steplist)) for index in steplist]
                    listtype = curclass.groups[keylist[0]].classtype
                    if listtype == 'dict':
                        nvdict['md'].__dict__[classtree[mod][0]].__dict__[classtree[mod][1]] = [OrderedDict() for i in range(max(1, len(curclass.groups)))]
                    else:
                        nvdict['md'].__dict__[classtree[mod][0]].__dict__[classtree[mod][1]] = [getattr(__import__(listtype), listtype)() for i in range(max(1, len(curclass.groups)))]
                        Tree = nvdict['md'].__dict__[classtree[mod][0]].__dict__[classtree[mod][1]][:]
                else:
                    nvdict['md'].__dict__[classtree[mod][0]].__dict__[classtree[mod][1]] = getattr(classtype[mod][1], classtype[mod][0])()
                    Tree = nvdict['md'].__dict__[classtree[mod][0]].__dict__[classtree[mod][1]]
            else:
                curclass = NCFile.groups[classtree[mod][0]]
                nvdict['md'].__dict__[mod] = getattr(classtype[mod][1], classtype[mod][0])()
                Tree = nvdict['md'].__dict__[classtree[mod][0]]
                #treating groups that are lists
            for i in range(0, max(1, len(curclass.groups))):
                if len(curclass.groups) > 0:
                    listclass = curclass.groups[keylist[i]]
                else:
                    listclass = curclass
                for var in listclass.variables:
                    #print("treating var {}".format(var))
                    if var not in ['errlog', 'outlog']:
                        varval = listclass.variables[str(var)]
                        vardim = varval.ndim
                        if vardim == 0:
                            if type(Tree) == list:
                                t = int(indexlist[i])
                                if listtype == 'dict':
                                    Tree[t][str(var)] = varval[0].data

                                else:
                                    Tree[t].__dict__[str(var)] = varval[0].data

                            else:
                                if str(varval[0]) == '':  #no value
                                    Tree.__dict__[str(var)] = []

                                elif varval[0] == 'True':  #treatin bool
                                    Tree.__dict__[str(var)] = True

                                elif varval[0] == 'False':  #treatin bool
                                    Tree.__dict__[str(var)] = False

                                else:
                                    Tree.__dict__[str(var)] = varval[0].item()
                        elif vardim == 1:
                            if varval.dtype == str:
                                if varval.shape[0] == 1:
                                    Tree.__dict__[str(var)] = [str(varval[0]), ]

                                elif 'True' in varval[:] or 'False' in varval[:]:
                                    Tree.__dict__[str(var)] = np.asarray([V == 'True' for V in varval[:]], dtype=bool)

                                else:
                                    Tree.__dict__[str(var)] = [str(vallue) for vallue in varval[:]]

                            else:
                                if type(Tree) == list:
                                    t = int(indexlist[i])
                                    if listtype == 'dict':
                                        Tree[t][str(var)] = varval[:].data

                                    else:
                                        Tree[t].__dict__[str(var)] = varval[:].data
                                else:
                                    try:
                                        #some thing specifically require a list
                                        mdtype = type(Tree.__dict__[str(var)])
                                    except KeyError:
                                        mdtype = float
                                    if mdtype == list:
                                        Tree.__dict__[str(var)] = [mdval for mdval in varval[:]]

                                    else:
                                        Tree.__dict__[str(var)] = varval[:].data

                        elif vardim == 2:
                            #dealling with dict
                            if varval.dtype == str:  #that is for toolkits wich needs to be ordered
                                if any(varval[:, 0] == 'toolkit'):  #toolkit definition have to be first
                                    Tree.__dict__[str(var)] = OrderedDict([('toolkit', str(varval[np.where(varval[:, 0] == 'toolkit')[0][0], 1]))])
                                    strings1 = [str(arg[0]) for arg in varval if arg[0] != 'toolkits']
                                    strings2 = [str(arg[1]) for arg in varval if arg[0] != 'toolkits']
                                    Tree.__dict__[str(var)].update(list(zip(strings1, strings2)))
                            else:
                                if type(Tree) == list:
                                    t = int(indexlist[i])
                                    if listtype == 'dict':
                                        Tree[t][str(var)] = varval[:, :].data
                                    else:
                                        Tree[t].__dict__[str(var)] = varval[:, :].data
                                else:
                                    Tree.__dict__[str(var)] = varval[:, :].data
                        elif vardim == 3:
                            if type(Tree) == list:
                                t = int(indexlist[i])
                                if listtype == 'dict':
                                    Tree[t][str(var)] = varval[:, :, :].data
                                else:
                                    Tree[t].__dict__[str(var)] = varval[:, :, :]
                            else:
                                Tree.__dict__[str(var)] = varval[:, :, :].data
                        else:
                            print('table dimension greater than 3 not implemented yet')
                for attr in listclass.ncattrs():
                    if attr != 'classtype':  #classtype is for treatment, don't get it back
                        if type(Tree) == list:
                            t = int(indexlist[i])
                            if listtype == 'dict':
                                Tree[t][str(attr).swapcase()] = str(listclass.getncattr(attr))
                            else:
                                Tree[t].__dict__[str(attr).swapcase()] = str(listclass.getncattr(attr))
                        else:
                            Tree.__dict__[str(attr).swapcase()] = str(listclass.getncattr(attr))
                            if listclass.getncattr(attr) == 'True':
                                Tree.__dict__[str(attr).swapcase()] = True
                            elif listclass.getncattr(attr) == 'False':
                                Tree.__dict__[str(attr).swapcase()] = False
        NCFile.close()
    if len(args) >= 2 and isinstance(args[1], str):  # (value)
        value = [nvdict[name] for name in args[1:]]
        return value

    elif len(args) == 2 and isinstance(args[1], list):  # ([values])
        value = [nvdict[name] for name in args[1]]
        return value

    elif (len(args) == 2 and isinstance(args[1], dict)) or (len(args) == 1):  # ({names:values})
        return nvdict


def netCDFread(filename):
    print(('Opening {} for reading '.format(filename)))
    NCData = Dataset(filename, 'r')
    class_dict = {}
    class_tree = {}

    for group in NCData.groups:
        if len(NCData.groups[group].groups) > 0:
            for subgroup in NCData.groups[group].groups:
                classe = str(group) + '.' + str(subgroup)
                class_dict[classe] = [str(getattr(NCData.groups[group].groups[subgroup], 'classtype')), ]
                if class_dict[classe][0] not in ['dict', 'list', 'cell']:
                    class_dict[classe].append(__import__(class_dict[classe][0]))
                class_tree[classe] = [group, subgroup]
        else:
            classe = str(group)
            try:
                class_dict[classe] = [str(getattr(NCData.groups[group], 'classtype')), ]
                if class_dict[classe][0] not in ['dict', 'list', 'cell']:
                    class_dict[classe].append(__import__(class_dict[classe][0]))
                    class_tree[classe] = [group, ]
            except AttributeError:
                print(('group {} is empty'.format(group)))
    NCData.close()
    return class_dict, class_tree
