r'''
    Module with all structures for defining rings with operators.

    Let `\sigma: R \rightarrow R` be an additive homomorphism, i.e., for all elements `r,s \in R`,
    the map satisfies `\sigma(r+s) = \sigma(r) + \sigma(s)`. We define the *ring* `R` *with operator*
    `\sigma` as the pair `(R, \sigma)`.

    Similarly, if we have a set of additive maps `\sigma_1,\ldots,\sigma_n : R \rightarrow R`.
    Then we define the *ring* `R` *with operators* `(\sigma_1,\ldots,\sigma_n)` as the tuple
    `(R, (\sigma_1,\ldots,\sigma_n))`.

    This module provides the framework to define this type of rings with as many operators as
    the user wants and we also provide a Wrapper class so we can extend existing ring structures that
    already exist in `SageMath <https://www.sagemath.org>`_.

    The factory :func:`RingWithOperator` allows the creation of these rings with operators and will determine
    automatically in which specified category a ring will belong. For example, we can create the differential
    ring `(\mathbb{Q}[x], \partial_x)` or the difference ring `(\mathbb{Q}[x], x \mapsto x + 1)` with the
    following code::

        sage: from dalgebra import *
        sage: dQx = RingWithOperators(QQ[x], lambda p : p.derivative())
        sage: sQx = RingWithOperators(QQ[x], lambda p : QQ[x](p)(x=QQ[x].gens()[0] + 1))

    Once the rings are created, we can create elements within the ring and apply the corresponding operator::

        sage: x = dQx(x)
        sage: x.operation()
        1
        sage: x = sQx(x)
        sage: x.operation()
        x + 1

    We can also create the same ring with both operators together::

        sage: dsQx = RingWithOperators(QQ[x], lambda p : p.derivative(), lambda p : QQ[x](p)(x=QQ[x].gens()[0] + 1))
        sage: x = dsQx(x)
        sage: x.operation(operation=0)
        1
        sage: x.operation(operation=1)
        x + 1

    However, these operators have no structure by themselves: `SageMath`_ is not able to distinguish the type
    of the operators if they are defined using lambda expressions or callables. This can be seen by the fact that
    the factory can not detect the equality on two identical rings::

        sage: dQx is RingWithOperators(QQ[x], lambda p : p.derivative())
        False

    To avoid this behavior, we can set the types by providing an optional list called ``types`` whose elements are
    strings with values:

    * ``homomorphism``: the operator is interpret as a homomorphism/shift/difference operator.
    * ``derivation``: the operator is considered as a derivation.
    * ``skew``: the operator is considered as a skew-derivation.
    * ``none``: the operator will only be considered as an additive Map without further structure.

    We can see that, when setting this value, the ring is detected to be equal::

        sage: dQx = RingWithOperators(QQ[x], lambda p : p.derivative(), types=["derivation"])
        sage: dQx is RingWithOperators(QQ[x], lambda p : p.derivative(), types=["derivation"])
        True
        sage: # Since we have one variable, the built-in `diff` also work
        sage: dQx is RingWithOperators(QQ[x], diff, types=["derivation"])
        True
        sage: # We can also use elements in the derivation module
        sage: dQx is RingWithOperators(QQ[x], QQ[x].derivation_module().gens()[0], types=["derivation"])
        True

    Also, we can detect this equality when adding operators sequentially instead of at once::

        sage: dsQx = RingWithOperators(QQ[x],
        ....:     lambda p : p.derivative(),
        ....:     lambda p : QQ[x](p)(x=QQ[x].gens()[0] + 1),
        ....:     types = ["derivation", "homomorphism"]
        ....: )
        sage: dsQx is RingWithOperators(dQx, lambda p : QQ[x](p)(x=QQ[x].gens()[0] + 1), types=["homomorphism"])
        True

    For specific types of operators as *derivations* or *homomorphism*, there are other functions where the ``types`` argument can be skipped
    taking the corresponding value by default::

        sage: dQx is DifferentialRing(QQ[x], lambda p : p.derivative())
        True
        sage: dsQx is DifferenceRing(DifferentialRing(QQ[x], lambda p : p.derivative()), lambda p : QQ[x](p)(x=QQ[x].gens()[0] + 1))
        True

    We can also have more complexes structures with different types of operators::

        sage: R.<x,y> = QQ[] # x is the usual variable, y is an exponential
        sage: dx, dy = R.derivation_module().gens(); d = dx + y*dy
        sage: DR = DifferentialRing(R, d)
        sage: # We add a special homomorphism where the two generators are squared but QQ is fixed
        sage: DSR = DifferenceRing(DR, R.Hom(R)([x^2, y^2]))
        sage: DSR.noperators()
        2
        sage: DSR.operator_types()
        ('derivation', 'homomorphism')

    We can see that these operator **do not commute**::

        sage: x = DSR(x); y = DSR(y)
        sage: x.difference().derivative()
        2*x
        sage: x.derivative().difference()
        1
        sage: y.difference().derivative()
        2*y^2
        sage: y.derivative().difference()
        y^2

    Finally, this module also allows the definition of skew-derivations for any ring. This requires the use
    of derivation modules with twist (see :sageref:`sage.rings.derivations <rings/sage/rings/derivation>`)::

        sage: R.<x,y> = QQ[]
        sage: s = R.Hom(R)([x-y, x+y])
        sage: td = R.derivation_module(twist=s)(x-y)
        sage: tR = RingWithOperators(R, s, td, types=["homomorphism", "skew"])
        sage: x,y = tR.gens()
        sage: (x*y).skew() == x.skew()*y + x.shift()*y.skew()
        True
        sage: (x*y).skew() == x.skew()*y.shift() + x*y.skew()
        True

    AUTHORS:

    - Antonio Jimenez-Pastor (:git:`GitHub <Antonio-JP>`)
'''

# ****************************************************************************
#  Copyright (C) 2023 Antonio Jimenez-Pastor <ajpa@cs.aau.dk>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#                  https://www.gnu.org/licenses/
# ****************************************************************************

from __future__ import annotations

from typing import Callable, Collection
from sage.all import ZZ, latex, Parent
from sage.categories.all import Morphism, Category, Rings, CommutativeRings, CommutativeAdditiveGroups
from sage.categories.morphism import IdentityMorphism, SetMorphism # pylint: disable=no-name-in-module
from sage.categories.pushout import ConstructionFunctor, pushout
from sage.misc.all import abstract_method, cached_method
from sage.rings.morphism import RingHomomorphism_im_gens # pylint: disable=no-name-in-module
from sage.rings.ring import Ring, CommutativeRing #pylint: disable=no-name-in-module
from sage.rings.derivation import RingDerivationModule
from sage.structure.element import parent, Element #pylint: disable=no-name-in-module
from sage.structure.factory import UniqueFactory #pylint: disable=no-name-in-module
from sage.symbolic.ring import SR #pylint: disable=no-name-in-module

_Rings = Rings.__classcall__(Rings)
_CommutativeRings = CommutativeRings.__classcall__(CommutativeRings)
_CommutativeAdditiveGroups = CommutativeAdditiveGroups.__classcall__(CommutativeAdditiveGroups)

