class EquationBlock(object):
    """
    Holds a bloc of equations.

    You can refer to existing equations by using the [] operator.
    """

    def __init__(self):
        self.Equations = {}

    def AddEquation(self, eqn):
        """
        Add an Equation object to the block.
        :param eqn: Equation
        :return: None
        """
        key = eqn.LeftHandSide
        self.Equations[key] = eqn

    def GetEquationList(self):
        """
        Return a sorted list of equations in the block.

        :return: list
        """
        out = list(self.Equations.keys())
        out.sort()
        return out

    def __getitem__(self, key):
        """
        Get an equation by the name of the variable on the left hand side.

        Usage
        eqn = equation_block['x']

        :param key: str
        :return: Equation
        """
        return self.Equations[key]

    def __contains__(self, key):
        """
        Is a variable in the EquationBlock
        :param key: str
        :return: bool
        """
        return key in self.Equations

    def ReplaceTokensFromLookup(self, lookup):
        for eq in self.Equations.values():
            eq.ReplaceTokensFromLookup(lookup) class Equation(object):
    """
    An equation object
    """
    def __init__(self, lhs, desc = '', rhs=()):
        """
        Create a new equation. Must pass the left hand side (variable name).

        May pass a list of terms; otherwise, starts out empty.

        May also pass strings as the right hand side; they will be parsed (with limited
        capabilities).

        >>> eq = Equation('x', desc='Variable x', rhs='y+2.')
        >>> print(str(eq))
        x=y+2. # Variable x

        You may even specify the right-hand side in the first variable
        >>> eq2 = Equation('x=z+2')
        >>> print(str(eq2))
        x=z+2

        Additionally, you can append a comment after a '#'
        >>> eq3 = Equation("e = m*c^2 # Einstein's thingy")
        >>> print(str(eq3))
        e=m*c^2 # Einstein's thingy

        :param lhs: str
        :param desc: str
        :param rhs: list
        """
        if '#' in lhs:
            new_lhs, new_desc = lhs.split('#', 1)
            new_desc.replace('#', '')
            new_lhs = new_lhs.strip()
            new_desc = new_desc.strip()
            self.__init__(new_lhs, new_desc, rhs)
            return
        if '=' in lhs:
            lhs, rhs = lhs.split('=', 1)
            lhs = lhs.strip()
            rhs.replace('=', '')
        self.LeftHandSide = lhs
        self.Description = desc
        if type(rhs) == str:
            rhs = Equation.ParseString(rhs)
        self.TermList = []
        for t in rhs:
            self.AddTerm(t)

    def GetRightHandSide(self):
        """
        Returns the string representation of the right hand side.
        :return: str
        """
        out = [str(s) for s in self.TermList]
        out = ''.join(out)
        if out.startswith('+'):
            out = out[1:]
        # If we have no terms, we are equal to zero.
        # This also happens if terms have a Constant = 0
        if out == '':
            out = '0.0'
        return out

    def RHS(self):
        """
        Convenience alias for GetRightHandSide()
        :return: str
        """
        return self.GetRightHandSide()

    def AddTerm(self, term):
        """
        Add a term to the equation. May be a string or Term object.
        :param term: Term
        :return: None
        """
        term = Term(term)
        if len(self.TermList) > 0:
            if term.IsBlob:
                raise LogicError('Cannot add a blob to non-empty equation')
        for other in self.TermList:
            if term.Term == other.Term:
                # Already exists; just add the constants together.
                other.Constant += term.Constant
                return
        # Otherwise, append
        self.TermList.append(term)

    def ReplaceTokensFromLookup(self, lookup):
        """
        Replace tokens based on lookup dict values.
        :param lookup: dict
        :return:
        """
        for t in self.TermList:
            t.ReplaceTokensFromLookup(lookup)

    @staticmethod
    def ParseString(eqn):
        """
        Parse a string, and return a list of Term objects.
        The first Term may be a "blob".
        :param eqn: str
        :return: list
        """
        is_a_term = True
        try:
            t = Term(eqn, is_blob=False)
        except LogicError:
            is_a_term = False
        except SyntaxError:
            is_a_term = False
        except NotImplementedError:
            is_a_term = False
        if is_a_term:
            # Is a single term; no need to parse...
            return [t,]
        return [Term(eqn, is_blob=True)]

    def __str__(self):
        """
        Return the string representatione of the entire equation
        :return: str
        """
        if self.Description == '':
            return '{0}={1}'.format(self.LeftHandSide, self.GetRightHandSide())
        else:
            return '{0}={1} # {2}'.format(self.LeftHandSide, self.GetRightHandSide(), self.Description) class Term(object):
    """
        Term object; only simple ones are supported by now.

        A simple term consists of a variable name plus a sign. The sign is typically +1. or -1.,
        but if the same variable is added to an equation, the sign is additive.

        May be a "blob," which are complex operations that are left untouched.

        Expected additions:
        [1] Constant term
        [2] one variable times (divided by) another.
    """
    def __init__(self, term, is_blob=False):
        """
        Pass in a string (or possibly another term object), and is parsed.

        If is_blob is True, we do not do any parsing (other than squeezing out internal spaces).

        An equation is allowed one "blob" term, which is the first term. It may be followed
        by non-blob terms.

        As parsing improves, terms can be peeled off of the "blob."

        :param term: str
        :param is_blob: False
        """
        if type(term) == Term:
            self.Constant = term.Constant
            self.Term = term.Term
            self.IsSimple = term.IsSimple
            # Ignore the is_blob input
            self.IsBlob = term.IsBlob
            return
        # Force to be a string; remove whitespace
        term_s = str(term).strip()
        # internal spaces do not matter
        term_s = term_s.replace(' ','')
        if is_blob:
            # If we are a "blob", don't do any parsing.
            self.Constant = 1.0
            self.Term = term_s
            self.IsSimple = True
            self.IsBlob = True
            return
        self.IsBlob = False
        # Rule #1: Eliminate '+' or '-' at front
        self.Constant = 1.0
        if term_s.startswith('+'):
            term_s = term_s[1:]
        elif term_s.startswith('-'):
            self.Constant = -1.0
            term_s = term_s[1:]
        # Rule #2: Allow matched "("
        if term_s.startswith('('):
            if not term_s.endswith(')'):
                raise SyntaxError('Term does not have matching ) - ' + str(term))
            # Remove brackets
            term_s = term_s[1:-1]
            # If we peeled the brackets, remove '+' or '-' again
            if term_s.startswith('+'):
                term_s = term_s[1:]
            elif term_s.startswith('-'):
                # Flip the sign
                self.Constant *= -1.0
                term_s = term_s[1:]
        # We now cannot have embedded '+' or '-' signs.
        if  '+' in term_s:
            raise LogicError('Term cannot contain interior "+" :' + str(term))
        if '-' in term_s:
            raise LogicError('Term cannot contain interior "-" :' + str(term))
        # Do we consist of anything besides a single name token?
        # If so, we are not simple.
        # (Will eventually allow for things like '2*x'.)
        if len(term_s) == 0:
            raise LogicError('Attempting to create an empty term object.')
        if is_python_3:
            g = tokenize.tokenize(BytesIO(term_s.encode('utf-8')).readline)  # tokenize the string
        else:  # pragma: no cover   [Do my coverage on Python 3]
            g = tokenize.generate_tokens(BytesIO(term_s.encode('utf-8')).readline)  # tokenize the string
        self.IsSimple = True
        g = tuple(g)
        if is_python_3:
            # Behaviour changed on me, so needed to clean up logic.
            # Remove "white space" tokens
            # Note: The change in behaviour happened between version 3.5 and 3.7. This new code
            # worked for me on version 3.5 as well as 3.7.
            cleaned = [x for x in g if x[0] not in (ENCODING, NEWLINE, ENDMARKER)]
            if len(cleaned) == 1:
                if cleaned[0][0] == NAME or cleaned[0][0] == NUMBER:
                    self.Term = term_s
                    return
            elif len(cleaned) == 3:
                # Allow
                if ((cleaned[0][0] in (NAME,NUMBER)) and (cleaned[1][0] == OP) and (cleaned[2][0] in (NAME, NUMBER))
                    and (cleaned[1][1] in ('*', '/'))):
                    self.Term = term_s
                    return
            # Puke otherwise
            raise NotImplementedError('Non-simple parsing not supported ' +
                                      '(Note: May fail for Python versions before 3.7 (?)): '+ term_s)
            #  Old code. Will delete once I know new code works...
            # if not g[0][0] == ENCODING:  # pragma: no cover
            #     raise LogicError('Internal error: tokenize behaviour changed')
            # if not g[-1][0] == ENDMARKER:  # pragma: no cover
            #     raise LogicError('Internal error: tokenize behaviour changed')
            #
            # if len(g) > 3:
            #     if len(g) == 4:
            #         # Behaviour changed in Version 3.7?
            #         if (not g[2][0] == NEWLINE) or (not g[1][0] == NAME):
            #             raise NotImplementedError('Non-simple parsing not done')
            #         self.Term = term_s
            #         return
            #     elif len(g) in (5, 6):
            #         # For some reason, getting a newline?
            #         # Allow variable*variable as a "simple" Variable.
            #         if g[1][0] == NAME and g[3][0] == NAME and g[2][0] == OP:
            #             if g[2][1] in ('*', '/'):
            #                 self.Term = term_s
            #                 return
            #     raise NotImplementedError('Non-simple parsing not done')
            #     # self.IsSimple = False
            # else:
            #     if not g[1][0] == NAME:
            #         raise NotImplementedError('Non-simple parsing not done')
            #         # self.IsSimple = False
            # self.Term = term_s
        else: # Python 2.7 # pragma: no cover
            # Missing the first term - augh
            if not g[-1][0] == ENDMARKER:  # pragma: no cover
                raise LogicError('Internal error: tokenize behaviour changed')
            if len(g) > 3:
                if len(g) == 4:
                    # Allow variable*variable as a "simple" Variable.
                    if g[0][0] == NAME and g[2][0] == NAME and g[1][0] == OP:
                        if g[1][1] in ('*', '/'):
                            self.Term = term_s
                            return
                raise NotImplementedError('Non-simple parsing not done')
                # self.IsSimple = False
            else:
                if not g[0][0] == NAME:
                    raise NotImplementedError('Non-simple parsing not done')
                    # self.IsSimple = False
            self.Term = term_s


    def __str__(self):
        """
        Convert the Term to a string
        :return: str
        """
        if self.IsBlob:
            return self.Term
        if not self.IsSimple: # pragma: no cover
            raise NotImplementedError('Not dealing with non-simple terms...')
        if self.Constant == 0.0:
            return ''
        if self.Constant == 1.0:
            lead = '+'
        elif self.Constant == -1:
            lead = '-'
        elif self.Constant > 0:
            lead = '+' + str(self.Constant) + '*'
        else:
            lead = str(self.Constant) + '*'
        return lead + self.Term

    def ReplaceTokensFromLookup(self, lookup):
        try:
            if self.IsBlob:
                self.Term = replace_token_from_lookup(self.Term, lookup)
                return
            if self.IsSimple:
                self.Term = replace_token_from_lookup(self.Term, lookup)
                #if self.Term in lookup:
                #    self.Term = lookup[self.Term]
                return
        except tokenize.TokenError:
            raise ValueError('Failure in parsing for token in: ("{0}", "{1}")'.format(self.Term, lookup))
        raise NotImplementedError('Need to handle non-simple terms somehow...') # pragma: no cover class EconomicObject(object):
    """
    EconomicObject class
    Base class for model entities. Gives them a unique numeric id to make debugging and
    some coding tasks easier. For example, if two EconomicObject's have a different ID, they are distinct,
    and so an EconomicObject can exclude itself when iterating through a list of Sector objects.

    Can create function for displaying, etc.
    """
    ID = 0

    def __init__(self, parent=None, code=''):
        self.ID = EconomicObject.ID
        EconomicObject.ID += 1
        self.Parent = parent
        self.Code = code
        self.LongName = ''
        Logger('EconomicObject Created: {0} ID = {1}', priority=1, data_to_format=(type(self),self.ID))

    def GetModel(self):
        """
        Get the base object of the hierarchy (which should be a Model).
        :return: Model
        """
        obj = self
        while obj.Parent is not None:
            obj = obj.Parent
        return obj

    def ShareParent(self, other):
        """
        Does this entity have the same parent?
        :param other:  EconomicObject
        :return: bool
        """
        return self.Parent == other.Parent class Model(EconomicObject):
    """
    Model class.

    All other entities live within a model.
    """

    def __init__(self):
        """
        Create the Model object. All the economic class objects (Country, Sector) live within
        a Model.

        Unlike other EconomicObject subclasses, the Model object is a base object, and has no parent.

        It is possible for two Model objects to coexist; there is no interaction between them
        (other than side effects from global Parameters).
        """
        EconomicObject.__init__(self)
        self.CountryList = []
        self.Exogenous = []
        self.InitialConditions = []
        self.FinalEquations = ''
        self.MaxTime = 100
        self.RegisteredCashFlows = []
        self.Aliases = {}
        self.TimeSeriesCutoff = None
        self.TimeSeriesSupressTimeZero = False
        self.EquationSolver = sfc_models.equation_solver.EquationSolver()
        self.GlobalVariables = []
        self.IncomeExclusions = []
        self.CurrencyZoneList = []
        self.State = 'Construction'
        self.DefaultCurrency = 'LOCAL'
        self.RunSteps = None
        self.ExternalSector = None
        self.FinalEquationBlock = EquationBlock()

    def main(self, base_file_name=None):  # pragma: no cover
        """
        Routine that does most of the work of model building. The model is build based upon
        the Sector objects that have been registered as children of this Model.

        The base_file_name is the base filename that is used for Logging operations; just used
        to call Logger.register_standard_logs(base_file_name). It is recommended that you call
        Logger.register_standard_logs() before main, so that Sector creation can be logged.

        The major operations:
        [1] Call GenerateEquations() on all Sectors. The fact that GenerateEquations() is only
         called now at the Sector level means that Sectors can be created independently, and the
        GenerateEquations() call which ties it to other sectors is only called after all other sectors
        are created. (There is at least one exception where Sectors have to be created in a specific
        order, which we want to avoid.)

        [2] Cleanup work: processing initial conditions, exogenous variables, replacing
            aliases in equations.

        [3] Tie the sector level equations into a single block of equations. (Currently strings, not
            Equation objects.)

        [4] The equations are passed to self.EquationSolver, and they are solved.

        The user can then use GetTimeSeries() to access the output time series (if they can be
        calculated.)

        :param base_file_name: str
        :return: None
        """
        self.State = 'Running'
        try:
            if base_file_name is not None:
                Logger.register_standard_logs(base_file_name)
            Logger('Starting Model main()')
            self._GenerateFullSectorCodes()
            self._GenerateEquations()
            self._FixAliases()
            self._GenerateRegisteredCashFlows()
            self._ProcessExogenous()
            self.FinalEquations = self._CreateFinalEquations()
            self.EquationSolver.ParseString(self.FinalEquations)
            self.EquationSolver.SolveEquation()
            self.LogInfo()
        except Warning as e:
            self.LogInfo(ex=e)
            print('Warning triggered: ' + str(e))
            return self.FinalEquations
        except Exception as e:
            self.LogInfo(ex=e)
            raise
        finally:
            self.State = 'Finished Running'
            Logger(self.EquationSolver.GenerateCSVtext(), 'timeseries')
            Logger.cleanup()
        return self.FinalEquations

    def _GetSteps(self): # pragma: no cover
        """
        This is experimental, for GUI use. Will integrate with main() later...
        :return:
        """
        if self.RunSteps is None:
            # Need to generate RunSteps
            self.RunSteps = []
            self.RunSteps.append({'Generate Sector Codes': self._GenerateFullSectorCodes})
            self.RunSteps.append({'Fix Aliases': self._FixAliases})
            self.RunSteps.append({'Generate Equations': self._GenerateEquationSteps})
            self.RunSteps.append({'Process Cash Flows': self._GenerateRegisteredCashFlows})
            self.RunSteps.append({'Process Exogenous': self._ProcessExogenous})
            self.RunSteps.append({'Fix Aliases (Pass #2)': self._FixAliases})
            self.RunSteps.append({'Final Equations': self._CreateFinalEquations})
            self.RunSteps.append({'Solve': self._FinalSteps})
            # self.EquationSolver.ParseString(self.FinalEquations)
            # self.EquationSolver.SolveEquation()
            # self.LogInfo()
            # self.State = 'Finished Running'
            # Logger(self.EquationSolver.GenerateCSVtext(), 'timeseries')
            # Logger.cleanup()
        out = [list(x.keys()) for x in self.RunSteps]
        for x in out:
            x.sort()
        return out

    def _GenerateEquationSteps(self): # pragma: no cover
        sector_list = self.GetSectors()
        for sec in sector_list:
            self.RunSteps[0][sec.FullCode] = sec._GenerateEquationsFrontEnd

    def _FinalSteps(self): # pragma: no cover
        self.EquationSolver.ParseString(self.FinalEquations)
        self.EquationSolver.SolveEquation()
        self.LogInfo()
        self.State = 'Finished Running'
        Logger(self.EquationSolver.GenerateCSVtext(), 'timeseries')
        Logger.cleanup()

    def _RunStep(self, command): # pragma: no cover
        if len(self.RunSteps) == 0:
            self.State = 'Finished Running'
            return
        self.State = 'Running'
        # Will throw KeyError if not in command list
        func = self.RunSteps[0].pop(command)
        try:
            func()
        except:
            self.RunSteps = []
            self.State = 'Finished Running'
            raise
        if len(self.RunSteps[0]) == 0:
            self.RunSteps.pop(0)
        if len(self.RunSteps) == 0:
            self.State = 'Finished Running'

    def _RunAllSteps(self):  # pragma: no cover
        while len(self.RunSteps) > 0:
            all_cmds = list(self.RunSteps[0].keys())
            self._RunStep(all_cmds[0])

    def AddExogenous(self, sector_fullcode, varname, value):
        """
        Add an exogenous variable to the model. Overwrites an existing variable definition.
        Need to use the full sector code.

        Exogenous variables are sepcified as time series, which are implemented as list variables ([1, 2, 3,...])
        The exogenous variable can either be specified as a string (which can be evaluated to a list), or else
        as a list object. The list object will be converted into a string representation using repr(), which
        means that it may be much longer than using something like '[20,] * 100'.

        At present, does not support the usage of specifying a constant value. For example value='20.' does
        not work, you need '[20.]*100.'

        :param sector_fullcode: str
        :param varname: str
        :param value: str
        :return:
        """
        Logger('Adding exogenous variable: {0} in {1}', priority=5,
               data_to_format=(varname, sector_fullcode))
        # If the user passes in a list or tuple, convert it to a string representation.
        if type(value) in (list, tuple):
            value = repr(value)
        self.Exogenous.append((sector_fullcode, varname, value))

    def AddInitialCondition(self, sector_fullcode, varname, value):
        """
        Set the initial condition for a variable. Need to use the full code of the sector -
        or its ID (int).

        :param sector_fullcode: str
        :param varname: str
        :param value: float
        :return:
        """
        Logger('Adding initial condition: {0} in {1}', priority=5,
               data_to_format=(varname, sector_fullcode))
        # Convert the "value" to a float, in case someone uses a string
        try:
            value = float(value)
        except:
            raise ValueError('The "value" parameter for initial conditions must be a float.')
        self.InitialConditions.append((sector_fullcode, varname, str(value)))

    def AddCashFlowIncomeExclusion(self, sector, cash_flow_name):
        """
        Specify that a cash flow is to be excluded from the list of flows that go into
        income.

        Examples of such flows within economics:
        - Household consumption of goods.
        - Business investment.
        - Dividend outflows for business (inflow is income for household).
        - (Financing flows - do not really appear?)

        Note that these exclusions are generally implemented at the sector code level.
        :param sector: Sector
        :param cash_flow_name: str
        :return: None
        """
        Logger('Registering cash flow exclusion: {0} for ID={1}', priority=5,
               data_to_format=(cash_flow_name, sector.ID))
        self.IncomeExclusions.append((sector, cash_flow_name))

    def _RegisterAlias(self, alias, sector, local_variable_name):
        """
        Used by Sector objects to register aliases for local variables.

        :param alias: str
        :param sector: Sector
        :param local_variable_name: str
        :return:
        """
        Logger('Registering alias {0} for {1} in ID={2}', priority=5,
               data_to_format=(alias, local_variable_name, sector.ID))
        self.Aliases[alias] = (sector, local_variable_name)

    def AddGlobalEquation(self, var, description, eqn):
        """
        Add a variable that is not associated with a sector.
        Typical example: 't'

        :param var: str
        :param description: str
        :param eqn: str
        :return: None
        """
        Logger('Registering global equation: {0} = {1}', priority=5, data_to_format=(var, eqn))
        self.GlobalVariables.append((var, eqn, description))

    def GetSectors(self):
        """
        Returns a list of Sector objects held within this Model.

        :return: list
        """
        out = []
        for cntry in self.CountryList:
            for sector in cntry.SectorList:
                out.append(sector)
        return out

    def GetTimeSeries(self, series, cutoff=None, group_of_series='main'):
        """
        Convenience function to retrieve time series from the EquationSolver.

        Use cutoff to truncate the length of the output.

        If self.TimeSeriesSupressZero is True, the first point is removed (the initial
        conditions period).

        :param group_of_series:
        :param series: str
        :param cutoff: int
        :return: list
        """
        if cutoff is None:
            cutoff = self.TimeSeriesCutoff
        try:
            series_holder = self.EquationSolver.TimeSeries
            if group_of_series == 'step': # pragma: no cover [The GUI tells us quickly if this breaks]
                series_holder = self.EquationSolver.TimeSeriesStepTrace
            elif group_of_series == 'initial': # pragma: no cover
                series_holder = self.EquationSolver.TimeSeriesInitialSteadyState
            if cutoff is None:
                val = series_holder[series]
            else:
                val = series_holder[series][0:(cutoff + 1)]
        except KeyError:
            raise KeyError('Time series "{0}" does not exist'.format(series))
        if self.TimeSeriesSupressTimeZero:
            val.pop(0)
        return val

    def _FixAliases(self):
        """
        Assign the proper names to variables in Sector objects (that were perviously aliases).
        :return:
        """
        Logger('Fixing aliases (Model._FixAliases)', priority=3)
        lookup = {}
        for alias in self.Aliases:
            sector, varname = self.Aliases[alias]
            lookup[alias] = sector.GetVariableName(varname)
        for sector in self.GetSectors():
            sector._ReplaceAliases(lookup)

    def LogInfo(self, generate_full_codes=True, ex=None):  # pragma: no cover
        """
        Write information to a file; if there is an exception, dump the trace.
        The log will normally generate the full sector codes; set generate_full_codes=False
        to leave the Model full codes untouched.

        :param generate_full_codes: bool
        :param ex: Exception
        :return:
        """
        # Not covered with unit tests [for now]. Output format will change a lot.
        if ex is not None:
            Logger('\nError or Warning raised:\n')
            try:
                traceback.print_exc(file=Logger.get_handle())
            except KeyError:
                # Log was not registered; do nothing!
                return
        Logger('-'*30)
        Logger('Starting LogInfo() Data Dump')
        Logger('-'*30)
        if generate_full_codes:
            self._GenerateFullSectorCodes()
            for c in self.CountryList:
                Logger('Country: Code= "%s" %s\n' % (c.Code, c.LongName))
                Logger('=' * 60 + '\n\n')
                for s in c.SectorList:
                    Logger(s.Dump() + '\n')
        Logger('Writing LogInfo to log="eqn"')
        Logger('\n\nFinal Equations:\n', log='eqn')
        Logger(self.FinalEquations + '\n', log='eqn')
        parser = EquationParser()
        parser.ParseString(self.FinalEquations)
        parser.EquationReduction()
        Logger('\n\nReduced Equations', log='eqn')
        Logger(parser.DumpEquations(), log='eqn')
        if ex is not None:
            Logger('\n\nError raised:\n')
            traceback.print_exc(file=Logger.get_handle())

    def _AddCountry(self, country):
        """
        Add a country to the list. This is called by the object constructore; users should
        not call this.

        :param country: Country
        :return: None
        """
        Logger('Adding Country: {0} ID={1}', data_to_format=(country.Code, country.ID))
        if country.Code in self:
            raise LogicError('Country with Code {0} already in Model'.format(country.Code))
        self.CountryList.append(country)
        self.DefaultCurrency = country.Currency
        czone = self._FitIntoCurrencyZone(country)
        country.CurrencyZone = czone

    def _FitIntoCurrencyZone(self, country):
        """
        Find whether Country fits into an existing CurrencyZone; if not,
        create a new one.
        :param country: Country
        :return: CurrencyZone
        """
        for czone in self.CurrencyZoneList:
            if country.Currency == czone.Currency:
                czone.CountryList.append(country)
                Logger('Fitting {0} into CurrencyZone {1}',
                       data_to_format=(country.LongName, czone.Currency))
                return czone
        Logger('Creating new currency zone {0}, adding {1} to it',
               data_to_format=(country.Currency, country.LongName))
        czone = CurrencyZone(self, country.Currency)
        czone.CountryList.append(country)
        self.CurrencyZoneList.append(czone)
        return czone

    def _GenerateFullSectorCodes(self):
        """
        Create full sector names (which is equal to '[country.Code]_[sector.Code]' - if there is more than one country.
        Equals the sector code otherwise.

        :return: None
        """
        Logger('Generating FullSector codes (Model._GenerateFullSectorCodes()', priority=3)
        add_country_code = len(self.CountryList) > 1
        for cntry in self.CountryList:
            for sector in cntry.SectorList:
                if add_country_code:
                    sector.FullCode = cntry.Code + '_' + sector.Code
                else:
                    sector.FullCode = sector.Code

    @staticmethod
    def GetSectorCodeWithCountry(sector):
        """
        Return the sector code including the country information.
        Need to use this if we want the FullCode before the model information
        is bound in main().

        We should not need to use this function often; generally only when we need
        FullCodes in constructors. For example, the multi-supply business sector
        needs the full codes of markets passed into the constructor.
        :param sector: Sector
        :return: str
        """
        return '{0}_{1}'.format(sector.Parent.Code, sector.Code)

    def RegisterCashFlow(self, source_sector, target_sector, amount_variable, is_income_source=True,
                         is_income_dest=True):
        """
        Register a cash flow between two sectors.

        The amount_variable is the name of the local variable within the source sector.

        Only allowed across currency zones if an ExternalSector object has been defined;
        otherwise throws a LogicError.

        The currency value of the amount_variable is assumed to be in the source currency.

        :param is_income_dest:
        :param is_income_source:
        :param source_sector: Sector
        :param target_sector: Sector
        :param amount_variable: str
        :return:
        """
        # if amount_variable not in source_sector.Equations:
        #     raise KeyError('Must define the variable that is the amount of the cash flow')
        Logger('Cash flow registered {0}: {1} -> {2}  [ID: {3} -> {4}]', priority=3,
               data_to_format=(amount_variable, source_sector.Code, target_sector.Code,
                               source_sector.ID, target_sector.ID))
        self.RegisteredCashFlows.append((source_sector, target_sector, amount_variable,
                                         is_income_source, is_income_dest))

    def _GenerateRegisteredCashFlows(self):
        """
        Create cash flows based on those previously registered.

        :return:
        """
        Logger('Model._GenerateRegisteredCashFlows()')
        Logger('Adding {0} cash flows to sectors', priority=3,
               data_to_format=(len(self.RegisteredCashFlows),))
        for source_sector, target_sector, amount_variable, is_income_source, is_income_dest in self.RegisteredCashFlows:
            is_cross_currency = source_sector.CurrencyZone != target_sector.CurrencyZone
            if is_cross_currency:
                if self.ExternalSector is None:
                    msg = """Only can have cross-currency flows if an ExternalSector object is created\nSource={0} Destination={1}""".format(
                        source_sector.FullCode, target_sector.FullCode)
                    raise LogicError(msg)
            full_variable_name = source_sector.GetVariableName(amount_variable)
            source_sector.AddCashFlow('-' + full_variable_name, eqn=None,
                                      is_income=is_income_source)
            if is_cross_currency:
                fx = self.ExternalSector['FX']
                fx._SendMoney(source_sector, full_variable_name)
                term = fx._ReceiveMoney(target_sector=target_sector, source_sector=source_sector,
                                        variable_name=full_variable_name)
            else:
                term = '+' + full_variable_name
            target_sector.AddCashFlow(term, eqn=None,
                                      is_income=is_income_dest)

    def LookupSector(self, fullcode):
        """
        Find a sector based on its FullCode.
        :param fullcode: str
        :return: Sector
        """
        for cntry in self.CountryList:
            try:
                s = cntry.LookupSector(fullcode, is_full_code=True)
                return s
            except KeyError:
                pass
        raise KeyError('Sector with FullCode does not exist: ' + fullcode)

    def _ProcessExogenous(self):
        """
        Handles the exogenous variables.

        :return: None
        """
        Logger('Processing {0} exogenous variables', priority=3, data_to_format=(len(self.Exogenous),))
        for sector_code, varname, eqn in self.Exogenous:
            if type(sector_code) is str:
                sector = self.LookupSector(sector_code)
            else:
                sector = sector_code
            if varname not in sector.EquationBlock.Equations:
                raise KeyError('Sector %s does not have variable %s' % (sector_code, varname))
            # Need to mark exogenous variables
            sector.SetEquationRightHandSide(varname, 'EXOGENOUS ' + eqn)

    def _GenerateInitialConditions(self):
        """
        Create block of equations for initial conditions.

        Validates that the variables exist.
        :return:
        """
        Logger('Generating {0} initial conditions', priority=3,
               data_to_format=(len(self.InitialConditions),))
        out = []
        for sector_code, varname, value in self.InitialConditions:
            sector = self.LookupSector(sector_code)
            if varname not in sector.EquationBlock.Equations:
                raise KeyError('Sector %s does not have variable %s' % (sector_code, varname))
            out.append(('%s(0)' % (sector.GetVariableName(varname),), value, 'Initial Condition'))
        return out

    def _GenerateEquations(self):
        """
        Call _GenerateEquations on all child Sector objects.

        :return:
        """
        Logger('Model._GenerateEquations()', priority=1)
        for cntry in self.CountryList:
            for sector in cntry.SectorList:
                Logger('Calling _GenerateEquations on {0}  ({1})', priority=3,
                       data_to_format=(sector.FullCode, type(sector)))
                sector._GenerateEquations()

    def DumpEquations(self):
        """
        Returns a string with basic information about the entities within this Model.
        Output is primitive, and aimed at debugging purposes. In other words, the format
        will change without warning.

        If you want information of a specific format, please create a specific reporting
        function for your needs.

        :return: str
        """
        out = ''
        for cntry in self.CountryList:
            for sector in cntry.SectorList:
                out += sector.Dump()
        return out

    def _CreateFinalEquations(self):
        """
        Create Final equations.

        Final output, which is a text block of equations
        :return: str
        """
        Logger('Model._CreateFinalEquations()')
        out = []
        for cntry in self.CountryList:
            for sector in cntry.SectorList:
                out.extend(sector._CreateFinalEquations())
        out.extend(self._GenerateInitialConditions())
        out.extend(self.GlobalVariables)
        if len(out) == 0:
            self.FinalEquations = ''
            raise Warning('There are no equations in the system.')
        # Build the FinalEquationBlock
        self.FinalEquationBlock = EquationBlock()
        for row in out:
            if 'EXOGENOUS' in row[1]:
                eq = Equation(row[0], desc=row[2], rhs=row[1].replace('EXOGENOUS', ''))
            else:
                eq = Equation(row[0], desc=row[2], rhs=row[1])
            self.FinalEquationBlock.AddEquation(eq)
        out = self._FinalEquationFormatting(out)
        self.FinalEquations = out
        return out

    def _FinalEquationFormatting(self, out):
        """
        Convert equation information in list into formatted strings.

        :param out: list
        :return:
        """
        Logger('_FinalEquationFormatting()', priority=5)
        endo = []
        exo = []
        for row in out:
            if 'EXOGENOUS' in row[1]:
                new_eqn = row[1].replace('EXOGENOUS', '')
                exo.append((row[0], new_eqn, row[2]))
            else:
                endo.append(row)
        max0 = max([len(x[0]) for x in out])
        max1 = max([len(x[1]) for x in out])
        formatter = '%<max0>s = %-<max1>s  # %s'
        formatter = formatter.replace('<max0>', str(max0))
        formatter = formatter.replace('<max1>', str(max1))
        endo = [formatter % x for x in endo]
        exo = [formatter % x for x in exo]
        s = '\n'.join(endo) + '\n\n# Exogenous Variables\n\n' + '\n'.join(exo)
        s += '\n\nMaxTime = {0}\nErr_Tolerance=1e-6'.format(self.MaxTime)
        return s

    def __getitem__(self, item):
        """
        Get a country using model[country_code] notation
        :param item: str
        :return: Country
        """
        for obj in self.CountryList:
            if item == obj.Code:
                return obj
        raise KeyError('Country {0} not in Model'.format(item))

    def __contains__(self, item):
        """
        Is a Country object (or string code) in a Model?
        :param item: str
        :return:
        """
        if type(item) is str:
            try:
                a = self[item]
                return True
            except KeyError:
                return False
        return item in self.CountryList class Country(EconomicObject):
    """
    Country class. Somewhat redundant in a closed economy model, but we need for multi-region models.
    """

    def __init__(self, model, code, long_name='', currency=None):
        """
        Create a Country object.

        Please note that the Currency cannot be changed by just changing the Currency data member. This
        is because the membership in a CurrencyZone object defines currency status, and not just the
        string value of Currency.

        TODO: Create a ChangeCurrency() method.

        :param model: Model
        :param code: str
        :param long_name: str
        :param currency: str
        """
        EconomicObject.__init__(self, model)
        self.Code = code
        if long_name == '':
            long_name = 'Country ' + code
        self.LongName = long_name
        self.SectorList = []
        self.CurrencyZone = None
        if currency is None:
            self.Currency = code
        else:
            self.Currency = currency
        model._AddCountry(self)


    def _AddSector(self, sector):
        """
        Add a sector to this country. This is called by the Sector constructore; users
        should not call directly.

        :param sector: Sector
        :return:
        """
        Logger('Adding Sector {0} To Country {1}', priority=1,
               data_to_format=(sector.Code, self.Code))
        if sector.Code in self:
            raise LogicError('Sector with Code {0} already in Country {1}'.format(
                sector.Code, self.Code))
        self.SectorList.append(sector)

    def LookupSector(self, code, is_full_code=False):
        """
        Get the sector object via code or fullcode or ID (if you pass an int).
        :param code: str
        :param is_full_code: bool
        :return: Sector
        """
        if type(code) is int:
            for s in self.SectorList:
                if code == s.ID:
                    return s
        elif is_full_code:
            for s in self.SectorList:
                if s.FullCode == code:
                    return s
        else:
            for s in self.SectorList:
                if s.Code == code:
                    return s
        raise KeyError('Sector does not exist - ' + str(code))

    def GetSectors(self):
        """
        What Sectors are in the Country?
        :return: list
        """
        return self.SectorList

    def __getitem__(self, item):
        """
        Get a Sector object using country[sector_code] notation.
        :param item: str
        :return:
        """
        for obj in self.SectorList:
            if obj.Code == item:
                return obj
        raise KeyError('Sector {0} does not exist in Country {1}'.format(item, self.Code))

    def __contains__(self, item):
        """
        Is a sector code (or Sector object) in this Country?
        :param item: str
        :return:
        """
        if type(item) == str:
            try:
                dummy = self[item]
                return True
            except KeyError:
                return False
        return item in self.SectorList class Logger(object):
    """
    Class to handle logging.

    Just use the constructor to write. Logging within code just looks like

    Logger("hello!")

    Logger("low level debugging", priority=9)

    If the priority > cutoff, ignore. Hence - level 1 are top level comments.

    To switch the file logged to, specify the log='<name>' parameter.
    The <name> is a short code used to identify the file. BY default, it equals
    'log' (main log).

    This allows us to centralise logging and file handling at a high level.
    """

    # Have a single file handle for all Logger objects
    log_file_handles = {}
    priority_cutoff = 10

    def __init__(self, txt, log='log', priority=1, data_to_format=None, endline=True):
        """
        Do a logging message.

        Usual usage:
        Logger('Hello World!')
        This writes "Hello World!" to the default log 'log'.

        We can add low-level logging by adding a priority (higher priority means
        less important. If that makes sense...)
        Logger('Low level message', priority=9)
        In this case, if we set the cutoff below 9 (default=10), the message will
        not be logged.

        We can direct messages to different files by using the log parameter.

        Logger('Start initial conditions calculation', log='init')
        will write the message to the log labelled 'init'.

        To start logging, need to use Logger.register_log() call. If
        a log is not registered, the Logger() call will do nothing.

        :param txt: str
        :param log: str
        :param priority: int
        :param data_to_format: tuple
        :param endline: bool
        """
        try:
            f = Logger.get_handle(log)
        except KeyError:
            # Was not registered, so just eat the message.
            return
        if priority > Logger.priority_cutoff:
            return
        if data_to_format is not None:
            txt = txt.format(*data_to_format)
        if priority < 1:
            priority = 1
        txt = (' ' * (priority-1)) + txt
        # Need to remove empty strings, or txt[-1] barfs
        if len(txt) == 0:
            if endline:
                txt = ' '
            else:
                return
        # Add an endline automatically (unless endline=False).
        if txt[-1] != '\n' and endline:
            txt += '\n'
        f.write(txt)

    @staticmethod
    def register_log(fname, log='log'):
        """
        Register a file for logging. The fname is the full file path to be
        written. The log parameter is the short code used to specify which log file
        is used.

        The file is only created (overwriting previous contents) when the file handle
        is accessed.

        :param fname: str
        :param log: str
        :return: None
        """
        if log in Logger.log_file_handles:
            raise ValueError('Log already registered: ' + log)
        Logger.log_file_handles[log] = fname

    @staticmethod
    def register_standard_logs(base_file_name): # pragma: no cover
        loginfo = {
            'log': base_file_name + '_log.txt',
            'timeseries': base_file_name + '_out.txt',
            'eqn': base_file_name + '_eqn.txt',
            'step': base_file_name + '_iteration.txt',
            'steadystate_0': base_file_name + '_steadystate.txt',
        }
        for logname in loginfo:
            try:
                Logger.register_log(loginfo[logname], log=logname)
            except ValueError:
                # If already registered, eat the error.
                pass

    @staticmethod
    def get_handle(log='log'):  # pragma: no cover
        """
        Get the file handle associated with the log.

        If the file was not yet opened, it is opened for writing.

        Throws a KeyError
        :param log: str
        :return:
        """
        f = Logger.log_file_handles[log]
        if type(f) is str:
            f = open(f, 'w')
            Logger.log_file_handles[log] = f
        return f

    @staticmethod
    def cleanup():
        for k, f in Logger.log_file_handles.items():
            if f is not None:
                if type(f) is not str:
                    f.close()
        Logger.log_file_handles = {} def replace_token_from_lookup(s, lookup):
    """
    Replace tokens using a lookup dictionary.

    >>> replace_token_from_lookup('y = x', {'y': 'H_y', 'x': 'H_x'})
    'H_y =H_x '

    :param s: str
    :param lookup: dict
    :return: str
    """
    result = []
    if is_python_3:
        g = tokenize.tokenize(BytesIO(s.encode('utf-8')).readline)  # tokenize the string
        for toknum, tokval, _, _, _ in g:
            if toknum == NAME and tokval in lookup:  # replace NAME tokens
                result.append((NAME, lookup[tokval]))
            else:
                result.append((toknum, tokval))
        return untokenize(result).decode('utf-8')
    else:  # pragma: no cover   [Do my coverage on Python 3]
        g = tokenize.generate_tokens(BytesIO(s.encode('utf-8')).readline)  # tokenize the string
        for toknum, tokval, _, _, _ in g:
            if toknum == NAME and tokval in lookup:  # replace NAME tokens
                result.append((NAME, lookup[tokval]))
            else:
                result.append((toknum, tokval))
        return str(untokenize(result).decode('utf-8')) class LogicError(ValueError):
    """
    Throw this when code logic is broken.
    """
    pass def create_equation_from_terms(terms):
    """
    Create a string equation (right hand side) from a list of terms.

    >>> create_equation_from_terms(['x','y'])
    'x+y'
    >>> create_equation_from_terms(['-x', 'y', '-z'])
    '-x+y-z'

    :param terms: list
    :return: str
    """
    if len(terms) == 0:
        return ''
    for i in range(0, len(terms)):
        term = terms[i].strip()
        if not term[0] in ('+', '-'):
            term = '+' + term
        terms[i] = term
    if terms[0][0] == '+':
        terms[0] = terms[0].replace('+', '')
    eqn = ''.join(terms)
    return eqn
