from collections import OrderedDict
from numpy import (int64, int32, int16, ndarray)
from scipy.sparse import issparse

from ..QuantumToolbox import operators as qOps #pylint: disable=relative-beyond-top-level
from ..QuantumToolbox import linearAlgebra as linAlg #pylint: disable=relative-beyond-top-level
from ..QuantumToolbox import states as qSta #pylint: disable=relative-beyond-top-level

from .base import addDecorator, _recurseIfList
from .baseClasses import qBaseSim, paramBoundBase, setAttr
#from quanguru.classes.exceptions import qSystemInitErrors, qCouplingInitErrors
from .QPro import freeEvolution

def _initStDec(_createAstate):
    def wrapper(obj, inp=None):
        if (issparse(inp) or isinstance(inp, ndarray)):
            if inp.shape[0] != obj.dimension:
                raise ValueError('Dimension mismatch')
            state = inp
        else:
            if inp is None:
                inp = obj.simulation._stateBase__initialStateInput.value

            if isinstance(obj.dimension, int):
                state = _createAstate(obj, inp)
            else:
                state = None
        return state
    return wrapper

def _computeDef(sys, state): # pylint: disable=unused-argument
    pass

def _calculateDef(sys): # pylint: disable=unused-argument
    pass

class genericQSys(qBaseSim):
    r"""
    Base class for both single (:class:`~qSystem`) and composite (:class:`~compQSystem`) quantum system classes.
    The ultimate goal is to make those two classes the same by combining them in here. Currently, a proxy
    :class:`~QuantumSystem` is introduced as a temporary solution.
    """
    label = 'genericQSys'
    #: (**class attribute**) number of instances created internally by the library
    _internalInstances: int = 0
    #: (**class attribute**) number of instances created explicitly by the user
    _externalInstances: int = 0
    #: (**class attribute**) number of total instances = _internalInstances + _externalInstances
    _instances: int = 0

    __slots__ = ['__unitary', '__dimension', '__dimsBefore', '__dimsAfter', '_inpCoef']

    def __init__(self, **kwargs):
        super().__init__()
        #: an internal :class:`~freeEvolution` protocol, this is the default evolution when a simulation is run.
        self.__unitary = freeEvolution(_internal=True)
        self._genericQSys__unitary.superSys = self # pylint: disable=no-member
        self._qBaseSim__simulation.addQSystems(subS=self, Protocol=self._freeEvol) # pylint: disable=no-member
        #: dimension of Hilbert space of the quantum system
        self.__dimension = None
        #: boolean to determine whether initialState inputs contains complex coefficients (the probability amplitudes)
        #: or the populations
        self._inpCoef = False
        self.__dimsBefore = 1
        self.__dimsAfter = 1
        self._named__setKwargs(**kwargs) # pylint: disable=no-member

    def __add__(self, other):
        if isinstance(self, compQSystem) and isinstance(other, qSystem):
            self.addSubSys(other)
            newComp = self
        elif ((isinstance(self, qSystem) and isinstance(other, qSystem)) or  # noqa: W504
              (isinstance(self, compQSystem) and isinstance(other, compQSystem))):
            newComp = compQSystem()
            # FIXME 'stepCount' getter creates problem with None defaults
            newComp.simulation._copyVals(self.simulation, ['totalTime', 'stepSize', 'delStates'])
            newComp.compute = _computeDef
            newComp.simulation.compute = _computeDef
            #newComp.calculate = _calculateDef
            #newComp.simulation.calculate = _calculateDef
            newComp.addSubSys(self)
            if other is self:
                newComp.addSubSys(other.copy())
            else:
                newComp.addSubSys(other)
        elif isinstance(self, qSystem) and isinstance(other, compQSystem):
            other.addSubSys(self)
            newComp = other
        elif isinstance(other, (float, int)):
            newComp = self
        return newComp

    def __sub__(self, other):
        self.removeSubSys(other, _exclude=[])
        return self

    def __rmul__(self, other):
        newComp = compQSystem()
        newComp.addSubSys(self)
        for _ in range(other - 1):
            newComp.addSubSys(self.copy())
        return newComp

    def copy(self, **kwargs): # pylint: disable=arguments-differ
        subSysList = []
        for sys in self.subSys.values():
            subSysList.append(sys.copy())

        if isinstance(self, qSystem):
            newSys = super().copy(dimension=self.dimension, terms=subSysList)
        elif isinstance(self, compQSystem):
            newSys = super().copy()
            for sys in subSysList:
                newSys.addSubSys(sys)

        if self.simulation._stateBase__initialStateInput._value is not None:
            newSys.initialState = self.simulation._stateBase__initialStateInput.value
        newSys._named__setKwargs(**kwargs)
        return newSys

    @property
    def ind(self):
        ind = 0
        if self.superSys is not None:
            ind += list(self.superSys.subSys.values()).index(self)
            if self.superSys.superSys is not None:
                ind += self.superSys.ind
        return ind

    @property
    def _dimsBefore(self):
        return self._genericQSys__dimsBefore if self._genericQSys__dimsBefore != 0 else 1

    @_dimsBefore.setter
    def _dimsBefore(self, val):
        if not isinstance(val, int):
            raise ValueError('?')
        oldVal = self._dimsBefore
        setAttr(self, '_genericQSys__dimsBefore', val)
        for sys in self.subSys.values():
            sys.delMatrices(_exclude=[]) # pylint: disable=protected-access
            if isinstance(sys, genericQSys):
                sys._dimsBefore = int((sys._dimsBefore*val)/oldVal)

    @property
    def _dimsAfter(self):
        return self._genericQSys__dimsAfter if self._genericQSys__dimsAfter != 0 else 1

    @_dimsAfter.setter
    def _dimsAfter(self, val):
        if not isinstance(val, int):
            raise ValueError('?')
        oldVal = self._dimsAfter
        setAttr(self, '_genericQSys__dimsAfter', val)
        for sys in self.subSys.values():
            sys.delMatrices(_exclude=[]) # pylint: disable=protected-access
            if isinstance(sys, genericQSys):
                sys._dimsAfter = int((sys._dimsAfter*val)/oldVal)

    @property
    def dimension(self):
        if self._genericQSys__dimension is None:
            try:
                dims = self.subSysDimensions
                self._genericQSys__dimension = 1 # pylint: disable=assigning-non-slot
                for val in dims:
                    self._genericQSys__dimension *= val # pylint: disable=assigning-non-slot
            except AttributeError:
                print(f'dimension? {self.name}')
        return self._genericQSys__dimension

    @property
    def _totalDim(self):
        return self.dimension * self._dimsBefore * self._dimsAfter#pylint:disable=E1101

    @property
    def _freeEvol(self):
        return self._genericQSys__unitary

    @property
    def unitary(self):
        unitary = self._genericQSys__unitary.unitary
        self._paramBoundBase__paramUpdated = False # pylint: disable=assigning-non-slot
        return unitary

    @qBaseSim.initialState.setter # pylint: disable=no-member
    def initialState(self, inp):
        if self.superSys is not None:
            self.superSys.simulation._stateBase__initialState._value = None
        self.simulation.initialState = inp # pylint: disable=no-member, protected-access
        if (isinstance(self, compQSystem) and isinstance(inp, list)):
            for ind, it in enumerate(inp):
                list(self.qSystems.values())[ind].initialState = it # pylint: disable=no-member

    def _constructMatrices(self):
        for sys in self.subSys.values():
            sys._constructMatrices() # pylint: disable=protected-access

    def addProtocol(self, protocol=None, system=None, protocolRemove=None):
        if system is None:
            system = self
        self.simulation.addProtocol(protocol=protocol, system=system, protocolRemove=protocolRemove)

    def _timeDependency(self, time=None):
        if time is None:
            time = self.simulation._currentTime
        for sys in self.subSys.values():
            sys._timeDependency(time)
        return time

class QuantumSystem(genericQSys):
    #: (**class attribute**) number of instances created internally by the library
    _internalInstances: int = 0
    #: (**class attribute**) number of instances created explicitly by the user
    _externalInstances: int = 0
    #: (**class attribute**) number of total instances = _internalInstances + _externalInstances
    _instances: int = 0
    def __new__(cls, sysType='composite', **kwargs):
        singleKeys = ['frequency', 'operator', 'order', 'dimension']
        for key in singleKeys:
            if key in kwargs.keys():
                sysType = 'single'

        if sysType == 'composite':
            newCls = compQSystem
        elif sysType == 'single':
            newCls = qSystem
        elif sysType == 'system coupling':
            newCls = qCoupling

        if newCls != cls:
            instance = newCls(**kwargs)
        return instance

    __slots__ = []