####################################################################################################
###
### DEFINING THE CATEGORY FOR RINGS WITH OPERATORS
###
####################################################################################################
class RingsWithOperators(Category):
    r'''
        Category for representing rings with operators.

        Let `\sigma: R \rightarrow R` be an additive homomorphism, i.e., for all elements `r,s \in R`,
        the map satisfies `\sigma(r+s) = \sigma(r) + \sigma(s)`. We define the *ring* `R` *with operator*
        `\sigma` as the pair `(R, \sigma)`.

        Similarly, if we have a set of additive maps `\sigma_1,\ldots,\sigma_n : R \rightarrow R`.
        Then we define the *ring* `R` *with operators* `(\sigma_1,\ldots,\sigma_n)` as the tuple
        `(R, (\sigma_1,\ldots,\sigma_n))`.

        This category defines the basic methods for these rings and their elements
    '''
    ## Defining a super-category
    def super_categories(self):
        return [_Rings]

    ## Defining methods for the Parent structures of this category
    class ParentMethods: #pylint: disable=no-member
        ##########################################################
        ### METHODS RELATED WITH THE OPERATORS
        ##########################################################
        ### 'generic'
        @abstract_method
        def operators(self) -> Collection[Morphism]:
            r'''
                Method to get the collection of operators that are defined over the ring.

                These operators are maps from ``self`` to ``self`` that compute the application
                of each operator over the elements of ``self``.
            '''
            raise NotImplementedError("Method 'operators' need to be implemented")

        def noperators(self) -> int:
            r'''
                Method to get the number of operators defined over a ring
            '''
            return len(self.operators())

        def operation(self, element : Element, operator : int = None) -> Element:
            r'''
                Method to apply an operator over an element.

                This method takes an element of ``self`` and applies one of the operators defined over ``self``
                over such element. This operator is given by its index, hence raising a :class:`IndexError` if
                the index is not in the valid range.

                INPUT:

                * ``element``: an element over the operator of this ring will be applied.
                * ``operator`` (`0` by default) the index of the operator that will be applied.

                OUTPUT:

                If the index is incorrect, an :class:`IndexError` is raised. Otherwise this method
                returns `f(x)` where `x` is the ``element`` and `f` is the operator defined by ``operator``.

                EXAMPLES::

                    sage: from dalgebra import *
                    sage: dQx = RingWithOperators(QQ[x], lambda p : p.derivative())
                    sage: sQx = RingWithOperators(QQ[x], lambda p : p(x=QQ[x].gens()[0] + 1))
                    sage: sdQx = RingWithOperators(QQ[x], lambda p : p(x=QQ[x].gens()[0] + 1), lambda p : p.derivative())
                    sage: p = QQ[x](x^3 - 3*x^2 + 3*x - 1)
                    sage: dQx.operation(p)
                    3*x^2 - 6*x + 3
                    sage: sQx.operation(p)
                    x^3
                    sage: sdQx.operation(p)
                    Traceback (most recent call last):
                    ...
                    IndexError: An index for the operation must be provided when having several operations
                    sage: sdQx.operation(p, 0)
                    x^3
                    sage: sdQx.operation(p, 1)
                    3*x^2 - 6*x + 3
                    sage: sdQx.operation(p, 2)
                    Traceback (most recent call last):
                    ...
                    IndexError: ... index out of range
            '''
            if operator is None and self.noperators() == 1: operator = 0
            elif operator is None: raise IndexError("An index for the operation must be provided when having several operations")
            return self.operators()[operator](element)

        @abstract_method
        def operator_types(self) -> tuple[str]:
            r'''
                Method to get the types of the operators.

                The only condition for `\sigma: R \rightarrow R` to be a valid operator is that it is
                an additive homomorphism. However, the behavior of `\sigma` with respect to the multiplication
                of `R` categorize `\sigma` into several possibilities:

                * "none": no condition is known over this method. This will disallow some extension operations.
                * "homomorphism": the map `\sigma` is an homomorphism, i.e., for all `r, s \in R` it satisfies
                  `\sigma(rs) = \sigma(r)\sigma(s)`.
                * "derivative": the map `\sigma` satisfies Leibniz rule, i.e., for all `r, s \in R` it satisfies
                  `\sigma(rs) = \sigma(r)s + r\sigma(s)`.
                * "skew": the map `\sigma` satisfies the skew-Leibniz rule, i.e., there is an homomorphism `\delta`
                  such for all `r, s \in R` it satisfies `\sigma(rs) = \sigma(r)s + \delta(r)\sigma(s)`.

                This method returns a tuple (sorted as the output of :func:`operators`) with the types of each of the
                operators.
            '''
            raise NotImplementedError("Method 'operator_types' need to be implemented")

        ### 'derivation'
        @cached_method
        def derivations(self) -> Collection[DerivationMap]:
            r'''
                Method to filter the derivations out of a ring with operators.

                Derivations are a particular type of operators. With this method we
                provide a similar interface as with the generic operators but just with
                derivation.

                Similarly, this class offers access to homomorphisms and skew derivations.

                When no derivation is declared for a ring, an empty tuple is returned.
            '''
            return tuple([operator for (operator, ttype) in zip(self.operators(), self.operator_types()) if ttype == "derivation"])

        def nderivations(self) -> int:
            r'''
                Method to get the number of derivations defined over a ring
            '''
            return len(self.derivations())

        def has_derivations(self) -> bool:
            r'''
                Method to know if there are derivations defined over the ring.
            '''
            return self.nderivations() > 0

        def is_differential(self) -> bool:
            r'''
                Method to check whether a ring is differential, i.e, all operators are derivations.
            '''
            return self.noperators() == self.nderivations()

        def derivative(self, element: Element, derivation: int = None) -> Element:
            r'''
                Method to apply a derivation over an element.

                This method applies a derivation over a given element in the same way an operator
                is applied by the method :func:`~RingsWithOperators.ParentMethods.operation`.
            '''
            if derivation is None and self.nderivations() == 1: derivation = 0
            elif derivation is None: raise IndexError("An index for the derivation must be provided when having several derivations")
            return self.derivations()[derivation](element)

        ### 'difference'
        @cached_method
        def differences(self) -> Collection[Morphism]:
            r'''
                Method to filter the differences out of a ring with operators.

                Differences are a particular type of operators. With this method we
                provide a similar interface as with the generic operators but just with
                difference.

                Similarly, this class offers access to derivations and skew derivations.

                When no difference is declared for a ring, an empty tuple is returned.
            '''
            return tuple([operator for (operator, ttype) in zip(self.operators(), self.operator_types()) if ttype == "homomorphism"])

        def ndifferences(self) -> int:
            r'''
                Method to get the number of differences defined over a ring
            '''
            return len(self.differences())

        def has_differences(self) -> bool:
            r'''
                Method to know if there are differences defined over the ring.
            '''
            return self.ndifferences() > 0

        def is_difference(self) -> bool:
            r'''
                Method to check whether a ring is difference, i.e, all operators are homomorphisms.
            '''
            return self.noperators() == self.ndifferences()

        def difference(self, element: Element, difference: int = None) -> Element:
            r'''
                Method to apply a difference over an element.

                This method applies a difference over a given element in the same way an operator
                is applied by the method :func:`~RingsWithOperators.ParentMethods.operation`.
            '''
            if difference is None and self.ndifferences() == 1: difference = 0
            elif difference is None: raise IndexError("An index for the difference must be provided when having several differences")
            return self.differences()[difference](element)

        def shift(self, element: Element, shift: int = None) -> Element:
            r'''
                Alias for :func:`~RingsWithOperators.ParentMethods.difference`.
            '''
            return self.difference(element, shift)

        ### 'skews'
        @cached_method
        def skews(self) -> Collection[Morphism]:
            r'''
                Method to filter the skew-derivations out of a ring with operators.

                Differences are a particular type of operators. With this method we
                provide a similar interface as with the generic operators but just with
                difference.

                Similarly, this class offers access to homomorphisms and derivations.

                When no skew-derivation is declared for a ring, an empty tuple is returned.
            '''
            return tuple([operator for (operator, ttype) in zip(self.operators(), self.operator_types()) if ttype == "skew"])

        def nskews(self) -> int:
            r'''
                Method to get the number of skew-derivations defined over a ring
            '''
            return len(self.skews())

        def has_skews(self) -> bool:
            r'''
                Method to know if there are skew-derivations defined over the ring.
            '''
            return self.ndifferences() > 0

        def is_skew(self) -> bool:
            r'''
                Method to check whether a ring is skewed, i.e, all operators are skew-derivations.
            '''
            return self.noperators() == self.nskews()

        def skew(self, element: Element, skew: int = None) -> Element:
            r'''
                Method to apply a skew-derivation over an element.

                This method applies a skew-derivation over a given element in the same way an operator
                is applied by the method :func:`~RingsWithOperators.ParentMethods.operation`.
            '''
            if skew is None and self.nskews() == 1: skew = 0
            elif skew is None: raise IndexError("An index for the skew must be provided when having several skews")
            return self.skews()[skew](element)

        ##########################################################
        ### OTHER METHODS
        ##########################################################
        @abstract_method
        def operator_ring(self) -> Ring:
            r'''
                Method to get the operator ring of ``self``.

                When we consider a ring with operators, we can always consider a new (usually non-commutative)
                ring where we extend ``self`` polynomially with all the operators and its elements represent
                new operators created from the operators defined over ``self``.

                This method return this new structure.
            '''
            raise NotImplementedError("Method 'operator_ring' need to be implemented")

        def operators_commute(self, op1: int, op2: int, points: int = 10, *args, **kwds) -> bool:
            r'''
                Method to check whether two operators of the ring commute.

                This method is not deterministic (meaning that it may return ``True`` even
                when the two operators do not fully commute) but it tries to check in a fix number
                of random elements if the two operators actually commute.

                It also try to see if the operators commute in the generators of the ring.

                INPUT:

                * ``op1``: index of the first operator to check.
                * ``op2``: index of the second operator to check.
                * ``points``: number of random points to be selected.
                * ``args``: arguments to be passed to the ``random_element`` method.
                * ``kwds``: arguments to be passed to the ``random_element`` method.

                OUTPUT:

                ``True`` if all the tests indicates the operators commute, ``False`` otherwise.
            '''
            op1 = self.operators()[op1]; op2 = self.operators()[op2]

            to_check = list(self.gens()); current = self.base()
            while current.ngens() > 0 and (not 1 in to_check):
                to_check.extend([self.element_class(self, el) for el in current.gens()])
                current = current.base()
            to_check.extend(self.random_element(*args, **kwds) for _ in range(points))

            return all(op1(op2(element)) == op2(op1(element)) for element in to_check)

        def all_operators_commute(self, points: int = 10, *args, **kwds):
            r'''
                Method to check whether all operators of the ring commute.

                This method is not deterministic (meaning that it may return ``True`` even
                when the two operators do not fully commute) but it tries to check in a fix number
                of random elements if the two operators actually commute.

                It also try to see if the operators commute in the generators of the ring.

                See :func:`operators_commute` for further information

                INPUT:

                * ``points``: number of random points to be selected.
                * ``args``: arguments to be passed to the ``random_element`` method.
                * ``kwds``: arguments to be passed to the ``random_element`` method.

                OUTPUT:

                ``True`` if all the tests indicates the operators commute, ``False`` otherwise.

                EXAMPLES::

                    sage: from dalgebra import *
                    sage: R.<x> = QQ[]; d = diff; s = R.Hom(R)(x+1)
                    sage: dsR = DifferenceRing(DifferentialRing(R, d), s)
                    sage: dsR.all_operators_commute()
                    True
                    sage: R.<x,y> = QQ[]
                    sage: dx,dy = R.derivation_module().gens(); d = dx + y*dy
                    sage: s = R.Hom(R)([x + 1, y^2])
                    sage: dsR = DifferenceRing(DifferentialRing(R, d), s)
                    sage: dsR.all_operators_commute()
                    False
            '''
            return all(
                self.operators_commute(i, j, points, *args, **kwds)
                for i in range(self.noperators())
                for j in range(i+1, self.noperators())
            )

        @abstract_method
        def constant_ring(self, operation: int = 0) -> Parent:
            r'''
                Method to obtain the constant ring of a given operation.

                The meaning of a ring of constants depends on the type of operator that
                we are considering:

                * "homomorphism": the elements that are fixed by the operator.
                * "derivation": the elements that goes to zero with the operator.
                * "skew": the elements that goes to zero with the operator.
                * "none": it makes no sense to talk about constant for these operators.
            '''
            raise NotImplementedError("Method 'constant_ring' not implemented")

    ## Defining methods for the Element structures of this category
    class ElementMethods: #pylint: disable=no-member
        ##########################################################
        ### APPLICATION METHODS
        ##########################################################
        def operation(self, operation : int = None, times : int = 1) -> Element:
            r'''
                Apply an operation to ``self`` a given amount of times.

                This method applies repeatedly an operation defined in the parent of ``self``.
                See :func:`~RingsWithOperators.ParentMethods.operation` for further information.
            '''
            if(not times in ZZ or times < 0):
                raise ValueError("The argument ``times`` must be a non-negative integer")

            if(times == 0):
                return self
            elif(times == 1):
                return self.parent().operation(self, operation)
            else:
                return self.parent().operation(self.operation(operation=operation, times=times-1), operation)

        def derivative(self, derivation: int = None, times: int = 1) -> Element:
            r'''
                Apply a derivation to ``self`` a given amount of times.

                This method applies repeatedly a derivation defined in the parent of ``self``.
                See :func:`~RingsWithOperators.ParentMethods.derivative` for further information.
            '''
            if(not times in ZZ or times < 0):
                raise ValueError("The argument ``times`` must be a non-negative integer")

            if(times == 0):
                return self
            elif(times == 1):
                return self.parent().derivative(self, derivation)
            else:
                return self.parent().derivative(self.derivative(derivation=derivation, times=times-1), derivation)

        def difference(self, difference: int = None, times: int = 1) -> Element:
            r'''
                Apply a difference to ``self`` a given amount of times.

                This method applies repeatedly a difference defined in the parent of ``self``.
                See :func:`~RingsWithOperators.ParentMethods.difference` for further information.
            '''
            if(not times in ZZ or times < 0):
                raise ValueError("The argument ``times`` must be a non-negative integer")

            if(times == 0):
                return self
            elif(times == 1):
                return self.parent().difference(self, difference)
            else:
                return self.parent().difference(self.difference(difference=difference, times=times-1), difference)

        def shift(self, shift: int = None, times: int = 1) -> Element:
            r'''
                Alias for :func:`~RingsWithOperators.ElementMethods.difference`.
            '''
            return self.difference(shift, times)

        def skew(self, skew: int = None, times: int = 1) -> Element:
            r'''
                Apply a skew-derivation to ``self`` a given amount of times.

                This method applies repeatedly a difference defined in the parent of ``self``.
                See :func:`~RingsWithOperators.ParentMethods.skew` for further information.
            '''
            if(not times in ZZ or times < 0):
                raise ValueError("The argument ``times`` must be a non-negative integer")

            if(times == 0):
                return self
            elif(times == 1):
                return self.parent().skew(self, skew)
            else:
                return self.parent().skew(self.skew(skew=skew, times=times-1), skew)

        ##########################################################
        ### BOOLEAN METHODS
        ##########################################################
        def is_constant(self, operation: int = 0):
            r'''
                Method to check whether an element is a constant with respect to one operator.

                INPUT:

                * ``operation``: index defining the operation we want to check.

                OUTPUT:

                A boolean value with ``True`` is the element is a constant (see
                :func:`~RingsWithOperators.ParentMethods.constant_ring` for further information
                on what is a constant depending on the type of operator).

                REMARK: this method do not require the implementation on :func:`~RingsWithOperators.ParentMethods.constant_ring`
                on its parent structure.

                EXAMPLES::

                    sage: from dalgebra import *
                    sage: R = DifferentialRing(QQ[x], diff)
                    sage: p = R(3)
                    sage: p.is_constant()
                    True
                    sage: p = R(x^3 - 3*x + 1)
                    sage: p.is_constant()
                    False

                Some interesting constants may arise unexpectedly when adding other derivations::

                    sage: R.<x,y> = QQ[]
                    sage: dx, dy = R.derivation_module().gens(); d = y*dx - x*dy
                    sage: dR = DifferentialRing(R, d)
                    sage: x,y = dR.gens()
                    sage: x.is_constant()
                    False
                    sage: y.is_constant()
                    False
                    sage: (x^2 + y^2).is_constant()
                    True
            '''
            ttype = self.parent().operator_types()[operation]
            if ttype == "homomorphism":
                result = self.operation(operation=operation) == self
            elif ttype in ("derivation", "skew"):
                result = self.operation(operation=operation) == self.parent().zero()
            else:
                raise ValueError(f"The operation {operation} has not a good type defined")

            return result

    # methods that all morphisms involving differential rings must implement
    class MorphismMethods:
        pass