from sfc_models.equation import EquationBlock, Equation, Term
from sfc_models.models import EconomicObject, Model, Country
from sfc_models.utils import Logger, replace_token_from_lookup, LogicError, create_equation_from_terms

        The is_absolute_weighting parameter is a placeholder; if set to true, asset demands are
        absolute. There is a TODO marking where code should be added.

        Note that weightings are (normally) from 0-1.
        :param asset_weighting_dict: dict
        :param residual_asset_code: str
        :param is_absolute_weighting: bool
        :return:
        """
        if is_absolute_weighting:
            # TODO: Implement absolute weightings.
            raise NotImplementedError('Absolute weightings not implemented')
        residual_weight = '1.0'
        if type(asset_weighting_dict) in (list, tuple):
            # Allow asset_weighting_dict to be a list of key: value pairs.
            tmp = dict()
            for code, eqn in asset_weighting_dict:
                tmp[code] = eqn
            asset_weighting_dict = tmp
        for code, weight_eqn in asset_weighting_dict.items():
            # Weight variable = 'WGT_{CODE}'
            weight = 'WGT_' + code
            self.AddVariable(weight, 'Asset weight for' + code, weight_eqn)
            self.AddVariable('DEM_' + code, 'Demand for asset ' + code, 'F * {0}'.format(weight))
            residual_weight += ' - ' + weight
        self.AddVariable('WGT_' + residual_asset_code, 'Asset weight for ' + residual_asset_code, residual_weight)
        self.AddVariable('DEM_' + residual_asset_code, 'Demand for asset ' + residual_asset_code,
                         'F * {0}'.format('WGT_' + residual_asset_code))

class Market(Sector):
    """
    Market Not really a sector, but keep it in the same list.
    """

    def __init__(self, country, code, long_name=''):
        if long_name == '':
            long_name = 'Market {0} in Country {1}'.format(code, country.Code)
        Sector.__init__(self, country, code, long_name, has_F=False)
        self.AddVariable('SUP_' + code, 'Supply for market ' + code, '')
        self.AddVariable('DEM_' + code, 'Demand for market ' + code, '')
        self.ResidualSupply = None
        self.OtherSuppliers = []

    def _SearchSupplier(self):
        """
        Find the sector that is a single supplier in a country.
        Throws a LogicError if more than one, or none.

        Need to set SupplyAllocation if you want to do something not
        covered by this default behaviour.

        :return: Sector
        """
        Logger('Market {0} searching Country {1} for a supplier', priority=3,
               data_to_format=(self.Code, self.Parent.Code))
        ret_value = None
        for sector in self.Parent.GetSectors():
            if sector.ID == self.ID:
                continue
            if 'SUP_' + self.Code in sector.EquationBlock.Equations:
                if ret_value is None:
                    ret_value = sector
                else:
                    raise LogicError('More than one supplier, must set SupplyAllocation: ' + self.Code)
        if ret_value is None:
            raise LogicError('No supplier: ' + self.Code)
        self.ResidualSupply = ret_value
        return ret_value

    def _GenerateEquations(self):
        """
        Generate the equations associated with this market.
        :return:
        """
        if self.ResidualSupply is None:
            supplier = self._SearchSupplier()
        self._GenerateTermsLowLevel('DEM', 'Demand')
        self._GenerateMultiSupply()

    def _GenerateTermsLowLevel(self, prefix, long_desc):
        """
        Generate the terms associated with this market, for supply and demand.

        TODO: This is now only called for the demand function; simplify to just refer
        to demand.

        :param prefix: str
        :param long_desc: str
        :return: None
        """
        Logger('Searching for demand for market {0}', priority=3, data_to_format=(self.FullCode,))
        if prefix not in ('SUP', 'DEM'):
            raise LogicError('Input to function must be "SUP" or "DEM"')
        # country = self.Parent
        short_name = prefix + '_' + self.Code
        long_name = prefix + '_' + self.FullCode
        self.AddVariable(short_name, long_desc + ' for Market ' + self.Code, '')
        term_list = []
        for s in self.CurrencyZone.GetSectors():
            if s.ID == self.ID:
                continue
            if self.ShareParent(s):
                var_name = short_name
            else:
                var_name = long_name
            try:
                term = s.GetVariableName(var_name)
            except KeyError:
                Logger('Variable {0} does not exist in {1}', priority=10,
                       data_to_format=(var_name, s.FullCode))
                continue
            term_list.append('+ ' + term)
            if prefix == 'SUP': # pragma: no cover
                # Since we assume that there is a single supplier, we can set the supply equation to
                # point to the equation in the market.
                s.AddCashFlow(var_name, self.GetVariableName(var_name), long_desc)
            else:
                # Must fill in demand equation in sectors.
                s.AddCashFlow('-' + var_name, '', long_desc)

        eqn = create_equation_from_terms(term_list)