class compQSystem(genericQSys):
    label = 'QuantumSystem'
    #: (**class attribute**) number of instances created internally by the library
    _internalInstances: int = 0
    #: (**class attribute**) number of instances created explicitly by the user
    _externalInstances: int = 0
    #: (**class attribute**) number of total instances = _internalInstances + _externalInstances
    _instances: int = 0

    __slots__ = ['__qCouplings', '__qSystems', 'couplingName']

    def __init__(self, **kwargs):
        if self.__class__.__name__ == 'compQSystem':
            compQSystem._externalInstances = qSystem._instances + compQSystem._instances
        super().__init__()
        self.__qCouplings = OrderedDict()
        self.__qSystems = OrderedDict()
        self.couplingName = None

        self._named__setKwargs(**kwargs) # pylint: disable=no-member

    def _timeDependency(self, time=None):
        time = super()._timeDependency(time=time)
        for coupling in self.qCouplings.values():
            coupling._timeDependency(time)

    @property
    def subSysDimensions(self):
        return [sys.dimension for sys in self.subSys.values()]

    @property
    def freeHam(self):
        ham = sum([val.totalHam for val in self.qSystems.values()])
        return ham

    @property
    def totalHam(self): # pylint: disable=invalid-overridden-method
        if ((self._paramUpdated) or (self._paramBoundBase__matrix is None)): # pylint: disable=no-member
            self._paramBoundBase__matrix = self.freeHam + self.couplingHam # pylint: disable=assigning-non-slot
            self._paramBoundBase__paramUpdated = False # pylint: disable=assigning-non-slot
        return self._paramBoundBase__matrix # pylint: disable=no-member

    @property
    def couplingHam(self):
        cham = sum([val.totalHam for val in self.qCouplings.values()])
        return cham

    @property
    def qSystems(self):
        return self._compQSystem__qSystems # pylint: disable=no-member

    @addDecorator
    def addSubSys(self, subSys, **kwargs): # pylint: disable=arguments-differ
        newSys = super().addSubSys(subSys, **kwargs)
        if isinstance(newSys, qCoupling):
            self._compQSystem__addCoupling(self._qBase__subSys.pop(newSys.name))  # pylint: disable=no-member
        elif isinstance(newSys, genericQSys):
            self._compQSystem__addSub(newSys)
        else:
            raise TypeError('?')
        newSys._paramBoundBase__paramBound[self.name] = self # pylint: disable=protected-access
        return newSys

    def createSubSys(self, subSysClass, **kwargs):
        return self.addSubSys(subSysClass, **kwargs)

    def __addSub(self, subSys):
        for subS in self._compQSystem__qSystems.values():
            subS._dimsAfter *= subSys.dimension
            subSys._dimsBefore *= subS.dimension

        if subSys._paramBoundBase__matrix is not None:
            for sys in subSys.subSys.values():
                sys._paramBoundBase__matrix = None
        # TODO big question here
        subSys.simulation._bound(self.simulation) # pylint: disable=protected-access
        self._compQSystem__qSystems[subSys.name] = subSys
        subSys.superSys = self
        return subSys

    @_recurseIfList
    def removeSubSys(self, subS, _exclude=[]):#pylint:disable=arguments-differ,dangerous-default-value,too-many-branches
        if isinstance(subS, str):
            subS = self.getByNameOrAlias(subS)
        couplings = list(self.qCouplings.values())
        for coupling in couplings:
            coupling.removeSubSys(subS, _exclude=_exclude)
            if len(coupling._qBase__subSys) == 0: # pylint: disable=protected-access
                self.qCouplings.pop(coupling.name)

        if subS in list(self.subSys.values()):
            for qS in self.subSys.values():
                qS.simulation._stateBase__initialState._value = None
                if qS.ind < subS.ind:
                    qS._dimsAfter = int(qS._dimsAfter/subS.dimension)
                elif qS.ind > subS.ind:
                    qS._dimsBefore = int(qS._dimsBefore/subS.dimension)
            self.qSystems.pop(subS.name)
            _exclude.append(self)
            super().removeSubSys(subS, _exclude=_exclude)
        elif subS in self.qCouplings.values():
            self.qCouplings.pop(subS.name)

        if self not in _exclude:
            _exclude.append(self)
            if ((self._dimsAfter != 1) or (self._dimsBefore != 1)):
                if self.ind < subS.superSys.ind:
                    self._dimsAfter = int(self._dimsAfter/subS.dimension)
                elif self.ind > subS.superSys.ind:
                    self._dimsBefore = int(self._dimsBefore/subS.dimension)

            for sys in self.subSys.values():
                sys.removeSubSys(subS, _exclude=_exclude)
                #_exclude.append(sys)

        if self.superSys is not None:
            self.superSys.removeSubSys(subS, _exclude=_exclude)
            _exclude.append(self.superSys)

        self.delMatrices(_exclude=[])
        self.simulation._stateBase__initialState._value = None
        self._genericQSys__dimension = None # pylint: disable=assigning-non-slot

    @property
    def qCouplings(self):
        return self._compQSystem__qCouplings

    def __addCoupling(self, couplingObj):
        self._compQSystem__qCouplings[couplingObj.name] = couplingObj
        couplingObj.superSys = self
        return couplingObj

    def createSysCoupling(self, *args, **kwargs):
        newCoupling = self.addSubSys(qCoupling, **kwargs)
        newCoupling.addTerm(*args)
        return newCoupling

    def addSysCoupling(self, couplingObj):
        self.addSubSys(couplingObj)

    @_initStDec
    def _createAstate(self, inp=None):
        if inp is None:
            inp = [qsys._createAstate() for qsys in self.subSys.values()]
        elif isinstance(inp, list):
            inp = [qsys._createAstate(inp[qsys.ind]) for qsys in self.subSys.values()]
        else:
            raise TypeError('?')
        return linAlg.tensorProd(*inp)

    def _constructMatrices(self):
        super()._constructMatrices()
        for sys in self.qCouplings.values():
            sys._constructMatrices() # pylint: disable=protected-access

    def updateDimension(self, qSys, newDimVal, oldDimVal=None, _exclude=[]):#pylint:disable=dangerous-default-value,too-many-branches
        # TODO can be combined with removeSubSys by a decorator or another method to simplfy both
        if oldDimVal is None:
            oldDimVal = qSys._genericQSys__dimension
        self._genericQSys__dimension = None # pylint: disable=assigning-non-slot
        if qSys in self.qSystems.values():
            _exclude.append(self)
            qSys._genericQSys__dimension = newDimVal
            ind = qSys.ind
            for qS in self.qSystems.values():
                if qS.ind < ind:
                    qS._dimsAfter = int((qS._dimsAfter*newDimVal)/oldDimVal)
                elif qS.ind > ind:
                    qS._dimsBefore = int((qS._dimsBefore*newDimVal)/oldDimVal)

            #if self.simulation._stateBase__initialStateInput.value is not None: # pylint: disable=no-member
            #    self.initialState = self.simulation._stateBase__initialStateInput.value # pylint: disable=no-member
            self._paramUpdated = True
            #self._constructMatrices()
            #for sys in self.subSys.values():
            #    if sys.simulation._stateBase__initialStateInput.value is not None:
            #        sys.initialState = sys.simulation._stateBase__initialStateInput.value

        if self not in _exclude:
            _exclude.append(self)
            if ((self._dimsAfter != 1) or (self._dimsBefore != 1)):
                if self.ind < qSys.superSys.ind:
                    self._dimsAfter = int((self._dimsAfter*newDimVal)/oldDimVal)
                elif self.ind > qSys.superSys.ind:
                    self._dimsBefore = int((self._dimsBefore*newDimVal)/oldDimVal)
            else:
                for sys in self.subSys.values():
                    if sys not in _exclude:
                        _exclude.append(sys)
                        if sys.ind < qSys.superSys.ind:
                            sys._dimsAfter = int((sys._dimsAfter*newDimVal)/oldDimVal)
                        elif sys.ind > qSys.superSys.ind:
                            sys._dimsBefore = int((sys._dimsBefore*newDimVal)/oldDimVal)

        if self.superSys is not None:
            self.superSys.updateDimension(qSys=qSys, newDimVal=newDimVal, oldDimVal=oldDimVal, _exclude=_exclude)
        self.delMatrices(_exclude=[])
        for c in self.qCouplings.values():
            c.delMatrices(_exclude=[])
        return qSys