_RingsWithOperators = RingsWithOperators.__classcall__(RingsWithOperators)
####################################################################################################
###
### DEFINING THE FACTORY FOR THE CREATION OF WRAPPED RINGS
###
####################################################################################################
class RingWithOperatorFactory(UniqueFactory):
    r'''
        Factory to create wrappers around existing rings.

        The :class:`RingsWithOperatorFactory` allows to create wrapper around existing rings
        with a predefined set of operators. For doing so, we have two possibilities:

        INPUT:

        * ``base``: a commutative ring to which we will add operators.
        * ``operators``: a list with operators that will be added to ``base``. It may be one of the following:
          - An additive callable: a :class:`AdditiveMap` will be created for it.
          - An additive homomorphism: a :class:`Morphism` with appropriate domain and codomain.
          - A ring homomorphism: a :class:`Morphism` in the appropriate Hom set.
          - A (skew)-derivation: an element of a module of (skew)-derivations. The corresponding :class:`SkewMap`
            will be created for it.
        * ``types`` (optional): if given, it must be a list with the corresponding types of the operators.
          We will use this information to create different types of :class:`Morphism`.

        SPECIAL CASES:

        If this is used over another wrapped ring, this Factory will create an extended version where the
        new operators are concatenated to the previous operators.

        OUTPUT:

        A :class:`RingWithOperators_Wrapper` with the new ring with operators.
    '''
    def create_key(self, base : CommutativeRing, *operators : Callable, **kwds):
        # checking the arguments
        if len(operators) < 1:
            raise ValueError("At least one operator must be given.")
        elif len(operators) == 1 and isinstance(operators[0], Collection):
            operators = operators[0]
        operators = list(operators)
        types = list(kwds.pop("types", len(operators)*["none"]))

        if isinstance(base, RingWithOperators_Wrapper):
            operators = list(base.construction()[0].operators) + operators
            types = list(base.construction()[0].types) + types
            base = base.wrapped

        # we convert the input into a common standard to create an appropriate key
        for (i, (operator, ttype)) in enumerate(zip(operators, types)):
            if ttype == "none":
                ## We decide the structure depending on the type of object
                if operator in base.Hom(base): # it is an homomorphism - we do nothing
                    types[i] = "homomorphism"
                    new_operator = operator
                elif isinstance(parent(operator), RingDerivationModule):
                    if operator.parent().twisting_morphism() is None: # derivation without twist
                        new_operator = DerivationMap(
                            base,
                            operator
                        )
                        types[i] = "derivation"
                    else:
                        new_operator = SkewMap(
                            base,
                            operator.parent().twisting_morphism(),
                            operator
                        )
                        types[i] = "skew"
                elif isinstance(operator, Callable):
                    new_operator = AdditiveMap(
                        base,
                        operator
                    )
                else:
                    raise TypeError(f"All operators must be callables. Found {operator}")
            elif ttype == "homomorphism":
                def hom_from_callable(base, func):
                    if base.ngens() > 0 and (not 1 in base.gens()):
                        base_map = hom_from_callable(base.base(), func)
                    else:
                        base_map = None
                    hom_set = base.Hom(base)
                    return hom_set([base(func(gen)) for gen in base.gens()], base_map = base_map)
                new_operator = hom_from_callable(base, operator)
            elif ttype == "derivation":
                der_module = base.derivation_module()
                new_operator = DerivationMap(
                    base,
                    sum((operator(base_gen)*der_gen for (base_gen,der_gen) in zip(base.gens(),der_module.gens())), der_module.zero())
                )
            elif ttype == "skew":
                if not isinstance(parent(operator), RingDerivationModule):
                    raise NotImplementedError("Building skew-derivation from callable not implemented")
                twist = operator.parent().twisting_morphism()
                twist = base.Hom(base).one() if twist is None else twist
                new_operator = SkewMap(base, twist, operator)

            if new_operator != operator:
                operators[i] = new_operator
        return tuple([base, tuple(operators), tuple(types)])

    def create_object(self, _, key):
        base, operators, types = key

        return RingWithOperators_Wrapper(base, *operators, types=types)
RingWithOperators = RingWithOperatorFactory("dalgebra.ring_w_operator.ring_w_operator.RingWithOperator")

def DifferentialRing(base : CommutativeRing, *operators : Callable):
    r'''
        Method that calls the :class:`RingWithOperatorFactory` with types always as "derivation".

        See documentation on :class:`RingWithOperatorFactory` for further information.
    '''
    # checking the arguments
    if len(operators) < 1:
        raise ValueError("At least one operator must be given.")
    elif len(operators) == 1 and isinstance(operators[0], Collection):
        operators = operators[0]

    return RingWithOperators(base, *operators, types=len(operators)*["derivation"])

def DifferenceRing(base: CommutativeRing, *operators : Callable):
    r'''
        Method that calls the :class:`RingWithOperatorFactory` with types always as "homomorphism".

        See documentation on :class:`RingWithOperatorFactory` for further information.
    '''
    # checking the arguments
    if len(operators) < 1:
        raise ValueError("At least one operator must be given.")
    elif len(operators) == 1 and isinstance(operators[0], Collection):
        operators = operators[0]

    return RingWithOperators(base, *operators, types=len(operators)*["homomorphism"])

