# This file is part of MarMOTS.
#
# MarMOTS is free software: you can redistribute it and/or modify it under the terms of the GNU Affero 
# General Public License as published by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# MarMOTS is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the 
# implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Affero General 
# Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License along with MarMOTS.  If not, 
# see <https://www.gnu.org/licenses/>.
#
# Copyright 2009, 2010, 2011, 2020 Jeremy Penner

import cPickle
import pdb
import traceback

def Decorator(dg):
    def DecoratorNew(*rgarg, **mparg):  
        def CallDecorator(o):
            return dg(o, *rgarg, **mparg) or o
        return CallDecorator
    return DecoratorNew

@Decorator
def Version(cls, vr):
    cls._version = vr

def RenameFrom(stName, namespace):
    def SetOldName(cls):
        namespace[stName] = cls
        return cls
    return SetOldName

class TransientSetter(object):
    def __init__(self, obj):
        self.obj = obj
        try:
            self._fWriteOld = obj._fWriteToPersistent
        except Exception:
            self._fWriteOld = True
    def __enter__(self):
        self.obj._fWriteToPersistent = False
        return self.obj
    def __exit__(self, type, value, traceback):
        self.obj._fWriteToPersistent = self._fWriteOld

@Version(1)
class TPrs(object):
    def __init__(self, *rgarg, **mparg):
        self._persistent = TPM()
        self._fWriteToPersistent = True
        self._versionCurr = self._version 
        self._fPersistent = True
        self.InitPersistent(*rgarg, **mparg)
        with self.SetTransiently():
            self.InitTransient()

    def InitPersistent(self, *rgarg, **mparg):
        pass

    def InitTransient(self):
        pass

    def SetTransiently(self):
        return TransientSetter(self)

    def MakeTransient(self, key, value):
        try:
            del self._persistent[key]
        except Exception:
            pass
        with self.SetTransiently():
            setattr(self, key, value)

    def FPersist(self):
        return self._fPersistent

    def DontPersist(self):
        self._fPersistent = False

    def UpgradeFrom(self, versionOld):
        pass

    def __getattr__(self, key):
        stPersistent = "no persistent yet"
        if key != "_persistent":
            try:
                return self._persistent[key]
            except Exception:
                stPersistent = str(self._persistent)
        raise AttributeError("Attribute '" + key + "' not found in object with type '" + type(self).__name__ + "':" + stPersistent + " __dict__:" + str(self.__dict__))

    __obj_setattr = object.__setattr__
    __obj_delattr = object.__delattr__
    def __setattr__(self, key, value):
        prop = getattr(self.__class__, key, None)
        if isinstance(prop, property) and prop.fset:
            prop.fset(self, value)
            return
        try:
            if self._fWriteToPersistent or key in self._persistent:
                self._persistent[key] = value
                try:
                    self.__obj_delattr(key)
                except AttributeError:
                    pass
                return
        except Exception: pass
        self.__obj_setattr(key, value)

    def __delattr__(self, key):
        try:
            del self._persistent[key]
        except Exception:
            self.__obj_delattr(key)

    def __reduce__(self):
        if self.FPersist():
            return (NewObjWithClass, (type(self),), self._persistent)
        else:
            return (NewObjWithClass, (None,))

    def __setstate__(self, state):
        self._persistent = state
        self._fWriteToPersistent = True
        if self._versionCurr != self._version:
            Odb.rgtprsToUpgrade.append(self)
        Odb.rgtprsToInit.append(self)
    def _Upgrade(self):
        self.UpgradeFrom(self._versionCurr)
        self._versionCurr = self._version

class Odb(object):
    @staticmethod
    def Save(tprs, fn):
        with open(fn, "w") as fl:
            cPickle.dump(tprs, fl)

    @classmethod
    def Load(cls, fn):
        cls.rgtprsToUpgrade = []
        cls.rgtprsToInit = []
        try:
            with open(fn, "r") as fl:
                tprs = cPickle.load(fl)
        except Exception as e:
            print "error unpickling:", e
            traceback.print_exc()
            return None
        for tprsToUpgrade in cls.rgtprsToUpgrade:
            tprsToUpgrade._Upgrade()
        for tprsToInit in cls.rgtprsToInit:
            with tprsToInit.SetTransiently():
                tprsToInit.InitTransient()
        if len(cls.rgtprsToUpgrade) > 0:
            fnUpgraded = fn + ".upgraded"
            Odb.Save(tprs, fnUpgraded)
            tprs = Odb.Load(fnUpgraded)
        cls.rgtprsToUpgrade = None
        cls.rgtprsToInit = None
        return tprs

clsList = list
clsDict = dict
clsSet = set
clsTuple = tuple

def MkTransparent(value):
    if type(value) == list:
        return TPL(value)
    elif type(value) == dict:
        return TPM(value)
    elif type(value) == set:
        return TPS(value)
    elif type(value) == tuple:
        return TPT(value)
    else:
        return value

def RgTransparent(rg):
    if rg == None: return []
    return [MkTransparent(x) for x in rg]

def MpTransparent(mp):
    if isinstance(mp, TPM): return mp
    mpNew = {}
    if mp != None:
        for k,v in mp.iteritems():
            mpNew[k] = MkTransparent(v)
    return mpNew

def FPersist(obj):
    return (not isinstance(obj, TPrs)) or obj.FPersist()