class termTimeDep(paramBoundBase):
    label = '_timeDep'
    #: (**class attribute**) number of instances created internally by the library
    _internalInstances: int = 0
    #: (**class attribute**) number of instances created explicitly by the user
    _externalInstances: int = 0
    #: (**class attribute**) number of total instances = _internalInstances + _externalInstances
    _instances: int = 0

    __slots__ = ['timeDependency', '__frequency', '__order', '__operator']

    def __init__(self, **kwargs):
        super().__init__()
        self.timeDependency = None
        self.__frequency = None
        self.__order = 1
        self.__operator = None
        self._named__setKwargs(**kwargs) # pylint: disable=no-member

    def copy(self, **kwargs):  # pylint: disable=arguments-differ
        newSys = super().copy(frequency=self.frequency, operator=self.operator, order=self.order, **kwargs)
        return newSys

    @property
    def operator(self):
        return self._termTimeDep__operator

    @operator.setter
    def operator(self, op):
        self._paramBoundBase__matrix = None # pylint: disable=assigning-non-slot
        setAttr(self, '_termTimeDep__operator', op)

    @property
    def order(self):
        return self._termTimeDep__order

    @order.setter
    def order(self, ordVal):
        setAttr(self, '_termTimeDep__order', ordVal)
        if self._paramBoundBase__matrix is not None: # pylint: disable=no-member
            self.freeMat = None

    @property
    def frequency(self):
        return self._termTimeDep__frequency

    @frequency.setter
    def frequency(self, freq):
        freq = 0 if freq == 0.0 else freq
        setAttr(self, '_termTimeDep__frequency', freq)

    def _constructMatrices(self):
        pass

    @property
    def totalHam(self):
        return self.frequency*self.freeMat

    @property
    def freeMat(self):
        #if ((self._paramBoundBase__matrix is None) or (self._paramUpdated)): # pylint: disable=no-member
        if self._paramBoundBase__matrix is None: # pylint: disable=no-member
            self.freeMat = None
            self._paramBoundBase__paramUpdated = False # pylint: disable=assigning-non-slot
        return self._paramBoundBase__matrix # pylint: disable=no-member

    @freeMat.setter
    def freeMat(self, qMat):
        if qMat is not None:
            self._paramBoundBase__matrix = qMat # pylint: disable=no-member, assigning-non-slot
        else:
            #if len(self._qBase__subSys) == 0: # pylint: disable=no-member
            #    raise ValueError('No operator is given for coupling Hamiltonian')
            #if self.operator is None:
            #    raise ValueError('No operator is given for free Hamiltonian')
            self._constructMatrices()

    def _timeDependency(self, time=None):
        if time is None:
            time = self.superSys.simulation._currentTime

        if callable(self.timeDependency):
            if hasattr(self, 'frequency'):
                self.frequency = self.timeDependency(self, time) # pylint: disable=assigning-non-slot,not-callable
            elif hasattr(self, 'couplingStrength'):
                self.couplingStrength = self.timeDependency(self, time) #pylint:disable=assigning-non-slot,not-callable

class term(termTimeDep):
    label = 'term'
    #: (**class attribute**) number of instances created internally by the library
    _internalInstances: int = 0
    #: (**class attribute**) number of instances created explicitly by the user
    _externalInstances: int = 0
    #: (**class attribute**) number of total instances = _internalInstances + _externalInstances
    _instances: int = 0

    __slots__ = []

    @paramBoundBase.superSys.setter
    def superSys(self, supSys):
        r"""
        Extends superSys setter to also add aliases to self.
        New aliases are (any name/alias of superSys) + Term + (number of terms)
        TODO What if there is already a superSys, and also alias list contains user given aliases as well.
        """
        paramBoundBase.superSys.fset(self, supSys) # pylint: disable=no-member
        termCount = len(self.superSys.subSys) if self in self.superSys.subSys.values() else len(self.superSys.subSys)+1 # pylint: disable=no-member,line-too-long # noqa: E501
        self.alias = [na+"Term"+str(termCount) for na in self.superSys.name._aliasClass__members()] # pylint: disable=no-member, protected-access,line-too-long # noqa: E501

    @property
    def _freeMatSimple(self):
        h = self._constructMatrices(dimsBefore=1, dimsAfter=1, setMat=False)
        return h

    def _constructMatrices(self, dimsBefore=None, dimsAfter=None, setMat=True): #pylint:disable=arguments-differ
        if dimsBefore is None:
            dimsBefore = self.superSys._dimsBefore # pylint: disable=no-member

        if dimsAfter is None:
            dimsAfter = self.superSys._dimsAfter # pylint: disable=no-member

        if not (isinstance(self.superSys.dimension, (int, int64, int32, int16)) and callable(self.operator)): # pylint: disable=no-member
            raise TypeError('?')

        dimension = self.superSys._genericQSys__dimension # pylint: disable=no-member
        if self.operator in [qOps.Jz, qOps.Jy, qOps.Jx, qOps.Jm, qOps.Jp, qOps.Js]:
            dimension = 0.5*(dimension-1)

        if self.operator not in [qOps.sigmam, qOps.sigmap, qOps.sigmax, qOps.sigmay, qOps.sigmaz]:
            mat = qOps.compositeOp(self.operator(dimension), #pylint:disable=assigning-non-slot
                                   dimsBefore, dimsAfter)**self.order
        else: # pylint: disable=bare-except
            mat = qOps.compositeOp( # pylint: disable=no-member, assigning-non-slot
                self.operator(), dimsBefore, dimsAfter)**self.order

        if setMat:
            self._paramBoundBase__matrix = mat #pylint:disable=assigning-non-slot
        return mat

class qSystem(genericQSys):
    label = 'QuantumSystem'
    #: (**class attribute**) number of instances created internally by the library
    _internalInstances: int = 0
    #: (**class attribute**) number of instances created explicitly by the user
    _externalInstances: int = 0
    #: (**class attribute**) number of total instances = _internalInstances + _externalInstances
    _instances: int = 0

    __slots__ = []
    #@qSystemInitErrors
    def __init__(self, **kwargs):
        if self.__class__.__name__ == 'qSystem':
            qSystem._externalInstances = qSystem._instances + compQSystem._instances
        super().__init__()
        qSysKwargs = ['terms', 'subSys', 'name', 'superSys', 'dimension']
        for key in qSysKwargs:
            val = kwargs.pop(key, None)
            if val is not None:
                setattr(self, key, val)
        self._named__setKwargs(**kwargs) # pylint: disable=no-member

        if len(self.subSys) == 0:
            self.addSubSys(term(superSys=self, **kwargs))

    # @genericQSys.name.setter #pylint: disable=no-member
    # def name(self, name):
    #     oldName = self.name
    #     genericQSys.name.fset(self, name) # pylint: disable=no-member
    #     for ii, sys in enumerate(self.subSys.values()):
    #         if sys.name == (oldName + 'term' + str(ii)):
    #             sys.name = self.superSys.name + 'term' + str(ii+1) # pylint: disable=no-member

    @genericQSys.dimension.setter # pylint: disable=no-member
    def dimension(self, newDimVal):
        if not isinstance(newDimVal, (int, int64, int32, int16)):
            raise ValueError('Dimension is not int')

        oldDimVal = self._genericQSys__dimension # pylint: disable=no-member

        for sys in self.subSys.values():
            sys.delMatrices(_exclude=[]) # pylint: disable=protected-access

        setAttr(self, '_genericQSys__dimension', newDimVal)
        # FIXME these should be called only if oldDim != newDim
        #if self.simulation._stateBase__initialStateInput.value is not None: # pylint: disable=protected-access
        #    self.initialState = self.simulation._stateBase__initialStateInput.value # pylint: disable=protected-access

        if isinstance(self.superSys, compQSystem):
            self.superSys.updateDimension(self, newDimVal, oldDimVal, _exclude=[]) # pylint: disable=no-member

    @property
    def totalHam(self): # pylint: disable=invalid-overridden-method
        if ((self._paramUpdated) or (self._paramBoundBase__matrix is None)): # pylint: disable=no-member
            h = sum([(obj.frequency * obj.freeMat) for obj in self.subSys.values()])
            self._paramBoundBase__matrix = h # pylint: disable=assigning-non-slot
            self._paramBoundBase__paramUpdated = False # pylint: disable=assigning-non-slot
        return self._paramBoundBase__matrix # pylint: disable=no-member

    @property
    def _totalHamSimple(self):
        return sum([(obj.frequency * obj._freeMatSimple) for obj in self.subSys.values()])#pylint:disable=protected-access

    @property
    def freeMat(self):
        return self.firstTerm.freeMat # pylint: disable=no-member

    @freeMat.setter
    def freeMat(self, qOpsFunc):
        if callable(qOpsFunc):
            self.firstTerm.operator = qOpsFunc
            self.firstTerm._constructMatrices() # pylint: disable=protected-access
        elif qOpsFunc is not None:
            self.firstTerm._paramBoundBase__matrix = qOpsFunc  # pylint: disable=assigning-non-slot
        else:
            if self.firstTerm.operator is None:
                raise ValueError('No operator is given for free Hamiltonian')
            self.firstTerm._constructMatrices() # pylint: disable=protected-access

    @property
    def operator(self):
        operators = [obj._termTimeDep__operator for obj in list(self.subSys.values())] # pylint: disable=protected-access
        return operators if len(operators) > 1 else operators[0]

    @operator.setter
    def operator(self, op):
        self.firstTerm.operator = op

    @property
    def frequency(self):
        #frequencies = [obj._termTimeDep__frequency for obj in list(self.subSys.values())] # pylint: disable=protected-access
        #return frequencies if len(frequencies) > 1 else frequencies[0]
        return self.firstTerm.frequency

    @frequency.setter
    def frequency(self, freq):
        self.firstTerm.frequency = freq

    @property
    def order(self):
        orders = [obj._termTimeDep__order for obj in list(self.subSys.values())] # pylint: disable=protected-access
        return orders if len(orders) > 1 else orders[0]

    @order.setter
    def order(self, ordVal):
        self.firstTerm.order = ordVal

    @property
    def firstTerm(self):
        return list(self.subSys.values())[0]

    @property
    def terms(self):
        qSys = list(self.subSys.values())
        return qSys if len(qSys) > 1 else qSys[0]

    @addDecorator
    def addSubSys(self, subS, **kwargs):
        if not isinstance(subS, term):
            raise TypeError('?')
        kwargs['superSys'] = self
        newS = super().addSubSys(subS, **kwargs)
        # FIXME use setAttr, check also for operator
        self._paramUpdated = True
        newS._paramBoundBase__paramBound[self.name] = self # pylint: disable=protected-access
        return subS

    @_recurseIfList
    def removeSubSys(self, subS, _exclude=[]): # pylint: disable=arguments-differ, dangerous-default-value
        if self not in _exclude:
            _exclude.append(self)
            if self.superSys is not None:
                self.superSys.removeSubSys(subS, _exclude=_exclude)
            super().removeSubSys(subS, _exclude=_exclude)

    @terms.setter
    def terms(self, subS):
        genericQSys.subSys.fset(self, subS) # pylint: disable=no-member
        for sys in self.subSys.values():
            sys.superSys = self

    def addTerm(self, operator, frequency=0, order=1):
        newTerm = self.addSubSys(term(operator=operator, frequency=frequency, order=order, superSys=self))
        return newTerm

    @_recurseIfList
    def removeTerm(self, termObj):
        self.removeSubSys(termObj, _exclude=[])

    @_initStDec
    def _createAstate(self, inp=None):
        if inp is None:
            raise ValueError(self.name + ' is not given an initial state')
        return qSta.superPos(self.dimension, inp, not self._inpCoef)