####################################################################################################
###
### DEFINING THE ELEMENT AND PARENT FOR WRAPPED RINGS
###
####################################################################################################
class RingWithOperators_WrapperElement(Element):
    def __init__(self, parent, element):
        if(not isinstance(parent, RingWithOperators_Wrapper)):
            raise TypeError("An element created from a non-wrapper parent")
        elif(not element in parent.wrapped):
            raise TypeError("An element outside the parent [%s] is requested" %parent)

        Element.__init__(self, parent=parent)
        self.wrapped = element

    # Arithmetic methods
    def _add_(self, x) -> RingWithOperators_WrapperElement:
        if parent(x) != self.parent(): # this should not happened
            x = self.parent().element_class(self.parent(), self.parent().base()(x))
        return self.parent().element_class(self.parent(), self.wrapped + x.wrapped)
    def _sub_(self, x) -> RingWithOperators_WrapperElement:
        if parent(x) != self.parent(): # this should not happened
            x = self.parent().element_class(self.parent(), self.parent().base()(x))
        return self.parent().element_class(self.parent(), self.wrapped - x.wrapped)
    def _neg_(self) -> RingWithOperators_WrapperElement:
        return self.parent().element_class(self.parent(), -self.wrapped)
    def _mul_(self, x) -> RingWithOperators_WrapperElement:
        if parent(x) != self.parent(): # this should not happened
            x = self.parent().element_class(self.parent(), self.parent().base()(x))
        return self.parent().element_class(self.parent(), self.wrapped * x.wrapped)
    def _rmul_(self, x) -> RingWithOperators_WrapperElement:
        if parent(x) != self.parent(): # this should not happened
            x = self.parent().element_class(self.parent(), self.parent().base()(x))
        return self.parent().element_class(self.parent(), self.wrapped * x.wrapped)
    def _lmul_(self, x) -> RingWithOperators_WrapperElement:
        if parent(x) != self.parent(): # this should not happened
            x = self.parent().element_class(self.parent(), self.parent().base()(x))
        return self.parent().element_class(self.parent(), self.wrapped * x.wrapped)
    def __pow__(self, n) -> RingWithOperators_WrapperElement:
        return self.parent().element_class(self.parent(), self.wrapped ** n)

    def __eq__(self, x) -> bool:
        if x is None: return False

        r = pushout(self.parent(), parent(x))
        if isinstance(r, RingWithOperators_Wrapper):
            return self.wrapped == r(x).wrapped
        return r(self) == r(x)

    def is_zero(self) -> bool:
        return self.wrapped == 0
    def is_one(self) -> bool:
        return self.wrapped == 1

    ## Other magic methods
    def __hash__(self) -> int:
        return hash(self.wrapped)
    def __str__(self) -> str:
        return str(self.wrapped)
    def __repr__(self) -> str:
        return repr(self.wrapped)
    def _latex_(self) -> str:
        return latex(self.wrapped)

class RingWithOperators_Wrapper(CommutativeRing):
    r'''
        Class for wrapping a Commutative ring and add operators over it.

        This class allows the user to translate a Commutative ring with some operations to
        the category of :class:`RingsWithOperators` preserving as many operations and properties
        of the original ring as possible, but adding the new functionality in the category.

        We do not recommend to use this class by itself. It should be created using the
        corresponding factory (see :class:`RingWithOperatorFactory` and its defined instance in
        ``dalgebra.ring_w_operator.ring_w_operator.RingWithOperators``).

        INPUT:

        * ``base``: the :class:`CommutativeRing` that will be wrapped.
        * ``operators``: a valid :class:`sage.categories.map.Map` to define an operator over ``self``.
        * ``types`` (optional): a list with the types (see :func:`RingsWithOperators.ParentMethods.operator_types`
          for further information). If nothing is given, the list will be automatically computed.
        * ``category`` (optional): argument from the category framework to allow further flexibility.
    '''
    Element = RingWithOperators_WrapperElement

    def __init__(self,
        base : CommutativeRing,
        *operators : Morphism | Collection[Morphism],
        types : Collection[str] = None,
        category = None
    ):
        #########################################################################################################
        ### CHECKING THE ARGUMENTS
        ### 'base'
        if not base in _CommutativeRings:
            raise TypeError("Only commutative rings can be wrapped as RingWithOperators")

        ### 'operators'
        if len(operators) == 1 and isinstance(operators[0], (list,tuple)):
            operators = operators[0]
        if any(not isinstance(operator, Morphism) for operator in operators):
            raise TypeError("All the given operators must be Maps")
        if any(operator.domain() != operator.codomain() or operator.domain() != base for operator in operators):
            raise TypeError("The operators must bu maps from and to the commutative ring given by 'base'")

        ### 'types'
        if types is None: # we compute the types using the maps
            types = []
            for operator in operators:
                if isinstance(operator, DerivationMap): types.append("derivation")
                elif isinstance(operator, SkewMap): types.append("skew")
                elif operator.category_for().is_subcategory(_CommutativeRings): types.append("homomorphism")
                else: types.append("none")
        else: # we check the operators behave as requested
            if not isinstance(types, (list, tuple)) or len(types) != len(operators):
                raise TypeError("The types must be a list of the same length of the operators")
            for operator, ttype in zip(operators, types):
                if ttype == "none":
                    if not operator.category_for().is_subcategory(_CommutativeAdditiveGroups):
                        raise ValueError(f"An operator invalid for type 'none' -> {operator}")
                elif ttype == "homomorphism":
                    if not operator.category_for().is_subcategory(_CommutativeRings):
                        raise ValueError(f"An operator invalid for type 'homomorphism' -> {operator}")
                elif ttype == "derivation":
                    if not isinstance(operator, DerivationMap):
                        raise ValueError(f"An operator invalid for type 'derivation' -> {operator}")
                elif ttype == "skew":
                    if not isinstance(operator, SkewMap):
                        raise ValueError(f"An operator invalid for type 'skew' -> {operator}")
                else:
                    raise ValueError(f"Invalid type provided -> {ttype}")

        self.__types = tuple(types)

        #########################################################################################################
        # CREATING CATEGORIES
        categories = [_RingsWithOperators, base.category()]
        if(isinstance(category, (list, tuple))):
            categories += list(category)
        elif(category != None):
            categories.append(category)

        #########################################################################################################
        ### CALLING THE SUPER AND ARRANGING SOME CONVERSIONS
        super().__init__(base.base(), category=tuple(categories))
        self.__wrapped = base

        # registering conversion to simpler structures
        current = self.base()
        morph = RingWithOperators_Wrapper_SimpleMorphism(self, current)
        current.register_conversion(morph)
        while(not(current.base() == current)):
            current = current.base()
            morph = RingWithOperators_Wrapper_SimpleMorphism(self, current)
            current.register_conversion(morph)

        #########################################################################################################
        ### CREATING THE NEW OPERATORS FOR THE CORRECT STRUCTURE
        self.__operators : tuple[WrappedMap] = tuple([WrappedMap(self, operator) for operator in operators])

    @property
    def wrapped(self) -> CommutativeRing: return self.__wrapped

    def operators(self) -> tuple[WrappedMap]: return self.__operators

    def operator_types(self) -> tuple[str]: return self.__types

    ## Coercion methods
    def _has_coerce_map_from(self, S) -> bool:
        r'''
            Return ``True`` if it is possible to have a coercion map from `S` to ``self``.
        '''
        if isinstance(S, RingWithOperators_Wrapper):
            return self.wrapped._has_coerce_map_from(S.wrapped) # the operators do not matter for coercing elements
        else:
            return self.wrapped._has_coerce_map_from(S)

    def _element_constructor_(self, x) -> RingWithOperators_WrapperElement:
        r'''
            Extended definition of :func:`_element_constructor_`.
        '''
        if x in SR:
            # conversion from symbolic ring --> using its string representation
            x = str(x)
        elif isinstance(parent(x), RingWithOperators_Wrapper):
            # conversion from other wrapped rings with operators --> we convert the element within
            x = x.wrapped
        p = self.wrapped._element_constructor_(x)
        return self.element_class(self, p)

    def _is_valid_homomorphism_(self, codomain, im_gens, base_map=None) -> bool:
        return self.wrapped._is_valid_homomorphism_(codomain, im_gens, base_map)

    def construction(self) -> RingWithOperatorsFunctor:
        return RingWithOperatorsFunctor([operator.function for operator in self.operators()], self.operator_types()), self.wrapped

    # Rings methods
    def characteristic(self) -> int:
        return self.wrapped.characteristic()

    def gens(self) -> tuple[RingWithOperators_WrapperElement]:
        return tuple([self.element_class(self, gen) for gen in self.wrapped.gens()])

    def ngens(self) -> int:
        return self.wrapped.ngens()

    ## Representation methods
    def __repr__(self) -> str:
        begin = "Differential " if self.is_differential() else "Difference " if self.is_difference() else ""
        return f"{begin}Ring [[{self.wrapped}], {repr(self.operators())}]"

    def __str__(self) -> str:
        return repr(self)

    def _latex_(self) -> str:
        return "".join((
            r"\left(",
            latex(self.wrapped),
            ", ",
            latex(self.operators()) if self.noperators() > 1 else latex(self.operators()[0]),
            r"\right)"
        ))

    ## Element generation
    def one(self) -> RingWithOperators_WrapperElement:
        r'''
            Return the one element in ``self``.

            EXAMPLES::

                sage: from dalgebra import *
                sage: R = RingWithOperators(QQ['x'], diff)
                sage: R.one()
                1
        '''
        return self.element_class(self, self.wrapped.one())

    def zero(self) -> RingWithOperators_WrapperElement:
        r'''
            Return the zero element in ``self``.

            EXAMPLES::

                sage: from dalgebra import *
                sage: R = RingWithOperators(QQ['x'], diff)
                sage: R.zero()
                0
        '''
        return self.element_class(self, self.wrapped.zero())

    def random_element(self,*args,**kwds) -> RingWithOperators_WrapperElement:
        r'''
            Creates a random element in this ring.

            This method creates a random element in the base infinite polynomial ring and
            cast it into an element of ``self``.
        '''
        p = self.wrapped.random_element(*args,**kwds)
        return self.element_class(self, p)

####################################################################################################
###
### DEFINING THE CONSTRUCTION FUNCTOR AND SIMPLE MORPHISM
###
####################################################################################################
class RingWithOperatorsFunctor(ConstructionFunctor):
    def __init__(self, operators: Collection[Morphism], types: Collection[str]):
        if len(operators) != len(types):
            raise ValueError("The length of the operators and types must coincide.")
        self.__operators = tuple(operators)
        self.__types = tuple(types)

        super().__init__(_CommutativeRings, _RingsWithOperators)

    ### Methods to implement
    def _coerce_into_domain(self, x: Element) -> Element:
        if not x in self.domain():
            raise TypeError(f"The object [{x}] is not an element of [{self.domain()}]")

    def _apply_functor(self, x):
        return RingWithOperators(x, self.__operators, self.__types)

    def _repr_(self):
        return f"RingWithOperators(*,{self.__operators}])"

    def __eq__(self, other):
        return self.__class__ == other.__class__ and self.__operators == other.__operators and self.__types == other.__types

    def merge(self, other):
        if isinstance(other, RingWithOperatorsFunctor):
            return RingWithOperatorsFunctor(self.__operators + other.__operators, self.__types + other.__types)
        else:
            raise NotImplementedError(f"{self} can only be merged with other RingWithOperatorsFunctor")

    @property
    def operators(self) -> Collection[Morphism]:  return self.__operators
    @property
    def types(self): return self.__types