def MkPersistable(obj):
    if FPersist(obj):
        return obj
    return None

def NewObjWithClass(cls):
    if cls == None:
        return None
    return cls.__new__(cls)

class TPL(clsList):
    __pl_setitem = clsList.__setitem__
    __pl_setslice = clsList.__setslice__
    __pl_iadd = clsList.__iadd__
    __pl_append = clsList.append
    __pl_insert = clsList.insert
    __pl_extend = clsList.extend

    def __init__(self, rg = None):
        clsList.__init__(self, RgTransparent(rg))

    def __setitem__(self, i, item):
        return self.__pl_setitem(i, MkTransparent(item))

    def __setslice__(self, i, j, other):
        return self.__pl_setslice(i, j, RgTransparent(other))

    def __iadd__(self, other):
        return self.__pl_iadd(RgTransparent(other))

    def append(self, item):
        return self.__pl_append(MkTransparent(item))

    def insert(self, i, item):
        return self.__pl_insert(i, MkTransparent(item))

    def extend(self, other):
        return self.__pl_extend(RgTransparent(other))

    def __reduce__(self):
        return (MkTransparent, ([MkPersistable(x) for x in self],))

# Strip weak references rather than replacing them with None
class TPLS(TPL):
    def __reduce__(self):
        return (MkTransparent, ([x for x in self if FPersist(x)],))

class TPM(clsDict):
    __pm_setitem = clsDict.__setitem__
    __pm_update = clsDict.update

    def __init__(self, *rgarg, **kwarg):
        clsDict.__init__(self, MpTransparent(dict(*rgarg, **kwarg)))

    def __setitem__(self, k, v):
        return self.__pm_setitem(k, MkTransparent(v))

    def update(self, other, **kwargs):
        return self.__pm_update(MpTransparent(other), **MpTransparent(kwargs))

    def __reduce__(self):
        mp = {}
        for k, v in self.iteritems():
            if FPersist(k):
                mp[k] = MkPersistable(v)
        return (MkTransparent, (mp,))

class TPS(clsSet):
    __ps_update = clsSet.update
    __ps_ior = clsSet.__ior__
    __ps_add = clsSet.add

    def __init__(self, rg = None):
        clsSet.__init__(self, RgTransparent(rg))

    def update(self, *rgother):
        return self.__ps_update(*RgTransparent(rgother))

    def __ior__(self, other):
        return self.__ps_ior(MkTransparent(other))

    def add(self, elem):
        return self.__ps_add(MkTransparent(elem))

    def __reduce__(self):
        return (MkTransparent, (set([x for x in self if FPersist(x)]),))

class TPT(clsTuple):
    def __new__(cls, tup = None):
        return clsTuple.__new__(cls, RgTransparent(tup))

    def __reduce__(self):
        return (MkTransparent, (tuple([MkPersistable(x) for x in self]),))

class Partial(object):
    def __init__(self, method, *args, **kwargs):
        self.method = method
        self.args = args
        self.kwargs = kwargs

    def __call__(self, *args, **kwargs):
        mergedkwargs = self.kwargs.copy()
        mergedkwargs.update(kwargs)
        return self.method(*(self.args + args), **mergedkwargs)

def _pickle_method(method):
    func_name = method.im_func.__name__
    obj = method.im_self
    cls = method.im_class
    return _unpickle_method, (func_name, obj, cls)

def _unpickle_method(func_name, obj, cls):
    for cls in cls.mro():
        try:
            func = cls.__dict__[func_name]
        except KeyError:
            pass
        else:
            break
    return func.__get__(obj, cls)

import copy_reg
import types
copy_reg.pickle(types.MethodType, _pickle_method, _unpickle_method)

def hashable(v):
    try:
        hash(v)
        return True
    except:
        return False

class StaticIndex(TPrs):
    def InitPersistent(self, fIncludeTemp = False):
        self.mpv_a_rge = {}
        self.rgv = set()
        self.fIncludeTemp = fIncludeTemp

    def FPersist(self):
        return False
    
    def AddEAV(self, e, a, v):
        if v not in self.mpv_a_rge:
            self.mpv_a_rge[v] = {}
        a_rge = self.mpv_a_rge[v]
        if a not in a_rge:
            a_rge[a] = []
        a_rge[a].append(e)

    def Populate(self, v, a = (), e = None):
        if hashable(v):
            if e:
                self.AddEAV(e, a, v)
            if v in self.rgv:
                return self
            self.rgv.add(v)
        if isinstance(v, TPrs):
            attrs = set(v._persistent.keys())
            if self.fIncludeTemp:
                attrs = attrs.union(set(v.__dict__.keys()))
            for attr in attrs:
                child = getattr(v, attr)
                self.Populate(child, (attr,), v)
        elif isinstance(v, clsList) or isinstance(v, clsTuple):
            for ichild, child in enumerate(v):
                self.Populate(child, a + (ichild,), e)
        elif isinstance(v, clsSet):
            for child in v:
                self.Populate(child, a, e)
        elif isinstance(v, clsDict):
            for k, child in v.items():
                self.Populate(child, a + (k,), e)
        return self
    
    def Rgclass(self):
        return set([v.__class__ for v in self.rgv])

    def RgvWithClass(self, cls):
        return [v for v in self.rgv if isinstance(v, cls)]
    
    def Rgref(self, v):
        return self.mpv_a_rge.get(v)