class Spin(qSystem): # pylint: disable=too-many-ancestors
    label = 'Spin'
    #: (**class attribute**) number of instances created internally by the library
    _internalInstances: int = 0
    #: (**class attribute**) number of instances created explicitly by the user
    _externalInstances: int = 0
    #: (**class attribute**) number of total instances = _internalInstances + _externalInstances
    _instances: int = 0

    __slots__ = ['__jValue']
    def __init__(self, **kwargs):
        super().__init__(terms=kwargs.pop('terms', None), subSys=kwargs.pop('subSys', None))
        self.operator = qOps.Jz
        self.__jValue = None
        self._named__setKwargs(**kwargs) # pylint: disable=no-member

    @property
    def jValue(self):
        return (self._genericQSys__dimension-1)/2 # pylint: disable=no-member

    @jValue.setter
    def jValue(self, value):
        self._Spin__jValue = value # pylint: disable=assigning-non-slot
        self.dimension = int((2*value) + 1)

class Qubit(Spin): # pylint: disable=too-many-ancestors
    label = 'Qubit'
    #: (**class attribute**) number of instances created internally by the library
    _internalInstances: int = 0
    #: (**class attribute**) number of instances created explicitly by the user
    _externalInstances: int = 0
    #: (**class attribute**) number of total instances = _internalInstances + _externalInstances
    _instances: int = 0

    __slots__ = []
    def __init__(self, **kwargs):
        super().__init__(terms=kwargs.pop('terms', None), subSys=kwargs.pop('subSys', None))
        kwargs['dimension'] = 2
        self.operator = qOps.Jz
        self._named__setKwargs(**kwargs) # pylint: disable=no-member

class Cavity(qSystem): # pylint: disable=too-many-ancestors
    label = 'Cavity'
    #: (**class attribute**) number of instances created internally by the library
    _internalInstances: int = 0
    #: (**class attribute**) number of instances created explicitly by the user
    _externalInstances: int = 0
    #: (**class attribute**) number of total instances = _internalInstances + _externalInstances
    _instances: int = 0

    __slots__ = []
    def __init__(self, **kwargs):
        super().__init__(terms=kwargs.pop('terms', None), subSys=kwargs.pop('subSys', None))
        self.operator = qOps.number
        self._named__setKwargs(**kwargs) # pylint: disable=no-member

class qCoupling(termTimeDep):
    label = 'qCoupling'
    #: (**class attribute**) number of instances created internally by the library
    _internalInstances: int = 0
    #: (**class attribute**) number of instances created explicitly by the user
    _externalInstances: int = 0
    #: (**class attribute**) number of total instances = _internalInstances + _externalInstances
    _instances: int = 0

    __slots__ = []

    #@qCouplingInitErrors
    def __init__(self, *args, **kwargs):
        super().__init__()
        self._named__setKwargs(**kwargs) # pylint: disable=no-member
        self.addTerm(*args)

    # TODO might define setters
    @property
    def couplingOperators(self):
        ops = []
        for co in self._qBase__subSys.values(): # pylint: disable=no-member
            ops.append(co[1])
        return ops

    @property
    def coupledSystems(self):
        ops = []
        for co in self._qBase__subSys.values(): # pylint: disable=no-member
            ops.append(co[0])
        return ops

    @property
    def couplingStrength(self):
        return self.frequency

    @couplingStrength.setter
    def couplingStrength(self, strength):
        self.frequency = strength

    def __coupOrdering(self, qts): # pylint: disable=no-self-use
        qts = sorted(qts, key=lambda x: x[0], reverse=False)
        oper = qts[0][1]
        for ops in range(len(qts)-1):
            oper = oper @ qts[ops+1][1]
        return oper

    def _constructMatrices(self):
        cMats = []
        for ind in range(len(self._qBase__subSys)): # pylint: disable=no-member
            qts = []
            for indx in range(len(list(self._qBase__subSys.values())[ind])): # pylint: disable=no-member
                sys = list(self._qBase__subSys.values())[ind][0][indx] # pylint: disable=no-member
                order = sys.ind
                oper = list(self._qBase__subSys.values())[ind][1][indx] # pylint: disable=no-member
                if oper in [qOps.sigmam, qOps.sigmap, qOps.sigmax, qOps.sigmay, qOps.sigmaz]:
                    cHam = qOps.compositeOp(oper(), sys._dimsBefore, sys._dimsAfter)
                else:
                    dimension = sys._genericQSys__dimension
                    if oper in [qOps.Jz, qOps.Jy, qOps.Jx, qOps.Jm, qOps.Jp, qOps.Js]:
                        dimension = 0.5*(dimension-1)
                    cHam = qOps.compositeOp(oper(dimension), sys._dimsBefore, sys._dimsAfter)
                ts = [order, cHam]
                qts.append(ts)
            cMats.append(self._qCoupling__coupOrdering(qts))
        #h = []
        #if ((self.couplingStrength != 0) or (self.couplingStrength is not None)):
        #    h = [self.couplingStrength * sum(cMats)]
        self._paramBoundBase__matrix = sum(cMats) # pylint: disable=assigning-non-slot
        return self._paramBoundBase__matrix # pylint: disable=no-member

    def __addTerm(self, count, ind, sys, *args):
        if callable(args[count][ind]):
            lo = len(self.subSys)
            self._qBase__subSys[str(lo)] = (sys, tuple(args[count])) # pylint: disable=no-member
            count += 1
            if count < len(args):
                count = self.__addTerm(count, ind, sys, *args)
        return count

    def addTerm(self, *args):
        counter = 0
        while counter in range(len(args)):
            # TODO write a generalisation for this one
            if isinstance(args[counter][0], qSystem):
                qSystems = args[counter]
                if callable(args[counter+1][1]):
                    #if tuple(args[counter + 1]) in self._qBase__subSys.keys(): # pylint: disable=no-member
                    #    print(tuple(args[counter + 1]), 'already exists')
                    lo = len(self.subSys)
                    self._qBase__subSys[str(lo)] = (qSystems, tuple(args[counter + 1])) # pylint: disable=no-member
                    counter += 2
                # TODO does not have to pass qSystem around
                if counter < len(args):
                    counter = self._qCoupling__addTerm(counter, 1, qSystems, *args)
        self._paramBoundBase__matrix = None # pylint: disable=assigning-non-slot
        return self

    @_recurseIfList
    def removeSysCoupling(self, sys):
        self.removeSubSys(sys, _exclude=[])

    @_recurseIfList
    def removeSubSys(self, subS, _exclude=[]): # pylint: disable=dangerous-default-value
        vals = self._qBase__subSys.values() # pylint: disable=no-member
        for ind, val in enumerate(vals):
            systs = val[0]
            if subS in systs:
                self._qBase__subSys.pop(str(ind)) # pylint: disable=no-member


'''
------------------------------------------------------------------------
Last updated 4/19/2018

This script produces the tables and figures for the Tax Function
Integration Paper

------------------------------------------------------------------------
'''

import pickle
import numpy as np
import scipy.optimize as opt
import matplotlib.pyplot as plt
import os
import xlsxwriter
import ogusa
from ogusa.utils import REFORM_DIR, BASELINE_DIR

# Read in tax function parameters
tax_func_params = {}
tax_func_list = ['DEP', 'DEP', 'DEP_totalinc', 'DEP_totalinc', 'GS',
                 'GS']
age_specific_list = [True, False, True, False, True, False]
guid_list = ['_DEP_TI_noAge', '_DEP_TI_Age', '_DEP_noAge', '_DEP_Age',
             '_GS_noAge', '_GS_Age']