class RingWithOperators_Wrapper_SimpleMorphism(Morphism):
    r'''
        Class representing maps to simpler rings.

        This map allows the coercion system to detect that some elements in a
        :class:`RingWithOperator_Wrapper` are included in simpler rings.
    '''
    def __init__(self, domain, codomain):
        super().__init__(domain, codomain)

    def _call_(self, p):
        return self.codomain()(p.wrapped)

####################################################################################################
###
### DEFINING THE REQUIRED MAPS FOR THIS MODULE
###
####################################################################################################
class AdditiveMap(SetMorphism):
    def __init__(self, domain : Parent, function : Callable):
        # We create the appropriate Hom set
        hom = domain.Hom(domain, category=_CommutativeAdditiveGroups)
        self.function = function
        super().__init__(hom, function)

    def __str__(self) -> str:
        return f"Additive Map [{repr(self)}]\n\t- From: {self.domain()}\n\t- To  : {self.codomain()}"

    def __repr__(self) -> str:
        return f"{repr(self.function)}"

    def _latex_(self) -> str:
        return latex(self.function)

    def __eq__(self, other) -> bool:
        return isinstance(other, AdditiveMap) and self.domain() == other.domain() and self.function == other.function

    def __hash__(self) -> int:
        return self.function.__hash__()

class SkewMap(AdditiveMap):
    def __init__(self, domain : Parent, twist : Morphism, function : Callable):
        # we check the input
        if not twist in domain.Hom(domain):
            raise TypeError("The twist for a skew derivation must be an homomorphism.")
        tw_der_module = domain.derivation_module(twist=twist)
        if not function in tw_der_module:
            raise TypeError("The function for a skew derivation must be in the corresponding module")
        self.twist = twist
        super().__init__(domain, function)

    def __str__(self) -> str:
        return f"Skew Derivation [{repr(self)}] over (({self.domain()}))"

class DerivationMap(SkewMap):
    def __init__(self, domain, function : Callable):
        super().__init__(domain, domain.Hom(domain).one(), function)

    def __str__(self) -> str:
        return f"Derivation [{repr(self)}] over (({self.domain()}))"

class WrappedMap(AdditiveMap):
    def __init__(self, domain : RingWithOperators_Wrapper, function : Morphism):
        if not isinstance(domain, RingWithOperators_Wrapper):
            raise TypeError("A WrappedMap can only be created for a 'RingWithOperators_Wrapper'")

        if function.domain() != domain.wrapped:
            raise ValueError(f"The map to be wrapped must have appropriate domain: ({domain.wrapped}) instead of ({function.domain()})")

        super().__init__(domain, lambda p : domain(function(domain(p).wrapped)))
        self.function = function

    def __repr__(self) -> str:
        if isinstance(self.function, RingHomomorphism_im_gens):
            im_gens = {v: im for (v,im) in zip(self.function.domain().gens(), self.function.im_gens())}
            return f"Hom({im_gens})"
        elif isinstance(self.function, IdentityMorphism):
            return "Id"
        else:
            return super().__repr__()

    def __str__(self) -> str:
        return f"Wrapped [{repr(self)}] over (({self.domain()}))"

    def _latex_(self) -> str:
        if isinstance(self.function, RingHomomorphism_im_gens):
            im_gens = {v: im for (v,im) in zip(self.function.domain().gens(), self.function.im_gens())}
            return r"\sigma\left(" + r", ".join(f"{latex(v)} \\mapsto {latex(im)}" for (v,im) in im_gens.items()) + r"\right)"
        elif isinstance(self.function, IdentityMorphism):
            return r"\text{id}"
        return super()._latex_()

__all__ = ["RingsWithOperators", "RingWithOperators", "DifferentialRing", "DifferenceRing"]


import torch as th
import pandas as pd
import numpy as np
import dgl

from bipartite_graph import BipartiteGraph

#######################
# user-item Subgraph Extraction
#######################

def map_newid(df, col):
    old_ids = df[col]
    old_id_uniq = old_ids.unique()

    id_dict = {old: new for new, old in enumerate(sorted(old_id_uniq))}
    new_ids = np.array([id_dict[x] for x in old_ids])

    return new_ids

def one_hot(idx, length):
    x = th.zeros([len(idx), length], dtype=th.int32)
    x[th.arange(len(idx)), idx] = 1.0
    return x

def get_neighbor_nodes_labels(u_node_idx, i_node_idx, graph,
                              hop=1, sample_ratio=1.0, max_nodes_per_hop=200):

    # 1. neighbor nodes sampling
    dist = 0
    u_nodes, i_nodes = th.unsqueeze(u_node_idx, 0), th.unsqueeze(i_node_idx, 0)
    u_dist, i_dist = th.tensor([0], dtype=th.long), th.tensor([0], dtype=th.long)
    u_visited, i_visited = th.unique(u_nodes), th.unique(i_nodes)
    u_fringe, i_fringe = th.unique(u_nodes), th.unique(i_nodes)

    for dist in range(1, hop+1):
        # sample neigh alternately
        # diff from original code : only use one-way edge (u-->i)
        u_fringe, i_fringe = graph.in_edges(i_fringe)[0], graph.out_edges(u_fringe)[1]
        u_fringe = th.from_numpy(np.setdiff1d(u_fringe.numpy(), u_visited.numpy()))
        i_fringe = th.from_numpy(np.setdiff1d(i_fringe.numpy(), i_visited.numpy()))
        u_visited = th.unique(th.cat([u_visited, u_fringe]))
        i_visited = th.unique(th.cat([i_visited, i_fringe]))

        if sample_ratio < 1.0:
            shuffled_idx = th.randperm(len(u_fringe))
            u_fringe = u_fringe[shuffled_idx[:int(sample_ratio*len(u_fringe))]]
            shuffled_idx = th.randperm(len(i_fringe))
            i_fringe = i_fringe[shuffled_idx[:int(sample_ratio*len(i_fringe))]]
        if max_nodes_per_hop is not None:
            if max_nodes_per_hop < len(u_fringe):
                shuffled_idx = th.randperm(len(u_fringe))
                u_fringe = u_fringe[shuffled_idx[:max_nodes_per_hop]]
            if max_nodes_per_hop < len(i_fringe):
                shuffled_idx = th.randperm(len(i_fringe))
                i_fringe = i_fringe[shuffled_idx[:max_nodes_per_hop]]
        if len(u_fringe) == 0 and len(i_fringe) == 0:
            break

        u_nodes = th.cat([u_nodes, u_fringe])
        i_nodes = th.cat([i_nodes, i_fringe])
        u_dist = th.cat([u_dist, th.full((len(u_fringe), ), dist,)])
        i_dist = th.cat([i_dist, th.full((len(i_fringe), ), dist,)])

    nodes = th.cat([u_nodes, i_nodes])

    # 2. node labeling
    # this labeling is based on hop from starting nodes
    u_node_labels = th.stack([x*2 for x in u_dist])
    v_node_labels = th.stack([x*2+1 for x in i_dist])
    node_labels = th.cat([u_node_labels, v_node_labels])

    return nodes, node_labels