tax_func_params_base = {}
tax_func_params_reform = {}
tpi_base = {}
tpi_reform = {}
ss_base = {}
ss_reform = {}
for guid in guid_list:
        # NOTE TAHT SWITHC BASELINE AND REFORM SINCE TCJA CURRENT LAW
        reform_path = os.path.join(BASELINE_DIR, guid,
                                   'TxFuncEst_baseline' + guid + '.pkl')
        base_path = os.path.join(REFORM_DIR, guid,
                                 'TxFuncEst_policy' + guid + '.pkl')
        tax_func_params_base[guid] = pickle.load(
            open(base_path, 'rb'), encoding='latin')
        tax_func_params_reform[guid] = pickle.load(
            open(reform_path, 'rb'), encoding='latin')
        tpi_reform_path = os.path.join(BASELINE_DIR, guid,
                                       'TPI/TPI_vars.pkl')
        tpi_base_path = os.path.join(REFORM_DIR, guid,
                                     'TPI/TPI_vars.pkl')
        tpi_base[guid] = pickle.load(
            open(tpi_base_path, 'rb'), encoding='latin')
        tpi_reform[guid] = pickle.load(
            open(tpi_reform_path, 'rb'), encoding='latin')
        ss_reform_path = os.path.join(BASELINE_DIR, guid,
                                      'SS/SS_vars.pkl')
        ss_base_path = os.path.join(REFORM_DIR, guid,
                                    'SS/SS_vars.pkl')
        ss_base[guid] = pickle.load(
            open(ss_base_path, 'rb'), encoding='latin')
        ss_reform[guid] = pickle.load(
            open(ss_reform_path, 'rb'), encoding='latin')

'''
------------------------------------------------------------------------
    Tables - all tables for paper saved to different worksheets in an
    Excel workbook
------------------------------------------------------------------------
'''
# open Excel workbook
workbook = xlsxwriter.Workbook('TFI_Tables.xlsx',
                               {'nan_inf_to_errors': True})

# Table 2: variation in phi by age for DEP function
# create list of list with table info
rate_labels = ['$ETR$', '$MTRx$', '$MTRy$']
rate_types = ['etr', 'mtrx', 'mtry']
table2 = []
table2.append([''])
table2[0].extend(('21 to 54', '55 to 65', '66 to 80', 'All ages'))
for i, label in enumerate(rate_labels):
    table2.append([label])
    table2[i + 1].extend((
        tax_func_params_base['_DEP_Age']['tfunc_' + rate_types[i] +
                                         '_params_S'][:34, 0, -1].mean(),
        tax_func_params_base['_DEP_Age']['tfunc_' + rate_types[i] +
                                         '_params_S'][34:45, 0, -1].mean(),
        tax_func_params_base['_DEP_Age']['tfunc_' + rate_types[i] +
                                         '_params_S'][45:60, 0, -1].mean(),
        tax_func_params_base['_DEP_Age']['tfunc_' + rate_types[i] +
                                         '_params_S'][:, 0, -1].mean()))
# save table of info to Excel
worksheet = workbook.add_worksheet('Table 2')
worksheet.merge_range('B1:D1', 'Age ranges')
for i, val in enumerate(table2):
    for j, val2 in enumerate(table2[i]):
        worksheet.write(i + 1, j, val2)

# Table 3: comparing std errors across tax functions
tax_func_labels = ['Ratio of polynomials, ETR',
                   'Ratio of polynomials, vary by age, ETR',
                   'Ratio of polynomials, vary by income source ETR',
                   'Ratio of polynomials, vary by age and income source ETR',
                   'Gouveia and Strauss (1994), ETR',
                   'Gouveia and Strauss (1994), vary by age, ETR']
table3 = []
table3.append([''])
table3[0].extend(('All ages', '21 to 54', '55 to 65', '66 to 80'))
for i, label in enumerate(tax_func_labels):
    table3.append([label])
    table3[i + 1].extend((
        tax_func_params_base[guid_list[i]]['tfunc_etr_sumsq'][:, 0].sum()
        / tax_func_params_base[guid_list[i]]['tfunc_etr_obs'][:, 0].sum(),
        tax_func_params_base[guid_list[i]]['tfunc_etr_sumsq'][:34, 0].sum()
        / tax_func_params_base[guid_list[i]]['tfunc_etr_obs'][:34, 0].sum(),
        tax_func_params_base[guid_list[i]]['tfunc_etr_sumsq'][34:45, 0].sum()
        / tax_func_params_base[guid_list[i]]['tfunc_etr_obs'][34:45, 0].sum(),
        tax_func_params_base[guid_list[i]]['tfunc_etr_sumsq'][45:60, 0].sum()
        / tax_func_params_base[guid_list[i]]['tfunc_etr_obs'][45:60, 0].sum()))
# save table of info to Excel
worksheet = workbook.add_worksheet('Table 3')
worksheet.merge_range('B1:E1', 'Age ranges')
for i, val in enumerate(table3):
    for j, val2 in enumerate(table3[i]):
        worksheet.write(i + 1, j, val2)

# Table 4: comparing SSE across tax functions
table4 = []
table4.append([''])
table4[0].extend(('All ages', '21 to 54', '55 to 65', '66 to 80'))
for i, label in enumerate(tax_func_labels):
    table4.append([label])
    table4[i + 1].extend((
        tax_func_params_base[guid_list[i]]['tfunc_etr_sumsq'][:, 0].mean(),
        tax_func_params_base[guid_list[i]]['tfunc_etr_sumsq'][:34, 0].mean(),
        tax_func_params_base[guid_list[i]]['tfunc_etr_sumsq'][34:45, 0].mean(),
        tax_func_params_base[guid_list[i]]['tfunc_etr_sumsq'][45:60, 0].mean()))
# save table of info to Excel
worksheet = workbook.add_worksheet('Table 4')
worksheet.merge_range('B1:E1', 'Age ranges')
for i, val in enumerate(table4):
    for j, val2 in enumerate(table4[i]):
        worksheet.write(i + 1, j, val2)

# Table 5: comparing observations across tax functions
table5 = []
table5.append([''])
table5[0].extend(('All ages', '21 to 54', '55 to 65', '66 to 80'))
for i, label in enumerate(tax_func_labels):
    table5.append([label])
    table5[i + 1].extend((
        tax_func_params_base[guid_list[i]]['tfunc_etr_obs'][:, 0].sum(),
        tax_func_params_base[guid_list[i]]['tfunc_etr_obs'][:34, 0].sum(),
        tax_func_params_base[guid_list[i]]['tfunc_etr_obs'][34:45, 0].sum(),
        tax_func_params_base[guid_list[i]]['tfunc_etr_obs'][45:60, 0].sum()))
# save table of info to Excel
worksheet = workbook.add_worksheet('Table 5')
worksheet.merge_range('B1:E1', 'Age ranges')
for i, val in enumerate(table5):
    for j, val2 in enumerate(table5[i]):
        worksheet.write(i + 1, j, val2)

# Table 6: parameter estimates from baseline and reform for DEP funcs
# report parameters from function for 42 year old in first year of window
param_names = ['$A$', '$B$', '$C$', '$D$', '$max_x$', '$min_x$',
               '$max_y$', '$min_y$', '$shift_x$', '$shift_y$',
               '$shift$', '$share$']
table6 = []
table6.append(['Parameter'])
table6[0].extend(tuple(rate_labels * 2))
for i, label in enumerate(param_names):
    table6.append([label])
    for j, rate in enumerate(rate_types):  # for baseline results
        table6[i + 1].append(
            tax_func_params_base['_DEP_Age']['tfunc_' + rate +
                                             '_params_S'][21, 0, i])
    for j, rate in enumerate(rate_types):  # for reform results
        table6[i + 1].append(
            tax_func_params_reform['_DEP_Age']['tfunc_' + rate +
                                               '_params_S'][21, 0, i])
table6.append(['Obs (N)'])
table6.append(['SSE'])
for j, rate in enumerate(rate_types):
    table6[-2].append(
        tax_func_params_base['_DEP_Age']['tfunc_' + rate +
                                         '_obs'][21, 0])
    table6[-1].append(
        tax_func_params_base['_DEP_Age']['tfunc_' + rate +
                                         '_sumsq'][21, 0])
for j, rate in enumerate(rate_types):
    table6[-2].append(
        tax_func_params_reform['_DEP_Age']['tfunc_' + rate +
                                           '_obs'][21, 0])
    table6[-1].append(
        tax_func_params_reform['_DEP_Age']['tfunc_' + rate +
                                           '_sumsq'][21, 0])
# save table of info to Excel
worksheet = workbook.add_worksheet('Table 6')
worksheet.merge_range('B1:D1', '2017 Law')
worksheet.merge_range('E1:G1', 'TCJA')
for i, val in enumerate(table6):
    for j, val2 in enumerate(table6[i]):
        worksheet.write(i + 1, j, val2)

# Table 8: GDP changes across tax functions
tax_func_labels = ['Ratio of polynomials, ETR',
                   'Ratio of polynomials, vary by age, ETR',
                   'Ratio of polynomials, vary by income source ETR',
                   'Ratio of polynomials, vary by age and income source ETR',
                   'Gouveia and Strauss (1994), ETR',
                   'Gouveia and Strauss (1994), vary by age, ETR']
table8 = []
for i, label in enumerate(tax_func_labels):
    table8.append([label])
    for y in range(10):
        table8[i].append(((tpi_reform[guid_list[i]]['Y'][y] -
                           tpi_base[guid_list[i]]['Y'][y])
                          / tpi_base[guid_list[i]]['Y'][y]))
    table8[i].extend((
        ((tpi_reform[guid_list[i]]['Y'][:11].sum() -
          tpi_base[guid_list[i]]['Y'][:11].sum())
         / tpi_base[guid_list[i]]['Y'][:11].sum()),
        ((ss_reform[guid_list[i]]['Yss'] -
          ss_base[guid_list[i]]['Yss'])
         / ss_base[guid_list[i]]['Yss'])))
# save table of info to Excel
worksheet = workbook.add_worksheet('Table 8')
worksheet.write(0, 0, 'Tax Function')
for y in range(2018, 2028):
    worksheet.write(0, y - 2017, str(y))
worksheet.write(0, 11, '2018-2027')
worksheet.write(0, 12, 'SS')
for i, val in enumerate(table8):
    for j, val2 in enumerate(table8[i]):
        worksheet.write(i + 1, j, val2)

# Table 9: Changes in all macro aggregates for DEP only
results_labels = ['GDP', 'Conusmption', 'Investment', 'Hours Worked',
                  'Avg. Wage', 'Interest Rate', 'Total Taxes']
var_names = ['Y', 'C', 'I', 'L', 'w', 'r', 'REVENUE']
ss_var_names = ['Yss', 'Css', 'Iss', 'Lss', 'wss', 'rss', 'revenue_ss']
table9 = []
for i, label in enumerate(results_labels):
    table9.append([label])
    for y in range(10):
        table9[i].append(((tpi_reform['_DEP_Age'][var_names[i]][y] -
                           tpi_base['_DEP_Age'][var_names[i]][y])
                          / tpi_base['_DEP_Age'][var_names[i]][y]))
    table9[i].extend((
        ((tpi_reform['_DEP_Age'][var_names[i]][:11].sum() -
          tpi_base['_DEP_Age'][var_names[i]][:11].sum())
         / tpi_base['_DEP_Age'][var_names[i]][:11].sum()),
        ((ss_reform['_DEP_Age'][ss_var_names[i]] -
          ss_base['_DEP_Age'][ss_var_names[i]])
         / ss_base['_DEP_Age'][ss_var_names[i]])))
# save table of info to Excel
worksheet = workbook.add_worksheet('Table 9')
worksheet.write(0, 0, 'Macroeconomic Variables')
for y in range(2018, 2028):
    worksheet.write(0, y - 2017, str(y))
worksheet.write(0, 11, '2018-2027')
worksheet.write(0, 12, 'SS')
for i, val in enumerate(table9):
    for j, val2 in enumerate(table9[i]):
        worksheet.write(i + 1, j, val2)

workbook.close()

'''
FIGURES
'''

"""
This version: 12 May 2018
Written by Kerk Phillips

This program fits tax functions for "Integrating Microsimulation Models of Tax
Policy into a DGE Macroeconomic Model: A Canonical Example," by DeBaker, Evans
and Phillips.

This file fits data from Tax Calc to three different functions, compares the
goodness-of-fit, and plots the functions against the data.  The data sample
can be restricted by the age and the capital income of the individuals.
"""

def GS(coeffs, *args):
    '''
    This is the functional from from Gouveia and Strass (1994) with an
    additional free parameter
    '''
    # unpack coefficients
    phi0, phi1, phi2 = coeffs
    # unpack data
    I, taxes, wgts = args
    # I = x+y
    errors = (taxes/I) - ((phi0*(I - (I**(-phi1) + phi2)**(-1/phi1)))/I)
    wsse = (wgts * (errors ** 2)).sum()
    print('GS SSE: ', wsse)
    print('coeffs = ', phi0, phi1, phi2)
    return wsse

'''
The functions below call those above and return the square-root of the sum of
squared errors (SSE) for each model
'''
# load data from pkl file then clean it up (like in txfunc.py)
micro_data = pickle.load(open("micro_data_2017Law.pkl", "rb"),
                         encoding='latin1')
data_orig = micro_data['2018']
data_orig['Total Labor Income'] = \
    (data_orig['Wage income'] +
     data_orig['SE income'])
data_orig['Effective Tax Rate'] = \
    (data_orig['Total tax liability'] /
     data_orig["Adjusted total income"])
data_orig["Total Capital Income"] = \
    (data_orig['Adjusted total income'] -
     data_orig['Total Labor Income'])
# use weighted avg for MTR labor - abs value because
# SE income may be negative
data_orig['MTR Labor'] = \
    (data_orig['MTR wage income'] * (data_orig['Wage income'] /
     (data_orig['Wage income'].abs() +
     data_orig['SE income'].abs())) +
     data_orig['MTR SE income'] *
     (data_orig['SE income'].abs() /
     (data_orig['Wage income'].abs() +
      data_orig['SE income'].abs())))
data = data_orig[['Age', 'MTR Labor', 'MTR capital income',
                  'Total Labor Income', 'Total Capital Income',
                  'Adjusted total income', 'Effective Tax Rate',
                  'Weights']]
# Clean up the data by dropping outliers
# drop all obs with ETR > 0.65
data_trnc = \
    data.drop(data[data['Effective Tax Rate'] > 0.65].index)
# drop all obs with ETR < -0.15
data_trnc = \
    data_trnc.drop(data_trnc[data_trnc['Effective Tax Rate']
                             < -0.15].index)
# drop all obs with ATI, TLI, TCI < $5
data_trnc = data_trnc[(data_trnc['Adjusted total income'] >= 5)
                      & (data_trnc['Total Labor Income'] >= 5) &
                      (data_trnc['Total Capital Income'] >= 5)]

# drop all obs with MTR on capital income > 10.99
data_trnc = \
    data_trnc.drop(data_trnc[data_trnc['MTR capital income']
                             > 0.99].index)
# drop all obs with MTR on capital income < -0.45
data_trnc = \
    data_trnc.drop(data_trnc[data_trnc['MTR capital income']
                             < -0.45].index)
# drop all obs with MTR on labor income > 10.99
data_trnc = data_trnc.drop(data_trnc[data_trnc['MTR Labor']
                                     > 0.99].index)
# drop all obs with MTR on labor income < -0.45
data_trnc = data_trnc.drop(data_trnc[data_trnc['MTR Labor']
                                     < -0.45].index)
data_baseline = data_trnc

# calculate total taxes
data_baseline['Taxes'] = (data_baseline['Effective Tax Rate'] *
                          data_baseline['Adjusted total income'])

# get time-series of interest  - for 43 year old in 2017
data_to_use = data_baseline[data_baseline['Age'] == 42].copy()
y = data_to_use['Total Capital Income'].values
x = data_to_use['Total Labor Income'].values
I = data_to_use['Adjusted total income'].values
taxes = data_to_use['Taxes'].values
wgts = data_to_use['Weights'].values
etr_data = data_to_use['Effective Tax Rate'].values
tx_objs = (I, taxes, wgts)

# # set bounds on coefficients
# Gbounds = ((0, None), (0, None), (0, None))
#
# # set starting guesses for coefficients
# Gguess = np.array([0.3745, 0.7525, 0.7368])
#
# # use minimizer to solve for coefficients
# # Gouveia and Strauss
# Gout = opt.minimize(GS, Gguess,
#     args=(tx_objs), method="L-BFGS-B", bounds=Gbounds, tol=1e-15)
# Gcoeffs = Gout.x
# G_SSE = Gout.fun
# print 'GS coeffs, Obs, SSE: ', Gcoeffs, len(taxes), G_SSE

# Plot tax function curves in 2D
# plot GS vs alt function against data for a given year/age
npts = len(I)  # points in grid
maxinc = 300000.  # upper end of grid
mininc = 0.  # lower end of grid
ygrid = np.zeros(npts)  # this is just a place holder to pass
xgrid = np.linspace(mininc, maxinc, npts)
datagrid = (ygrid, xgrid)  # tuple to pass to functions
Gcoeffs = np.zeros((3,))
Gcoeffs[0] = (tax_func_params_base['_GS_Age']['tfunc_etr_params_S'][21, 0, 0])
Gcoeffs[1] = (tax_func_params_base['_GS_Age']['tfunc_etr_params_S'][21, 0, 1])
Gcoeffs[2] = (tax_func_params_base['_GS_Age']['tfunc_etr_params_S'][21, 0, 2])
# put in terms of ETR by dividing by xgrid
Gyfit = ((Gcoeffs[0] * (xgrid - (xgrid ** (-Gcoeffs[1]) + Gcoeffs[2]) **
                        (-1 / Gcoeffs[1]))) / xgrid)