def subgraph_extraction_labeling(u_node_idx, i_node_idx, graph,
                                 hop=1, sample_ratio=1.0, max_nodes_per_hop=200,):

    # extract the h-hop enclosing subgraph nodes around link 'ind'
    nodes, node_labels = get_neighbor_nodes_labels(u_node_idx=u_node_idx, i_node_idx=i_node_idx, graph=graph,
                                                  hop=hop, sample_ratio=sample_ratio, max_nodes_per_hop=max_nodes_per_hop)

    subgraph = dgl.node_subgraph(graph, nodes, store_ids=True)

    subgraph.ndata['nlabel'] = one_hot(node_labels, (hop+1)*2)
    subgraph.ndata['x'] = subgraph.ndata['nlabel']

    # set edge mask to zero as to remove links between target nodes in training process
    subgraph.edata['edge_mask'] = th.ones(subgraph.number_of_edges(), dtype=th.int32)
    su = subgraph.nodes()[subgraph.ndata[dgl.NID]==u_node_idx]
    si = subgraph.nodes()[subgraph.ndata[dgl.NID]==i_node_idx]
    _, _, target_edges = subgraph.edge_ids([su, si], [si, su], return_uv=True)
    subgraph.edata['edge_mask'][target_edges.to(th.long)] = 0

    # mask target edge label
    subgraph.edata['label'][target_edges.to(th.long)] = 0.0

    # timestamp normalization
    # compute ts diff from target edge & min-max normalization
    n = subgraph.edata['ts'].shape[0]
    timestamps = subgraph.edata['ts'][:n//2]
    standard_ts = timestamps[target_edges.to(th.long)[0]]
    timestamps = th.abs(timestamps - standard_ts.item())
    timestamps = 1 - (timestamps - th.min(timestamps)) / (th.max(timestamps)-th.min(timestamps) + 1e-5)
    subgraph.edata['ts'] = th.cat([timestamps, timestamps], dim=0) + 1e-5

    return subgraph

#######################
# Ego-graph Extraction
#######################
def get_egograph_neighbor(center_node_idx:int, graph:dgl.DGLGraph,
                          hop=1, max_nodes_per_hop=20):

    # 1. neighbor nodes sampling
    node_dist = th.tensor([0], dtype=th.long)
    visited_nodes = th.tensor([center_node_idx], dtype=th.long)
    nodes = th.tensor([center_node_idx], dtype=th.long)
    fringe = th.tensor([center_node_idx], dtype=th.long)

    for dist in range(1, hop+1):
        fringe = graph.in_edges(fringe)[0]
        fringe = th.from_numpy(np.setdiff1d(fringe.numpy(), visited_nodes.numpy()))
        visited_nodes = th.unique(th.cat([visited_nodes, fringe]))

        if max_nodes_per_hop < len(fringe):
            shuffled_idx = th.randperm(len(fringe))
            fringe = fringe[shuffled_idx[:max_nodes_per_hop]]

        if len(fringe) == 0 :
            break

        nodes = th.cat([nodes, fringe])
        node_dist = th.cat([node_dist, th.full((len(fringe),), dist,)])

    # 2. node labeling
    # this labeling is based on hop from starting nodes
    node_labels = th.stack([x for x in node_dist])

    return nodes, node_labels

def egograph_extraction(node_idx, graph,
                        hop=1, max_nodes_per_hop=10,):

    # extract the h-hop enclosing subgraph nodes around link 'ind'
    nodes, node_labels = get_egograph_neighbor(center_node_idx=node_idx, graph=graph,
                                               hop=hop, max_nodes_per_hop=max_nodes_per_hop)

    subgraph = dgl.node_subgraph(graph, nodes, store_ids=True)

    subgraph.ndata['nlabel'] = one_hot(node_labels, hop+1)
    subgraph.ndata['x'] = subgraph.ndata['nlabel']
    return subgraph

#######################
# Ego graph Dataset
#######################

class EgoGraphDataset(th.utils.data.Dataset):
    def __init__(self, graph,
                hop=2, max_nodes_per_hop=5):

        self.nodes = graph.nodes()
        self.graph = graph

        self.hop = hop
        self.max_nodes_per_hop = max_nodes_per_hop

    def __len__(self):
        return len(self.nodes)

    def __getitem__(self, idx):
        node_idx = self.nodes[idx]
        ego_graph = egograph_extraction(node_idx, self.graph, hop=self.hop, max_nodes_per_hop=self.max_nodes_per_hop)

        return ego_graph

def collate_data(data):
    g_list = data
    g = dgl.batch(g_list)
    return g

""" Dataset for classifier"""
class AmazonDataset(th.utils.data.Dataset):
    def __init__(self, df:pd.DataFrame, embeds, start=0, end=-1):
        df = df.reset_index()
        self.labels = df['rating'] #pre processed 0~4
        self.users = df['user_id']
        self.items = df['item_id'] + max(self.users) + 1
        self.embeds = embeds[start:]

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        uid = self.users[idx]
        iid = self.items[idx]
        u_emb = self.embeds[uid]
        i_emb = self.embeds[iid]
        label = self.labels[idx]
        return uid, iid, u_emb, i_emb, label


"""
FILE: kernelregression.py
LAST MODIFIED: 24-12-2015
DESCRIPTION: Module for Gaussian kernel regression

===============================================================================
This file is part of GIAS2. (https://bitbucket.org/jangle/gias2)

This Source Code Form is subject to the terms of the Mozilla Public
License, v. 2.0. If a copy of the MPL was not distributed with this
file, You can obtain one at http://mozilla.org/MPL/2.0/.
===============================================================================
"""
import logging

import numpy as np

log = logging.getLogger(__name__)

def _gaussianKernel(T, t, s):
    c1 = 1.0 / (s * np.sqrt(2 * np.pi))
    c2 = np.exp(-(T - t) ** 2.0 / (2 * s * s))
    return c1 * c2

def _gaussianKernelMax(s):
    return 1.0 / (s * np.sqrt(2 * np.pi))

def _weightedMean(x, w):
    # print x
    # print x.shape
    # print w.shape
    return (x * w).sum(-1) / w.sum()

def _weightedSD(x, w):
    u = _weightedMean(x, w)
    d = ((x.T - u).T) ** 2.0
    return np.sqrt((d * w).sum(-1) / w.sum())

class KernelRegressor(object):
    sigmaeps = 1e-6

    def __init__(self, k=2, sigma0=1.0, wmin=0.35):
        self.k = k
        self.sigma0 = sigma0
        self.wmin = wmin  # kernel weight cutoff

        self.nTarg = None  # target number of obs at each time point
        self.sigmas = None
        self.n = None
        self.y = None
        self.ytmeans = None
        self.ytSDs = None
        self.ytweights = None
        self.ytinds = None
        self.x = None
        self.xsamples = None
        self.xt = None
        self.xtn = None
        self.xtnTarg = None
        self.xmin = None
        self.xmax = None

    def fit(self, x, y, xmin, xmax, xsamples):
        """
        x is a 1-d array, the independent variable e.g. time

        if y is multivariate, each observation should be a column, so variables
        in rows.
        """
        self.x = x
        self.y = y
        self.xmin = xmin
        self.xmax = xmax
        self.sigmaeps = 1e-3 * (xmax - xmin)
        self.xsamples = xsamples
        self.xt = np.linspace(self.xmin, self.xmax, self.xsamples)

        self._fitInit()
        self._optWidths()

        return self.ytmeans

    def _fitInit(self):

        # initialise kernel width at each time point
        self.sigmas = np.ones(self.xsamples) * self.sigma0

        # calc number of obs at each y sampling (yt[i])
        self.xtn = []
        self.ytmeans = []
        self.ytweights = []
        self.ytSDs = []
        self.ytinds = []
        for t, s in zip(self.xt, self.sigmas):
            ty, tw, tyi = self._getKernelY(t, s)
            self.xtn.append(len(tw))
            self.ytweights.append(tw)
            self.ytinds.append(tyi)
            self.ytmeans.append(_weightedMean(ty, tw))
            self.ytSDs.append(_weightedSD(ty, tw))

        self.ytmeans = np.array(self.ytmeans)

        # calculate target number obs per time point
        self.xtnTarg = np.median(self.xtn)

    def _optWidths(self):
        for i in range(self.xsamples):
            tyi = None
            change = 0
            if self.xtn[i] > (self.xtnTarg + self.k):
                while self.xtn[i] > (self.xtnTarg + self.k):
                    self.sigmas[i] -= self.sigmaeps
                    ty, tw, tyi = self._getKernelY(self.xt[i], self.sigmas[i])
                    self.xtn[i] = len(tw)
                change = 1

            if self.xtn[i] < (self.xtnTarg - self.k):
                while self.xtn[i] < (self.xtnTarg - self.k):
                    self.sigmas[i] += self.sigmaeps
                    ty, tw, tyi = self._getKernelY(self.xt[i], self.sigmas[i])
                    self.xtn[i] = len(tw)
                change = 1

            # calculated weighted mean
            # if tyi is not None:
            #   print self.xt[i], tyi
            # else:
            #   print self.xt[i], 'no opt'
            if change:
                self.ytmeans[i] = _weightedMean(ty, tw)
                self.ytSDs[i] = _weightedSD(ty, tw)
                self.ytweights[i] = tw
                self.ytinds[i] = tyi

    def _getKernelY(self, t, s):
        xw = _gaussianKernel(self.x, t, s)
        validMask = xw >= (self.wmin * _gaussianKernelMax(s))
        validW = xw[validMask]
        validY = self.y[:, validMask]
        return validY, validW, np.where(validMask)[0]

def test():
    ndata = 500
    x = 3.0 * np.random.rand(ndata)
    y = np.sin(x) + 0.1 * (np.random.rand(ndata) - 0.5)

    r = KernelRegressor(k=2, sigma0=0.05, wmin=0.35)
    r.fit(x, y, 0.1, 2.9, 10)

    log.debug('xt      sigma     n     y_mean')
    for xt, s, n, ym in zip(r.xt, r.sigmas, r.xtn, r.ytmeans):
        log.debug(('{:4.3f}   {:4.3f}   {:4d}    {:4.3f}').format(xt, s, n, ym))

    import matplotlib.pyplot as plt
    f = plt.figure()
    plt.scatter(x, y)
    plt.plot(r.xt, r.ytmeans, 'r')
    plt.plot(r.xt, r.ytmeans + r.ytSDs, 'g--')
    plt.plot(r.xt, r.ytmeans - r.ytSDs, 'g--')
    plt.show()

if __name__ == '__main__':
    test()


#coding: utf-8
__author__ = "Lário dos Santos Diniz"

from django.contrib import admin

from .models import (RPGSystem)

class RPGSystemAdmin(admin.ModelAdmin):

    list_display = ['name', 'description', 'site']
    search_fields = ['name', 'description', 'site']

admin.site.register(RPGSystem, RPGSystemAdmin)


# Import required libraries
import pandas as pd
from sqlalchemy import create_engine

# Load data from source into a Pandas dataframe
df = pd.read_csv('source_data.csv')

# Perform data transformation and cleaning
df = df.dropna()
df['column_name'] = df['column_name'].str.upper()

# Load data into a SQL database
engine = create_engine('sqlite:///data.db')
df.to_sql('table_name', engine, if_exists='replace')

# Perform data reconciliation
query = '''
SELECT COUNT(*)
FROM table_name
'''
count = pd.read_sql_query(query, engine)['COUNT(*)'][0]

# Log the reconciliation result
if count == df.shape[0]:
    print('Data integration successful')
else:
    print('Data integration failed')

# This is just a basic example to illustrate the steps involved in a data integration workflow.
# ou can add more complex transformations, data quality checks, and error handling to suit your specific requirements.


import random
import sys

import numpy
import torch
import pygad
import pygad.torchga
from nn import create_ga, create_network
import math

class Gym:

    def __init__(self, w, h, left_ai, right_ai):
        self.turn_i = 0
        self.w = w
        self.h = h
        self.left_ai = left_ai
        self.right_ai = right_ai
        self.left_units = []
        self.right_units = []

    def prepare(self):
        i = 0
        for unit in self.left_units:
            unit.set_position(0, i)
            unit.side = True
            i += random.randint(1, 2)
        i = 0
        for unit in self.right_units:
            unit.set_position(self.w - 1, i)
            unit.side = False
            i += random.randint(1, 2)

    def turn(self, represent, dprint = True):
        # for now - left to right foreach

        states = []

        if len(self.left_units) == 0:
            self.left_ai.score -= 100
            self.right_ai.score += 100
        elif len(self.right_units) == 0:
            self.left_ai.score += 100
            self.right_ai.score -= 100

        for unit in self.left_units:
            self.action(
                self.left_ai,
                unit,
                self.left_ai.decide(unit,
                                    self)
            )
            if represent:
                states.append(self.represent(dprint))

        for unit in self.right_units:
            self.action(
                self.right_ai,
                unit,
                self.right_ai.decide(unit,
                                     self)
            )
            if represent:
                states.append(self.represent(dprint))
        self.turn_i += 1
        return states

    def action(self, ai, unit, act):
        act = numpy.array(act)
        pos = numpy.array([unit.x, unit.y])

        new_pos = (act - pos)
        x_diff = round(new_pos[0])
        y_diff = round(new_pos[1])

        new_x = unit.x + x_diff
        new_y = unit.y + y_diff
        dist = math.sqrt(x_diff ** 2 + y_diff ** 2)

        if self.is_out_of_bounds(new_x, new_y) or dist > unit.speed:
            ai.score -= 1
        elif self.is_occupied(new_x, new_y) and self.get_unit_at(new_y, new_y) != unit:
            ai.score -= 1
        else:
            unit.set_position(new_x, new_y)
            e = self.get_adjacent_enemy(new_x, new_y, unit.side)
            if e is not None:
                dmg = unit.damage * unit.count
                fh = (e.count - 1) * e.max_health + e.health
                fh -= dmg
                new_c = fh // e.max_health
                new_h = fh  % e.max_health
                e.count = new_c
                e.health = new_h
                if new_c <= 0:
                    [self.right_units, self.left_units][e.side].remove(e)
                    ai.score += 50
                ai.score += dmg
                opp_ai = self.right_ai if self.left_ai == ai else self.left_ai
                opp_ai.score -= dmg

            # here we check if enemy adjacent

    def get_adjacent_enemy(self, x, y, side):
        adj_mat = [
            [-1, 0],  # L
            [-1, -1],  # TL
            [0, -1],  # T
            [1, -1],  # TR
            [1, 0],  # R
            [1, 1],  # BR
            [0, 1],  # B
            [-1, 1],  # BL
        ]

        for adj in adj_mat:
            pos = numpy.array([x, y]) + adj

            if self.is_out_of_bounds(*pos):
                continue

            if self.is_occupied(*pos):
                u = self.get_unit_at(*pos)
                if u.side != side:
                    return u

        return None

    def is_out_of_bounds(self, x, y):
        return x >= self.w or y >= self.h or x < 0 or y < 0

    def is_occupied(self, x, y):
        return self.get_unit_at(x, y) is not None

    def get_unit_at(self, x, y):
        for unit in self.left_units + self.right_units:
            if unit.x == x and unit.y == y:
                return unit
        return None

    def represent(self, dprint = True):
        table = []
        for x in range(self.w):
            arr = []
            table.append(arr)
            for y in range(self.h):
                arr.append('-')

        # default, 1st unit, 2nd unit in team
        colors = ['\033[91m', '\033[94m', '\033[92m', '\033[93m', '\033[95m', '\033[96m']

        i = 1
        for unit in self.left_units:
            table[unit.y][unit.x] = colors[i] + '▣' + colors[0]
            i += 1
        i = 1

        for unit in self.right_units:
            table[unit.y][unit.x] = colors[i] + '△' + colors[0]
            i += 1

        t = ''
        if dprint:
            print('State, turn ', self.turn_i)
        for x in range(self.w):
            for y in range(self.h):
                if dprint:
                    print(table[x][y], end='  ')
                t += table[x][y] + '  '
            if dprint:
                print()
            t += '\n'
        return t

class Unit:

    def __init__(self, speed, damage, health, count):
        self.speed = speed
        self.damage = damage
        self.count = count
        self.max_health = health
        self.health = health
        self.side = True
        self.x = -1
        self.y = -1

    def set_position(self, x, y):
        self.x = x
        self.y = y

class AI:
    def __init__(self, model):
        self.model = model
        self.score = 0

    def get_input(self, unit, gym):
        unit_table = numpy.zeros(8 * 8).reshape((8, 8))
        for u in gym.left_units + gym.right_units:
            # index = u.x + u.y * gym.w
            if unit.side == u.side:
                unit_table[u.y][u.x] = -1
            else:
                unit_table[u.y][u.x] = 1

        count_table = numpy.zeros(8 * 8).reshape((8, 8))
        for u in gym.left_units + gym.right_units:
            # index = u.x + u.y * gym.w
            count_table[u.y][u.x] = u.count

        return torch.tensor([[
            unit_table,
            count_table
        ]]).float()

    def decide(self, unit, gym):
        # inputs
        input = self.get_input(unit, gym)
        # outputs
        # 1x64 array of weights
        decision = self.model.model(torch.tensor(input).float())
        # we need to find in range best spot
        bw = -1
        bwi = -1

        for (idx, weight) in enumerate(decision):
            px, py = idx % 8, idx // 8
            dist = math.sqrt((px - unit.x) ** 2 + (py - unit.y) ** 2)
            if weight > bw and dist <= unit.speed:
                bw = weight
                bwi = idx

        return [bwi % 8, bwi // 8]

def create_gym(left_ai, right_ai):
    g = Gym(8, 8, left_ai, right_ai)
    speed = 2
    damage = 15
    health = 20

    min, max = 3,3

    # 5 of 1's 1 of 15
    g.left_units.append(
        Unit(speed,damage, health, 15)
    )
    for i in range(random.randint(min, max)):
        g.left_units.append(
            Unit(speed, damage, health, 1)
        )
    g.right_units.append(
        Unit(speed, damage, health, 15)
    )
    for i in range(random.randint(min, max)):
        g.right_units.append(
            Unit(speed, damage, health, 1)
        )
    g.prepare()
    return g

def save_model(model, filename):
    torch.save(model, filename)

class GymTester:

    def run_simulation(self, ai_to_test, solution, side=False, rep=False):
        # get weights from current model
        model_weights_dict = pygad.torchga.model_weights_as_dict(model=ai_to_test.model, weights_vector=solution)
        # Use the current solution as the model parameters.
        ai_to_test.model.load_state_dict(model_weights_dict)
        # create gym where left is tested model and right is previous generation

        ais = [AI(ai_to_test), AI(ai_to_test) if self.previous_best_model is None else AI(self.previous_best_model)]
        if side:
            ais = ais[::-1] # reverse so tested ai becomes second(right side)
        g = create_gym(ais[0], ais[1])
        for i in range(15):
            g.turn(rep)

        return g

    def callback_gen(self, ga_instance):
        model = self.previous_best_model.model
        weights = pygad.torchga.model_weights_as_dict(model=model, weights_vector=self.ga_instance.best_solution()[0])
        self.previous_best_model.model.load_state_dict(weights)
        print("%d / %d Fitness of the best solution :" % (ga_instance.generations_completed, self.iterations), ga_instance.best_solution()[1])

    def fitness_func(self, solution, solution_idx):
        # random_side = bool(random.getrandbits(1))
        random_side = False
        g = self.run_simulation(self.training_ga, solution, side=random_side, rep = False)

        score = g.left_ai.score if random_side == False else g.right_ai.score
        return score

    def __init__(self, iterations_to_run):
        self.iterations = iterations_to_run
        self.input_size = 8 * 8
        self.training_ga, self.ga_instance = create_ga(self.input_size, iterations_to_run,
                                                       lambda x, y: self.fitness_func(x, y), 10,
                                                       callback_gen=lambda instance: self.callback_gen(instance))
        self.previous_best_model = create_network(self.input_size, 64, 10)
        pass

    def run(self):
        self.ga_instance.run()

        solution, solution_fitness, solution_idx = self.ga_instance.best_solution()

        save_model(self.training_ga.model, 'models/%d_%d.pt' % (self.iterations, solution_fitness))

        self.previous_best_model = self.training_ga
        self.run_simulation(self.training_ga, solution, side=False, rep=True)

        # self.ga_instance.plot_result(title="PyGAD & PyTorch - Iteration vs. Fitness", linewidth=4)

def main(iters):
    GymTester(iters).run()

if __name__ == "__main__":
   # stuff only to run when not called via 'import' here
   iters = int(sys.argv[1])
   main(iters)


import csv
import models
from operator import attrgetter
import statistics

def filter_data(columns, row):
    ENERGYSTARScore = row[columns.index('ENERGYSTARScore')]
    if ENERGYSTARScore == '':
        return False

    YearBuilt = row[columns.index('YearBuilt')]
    if int(YearBuilt) < 1920:
        return False

    BuildingName = row[columns.index('BuildingName')]
    if 'CENTER' in BuildingName:
        return False

    return True

def parse_file(filename):
    print('Starting parsing file:')
    list_of_buildings = []
    # read file
    with open(filename) as file_raw:
        print(f'Opening file {filename}')
        data_raw = csv.reader(file_raw, delimiter=',')
        print('Iterating over each line in the file:')
        firstline = True
        columns = []
        for row in data_raw:
            if firstline:
                columns = row
                firstline = False
                continue
            # filter data
            if filter_data(columns, row) == False:
                continue

            OSEBuildingID = row[columns.index('OSEBuildingID')]
            # Question 1
            BuildingName = row[columns.index('BuildingName')]
            NumberofFloors = int(row[columns.index('NumberofFloors')])
            # Question 2
            ENERGYSTARScore = int(row[columns.index('ENERGYSTARScore')])
            SiteEUI = float(row[columns.index('SiteEUI(kBtu/sf)')])
            NaturalGas = row[columns.index('NaturalGas(kBtu)')]
            # Question 3
            Neighborhood = row[columns.index('Neighborhood')]
            Electricity = int(row[columns.index('Electricity(kBtu)')])
            # Question 4
            PropertyGFABuildings = int(
                row[columns.index('PropertyGFABuilding(s)')])
            PrimaryPropertyType = row[columns.index('PrimaryPropertyType')]
            LargestPropertyUseType = row[columns.index(
                'LargestPropertyUseType')]
            SecondLargestPropertyUseType = row[columns.index(
                'SecondLargestPropertyUseType')]
            ThirdLargestPropertyUseType = row[columns.index(
                'ThirdLargestPropertyUseType')]
            PropertyTypes = [
                PrimaryPropertyType,
                LargestPropertyUseType,
                SecondLargestPropertyUseType,
                ThirdLargestPropertyUseType]

            # instance a Building
            building = models.Building(
                OSEBuildingID,
                BuildingName,
                NumberofFloors,
                ENERGYSTARScore,
                SiteEUI,
                NaturalGas,
                Neighborhood,
                Electricity,
                PropertyGFABuildings,
                PropertyTypes)

            list_of_buildings.append(building)

    print('Finish parsing file and filter data.')
    print(f'list_of_buildings length is {len(list_of_buildings)}')
    return list_of_buildings

def find_buildings_with_largest_number_of_floors(
        list_of_buildings, max_number_of_floors):
    list_of_buildings_with_largest_number_of_floors = []
    for building in list_of_buildings:
        if building.NumberofFloors == max_number_of_floors:
            list_of_buildings_with_largest_number_of_floors.append(building)
            continue
        if building.NumberofFloors > max_number_of_floors:
            max_number_of_floors = building.NumberofFloors
            list_of_buildings_with_largest_number_of_floors.clear()
            list_of_buildings_with_largest_number_of_floors.append(building)
            continue
    return list_of_buildings_with_largest_number_of_floors

def question_1(list_of_buildings):
    print('Solving question 1:')
    max_number_of_floors = 0
    list_of_buildings_with_largest_number_of_floors = find_buildings_with_largest_number_of_floors(
        list_of_buildings, max_number_of_floors)
    print('Building(s) with the largest number of floors:')
    for building in list_of_buildings_with_largest_number_of_floors:
        print(
            'Id {0}, name {1}, with {2} floors.'.format(
                building.OSEBuildingID,
                building.BuildingName,
                building.NumberofFloors))

def question_2(list_of_buildings):
    print('Solving question 2:')
    energy_star_score_threshold = 97
    nums_of_building_energy = sum(
        building.ENERGYSTARScore >= energy_star_score_threshold for building in list_of_buildings)
    print(
        'The number of buildings with ENERGYSTARScore of at least {0} is {1}.'.format(
            energy_star_score_threshold,
            nums_of_building_energy))

def filter_natural_gas(list_of_buildings):
    print('Filtering buildings that used natural gas:')
    list_of_buildings_used_naturalgas = []
    for building in list_of_buildings:
        if building.NaturalGas:
            list_of_buildings_used_naturalgas.append(building)
    return list_of_buildings_used_naturalgas

def question_3(list_of_buildings):
    print('Solving question 3:')
    list_of_buildings_used_naturalgas = filter_natural_gas(list_of_buildings)
    list_of_gas_used = []
    for building in list_of_buildings_used_naturalgas:
        list_of_gas_used.append(building.SiteEUI)
    median_SiteEUI = statistics.median(list_of_gas_used)

    print(
        f'The median of SiteEUI among buildings using natural gas is {median_SiteEUI} kBtu.')

    # list_of_gas_used.sort()
    # n = len(list_of_gas_used)
    # if n % 2 != 0:
    #     median_SiteEUI = list_of_gas_used[int(n / 2)]
    # else:
    #     median_SiteEUI = (list_of_gas_used[int(n / 2)] + list_of_gas_used[int(n / 2 - 1)]) / 2
    # print(
    #     f'The median of SiteEUI among buildings using natural gas is {median_SiteEUI}.')

def filter_neighborhood(list_of_buildings, neighborhood):
    # neighborhood = 'BALLARD'
    list_of_buildings_in_neighborhood = []
    for building in list_of_buildings:
        if building.Neighborhood == neighborhood:
            list_of_buildings_in_neighborhood.append(building)
    return list_of_buildings_in_neighborhood

def get_threshold_building(filename, threshold_building_name):
    with open(filename) as file_raw:
        data_raw = csv.reader(file_raw, delimiter=',')
        firstline = True
        columns = []
        for row in data_raw:
            if firstline:
                columns = row
                firstline = False
                continue

            if threshold_building_name in row:
                threshold_building_electricity = int(
                    row[columns.index('Electricity(kBtu)')])
                print('Building {0} used the amount of Electricity of {1} kBtu.'.format(
                    threshold_building_name, threshold_building_electricity))
                return threshold_building_electricity

def filter_electricity_use(list_of_buildings, threshold_building_electricity):
    list_of_building_used_more_electricity = []
    for building in list_of_buildings:
        if building.Electricity > threshold_building_electricity:
            list_of_building_used_more_electricity.append(building)
    return list_of_building_used_more_electricity

def question_4(filename, list_of_buildings):
    print('Solving question 4:')
    neighborhood = 'BALLARD'
    threshold_building_name = 'BIOMED FAIRVIEW RESEARCH CENTER'

    threshold_building_electricity = get_threshold_building(
        filename, threshold_building_name)
    list_of_buildings_in_neighborhood = filter_neighborhood(
        list_of_buildings, neighborhood)
    list_of_buildings_used_more_electricity = filter_electricity_use(
        list_of_buildings_in_neighborhood, threshold_building_electricity)

    for building in list_of_buildings_used_more_electricity:
        print(
            'Building id {0}, name {1} used the amount of Electricity of {2} kBtu.'.format(
                building.OSEBuildingID,
                building.BuildingName,
                building.Electricity))

def filter_property_type_block_list(list_of_buildings, block_list):
    print('Filtering buildings that are not offices or hospitals:')
    list_of_buildings_not_block_listed = []
    for building in list_of_buildings:
        chosen = True
        for property_type in building.PropertyTypes:
            if not chosen:
                continue
            for block_type in block_list:
                if block_type in property_type:
                    chosen = False
                    break
        if chosen:
            list_of_buildings_not_block_listed.append(building)
    return list_of_buildings_not_block_listed

def filter_property_floor_area_building(
        list_of_buildings, floor_area_threshold):
    list_of_large_buildings = []
    for building in list_of_buildings:
        if building.PropertyGFABuildings > floor_area_threshold:
            list_of_large_buildings.append(building)
    return list_of_large_buildings

def question_5(list_of_buildings):
    print('Solving question 5:')
    football_field_area = 57600  # square feet
    floor_area_threshold = 15 * football_field_area
    print(f'Threshold property area is {floor_area_threshold} square feet.')

    block_list = ['Hospital', 'Office']
    list_of_buildings_not_offices_or_hospitals = filter_property_type_block_list(
        list_of_buildings, block_list)

    list_of_large_buildings = filter_property_floor_area_building(
        list_of_buildings_not_offices_or_hospitals, floor_area_threshold)

    for building in list_of_large_buildings:
        print('Building id {0}, name {1}, total floor area for buildings {2} square feet.'.format(
            building.OSEBuildingID, building.BuildingName, building.PropertyGFABuildings
        ))

def ada_data_challenge(filename):
    list_of_buildings = parse_file(filename)
    question_1(list_of_buildings)
    question_2(list_of_buildings)
    question_3(list_of_buildings)
    question_4(filename, list_of_buildings)
    question_5(list_of_buildings)


import serial
import matplotlib
matplotlib.use('TkAgg') # MUST BE CALLED BEFORE IMPORTING plt
from matplotlib import pyplot as plt
import queue
import threading
import animation
import seaborn as sns
import numpy as np
import time

class ArduinoReader(threading.Thread):
    def __init__(self, stop_event, sig, serport):
        threading.Thread.__init__(self)
        self.stopped = stop_event
        self.signal = sig
        self.pixalarray = [[] for _ in range(8)]
        self.pixalarray0 = [[] for _ in range(8)]
        port = serport
        # self.s = serial.Serial(port, 9600, timeout=1, rtscts=True, dsrdtr=True)
        self.s = serial.Serial(port, 115200, timeout=0.1, rtscts=True, dsrdtr=True)
        if not self.s.isOpen():
            self.s.open()
        print("connected: ", self.s)

    def run(self):
        while not self.stopped.is_set():
            # print(self.s.readline().rstrip())
            try:
                dstr = str(self.s.readline())
                # print(dstr.find('Frame'))
                if dstr.find('Frame') > 0:
                    data_d1 = str(self.s.readline()).split(',')
                    self.pixalarray = [float(x) for x in data_d1[1:-1]]
                    data_d0 = str(self.s.readline()).split(',')
                    self.pixalarray0 = [float(x) for x in data_d0[1:-1]]
                    if len(self.pixalarray) == 64 and len(self.pixalarray0) == 64:
                        self.signal.put([self.pixalarray, self.pixalarray0])
            except:
                continue
        self.clean()

    def get_signal(self):
        return self.signal

    def clean(self):
        # self.s.cancel_read()
        while self.s.isOpen():
            self.s.close()
            # print('the serial port is open? {}'.format(self.s.isOpen()))

def colorscale(data, minc, maxc):
    data_scale = 256 * (data - minc) / (maxc - minc)
    if data_scale < 0:
        data_scale = 0
    elif data_scale > 255:
        data_scale = 255
    print(data, data_scale)
    return int(data_scale)

if __name__ == '__main__':
    try:
        q = queue.Queue()
        stop_event = threading.Event()
        data_reader = ArduinoReader(stop_event, q, 'COM3')
        data_reader.start()
        fig, (ax1, ax0) = plt.subplots(1, 2)
        im1 = ax1.imshow(np.random.uniform(low=22, high=32, size=(8, 8)), vmin=22, vmax=32, cmap='jet', interpolation='lanczos')
        im0 = ax0.imshow(np.random.uniform(low=22, high=32, size=(8,8)), vmin = 22, vmax = 32, cmap='jet', interpolation='lanczos')
        plt.tight_layout()
        plt.ion()
        while True:
            [frame1, frame0] = q.get()
            im1.set_array(np.reshape(frame1, (8, 8)))
            im0.set_array(np.reshape(frame0, (8, 8)))
            # plt.draw()
            plt.pause(0.001)
        plt.ioff()
        plt.show()
    finally:
        stop_event.set()
        data_reader.clean()
        data_reader.join()