A = (tax_func_params_base['_DEP_Age']['tfunc_etr_params_S'][21, 0, 0])
B = (tax_func_params_base['_DEP_Age']['tfunc_etr_params_S'][21, 0, 1])
C = (tax_func_params_base['_DEP_Age']['tfunc_etr_params_S'][21, 0, 2])
D = (tax_func_params_base['_DEP_Age']['tfunc_etr_params_S'][21, 0, 3])
max_x = (tax_func_params_base['_DEP_Age']['tfunc_etr_params_S'][21, 0, 4])
min_x = (tax_func_params_base['_DEP_Age']['tfunc_etr_params_S'][21, 0, 5])
max_y = (tax_func_params_base['_DEP_Age']['tfunc_etr_params_S'][21, 0, 6])
min_y = (tax_func_params_base['_DEP_Age']['tfunc_etr_params_S'][21, 0, 7])
shift_x = (tax_func_params_base['_DEP_Age']['tfunc_etr_params_S'][21, 0, 8])
shift_y = (tax_func_params_base['_DEP_Age']['tfunc_etr_params_S'][21, 0, 9])
shift = (tax_func_params_base['_DEP_Age']['tfunc_etr_params_S'][21, 0, 10])
share = (tax_func_params_base['_DEP_Age']['tfunc_etr_params_S'][21, 0, 11])
print('Main spec params = ',
      tax_func_params_base['_DEP_Age']['tfunc_etr_params_S'][21, 0, :])

X = xgrid.copy()
Y = xgrid.copy()*0.0
X2 = X ** 2
Y2 = Y ** 2
tau_x = (((max_x - min_x) * (A * X2 + B * X) / (A * X2 + B * X + 1))
         + min_x)
tau_y = (((max_y - min_y) * (C * Y2 + D * Y) / (C * Y2 + D * Y + 1))
         + min_y)
DEP_fit_nocapital = ((((tau_x + shift_x) ** share) * ((tau_y + shift_y)
                                                      ** (1 - share))) +
                     shift)

X = xgrid.copy() * 0.7
Y = xgrid.copy() * 0.3
X2 = X ** 2
Y2 = Y ** 2
tau_x = (((max_x - min_x) * (A * X2 + B * X) / (A * X2 + B * X + 1))
         + min_x)
tau_y = (((max_y - min_y) * (C * Y2 + D * Y) / (C * Y2 + D * Y + 1))
         + min_y)
DEP_fit_nolabor = (((tau_x + shift_x) ** share) * ((tau_y + shift_y) **
                                                   (1 - share))) + shift

# plot data
# plt.plot(I, etr_data, '.b', ms = 1, label='data')
plt.plot(I, etr_data, '.b', label='data')
plt.plot(xgrid, Gyfit, 'm-', lw=1, label='GS')
plt.plot(xgrid, DEP_fit_nocapital, '-', color='orange', lw=1,
         label='DEP, no y')
plt.plot(xgrid, DEP_fit_nolabor, '-', color='green', lw=1,
         label='DEP, mixed')
plt.xlabel('Total Income')
plt.ylabel('Effective Tax Rate')
# set axes range
plt.xlim(0, 100000)
plt.ylim(-0.2, 0.32)
plt.legend(loc='lower right')
plt.suptitle('Age = 43, Year = 2018')
# save high quality version to external file
plt.savefig('Compare_ETR_functions.png')
plt.show()

plt.plot(I, etr_data, '.b', label='data')
plt.plot(xgrid, DEP_fit_nocapital, '-', color='orange', lw=1,
         label='DEP, no y')
plt.plot(xgrid, DEP_fit_nolabor, '-', color='green', lw=1,
         label='DEP, mixed')
plt.xlabel('Total Income')
plt.ylabel('Effective Tax Rate')
# set axes range
plt.xlim(0, 100000)
plt.ylim(-0.2, 0.32)
plt.legend(loc='lower right')
plt.suptitle('Age = 43, Year = 2018')
# save high quality version to external file
plt.savefig('Compare_DEP_ETR_functions.png')
plt.show()


#!/usr/bin/env python

"""capture_stats.py: Sensitivity test experiments python program for Doppler shift"""

__author__ = "Chakraborty, S."
__copyright__ = "Copyright 2021, SuperDARN@VT"
__credits__ = []
__license__ = "MIT"
__version__ = "1.0."
__maintainer__ = "Chakraborty, S."
__email__ = "shibaji7@vt.edu"
__status__ = "Research"

import os
import sys
sys.path.append("../sd/")
import datetime as dt
import pandas as pd
import argparse
from dateutil import parser as dparser
import numpy as np
from netCDF4 import Dataset
import time
import glob
from scipy.integrate import trapz
from scipy import signal
from scipy.io import loadmat

from pysolar.solar import get_altitude
import plotlib

def calculate_sza(d, lats, lons, alt=300):
    d = d.replace(tzinfo=dt.timezone.utc)
    szas = []
    for la, lo in zip(lats, lons):
        szas.append(90. - get_altitude(la, lo, d))
    return szas

INT_F = 300
def get_freq(dn, rad):
    f = 12.
    fname = "../data/op/{dn}/waccmx/sd_{rad}_data.csv.gz".format(dn=dn.strftime("%Y.%m.%d.%H.%M"),rad=rad)
    if os.path.exists(fname):
        os.system("gzip -d " + fname)
        du = pd.read_csv(fname.replace(".gz", ""))
        os.system("gzip " + fname.replace(".gz", ""))
        if len(du) > 0: f = np.median(du.tfreq)/1e3
    return f

def get_vdeta(d, rlim, freq):
    d = d[(d.height>=rlim[0]) & (d.height<rlim[1])]
    f = trapz(signal.resample(d.dop,INT_F))
    v = (0.5 * f * 3e8 / (freq * 1e6))
    return v

def _estimate_dop_delh_(x, y, freq, phi=0):
    dh = (np.max(x.height) - np.max(y.height)) * 1000.
    xf = (-2.*freq*1e6/3e8) * (dh/60.) * np.cos(np.deg2rad(phi))
    xd = 0.5 * xf * 3e8 / (freq * 1e6)
    return xd

sim_fname = "finescale_simulate_total_{kind}.csv"
t, T = True, True
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-v", "--verbose", action="store_false", help="Increase output verbosity (default True)")
    parser.add_argument("-pl", "--plot", action="store_true", help="Analyze sensitivity (default False)")
    parser.add_argument("-t", "--type", default="A", help="Flare type (A/X/M)")
    parser.add_argument("-n", "--number", default=-1, type=int, help="Number of radar event")
    d_reg, e_reg, f_reg = [60,90], [100,130], [140,300]
    args = parser.parse_args()
    model = "waccmx"
    vD, vE, vF, vFh, vT, SZA, E, F = [], [], [], [], [], [], [], []
    events = pd.read_csv("op/radar_event_list.csv", parse_dates=["date"])
    T, kind = 1, "A"
    if args.type == "M" or args.type == "X":
        events = events[events.type.str.contains(args.type)]
        T, kind = 0, args.type
    if args.number >= 0:
        events = events.iloc[[args.number]]
        T, kind = 0, events.rad.tolist()[0] + "_" + events.date.tolist()[0].strftime("%Y-%m-%d-%H-%M")
    if not args.plot:
        ix = 0
        for d, r in zip(events.date.tolist(), events.rad.tolist()):
            print("Events - ", r, d)
            if ix >= T:
                freq = get_freq(d, r)
                for bm in range(24):
                    bearing_file = "../data/op/%s/waccmx/%s/bm.%02d/bearing.mat"%(d.strftime("%Y.%m.%d.%H.%M"), r, bm)
                    print(bearing_file)
                    obj = loadmat(bearing_file)
                    lat, lon = obj["lat"], obj["lon"]
                    sza = np.mean(calculate_sza(d, lat, lon))
                    dic = "../data/op/{dn}/waccmx/{r}/bm.{bm}/".format(r=r,bm="%02d"%bm,dn=d.strftime("%Y.%m.%d.%H.%M"))
                    for i in range(18,19):
                        i_start, i_end = 16, 30
                        for elv in np.linspace(i_start, i_end, (i_end-i_start)*2 + 1):
                            if elv.is_integer(): fname = dic + "ti(%02d)_elv(%d)_f.csv"%(i,elv)
                            else: fname = dic + "ti(%02d)_elv(%.1f)_f.csv"%(i,elv)
                            files = glob.glob(fname)
                            for f in files:
                                try:
                                    T = False
                                    ff = pd.read_csv(f)
                                    vd = np.abs(get_vdeta(ff, d_reg, freq))
                                    ve = np.abs(get_vdeta(ff, e_reg, freq))
                                    vf = np.abs(get_vdeta(ff, f_reg, freq))
                                    bf = pd.read_csv(f.replace("_f.","_d."))
                                    vfh = np.abs(_estimate_dop_delh_(ff,bf,freq))
                                    vD.append(vd)
                                    vE.append(ve)
                                    vF.append(vf)
                                    vFh.append(vfh)
                                    vT.append(vd+ve+vf+vfh)
                                    SZA.append(sza)
                                    E.append(elv)
                                    F.append(freq)
                                except: pass
            #if t and T: break
            ix += 1
        x = pd.DataFrame()
        x["vD"], x["vE"], x["vF"], x["vFh"], x["vT"], x["sza"], x["elv"], x["freq"] = vD, vE, vF, vFh, vT, SZA, E, F
        x = x.round(3)
        x.to_csv("op/" + sim_fname.format(kind=kind), index=False)
    else:
        print(" Compile plots...")
        print("op/"+sim_fname.format(kind=kind))
        x0 = pd.read_csv("op/"+sim_fname.format(kind=kind))
        x0 = x0[(np.abs(x0.vT)>30) & (np.abs(x0.vT)<300)]
        vd0, ve0, vf0 = np.array(x0.vD/x0.vT), np.array(x0.vE/x0.vT), np.array((x0.vF + x0.vFh)/x0.vT)
        vdn0, vdh0 = np.array((x0.vD + x0.vE + x0.vF)/x0.vT), np.array(x0.vFh/x0.vT)
        x = pd.read_csv("op/finescale_simulate.csv")
        x = x[(np.abs(x.vT)>30) & (np.abs(x.vT)<300)]
        vd, ve, vf = np.array(x.vD/x.vT), np.array(x.vE/x.vT), np.array((x.vF + x.vFh)/x.vT)
        vdn, vdh = np.array((x.vD + x.vE + x.vF)/x.vT), np.array(x.vFh/x.vT)
        plotlib.plot_mastogram(vd, ve, vf, vdn, vdh, vd0, ve0, vf0, vdn0, vdh0)
    if os.path.exists("__pycache__/"):
        os.system("rm -rf __pycache__/")
        os.system("rm -rf py*.log")


#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
vslab2d
=======

A script for displaying and saving density distribution generated by
scft_confined/SlabXX2d, where XX stands for the polymer model. For A-B diblock
copolymer, XX=AB.

Copyright (C) 2013 Yi-Xin Liu (lyx@fudan.edu.cn)

'''

import argparse
import os
import glob
import json
from ConfigParser import SafeConfigParser

import numpy as np
from scipy.io import loadmat, savemat
import matplotlib.pylab as plt

from chebpy import cheb_barycentric_matrix

from scftpy import scft_contourf, SCFTConfig

parser = argparse.ArgumentParser(description='vis_slab2d options')

parser.add_argument('-b', '--batch',
                    action='store_true',
                    help='If present or True, perform batch mode.')
parser.add_argument('-r', '--path',
                    default='.',
                    help='Path to be processed, for batch mode use.')
parser.add_argument('-p', '--param_file',
                    default='param.ini',
                    help='SCFT configuration file, *.ini')
parser.add_argument('-d', '--data_file',
                    default='scft_out',
                    help='SCFT generated data file, *.mat')
parser.add_argument('-s', '--save',
                    action='store_true',
                    help='If present or True, save figure.')
args = parser.parse_args()

def batch_vis_slab2d(path='.', param='param.ini', data='scft_out'):
    '''
    Batch mode vis_slab2d, for all directory in the <path>.
    Note: ONLY directories in the path will be processed.
    The input var <data> is the main part of the whole data file name. The
    suffix is in the form '_XXXX', where XXXX is the max number of time steps.
    Then the full data file name is 'scft_out_XXXX.mat', the '.mat' can be
    ignored.
    Generated figures are stored in the same directory as its data file.
    Other data, such as H, are stored in the parent path as 'data.mat'.
    '''
    is_save = True  # Do not show figure in the batch mode

    var = []
    F = []

    for f in os.listdir(path):
        p = os.path.join(path, f)  # path
        if os.path.isdir(p):
            pt = os.path.join(p, data+'_*.mat')  # path to be globbed
            datafiles = glob.glob(pt)
            fnames = [os.path.basename(x) for x in datafiles]
            data_name = get_final_datafile(fnames)
            if data_name == '':
                print p, ' data file missing.'
                continue
            pfile = os.path.join(p, param)
            if not os.path.exists(pfile):
                print p, ' configuration file missing.'
                continue
            dfile = os.path.join(p, data_name)
            vis_slab2d(pfile, dfile, is_save)
            print pfile, dfile
            v = get_var(pfile)
            var.append(v)
            mat = loadmat(dfile)
            F.append(mat['F'][-1, 0])
            print v, mat['F'][-1, 0]

    savemat(os.path.join(path, 'data'), {'v': var, 'F': F})

def get_var(param_file):
    '''
        Get the main batch variable and its current value.
    '''
    cfg = SafeConfigParser(allow_no_value=True)
    cfg.optionxform = str
    cfg.read(param_file)
    section = cfg.get('Batch', 'section')
    # name list of the batch variable
    batch_var = json.loads(cfg.get('Batch', 'var'))
    var_name = batch_var[0]  # the main batch variable is the first one
    if (var_name == 'BC_coefficients_left'
            or var_name == 'BC_coefficients_right'):
        bc = json.loads(cfg.get(section, var_name))
        var = bc[1]
    else:
        var = cfg.getfloat(section, var_name)
    return var

def get_final_datafile(namelist):
    '''
    Each name has the form 'scft_out_XXXX.mat', where XXXX is a number.
    '''
    data = ''
    num = 0  # a number to be compared
    for f in namelist:
        name, ext = os.path.splitext(f)  # split into 'scft_out_XXXX', '.mat'
        fragments = name.split('_')  # split into 'scft', 'out', 'XXXX'
        n = int(fragments[-1])
        if n > num:
            num = n
            data = name
    return data

def vis_slab2d(param='param.ini', data='scft_out', is_save=False):
    '''
    Visualize 2D data generated by DiskXX, here XX represents the polymer model, e.g. XX = AB stands for A-B diblock copolymers.
    '''
    is_show = not is_save
    path = os.path.dirname(data)

    config = SCFTConfig.from_file(param)
    Nx, Ny = config.grid.Lx, config.grid.Ly
    Lx = config.uc.a
    Ly = config.uc.b
    print Nx, Ny, Lx, Ly
    Nxp = Nx + 1
    Nyp = 2 * Ny

    mat = loadmat(data)
    phiA = mat['phiA']
    phiB = mat['phiB']
    #phiAB = phiA - phiB

    if not (Nx, Ny) == phiA.shape:
        raise 'Data file does not match param file.'

    # Periodic in x direction, Fourier
    xxp = np.linspace(0, Lx, Nxp)
    # Non-periodic in y direction, Chebyshev
    #ii = np.arange(Ny)
    #yy = np.cos(np.pi * ii / (Ny-1))  # rr [-1, 1]
    yyp = np.linspace(0, Ly, Nyp)
    yp, xp = np.meshgrid(yyp, xxp)

    phiAp = np.zeros([Nxp, Ny])
    phiBp = np.zeros([Nxp, Ny])
    phiAp[:-1, :] = phiA
    phiAp[-1, :] = phiA[0, :]
    phiBp[:-1, :] = phiB
    phiBp[-1, :] = phiB[0, :]
    phiAp = cheb_interp2d_y(phiAp, yyp)
    phiBp = cheb_interp2d_y(phiBp, yyp)
    phiABp = phiBp - phiAp
    if is_show:
        plt.plot(yyp, phiAp[Nxp/2, :])
        plt.plot(yyp, phiBp[Nxp/2, :])
        plt.plot(yyp, phiABp[Nxp/2, :])
        #scft_contourf(xp, yp, phiAp, show_cbar=True)
        #scft_contourf(xp, yp, phiBp, show_cbar=True)
        #scft_contourf(xp, yp, phiABp, show_cbar=True)
        scft_contourf(xp, yp, phiAp)
        scft_contourf(xp, yp, phiBp)
        scft_contourf(xp, yp, phiABp)
        plt.show()
    if is_save:
        figA = os.path.join(path, 'phiA.png')
        figB = os.path.join(path, 'phiB.png')
        figAB = os.path.join(path, 'phiAB.png')
        scft_contourf(xp, yp, phiAp)
        plt.savefig(figA)
        scft_contourf(xp, yp, phiBp)
        plt.savefig(figB)
        scft_contourf(xp, yp, phiABp)
        plt.savefig(figAB)

def cheb_interp2d_y(u, vy):
    '''
    Use chebyshev interpolation for the last dimension of Cartesian coordinates
    (x, y).
    u(x, y): source data
    vy: vector to be interpolated, size is Nyp.
    '''
    Nx, Ny = u.shape
    Nyp = vy.size
    uout = np.zeros([Nx, Nyp])
    vyp = np.linspace(-1, 1, Nyp)
    T = cheb_barycentric_matrix(vyp, Ny-1)
    for i in xrange(Nx):
        uout[i] = np.dot(T, u[i])
    return uout

if __name__ == '__main__':
    if args.batch:
        batch_vis_slab2d(args.path, args.param_file, args.data_file)
    else:
        vis_slab2d(args.param_file, args.data_file, args.save)