import logging
logging.basicConfig(format='%(levelname)s:%(message)s', level=logging.WARN)

class Space:
    left_trans = dict(N="W", E="N", S="E", W="S")
    move_xy = dict(N=(0, 1), E=(1, 0), S=(0, -1), W=(-1, 0))

    def to_left(self, dir):
        return self.left_trans[dir]

    def to_back(self, dir):
        return self.left_trans[self.left_trans[dir]]

    def to_right(self, dir):
        return self.left_trans[self.left_trans[self.left_trans[dir]]]

    def neighbour(self, pos, dir):
        return(pos[0]+self.move_xy[dir][0],
               pos[1]+self.move_xy[dir][1])

class Factory(Space):
    """The factory contains the static layout of the floor as well as the agents.
       It can add new features to places, but this should only be done in the
       beginning.
       Each field is initially a "FreeElement" surrounded by four "OpenBorder".
       These are stored in a dict "floor" indexed by positions, where at each
       position there are five slots: "N", "S", "W", "E" for the borders
       and "P" for the place.
       Apart from that the factory senses collisions between agents and
       dispatches the methods call to apply the different factory elements.
    """

    def __init__(self, cols=5, rows=5, installs=()):
        """Inserts objects in 'installs' according to their position in a dict
        """
        self.rows = rows
        self.cols = cols
        self.agents = []
        self.beams = []
        self.step = 0
        self.reg_phase = 0
        self.floor = dict()  # floor with coordinates
        self.init_floor(cols, rows, installs)

    def init_floor(self, cols, rows, installs):
        """Insert borders, free elements and all objects to be installed
        """
        for x in range(1, cols+1):
            for y in range(1, rows+1):
                self.floor[(x, y)] = dict()
                self.install(FreeElement(x, y))  # might be overwritten later!
                for slot in ("N", "E", "S", "W"):
                    fx = ((slot == "N" and y == self.rows) or
                          (slot == "S" and y == 1) or
                          (slot == "E" and x == self.cols) or
                          (slot == "W" and x == 1))
                    self.install(OpenBorder(x, y, dir=slot, factoryexit=fx))
        for obj in installs:
            self.install(obj)

    def install(self, obj):
        """Installs an object in the right slot and inserts a pointer
           back to factory in each object!
        """

        obj.factory = self
        if isinstance(obj, MoveableThing):
            self.agents.append(obj)
        elif isinstance(obj, BorderElement):
            self.floor[obj.pos][obj.dir] = obj
        else:
            self.floor[obj.pos]["P"] = obj

    def occupied(self, pos, virtual=False):
        """Checks for agents in this field and returns them.
           if virtual is True, also the virtual agents are returned"""
        result = []
        if pos:
            for a in self.agents:
                if a.pos == pos and (not a.virtual or virtual):
                    result.append(a)
        return result

    def collision(self, agent):
        """Returns all agents that collide with 'agent'.
        """
        if agent.pos and self.floor[agent.pos]['P'].nocollisions:
            return []  # in Pits, nocollisions is True!
        result = self.occupied(agent.pos)
        if agent in result:
            result.remove(agent)
        return result

    def push_conflict(self, pusher):
        """Checks whether there is another pusher affecting the same cells.
           If so, we set the conflict flag in both pushers.
        """
        conflict_detected = False
        for pos in self.floor:
            obj = self.floor[pos]["P"]
            try:
                if obj != pusher and obj.marked & pusher.marked != set():
                    logging.info("Pusher conflict: %s and %s" %
                                 (obj.pos, pusher.pos))
                    obj.conflict = True
                    pusher.conflict = True
                    conflict_detected = True
            except Exception:
                pass
        if conflict_detected:
            logging.info("Pusher stopped at %s" % (pusher.pos,))
        return conflict_detected

    def apply(self):
        """Apply all factory elements. Order is important!"""
        for pos in self.floor:
            self.floor[pos]['P'].apply_element()
        for a in self.agents:  # consider robots as factory elements
            a.apply_element()
        for b in self.beams:  # finally consider all shot beams
            b.apply_element()
        self.beams = []  # remove all beams after one round of shots

    def exec_reg_phase(self, reg_phase, cmdlist):
        """Execute one register phase."""
        self.reg_phase = reg_phase
        logging.info("*** Starting register phase %s" % reg_phase)
        self.step = 0
        for cmd in cmdlist:
            for repitition in range(cmd[0]):
                cmd[1]()
                self.apply()  # only pits will be considered here (step==0)
                self.resolve_conflicts()
        for self.step in range(1, 9):
            self.apply()
            self.resolve_conflicts()

    def resolve_conflicts(self):
        """Resolve all conflicts created by collosions.
           Simply send all agents a "resolve" message!
           Finally, all agents marked as killed will die.
        """
        for a in self.agents:
            a.resolve()
        for a in self.agents:
            a.reset()
        for a in self.agents:
            if a.killed:
                a.die()

class Thing(Space):
    """Each thing has a position on the floor

    """
    def __init__(self, x, y, factory=None, **kw):
        self.factory = factory
        self.pos = (x, y)

    def __str__(self):
        try:
            return self.name.upper()
        except AttributeError:
            return self.__class__.__name__.upper()

class OrientedThing(Thing):
    """Anything oriented using cardinal directions (N, E, S, W)"""

    def __init__(self, x, y, dir="N", **kw):
        super().__init__(x, y, **kw)
        self.dir = dir

class MoveableThing(OrientedThing):
    """Anything moveable"""

    def __init__(self, x, y, dir="N", **kw):
        super().__init__(x, y, dir, **kw)
        self.lastconf = None  # last configuration, i.e. pos and dir
        self.pushmarker = None  # last pusher element which marked fields

    def move(self, dir):
        """Passive or active movement to a neighbouring field.
           Is blocked immediately by walls.
           Other robots are pushed.
           Collisions are handled as in the transport case, i.e.,
           resolve is called afterwards.
        """
        logging.info("move: move of %s from %s in direction %s" %
                     (self, self.pos, dir))
        self.lastconf = (self.pos, self.dir)
        # try to move into direction dir
        if not self.factory.floor[self.pos][dir].leavecell(self):
            # could not leave cell
            logging.info("move: %s cannot move because of a wall" % self)
            self.lastconf = None
            return
        colliders = self.factory.collision(self)
        for collider in colliders:  # if collision with another robot, push
            collider.move(dir)

    def transport(self, dir):
        """Transport on a conveyor belt without collision checks
        After the parallel move, somebody has to call the resolution method
        """
        if self.lastconf:  # has already been moved in parallel
            return False
        self.lastconf = (self.pos, self.dir)
        if self.factory.floor[self.pos][dir].leavecell(self):
            logging.info("transport: %s to %s" % (self, self.pos))
            return True
        return False

    def mark(self, dir, pusher):
        """Marks current field and then propagates to neighbour.
           Marks also the agent with the pusher who tried to push.
        """
        self.pushmarker = pusher
        marked = {self.pos}
        if self.factory.floor[self.pos][dir].leavecell(self, onlycheck=True):
            neighbour_cell = self.neighbour(self.pos, dir)
            if (neighbour_cell in self.factory.floor and
               not self.factory.floor[neighbour_cell]['P'].nocollisions):
                marked |= {neighbour_cell}
                for a in self.factory.occupied(neighbour_cell):
                    marked |= a.mark(dir, pusher)
        return marked

    def resolve(self):
        """Is called for each robot after every robot has been
           transported or moved
        """
        collider = self.factory.collision(self)
        if collider:
            for a in collider + [self]:
                a.retract()
        if self.pushmarker and self.pushmarker.conflict:
            self.retract()

    def retract(self):
        """There was a collision or confilct, so we have to retract the robot
           to the last position
        """
        if self.lastconf:
            self.pos = self.lastconf[0]
            self.dir = self.lastconf[1]
            self.lastconf = None
            self.killed = False
            logging.info("retract: send %s back to %s" % (self, self.pos))
            for a in self.factory.collision(self):
                a.retract()

    def reset(self):
        self.lastconf = None
        if self.pushmarker:
            self.pushmarker.conflict = False
            self.pushmarker.marked = set()
            self.pushmarker = None

class FlyingThing(OrientedThing):

    def fly(self, shootby=None):
        """Movement until we hit an obstacle; shootby is the agent
           shooting the object (not leading to a collision or to damage!)
        """
        if not self.pos:
            return  # left factory!
        logging.info("fly: %s wants to fly from %s in direction %s" %
                     (self, self.pos, self.dir))
        occ = self.factory.occupied(self.pos, virtual=True)
        if (occ and shootby not in occ):
            logging.info("fly: %s cannot move because of an obstacle" % self)
        elif not self.factory.floor[self.pos][self.dir].leavecell(self):
            # could not leave cell
            logging.info("fly: %s cannot move because of a wall" % self)
        else:
            self.fly()

class TurnableThing(OrientedThing):

    def rotate_left(self, *rest):
        self.dir = self.to_left(self.dir)
        logging.info("rotate_left: %s facing now %s" %
                     (self, self.dir))

    def u_turn(self, *rest):
        self.dir = self.to_back(self.dir)
        logging.info("u_turn: %s facing now %s" %
                     (self, self.dir))

    def rotate_right(self, *rest):
        self.dir = self.to_right(self.dir)
        logging.info("rotate_right: %s facing now %s" %
                     (self, self.dir))

class AffectingElement(Thing):

    # kill agents and other things
    def kill(self, obj, message):
        logging.info(message % (obj, obj.pos))
        obj.pos = None
        obj.killed = True

    # damage an agent
    def damage(self, agent, message):
        logging.info(message % (agent, agent.pos))
        agent.damage += 1
        if agent.damage >= 10:
            self.kill(agent, "%s killed by too much damage at %s")

class FactoryElement(AffectingElement):
    """This is a factory element. You need to specify in which register phase
       and which element step the element will be active. This can be done on a
       per class or per instance base. The main methods to specify are the
       acton and act methods.

       Element steps can be:
       0: the robots move
       1: express conveyor belts move one step
       2: express conveyor belts and normal conveyor belts move one step
       3: pushers push
       4: gears turn
       5: crushers crush
       6: lasers fire (installed and robot lasers)
       7: one way portals
       8: repair and checkpoint actions
    """
    nocollisions = False
    active_steps = {}

    def __init__(self, x, y, reg_phases=None, **kw):
        if (reg_phases):
            self.active_reg_phases = reg_phases
        else:
            self.active_reg_phases = {1, 2, 3, 4, 5}
        super().__init__(x, y, **kw)

    # apply an element in isolation (e.g., shoot laser beam) and
    # on all agents if all preconditions are met
    def apply_element(self):
        if (self.factory.step in self.active_steps and
           self.factory.reg_phase in self.active_reg_phases):
            # go over all agents in this cell:
            for agent in self.factory.occupied(self.pos, virtual=True):
                self.acton(agent)
            self.act()

    # apply in isolation
    def act(self):
        pass

    # act on one agent
    def acton(self, agent):
        pass

    # do something on arrival from direction dir (for turning conveyors)
    def on_arrival(self, agent, dir):
        pass

class Beam(FlyingThing, FactoryElement):

    active_steps = {6}

    def __init__(self, x, y, shootby=None, **kw):
        super().__init__(x, y, **kw)
        self.shootby = shootby
        self.factory.beams.append(self)
        if shootby:
            logging.info("%s shot by %s at %s, direction %s" %
                         (self, self.shootby, self.pos, self.dir))
        else:
            logging.info("%s shot at %s, direction %s" %
                         (self, self.pos, self.dir))
        self.fly(shootby=shootby)

    def acton(self, agent):
        if not self.pos or self.shootby is agent:
            return
        self.damage(agent, "%s got damaged by laserbeam at %s")

class Robot(MoveableThing, TurnableThing, FactoryElement):

    active_steps = {6}

    def __init__(self, x, y, dir="N",
                 name="Anonymous", **kw):
        super().__init__(x, y, dir, **kw)
        self.name = name
        self.damage = 0
        self.lives = 3
        self.alive = True
        self.virtual = False
        self.killed = False
        logging.info("init: %s starts at %s with orientation %s" %
                     (self, self.pos, self.dir))

    def die(self):
        logging.info("%s dies" % self)
        self.alive = False
        self.lives -= 1
        self.damage = 0
        self.killed = False

    def go(self):
        logging.info("GO command: %s" % self)
        self.onestep(True)

    def backup(self):
        logging.info("BACKUP command: %s" % self)
        self.onestep(False)

    def onestep(self, forward):
        """active execution of one step (forward or backward)"""
        forb = "forward"
        if not forward:
            forb = "backward"
        logging.info("onestep: %s wants to make 1 step %s (dir=%s)" %
                     (self, forb, self.dir))
        if not self.alive:
            logging.info("onestep: %s is dead and cannot move" % self)
            return
        if forward:
            self.move(self.dir)
        else:
            self.move(self.to_back(self.dir))

    # And now the robot acts as a factory element:
    def act(self):
        if self.alive and not self.virtual:
            Beam(self.pos[0], self.pos[1], shootby=self,
                 factory=self.factory, dir=self.dir)

class FreeElement(FactoryElement):
    """ Nothing happens to anybody on free space

    >>> f = Factory(2, 2, installs=(Robot(1, 1, dir="S", name="T"),))
    >>> f.exec_reg_phase(1,())
    >>> f.agents[0].pos, f.agents[0].dir
    ((1, 1), 'S')

    """
    pass

class Pit(FactoryElement):
    """ A pit kills the agent

    >>> t = Robot(2,1, dir = 'N', name='T')
    >>> f = Factory(2, 2, installs=(Pit(2,2), t))
    >>> f.exec_reg_phase(1, [(2, t.go)])
    >>> t.pos, t.dir
    (None, 'N')

    """
    nocollisions = True
    active_steps = {0, 1, 2, 3, 4, 5, 6, 7, 8}  # also active when robots move!

    def acton(self, agent):
        self.kill(agent, "%s was killed in pit at %s")

class BorderElement(OrientedThing, AffectingElement):
    """Border elements are a bit different
       They interact through leavecell() and entercall() methods
       with moving and flying things.
       """

    def leavecell(self, thing, onlcheck=False):
        """Try to leave cell through this border (containing direction).
        Change thing.pos (if not onlycheck) and return True if successful,
        otherwise return False
        """
        raise NotImplementedError("leavecall must be defined")

    def entercell(self, thing, onlycheck=False):
        """Try to enter cell through this border. Change thing.pos
        (if not onlycheck) and return True
        if successful. Otherwise return false.
        """
        raise NotImplementedError("entercall must be defined")

class OpenBorder(BorderElement):
    """Open border is the default for all cells. Leaving the factory leads to
       a fast death
       >>> r = Robot(1, 1, dir="E", name="r")
       >>> f = Factory(2, 2, installs=(r,))
       >>> f.exec_reg_phase(1, [(1, r.go)])
       >>> r.pos
       (2, 1)
       >>> f.exec_reg_phase(2, [(2, r.go)])
       >>> print(r.pos)
       None

       """

    def __init__(self, x, y, factoryexit=False, **kw):
        super().__init__(x, y, **kw)
        self.exit = factoryexit

    def leavecell(self, thing, onlycheck=False):
        if self.exit:
            if not onlycheck:
                self.kill(thing, "%s was killed leaving factory at %s")
            return True
        else:
            return self.factory.floor[self.neighbour(
                self.pos, self.dir)][self.to_back(self.dir)].\
                entercell(thing, onlycheck)

    def entercell(self, thing, onlycheck=False):
        if not onlycheck:
            thing.pos = self.pos
        return True

class Wall(BorderElement):

    def leavecell(self, thing, onlycheck=False):
        if not onlycheck:
            logging.info("%s bumped into a wall and does not leave %s" %
                         (thing, thing.pos))
        return False

    def entercell(self, thing, onlycheck=False):
        if not onlycheck:
            logging.info("%s bumped into a wall and is back at %s" %
                         (thing, thing.pos))
        return False

class Pusher(FactoryElement, OrientedThing):
    """Pusher pushes an agent and all others in line one step
       A wall opposite to push direction should be installed!

       In order to find out whether pushing doesn't lead to
       collisions with robots pushed by other pushers, we first
       need to find out which fields we are going to affect.

       >>> f = Factory(2, 2, installs=(Robot(1, 1, dir="S", name="T"),))
       >>> f.install(Pusher(1, 1, dir='E', reg_phases={2, 5}))
       >>> f.install(Robot(2, 1, dir="E", name="X"))
       >>> f.exec_reg_phase(1,())
       >>> f.agents[0].pos, f.agents[0].dir
       ((1, 1), 'S')
       >>> f.exec_reg_phase(2,())
       >>> f.agents[0].pos, f.agents[1].pos
       ((2, 1), None)
       >>> r12 = Robot(1,2,name='R(1,2)')
       >>> r21 = Robot(2,1,name='R(2,1)')
       >>> r22 = Robot(2,2,name='R(2,2)')
       >>> r23 = Robot(2,3,name='R(2,3)')
       >>> r33 = Robot(3,3,name='R(3,3)')
       >>> f = Factory(3, 3, installs=(r12, r21, r22, r23, r33))
       >>> f.install(Pusher(1,2,dir='E',reg_phases={1}))
       >>> f.install(Pusher(2,1,dir='N',reg_phases={1,2,4}))
       >>> f.install(Pusher(3,3,dir='W',reg_phases={1,2,3}))
       >>> f.exec_reg_phase(1,())
       >>> print(r12.pos, r21.pos, r22.pos, r23.pos, r33.pos)
       (1, 2) (2, 1) (2, 2) (2, 3) (3, 3)
       >>> f.exec_reg_phase(2,())
       >>> print(r12.pos, r21.pos, r22.pos, r23.pos, r33.pos)
       (1, 2) (2, 1) (2, 2) (2, 3) (3, 3)
       >>> f.exec_reg_phase(3,())
       >>> print(r12.pos, r21.pos, r22.pos, r23.pos, r33.pos)
       (1, 2) (2, 1) (2, 2) (1, 3) (2, 3)
       >>> f.exec_reg_phase(4,())
       >>> print(r12.pos, r21.pos, r22.pos, r23.pos, r33.pos)
       (1, 2) (2, 2) (2, 3) (1, 3) None

    """

    active_steps = {3}

    def __init__(self, x, y, **kw):
        super().__init__(x, y, **kw)
        self.marked = set()  # set of potentially affected positions
        self.conflict = False  # set when a conflict has been detected

    def acton(self, agent):
        logging.info("PUSHER: try to push %s at %s into dir %s" %
                     (agent, self.pos, self.dir))
        # mark all potentially affected agents with pusher (and remember)
        self.marked = agent.mark(self.dir, self)
        if not self.factory.push_conflict(self):
            agent.move(self.dir)

class Gear(FactoryElement):
    """Gear turns robot on the spot

       >>> f = Factory(2, 2, installs=(Robot(1, 1, dir="S", naem="T"),\
           Gear(1, 1)))
       >>> f.exec_reg_phase(1,())
       >>> f.agents[0].pos, f.agents[0].dir
       ((1, 1), 'W')

    """
    active_steps = {4}

    def __init__(self, x, y, clockwise=True, **kw):
        super().__init__(x, y, **kw)
        self.clockwise = clockwise

    def acton(self, agent):
        if self.clockwise:
            logging.info("GEAR: turn %s clockwise at %s" % (agent, self.pos))
            agent.rotate_right()
        else:
            logging.info("GEAR: turn %s counter-clockwise at %s" %
                         (agent, self.pos))
            agent.rotate_left()

class Conveyor(FactoryElement, OrientedThing):
    """Conveyor transport robots one step

       >>> f = Factory(2, 2, installs=(Robot(1, 1, dir="S", name="T"),\
           Conveyor(1, 1, dir="N")))
       >>> f.exec_reg_phase(1,())
       >>> f.agents[0].pos, f.agents[0].dir
       ((1, 2), 'S')
       >>> f.install(Robot(1, 1, dir="W", name="X"))
       >>> f.exec_reg_phase(1,())
       >>> f.agents[0].pos, f.agents[0].dir, f.agents[1].pos, f.agents[1].dir
       ((1, 2), 'S', (1, 1), 'W')
       >>> f = Factory(2, 3, installs=(Conveyor(1,2,dir='E'),\
           Conveyor(2,1,dir='N')))
       >>> t = Robot(1,2,dir='S',name='T'); f.install(t)
       >>> x = Robot(2,1,dir='E',name='X'); f.install(x)
       >>> f.exec_reg_phase(1, [])
       >>> t.pos, x.pos
       ((1, 2), (2, 1))
       >>> f.install(Pit(2, 2))
       >>> f.exec_reg_phase(2, [])
       >>> t.pos, x.pos
       (None, None)

    """
    active_steps = {2}

    def acton(self, agent):
        if agent.transport(self.dir) and agent.pos:
            self.factory.floor[agent.pos]['P'].on_arrival(agent, self.dir)

class FastConveyor(Conveyor):

    active_steps = {1, 2}

class TurningConveyor(Conveyor):

    """Note: The turning element transports things normally. Only
       when another conveyor moves something on us, then this sending
       element has to call the on_arrival method!

       >>> f = Factory(5, 5, installs=())
       >>> f.install(Conveyor(2, 1, dir='N'))
       >>> f.install(TurningConveyor(2, 2, dir='W'))
       >>> f.install(Robot(2, 1, dir="E", name='Y'))
       >>> f.exec_reg_phase(2,())
       >>> f.agents[0].pos, f.agents[0].dir
       ((2, 2), 'N')

    """

    def on_arrival(self, agent, dir):  # called only by ordinary conveyor
        logging.info("%s arrives on %s with dir %s" % (agent, self, dir))
        if self.dir == self.to_right(dir):
            logging.info("%s: rotate %s right" % (self, agent))
            agent.rotate_right()
        if self.dir == self.to_left(dir):
            logging.info("%s: rotate %s left" % (self, agent))
            agent.rotate_left()

class FastTurningConveyor(FastConveyor, TurningConveyor):
    """Everything is handled by the base classes
    """
    pass

class Crusher(FactoryElement):
    """Crushers destroy robots, when the crushers are active.

    >>> r = Robot(1, 2, dir="S", name="T")
    >>> f = Factory(2, 2, installs=(r,))
    >>> f.install(Crusher(1, 1, reg_phases={3, 4}))
    >>> f.exec_reg_phase(1,[(1, r.go)])
    >>> r.pos, r.alive
    ((1, 1), True)
    >>> f.exec_reg_phase(2,())
    >>> r.pos, r.alive
    ((1, 1), True)
    >>> f.exec_reg_phase(3,())
    >>> r.pos, r.alive
    (None, False)

    """
    active_steps = {5}

    def acton(self, agent):
        self.kill(agent, "%s is killed by crusher at %s")

class Laser(FactoryElement):
    """This is an industrial strenght laser!
    >>> r = Robot(3,3, dir='S',name='R')
    >>> f = Factory(3, 3, installs=(r, Laser(3, 1, dir='N', beams=3)))
    >>> f.exec_reg_phase(1, ())
    >>> r.damage
    3

    """
    active_steps = {6}

    def __init__(self, x, y, beams=1, **kw):
        self.beams = beams
        super().__init__(x, y, **kw)

    def act(self):
        for b in range(self.beams):
            Beam(self.pos[0], self.pos[1], factory=self.factory)

class RepairSite():  # TODO: Insert appropriate superclass(es)
    """Repair fields reduce the damage of a robot standing on it

    >>> r = Robot(1, 4, dir='S', name='T')
    >>> f = Factory(1, 4, installs=(r,))
    >>> f.install(RepairSite(1, 3, reducedamage=1))
    >>> f.install(RepairSite(1, 2, reducedamage=2))
    >>> f.install(RepairSite(1, 1, reducedamage=2))
    >>> r.damage = 4
    >>> r.pos, r.damage
    ((1, 4), 4)
    >>> f.exec_reg_phase(1,[(1, r.go)])
    >>> r.pos, r.damage
    ((1, 3), 3)
    >>> f.exec_reg_phase(1,[(1, r.go)])
    >>> r.pos, r.damage
    ((1, 2), 1)
    >>> f.exec_reg_phase(1,[(1, r.go)])
    >>> r.pos, r.damage
    ((1, 1), 0)

    """
    # TODO: Insert your code here!
    pass

class OneWayPortal():  # TODO: Insert appropriate superclass(es)
    """OneWayPortal are elements that teleport robots to a target place
       become active in the element move phase 6

    >>> r = Robot(1, 1, dir='E', name='Rob')
    >>> t = Robot(3, 3, dir='S', name='Tom')
    >>> f = Factory(3, 3, installs=(r, t))
    >>> f.install(OneWayPortal(1, 1, target=(3, 3)))
    >>> f.install(Wall(3, 2, dir='N'))
    >>> f.exec_reg_phase(1, ())
    >>> r.pos, t.pos
    ((1, 1), (3, 3))
    >>> f.install(OneWayPortal(3, 3, target=(3, 2)))
    >>> f.exec_reg_phase(2, ())
    >>> r.pos, t.pos
    ((3, 3), (3, 2))

    """
    # TODO: Insert your code here!
    pass

if __name__ == "__main__":
    import doctest
    logging.getLogger().setLevel(logging.INFO)
    doctest.testmod()


#Python Midterm Review
#Syuta Sugawara
"""
1.
Given an 8-bit binary string, (eg. "10101010") convert
it to decimal (base 10).  Use Python bit operations.
"""

def convert_bit_decimal(bit):
    bin_li = list(bit)
    bin_li.reverse()
    decimal_number = 0
    for x in range(0, len(bin_li)):
        decimal_number += int(bin_li[x]) << x
    return decimal_number

bit=input("Enter 8digit:")
print(convert_bit_decimal(bit))

"""
2.
Show how to print out a list in reverse order -
example list( [1,2,3,4,5,6,7,8,9,10] ).
Do not use the list.reverse() function.
"""

li=[1,2,3,4,5,6,7,8,9,10]
print(li[::-1])
for i in li[::-1]:
   print(i)

"""
3.
Write a function to return a MegaMillions lottery ticket:
5 numbers range from 1 to 70 and 1 number ranges from 1 to 25.

"""

def MegaMillion():

      import random
      MegaMillions=[]
      while len(MegaMillions)<=5:
            a_num=random.randint(1,70)
            if a_num not in MegaMillions:
                  MegaMillions.append(a_num)

      while len(MegaMillions)<=6:
            b_num=random.randint(1,25)
            if b_num not in MegaMillions:
                  MegaMillions.append(b_num)

      print(MegaMillions)

MegaMillion()


import requests
def getHTMLText(url,kv):
    try:
        hd = {'user_agent':'Mozilla/5.0'}
        r = requests.get(url,headers = hd,params = kv)
        r.raise_for_status() #如果不是200，引发HTTPError异常
        r.encoding = r.apparent_encoding
        return r.text
    except:
        return "异常"

if __name__ == "__main__" :
    keyword = input('请输入搜索关键词')
    a = input('使用百度请按1，使用360搜索请按2')
    if a == '1':
        kv = {'wd=':keyword}
        url = "https://www.baidu.com/s"
    if a ==  '2':
        kv ={'q':keyword}
        url = "https://www.so.com/s"

    print(getHTMLText(url,kv)[:])

第1单元/2实例3：百度360搜索关键词提交.py

2实例3：百度360搜索关键词提交.py


from django.urls import path
from .views import (
    SiteSettingsView,
    PageListView,
    PageDetailView,
    NewsListView,
    NewsByCategoryListView,
    NewsByDateListView,
    NewsDetailView,
    ContactDetailsView,
    ContactDetailsMessageView,
)

urlpatterns = [
    path("site-settings/", SiteSettingsView.as_view(), name="site-settings"),
    path("pages/", PageListView.as_view(), name="pages"),
    path("pages/<slug:slug>/", PageDetailView.as_view(), name="page-detail"),
    path("news/", NewsListView.as_view(), name="news"),
    path("news/category/", NewsByCategoryListView.as_view(), name="news-by-category"),
    path("news/date/", NewsByDateListView.as_view(), name="news-by-date"),
    path("news/<slug:slug>/", NewsDetailView.as_view(), name="news-detail"),
    path("contact-details/", ContactDetailsView.as_view(), name="contact-details"),
    path("contact-details/message/", ContactDetailsMessageView.as_view(), name="contact-details-message"),
]


import numpy as np

class EarlyStopper:
    def __init__(self, tolerance=7, min_delta=1e-4):
        self.tolerance = tolerance
        self.min_delta = float(min_delta)
        self.counter = 0
        self.best_score = None
        self.early_stop = False

    def __call__(self, val_loss: float):
        if self.best_score is None:
            self.best_score = val_loss
        elif np.abs(val_loss - self.best_score) < self.min_delta:
            self.counter += 1
            print(f'Early stopping counter: {self.counter} out of {self.tolerance}')
            if self.counter >= self.tolerance:
                print("Early stopping")
                self.early_stop = True

        else:
            self.best_score = val_loss
            self.counter = 0

        return self.early_stop


class UnionFind:
    def __init__(self,n):
        self.parent = [i for i in range(n)]
        self.rank = [0]*n

    def find(self,a):
        if self.parent[a]==a:
            return a
        self.find(self.parent[a])
        return self.parent[a]

    def union(self,a, b):
        parA = self.find(a)
        parB = self.find(b)
        if self.rank[parA]==self.rank[parB]:
            self.parent[parB] = parA
            self.rank[parA]+=1
        elif self.rank[parA] < self.rank[parB]:
            self.parent[parA] = parB
        else:
            self.parent[parB] = parA

class Edge:
    def __init__(self, s, d, w):
        #0--5---1
        self.src = s
        self.dest = d
        self.wt = w

def createGraph():
    graph = []
    graph.append(Edge(0, 1, 10))
    graph.append(Edge(0, 2, 15))
    graph.append(Edge(0, 3, 30))
    graph.append(Edge(1, 3, 40))
    graph.append(Edge(2, 3, 50))

    return graph

def kruskals(g):
    mstCost = 0
    mst = []
    count = 0
    g.sort(key=lambda x: x.wt)

    uf = UnionFind(len(g))
    while(count<len(g)-1):

        r1 = uf.find(g[count].src)
        r2 = uf.find(g[count].dest)

        if r1!=r2:
            uf.union(g[count].src, g[count].dest)
            mstCost += g[count].wt
            mst.append((g[count].src, g[count].dest, g[count].wt))
        count=count+1
    printMST(mst)
    print(mstCost)

def printMST(mst):
    for e in mst:
        print(f"Edge: {e[0]}->{e[1]} Weight:{e[2]}")

n=5
g = createGraph()
kruskals(g)
#TC - O(ElogE + ElogV)
#SC - O(|E| + |V|)


import gi
gi.require_version('Gtk', '3.0')
from gi.repository import Gtk
from matplotlib.figure import Figure
from matplotlib.backends.backend_gtk3cairo import FigureCanvasGTK3Cairo as FigureGTKCanvas
from matplotlib.backends.backend_gtk3 import NavigationToolbar2GTK3 as NavigationToolbar
from matplotlib.patches import Rectangle
import matplotlib.image as image
import numpy as np

class MatPlotLibGTKWindow(Gtk.Window):
    """An abstract base class for a matplotlib drawing widget.

    Extend Gtk.Window, and have a matplotlib.Figure attribute as an handle
    to the matplotlib API. This window have a Gtk.DrawingArea on wich
    a gtk3cairo Matplotlib backend make the plot.

    """

    def __init__(self, title="", application=None):
        Gtk.Window.__init__(self, application=application)

        # matplotlib figure and canvas
        self.figure = Figure(dpi=36)
        self.gtk_drawing_area = FigureGTKCanvas(self.figure)

        # container and layout
        box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL)
        self.add(box)

        self.sw = Gtk.ScrolledWindow()
        self.sw.add_with_viewport(self.gtk_drawing_area)
        box.pack_start(self.sw, True, True, 0)

        toolbar = NavigationToolbar(self.gtk_drawing_area, self)
        box.pack_start(toolbar, False, True, 0)

        # window config
        self.set_title(title)
        self.set_default_size(600, 600)
        self.set_position(Gtk.WindowPosition.CENTER)

    def prepareFigure(self):
        raise NotImplementedError

    def draw(self):
        pass
        #self.figure.show()

class CellViewer(MatPlotLibGTKWindow):
    """A Gtk Window to view an halftone cell."""

    def __init__(self, cell, application=None):
        if cell is None :
            raise AttributeError("cell should not be None")
        MatPlotLibGTKWindow.__init__(self, title="Cell view", application=application)
        self.cell = cell

        self.prepareFigure()
        self.plotWhiteningOrder()
        #self.draw()
        self.show_all()

    def prepareFigure(self):
        """Draw axes, ticks, and pixels shape of the cell"""
        # axes
        f_axes = self.figure.gca(xlim=(-1.5,1.5), ylim=(-1.5,1.5))
        f_axes.spines['bottom'].set_position(('data',-1.5))
        f_axes.spines['left'].set_position(('data',-1.5))
        f_axes.spines['left'].set_color('black')

        # x-ticks
        cn = self.cell.coordinates
        xt = [p.x for p in cn[0:self.cell.width]]
        #f_axes.set_xticks(xt)

        # y-ticks
        yt = []
        for i in range(0, len(cn), self.cell.width):
            yt.append(cn[i].y)
        #f_axes.set_yticks(yt)

        f_axes.grid(True)

        # draw pixel limit
#        ix = (xt[1] - xt[0]) / 2
#        iy = (yt[1] - yt[0]) / 2
#        xt.append(xt[-1:][0] + (2 * ix))
#        f_axes.vlines([x - ix for x in xt], -1 - iy, 1 + iy)
#        yt.append(yt[-1:][0] + (2 * iy))
#        f_axes.hlines([y - iy for y in yt], -1 - ix, 1 + ix)

    def plotWhiteningOrder(self):
        arr = np.zeros((self.cell.width, self.cell.height), dtype=np.uint8)
        for p in self.cell.whiteningOrder:
            arr[p.x][p.y] = 1 - (p.w / 255)
        self.figure.figimage(arr)

#        f_axes = self.figure.gca()
#        cwo = self.cell.whiteningOrder
#        cn = self.cell.coordinates
#
#        dcw = cn[1].x - cn[0].x
#        dch = cn[self.cell.width].y - cn[0].y
#        ix = dcw / 2
#        iy = dch / 2
#
#        for d in cwo:
#            np = cn[d.x + self.cell.width * d.y]
#            col = str(1 - (d.w / 255))
#            r = Rectangle((np.x - ix, np.y - iy), dcw, dch, color=col)
#            f_axes.add_patch(r)

    def plotCellData(self):
        pass


import json
import unittest
from graph import Graph

class GraphTests(unittest.TestCase):

	def setUp(self):
		# create graph from JSON file
		self.g = Graph('model/data/data.json')
		# read JSON file into Python dictionary
		with open('model/data/data.json') as f:
			self.json_data = json.load(f)

	# test that Graph constructor correctly built graph from JSON file
	def test_json_to_graph(self):
		# each movie and actor in JSON data should have a node in graph
		self.assertEqual(len(self.json_data[1]), len(self.g.movie_vertices))
		self.assertEqual(len(self.json_data[0]), len(self.g.actor_vertices))
		# test that edge weight are properly set to (age * gross) and edges all built
		bruce_willis = None
		blind_date = None
		for actor in self.g.actor_vertices:
			if actor.name == 'Bruce Willis':
				bruce_willis = actor
		for movie in self.g.movie_vertices:
			if 'Blind Date' in movie.name:
				blind_date = movie
		# christian bale has 10 movies in test set
		self.assertEqual(len(bruce_willis.neighbors), 91)
		# both edge directions should have correct weight
		self.assertIn(bruce_willis.age * blind_date.gross, bruce_willis.neighbors.values())
		self.assertIn(bruce_willis.age * blind_date.gross, blind_date.neighbors.values())

	# test graph is properly converted back to JSON
	def test_graph_to_json(self):
		# convert graph to json
		d = self.g.to_json()
		# check that dictionaries are the same (order doesn't matter)
		self.assertItemsEqual(self.json_data, d)

	# test small number oldest X
	def test_oldest_X_small(self):
		oldest = self.g.get_oldest_X_actors(2)
		# oldest person(s) in test set is 94
		self.assertEqual(oldest[1][1], 94)

	# test large number oldest X
	def test_oldest_X_large(self):
		# choose number larger than test set
		oldest = self.g.get_oldest_X_actors(100000)
		# oldest person(s) in test set is 94
		self.assertEqual(oldest[-1][1], 94)
		# younger person in test set is -1?
		self.assertEqual(oldest[0][1], -1)

	# test a year with small number of movies
	def test_movies_from_year_small(self):
		movies = self.g.get_movies_from_year(2017)
		# only one movie from 2017 in test set
		self.assertEqual(len(movies), 1)
		self.assertEqual(movies[0], 'The Bye Bye Man')

	# test a year that hasn't happened
	def test_movies_from_year_invalid(self):
		# no movies from 2018 exist yet
		movies = self.g.get_movies_from_year(2018)
		self.assertEqual(len(movies), 0)

	# test a year with large number of movies
	def test_movies_from_year_large(self):
		movies = self.g.get_movies_from_year(2015)
		# 28 movies from 2015 in test set
		self.assertEqual(len(movies), 3)
		# one of these movies should be Vice
		self.assertIn('Vice', movies)

	# test a year with small number of movies
	def test_actors_from_year_small(self):
		actors = self.g.get_actors_from_year(2017)
		# 1 actor in test set from John Wick: Chapter 2 (only 2017 movie we have)
		self.assertEqual(len(actors), 1)
		# Keanu Reeves should be the actor
		self.assertIn('Faye Dunaway', actors.pop())

	# look for actors in invalid year
	def test_actors_from_year_invalid(self):
		# no movies from 2018 exist yet
		actors = self.g.get_actors_from_year(2018)
		self.assertEqual(len(actors), 0)

	# look for actors in year with a lot of movies
	def test_actors_from_year_large(self):
		actor_names = self.g.get_actors_from_year(2015)
		# 10 actors in test set from 2015 movies
		self.assertEqual(len(actor_names), 10)
		# two of these actors should be Danny McBride and Bruce Willis
		self.assertIn('Bruce Willis', actor_names)
		self.assertIn('Danny McBride', actor_names)

	# check top grossing with small arg
	def test_top_grossing_actors_small(self):
		top_actors = self.g.get_top_X_grossing_actors(1)
		# top grossing actor in test set is Morgan Freeman with ~3.2 billion gross
		self.assertEqual('Bruce Willis', top_actors[0][0])
		self.assertEqual(562709189, top_actors[0][1])

	# check top grossing with invalid arg
	def test_top_grossing_actors_invalid_param(self):
		# invalid param should not cause an error
		top_actors = self.g.get_top_X_grossing_actors(-1)
		self.assertEqual(len(top_actors), 0)

	# test top grossing with large arg
	def test_top_grossing_actors_large(self):
		# choose number larger than set
		top_actors = self.g.get_top_X_grossing_actors(100000)
		# top grossing actor in test set is Bruce Willis
		self.assertEqual('Bruce Willis', top_actors[308][0])
		self.assertEqual(562709189, top_actors[308][1])
		# second top grossing actor is Faye Dunaway
		self.assertEqual('Faye Dunaway', top_actors[307][0])
		self.assertEqual(515893034, top_actors[307][1])
		# several bottom grossing actors, all with 0
		self.assertEqual(0, top_actors[0][1])

	# test DFS traversal
	def test_dfs_traversal(self):
		self.g.dfs_traversal()
		# all nodes should be visited now
		for actor_node in self.g.actor_vertices:
			self.assertTrue(actor_node.visited)
		for movie_node in self.g.movie_vertices:
			self.assertTrue(movie_node.visited)

	# Test BFS traversal w/ a depth 1
	def test_bfs_d1(self):
		for actor in self.g.actor_vertices:
			if actor.name == 'Bruce Willis':
				bruce_willis = actor
			if actor.name == 'Kim Basinger':
				kim_bas = actor
		i = self.g.bfs(bruce_willis, kim_bas.name)
		# shortest path should have length 2
		self.assertEqual(i, 1)

	# Test BFS traversal w/ a larger depth
	def test_bfs_d2(self):
		for actor in self.g.actor_vertices:
			if actor.name == 'Kirstie Alley':
				bruce_willis = actor
			if actor.name == 'John Pankow':
				other = actor
		i = self.g.bfs(bruce_willis, other.name)
		# shortest path should have length 2
		self.assertEqual(i, 2)
		# lengths should be reciprocal
		j = self.g.bfs(other, bruce_willis.name)
		self.assertEqual(j,2)

	# Test nodes not reachable
	def test_bfs_unreachabel(self):
		for actor in self.g.actor_vertices:
			if actor.name == 'Bruce Willis':
				bruce_willis = actor
		i = self.g.bfs(bruce_willis, 'Not an Actor')
		# -1 should be returned
		self.assertEqual(i, -1)

	# Test BFS traversal w/ a larger depth
	def test_bfs_d3(self):
		for actor in self.g.actor_vertices:
			if actor.name == 'Danny Aiello':
				first = actor
			if actor.name == 'Kirstie Alley':
				other = actor
		i = self.g.bfs(first, other.name)
		# shortest path should have length 2
		self.assertEqual(i, 2)
		# lengths should be reciprocal
		j = self.g.bfs(other, first.name)
		self.assertEqual(j, 2)

	# Test top 5 hub actors
	def test_hub_actors(self):
		# select top 5 hub actors
		hub_actors = self.g.get_hub_actors(5)
		# top actor should be Bruce willis
		self.assertEqual(hub_actors[-1][0], 'Bruce Willis')
		self.assertEqual(hub_actors[-1][1], 305)
		# second to top is Jack Warden
		self.assertEqual(hub_actors[-2][0], 'Jack Warden')
		self.assertEqual(hub_actors[-2][1], 39)

	# Test getting hub actors with larger parameter
	def test_hub_actors_large(self):
		# select top 5 hub actors
		hub_actors = self.g.get_hub_actors(500)
		# top actor should still be Bruce willis
		self.assertEqual(hub_actors[-1][0], 'Bruce Willis')
		self.assertEqual(hub_actors[-1][1], 305)
		# Multiple bottom individuals with no connections
		self.assertEqual(hub_actors[0][1], 0)

	# Test hub actors with invalid param
	def test_hub_actors_invalid(self):
		hub_actors = self.g.get_hub_actors(-15)
		self.assertEqual(hub_actors, [])

	# Test gross for age group
	def test_get_gross_for_age_group(self):
		gross = self.g.get_gross_for_age_group(60, 69)
		self.assertEqual(gross, 1005773688)
		gross2 = self.g.get_gross_for_age_group(20, 29)
		self.assertEqual(gross2, 21767523)

	# Test gross for invalid age group
	def test_get_gross_for_age_group_invalid(self):
		# start age larger than end
		gross = self.g.get_gross_for_age_group(69, 60)
		self.assertEqual(gross, 0)
		# no actors in age group
		gross = self.g.get_gross_for_age_group(100, 150)
		self.assertEqual(gross, 0)

if __name__ == '__main__':
	unittest.main()


import random
import logging
import time
import argparse
import ConfigParser
from datetime import datetime

parser = argparse.ArgumentParser(description='Random Log Generator.')
parser.add_argument('-i', '--iterations', type=int, help='Number of iterations', required=True)
parser.add_argument('-f', '--framework', type=str, default='lumen', help='Name of framework')
parser.add_argument('-s', '--sleep', type=int, default=0, help='Delay for the log to be printed')
args = parser.parse_args()

config = ConfigParser.ConfigParser()
config.readfp(open('loggenerator.conf'))
msg = config.get(args.framework,"msg").split(',')

logging.basicConfig(
		filename='%s.log'%(str(datetime.now().date())), \
		level=logging.DEBUG,\
		format='%(message)s')

for i in range (args.iterations):
	logging.info('[%s] '%(str(datetime.now().strftime('%Y-%m-%d %H:%M:%S')))+msg[random.randint(0, len(msg)-1)])
	time.sleep(args.sleep)


import torch.nn as nn
import torch.nn.functional as F
import torch

class DEQCpResidualBlockAdjacentBlock(nn.Module):
    def __init__(self, dilations, in_channels, out_channels, kernel_size, block_id):
        super(DEQCpResidualBlockAdjacentBlock, self).__init__()
        self.kernel_size = kernel_size
        self.dilation = dilations
        self.block_id = block_id

        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=(1, kernel_size),
                                   padding=0, dilation=dilations[0], bias=True)
        self.rel1 = nn.ReLU()
        self.conv2 = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=(1, kernel_size),
                                   padding=0, dilation=2*dilations[0], bias=True)
        self.rel2 = nn.ReLU()
        self.init_weights()

    def init_weights(self):
        bound = 0.01
        self.conv1.weight.data.normal_(0, bound)
        self.conv1.bias.data.normal_(0, bound)
        self.conv2.weight.data.normal_(0, bound)
        self.conv2.bias.data.normal_(0, bound)

    def copy(self, func):
        self.conv1.weight.data = func.conv1.weight.data.clone()
        self.conv1.bias.data = func.conv1.bias.data.clone()
        self.conv2.weight.data = func.conv2.weight.data.clone()
        self.conv2.bias.data = func.conv2.bias.data.clone()

    def forward_conv(self, conv, x, dilation):
        padding = [(self.kernel_size - 1) * dilation, 0, 0, 0, 0, 0]
        padded = torch.nn.functional.pad(x, padding)
        input_expanded = padded.unsqueeze(dim=2)
        conv.dilation = dilation
        y = conv(input_expanded)
        y = y.squeeze(dim=2)
        return y

    # dilations for deq is a tupple eg(1,4) so we have 4 convolutions or 2 residual blocks -> 1,2,4,8
    def forward(self, z1ss, uss, *kwargs):
        # pytorch expect BCHW input that is why we permute and add height in dim 2

        y1 = self.forward_conv(self.conv1, z1ss, self.dilation[0])
        y1 = self.layer_norm(y1)
        y1 = self.rel1(y1)

        y1 = self.forward_conv(self.conv2, y1, 2 * self.dilation[0])
        y1 = self.layer_norm(y1)
        y1 = self.rel2(y1)

        y1 = uss + y1

        y2 = y1
        y2 = self.forward_conv(self.conv1, y2, self.dilation[1])
        y2 = self.layer_norm(y2)
        y2 = self.rel1(y2)

        y2 = self.forward_conv(self.conv2, y2, 2 * self.dilation[1])
        y2 = self.layer_norm(y2)
        y2 = self.rel2(y2)

        return uss + y2

    def layer_norm(self, x, epsilon=1e-8):
        shape = x.shape
        x = x.permute(0, 2, 1)
        mean = x.mean(dim=len(shape) - 1, keepdim=True)
        variance = x.var(dim=len(shape) - 1, keepdim=True)

        x = (x - mean) / torch.sqrt(variance + epsilon)
        y = x.permute(0, 2, 1)
        return y


import os
from pyglet.gl import *

from .player import Player
from .platform import Platform

class mainMenu():
	def __init__(self, main_window):
		self.main_window = main_window

		self.main_batch = pyglet.graphics.Batch()
		depressed = pyglet.image.load(f'{os.getcwd()}/resources/play_game_button.png')
		pressed = pyglet.image.load(f'{os.getcwd()}/resources/play_game_button.png')

		self.frame = pyglet.gui.Frame(self.main_window, order=4)
		self.play_button = pyglet.gui.ToggleButton(100, 400, pressed=pressed, depressed=depressed, batch=self.main_batch)
		self.play_button.set_handler('on_toggle', self.play_game)
		self.frame.add_widget(self.play_button)

	def play_game(self, value):
		self.main_window.view = GameView(self.main_window)

	def render(self):
		glClearColor(25/255, 94/255, 131/255, 0/255)
		self.main_batch.draw()

class GameView():
	def __init__(self, main_window):
		self.main_window = main_window

		self.platform = Platform(self)
		self.player = Player(self)

	def render(self):
		glClearColor(94/255, 25/255, 131/255, 0/255)
		self.platform.render()
		self.player.render()


class Queue:
    def __init__(self):
        self.items=[]

    def push(self,val):
        self.items.append(val)

    def pop(self):
        front=0
        if(self.empty()==0):
            front=self.items[0]
            self.items=self.items[1:]
            return front
        else:
            return -1

    def empty(self):
        if(len(self.items)==0):
            return 1
        else:
            return 0

import sys

q=Queue()
n=int(input())

for i in range(n):
    command=sys.stdin.readline().split()

    if(command[0]=='push'):
        q.push(command[1])

    elif(command[0]=='pop'):
        print(q.pop())

    elif(command[0]=='size'):
        print(len(q.items))

    elif(command[0]=='empty'):
        print(q.empty())

    elif(command[0]=='front'):
        if(q.empty()==1):
            print(-1)
        else:
            print(q.items[0])

    elif(command[0]=='back'):
        if(q.empty()==1):
            print(-1)
        else:
            print(q.items[-1])


import random
import re
from sys import get_coroutine_origin_tracking_depth
from sys import exit
random.seed(101)
import matplotlib.pyplot as plt
import math
import matplotlib.patches as mpatches
#from scipy.linalg import svd
import itertools
import torch
import time
import numpy as np
from tqdm import tqdm
from evaluator import ProxyEvaluator
import collections
import os
from data import Data
from parse import parse_args
# need change 1
from model import CausE, IPS, LGN, MACR, INFONCE, INFONCE_linear, AdvInfoNCE, SAMREG, BC_LOSS, SimpleX, SGL, NCL, XSimGCL, PID, NBC_LOSS
from torch.utils.data import Dataset, DataLoader
from collect_log import read_log
import torch.nn.functional as F

def merge_user_list(user_lists):
    out = collections.defaultdict(list)
    # Loop over each user list
    for user_list in user_lists:
        # Loop over each user in the user list
        for key, item in user_list.items():
            out[key] = out[key] + item
    return out

def merge_user_list_no_dup(user_lists):
    out = collections.defaultdict(list)
    for user_list in user_lists:
        for key, item in user_list.items():
            out[key] = out[key] + item

    for key in out.keys():
        out[key]=list(set(out[key]))
    return out

def save_checkpoint(model, epoch, checkpoint_dir, buffer, max_to_keep=10):
    state = {
        'epoch': epoch,
        'state_dict': model.state_dict(),
    }

    filename = os.path.join(checkpoint_dir, 'epoch={}.checkpoint.pth.tar'.format(epoch))
    torch.save(state, filename)
    buffer.append(filename)
    if len(buffer)>max_to_keep:
        os.remove(buffer[0])
        del(buffer[0])

    return buffer

def restore_checkpoint(model, checkpoint_dir, device, force=False, pretrain=False):
    """
    If a checkpoint exists, restores the PyTorch model from the checkpoint.
    Returns the model and the current epoch.
    """
    cp_files = [file_ for file_ in os.listdir(checkpoint_dir)
                if file_.startswith('epoch=') and file_.endswith('.checkpoint.pth.tar')]

    if not cp_files:
        print('No saved model parameters found')
        if force:
            raise Exception("Checkpoint not found")
        else:
            return model, 0,

    epoch_list = []

    regex = re.compile(r'\d+')

    for cp in cp_files:
        epoch_list.append([int(x) for x in regex.findall(cp)][0])

    epoch = max(epoch_list)

    if not force:
        print("Which epoch to load from? Choose in range [0, {})."
              .format(epoch), "Enter 0 to train from scratch.")
        print(">> ", end = '')
        # inp_epoch = int(input())
        inp_epoch = epoch
        if inp_epoch not in range(epoch + 1):
            raise Exception("Invalid epoch number")
        if inp_epoch == 0:
            print("Checkpoint not loaded")
            clear_checkpoint(checkpoint_dir)
            return model, 0,
    else:
        print("Which epoch to load from? Choose in range [0, {}).".format(epoch))
        inp_epoch = int(input())
        if inp_epoch not in range(0, epoch):
            raise Exception("Invalid epoch number")

    filename = os.path.join(checkpoint_dir,
                            'epoch={}.checkpoint.pth.tar'.format(inp_epoch))

    print("Loading from checkpoint {}?".format(filename))

    checkpoint = torch.load(filename, map_location = str(device))

    try:
        if pretrain:
            model.load_state_dict(checkpoint['state_dict'], strict=False)
        else:
            model.load_state_dict(checkpoint['state_dict'])
        print("=> Successfully restored checkpoint (trained for {} epochs)"
              .format(checkpoint['epoch']))
    except:
        print("=> Checkpoint not successfully restored")
        raise

    return model, inp_epoch

def restore_best_checkpoint(epoch, model, checkpoint_dir, device):
    """
    Restore the best performance checkpoint
    """
    cp_files = [file_ for file_ in os.listdir(checkpoint_dir)
                if file_.startswith('epoch=') and file_.endswith('.checkpoint.pth.tar')]

    filename = os.path.join(checkpoint_dir,
                            'epoch={}.checkpoint.pth.tar'.format(epoch))

    print("Loading from checkpoint {}?".format(filename))

    checkpoint = torch.load(filename, map_location = str(device))

    model.load_state_dict(checkpoint['state_dict'])
    print("=> Successfully restored checkpoint (trained for {} epochs)"
          .format(checkpoint['epoch']))

    return model

def clear_checkpoint(checkpoint_dir):
    filelist = [f for f in os.listdir(checkpoint_dir) if f.endswith(".pth.tar")]
    for f in filelist:
        os.remove(os.path.join(checkpoint_dir, f))

    print("Checkpoint successfully removed")

def evaluation(args, data, model, epoch, base_path, evaluator, name="valid"):
    # Evaluate with given evaluator

    ret, _ = evaluator.evaluate(model)

    n_ret = {"recall": ret[1], "hit_ratio": ret[5], "precision": ret[0], "ndcg": ret[3], "mrr":ret[4], "map":ret[2]}

    perf_str = name+':{}'.format(n_ret)
    print(perf_str)
    with open(base_path + 'stats_{}.txt'.format(args.saveID), 'a') as f:
        f.write(perf_str + "\n")
    # Check if need to early stop (on validation)
    is_best=False
    early_stop=False
    if name=="valid":
        if ret[1] > data.best_valid_recall:
            data.best_valid_epoch = epoch
            data.best_valid_recall = ret[1]
            data.patience = 0
            is_best=True
        else:
            data.patience += 1
            if data.patience >= args.patience:
                print_str = "The best performance epoch is % d " % data.best_valid_epoch
                print(print_str)
                early_stop=True

    return is_best, early_stop

def Item_pop(args, data, model):

    for K in range(5):

        eval_pop = ProxyEvaluator(data, data.train_user_list, data.pop_dict_list[K], top_k=[(K+1)*10],
                                   dump_dict=merge_user_list([data.train_user_list, data.valid_user_list]))

        ret, _ = eval_pop.evaluate(model)

        print_str = "Overlap for K = % d is % f" % ( (K+1)*10, ret[1] )

        print(print_str)

        with open('stats_{}.txt'.format(args.saveID), 'a') as f:
            f.write(print_str + "\n")

def ensureDir(dir_path):

    if not os.path.exists(dir_path):
        os.makedirs(dir_path)

def split_grp_view(data,grp_idx):
    n=len(grp_view)
    split_data=[{} for _ in range(n)]

    for key,item in data.items():
        for it in item:
            if key not in split_data[grp_idx[it]].keys():
                split_data[grp_idx[it]][key]=[]
            split_data[grp_idx[it]][key].append(it)
    return split_data

def checktensor(tensor):
    t=tensor.detach().cpu().numpy()
    if np.max(np.isnan(t)):
        idx=np.argmax(np.isnan(t))
        return idx
    else:
        return -1

def get_rotation_matrix(axis, theta):
    """
    Find the rotation matrix associated with counterclockwise rotation
    about the given axis by theta radians.
    Credit: http://stackoverflow.com/users/190597/unutbu

    Args:
        axis (list): rotation axis of the form [x, y, z]
        theta (float): rotational angle in radians

    Returns:
        array. Rotation matrix.
    """

    axis = np.asarray(axis)
    theta = np.asarray(theta)
    axis = axis/math.sqrt(np.dot(axis, axis))
    a = math.cos(theta/2.0)
    b, c, d = -axis*math.sin(theta/2.0)
    aa, bb, cc, dd = a*a, b*b, c*c, d*d
    bc, ad, ac, ab, bd, cd = b*c, a*d, a*c, a*b, b*d, c*d
    return np.array([[aa+bb-cc-dd, 2*(bc+ad), 2*(bd-ac)],
                     [2*(bc-ad), aa+cc-bb-dd, 2*(cd+ab)],
                     [2*(bd+ac), 2*(cd-ab), aa+dd-bb-cc]])

grads = {}
def save_grad(name):
    def hook(grad):
        torch.clamp(grad, -1, 1)
        grads[name] = grad
    return hook

"""
It takes in the embeddings of users and items, the popularity of users and items, and the data, and
plots the embeddings of users and items in 3D space, and plots the angular distribution of the
embeddings of items.

The embeddings of users and items are plotted in 3D space, and the angular distribution of the
embeddings of items are plotted in 2D space.

:param items: the item embeddings
:param users: the user embeddings
:param data: the data object
:param p_item: the popularity of each item
:param p_user: the popularity of each user
:param name: the name of the dataset
"""
def visulization(items,users,data,p_item,p_user,name):
    test_ood_user_list=data.test_ood_user_list
    test_id_user_list=data.test_id_user_list
    train_user_list=data.train_user_list

    def split_grp_view(data,grp_idx):
        n=len(grp_view)
        split_data=[collections.defaultdict(list) for _ in range(n)]

        for key,item in data.items():
            for it in item:
                if key not in split_data[grp_idx[it]].keys():
                    split_data[grp_idx[it]][key]=[]
                split_data[grp_idx[it]][key].append(it)
        return split_data

    pop_sorted=np.sort(p_item)
    n_items=p_item.shape[0]

    n_groups=3
    grp_view=[]
    for grp in range(n_groups):
        split=int((n_items-1)*(grp+1)/n_groups)
        grp_view.append(pop_sorted[split])
    #print("group_view:",grp_view)
    idx=np.searchsorted(grp_view,p_item)

    pop_group=[[] for _ in range(n_groups)]
    for i in range(n_items):
        pop_group[idx[i]].append(i)

    eval_test_ood_split=split_grp_view(test_ood_user_list,idx)
    eval_test_id_split=split_grp_view(test_id_user_list,idx)
    eval_train_split=split_grp_view(train_user_list,idx)

    pop_users=p_user.tolist()

    u_pop_sorted=np.sort(p_user)
    print(u_pop_sorted[-10:])

    fig = plt.figure(constrained_layout=True,figsize=(12,6))

    def plot_embed(ax1,ax2,idx):
        u, v = np.mgrid[0:2*np.pi:20j, 0:2*np.pi:20j]
        x1 = np.cos(u)*np.sin(v)
        y1 = np.sin(u)*np.sin(v)
        z1 = np.cos(v)
        ax1.plot_wireframe(x1, y1, z1, color="0.5",linewidth=0.1)
        user_idx=pop_users.index(idx)
        m_user=users[user_idx]
        target=np.array([1,-1,1])
        r_theta=np.arccos(np.dot(m_user,target)/(np.linalg.norm(m_user)*np.linalg.norm(target)))
        axis=np.cross(m_user,target)
        R=get_rotation_matrix(axis,r_theta)
        grp_theta=[]
        grp_r=[]
        sizes=[10,10,10]

        cmap_b = 'b'
        cmap_r = 'r'
        cmaps=[cmap_b,cmap_r]

        norm = plt.Normalize(vmin=-3, vmax=3)

        all_sampled=set([])
        all_pos=set([])
        for i,grp in enumerate(pop_group):
            sampled_group=set(np.random.choice(np.array(grp),50,replace=False).tolist())
            if user_idx in eval_test_id_split[i].keys():
                for item in eval_test_id_split[i][user_idx]:
                    sampled_group.add(item)
                    all_pos.add(item)
            for item in eval_train_split[i][user_idx]:
                sampled_group.add(item)
                all_pos.add(item)
            if user_idx in eval_test_ood_split[i].keys():
                for item in eval_test_ood_split[i][user_idx]:
                    sampled_group.add(item)
                    all_pos.add(item)

            all_sampled=all_sampled.union(sampled_group)

        all_neg=all_sampled.difference(all_pos)
        #print(all_neg)
        all_pos=np.array(list(all_pos),dtype=int)
        all_neg=np.array(list(all_neg),dtype=int)
        nor = plt.Normalize(vmin=-3, vmax=3)
        r=np.linalg.norm(target)

        lab=["neg","pos"]
        for i,idx in enumerate([all_neg,all_pos]):
            g_item=items[idx]
            g_item=np.matmul(g_item,R.T)
            norm=np.linalg.norm(g_item,axis=1)
            x=g_item[:,0]/norm#*r
            y=g_item[:,1]/norm#*r
            z=g_item[:,2]/norm#*r

            for j in range(len(idx)):
                ax1.plot([0,g_item[j][0]/norm[j]],[0,g_item[j][1]/norm[j]],[0,g_item[j][2]/norm[j]],color = cmaps[i],alpha=0.1)

            ax1.scatter(x, y, z, c = cmaps[i], marker =".",s=10,label=lab[i])

        #print("V^{T}",V_transpose)

        ax1.scatter(target[0]/r, target[1]/r, target[2]/r, c = 'g', marker ="*",s=120,label="user")
        ax1.plot([0,target[0]/r],[0,target[1]/r],[0,target[2]/r],color = 'g',alpha=0.1)
        ax1.legend()

        all_items=set([i for i in range(n_items)])
        all_neg=all_items.difference(all_pos)
        all_neg=np.array(list(all_neg),dtype=int)

        grp=["(neg):","(pos):"]
        txt=""

        for i,idx in enumerate([all_neg,all_pos]):
            g_item=items[idx]
            g_item=np.matmul(g_item,R.T)
            norm=np.linalg.norm(g_item,axis=1)
            cos=np.arccos(np.matmul(target,g_item.T)/norm/r)
            me=float(np.mean(cos))
            me=round(me,3)
            if i==1:
                txt="mean angle"+grp[i]+str(me)+"\n"+txt
            else:
                txt="mean angle"+grp[i]+str(me)+txt

            ax2.hist(cos,50,range=[0,np.pi],color=cmaps[i],weights=np.zeros_like(cos) + 1. / cos.size,edgecolor='black',alpha=0.6)
            mi_x,ma_x=ax2.get_xlim()
            mi_y,ma_y=ax2.get_ylim()
        ax2.text(mi_x+(ma_x-mi_x)*0.45, mi_y+(ma_y-mi_y)*0.75,txt , style ='italic')
        red_patch = mpatches.Patch(color='red', alpha=0.6, label='pos')
        blue_patch = mpatches.Patch(color='blue', alpha=0.6,label='neg')
        ax2.legend(handles=[red_patch,blue_patch])

    pops=[205,30,10]

    fig = plt.figure(figsize=(6,8),constrained_layout=True)
    gs = fig.add_gridspec(3, 2)
    ax1 = fig.add_subplot(gs[0:2, 0:2],projection='3d')
    ax2 = fig.add_subplot(gs[2,0:2])
    ax1.set_xticks([-1,-0.5,0,0.5,1])
    ax1.set_yticks([-1,-0.5,0,0.5,1])
    ax1.set_zticks([-1,-0.5,0,0.5,1])
    ax1.grid(False)

    plot_embed(ax1,ax2,pops[0])
    #ax1.set_title("High Pop User(p=205)")
    #ax2.set_title("Angular Distribution(High Pop)")

    plt.savefig(name+"high_pop_"+str(pops[0])+".png",bbox_inches='tight')
    plt.close()

    fig = plt.figure(figsize=(6,8),constrained_layout=True)
    gs = fig.add_gridspec(3, 2)
    ax3 = fig.add_subplot(gs[0:2, 0:2],projection='3d')
    ax4 = fig.add_subplot(gs[2,0:2])
    ax3.set_xticks([-1,-0.5,0,0.5,1])
    ax3.set_yticks([-1,-0.5,0,0.5,1])
    ax3.set_zticks([-1,-0.5,0,0.5,1])
    ax3.grid(False)
    plot_embed(ax3,ax4,pops[1])
    #ax3.set_title("Mid Pop User(p=30)")
    #ax4.set_title("Angular Distribution(Mid Pop)")

    plt.savefig(name+"mid_pop_"+str(pops[1])+".png",bbox_inches='tight')
    plt.close()

    fig = plt.figure(figsize=(6,8),constrained_layout=True)
    gs = fig.add_gridspec(3, 2)
    ax5 = fig.add_subplot(gs[0:2, 0:2],projection='3d')
    ax6 = fig.add_subplot(gs[2,0:2])
    ax5.set_xticks([-1,-0.5,0,0.5,1])
    ax5.set_yticks([-1,-0.5,0,0.5,1])
    ax5.set_zticks([-1,-0.5,0,0.5,1])
    ax5.grid(False)
    plot_embed(ax5,ax6,pops[2])
    #ax5.set_title("Low Pop User(p=10)")
    #ax6.set_title("Angular Distribution(Low Pop)")

    plt.savefig(name+"low_pop_"+str(pops[2])+".png",bbox_inches='tight')
    plt.close()

def seed_torch(seed=101):
	random.seed(seed)
	os.environ['PYTHONHASHSEED'] = str(seed) # 为了禁止hash随机化，使得实验可复现
	np.random.seed(seed)
	torch.manual_seed(seed)
	torch.cuda.manual_seed(seed)
	torch.cuda.manual_seed_all(seed) # if you are using multi-GPU.
	torch.backends.cudnn.benchmark = False
	torch.backends.cudnn.deterministic = True

def align_loss(x, y, alpha=2):
    return (x - y).norm(p=2, dim=1).pow(alpha).mean()

def uniform_loss(x, t=2):
    return torch.pdist(x, p=2).pow(2).mul(-t).exp().mean().log()

if __name__ == '__main__':

    start = time.time()

    args = parse_args()
    seed_torch(args.seed)
    print(args)
    data = Data(args)
    data.load_data()
    device="cuda:"+str(args.cuda)
    device = torch.device(args.cuda)
    saveID = args.saveID
    saveID += str(args.dsc)
    #@ 保存的目录
    # need change 2
    if args.modeltype == "INFONCE":
        saveID += "_INFONCE_" + "n_layers=" + str(args.n_layers) + "tau=" + str(args.tau) + "batch_size=" + str(args.batch_size) + "neg_sample=" + str(args.neg_sample) + "lr=" + str(args.lr)
    if args.modeltype == "INFONCE_linear":
        saveID += "_INFONCE_linear_" + "n_layers=" + str(args.n_layers) + "tau=" + str(args.tau) + "batch_size=" + str(args.batch_size) + "neg_sample=" + str(args.neg_sample) + "lr=" + str(args.lr)

    if args.modeltype == "BC_LOSS":
        saveID += "n_layers=" + str(args.n_layers) + "tau1=" + str(args.tau1) + "tau2=" + str(args.tau2) + "w=" + str(args.w_lambda)

    if args.modeltype == "NBC_Loss":
        saveID += "n_layers=" + str(args.n_layers) + "tau1=" + str(args.tau1) + "tau2=" + str(args.tau2) + "w=" + str(args.w_lambda) + "k_score=" + str(args.k_score)

    if args.modeltype == "SimpleX":
        saveID += "_SimpleX_" + "n_layers=" + str(args.n_layers) + "tau=" + str(args.tau) + "batch_size=" + str(args.batch_size) + "neg_sample=" + str(args.neg_sample) + "lr=" + str(args.lr) + "w=" + str(args.w_neg) + "m=" + str(args.neg_margin)

    if args.modeltype == "XSimGCL":
        saveID += "_XSimGCL_" + "n_layers=" + str(args.n_layers) + "batch_size=" + str(args.batch_size) + "lr=" + str(args.lr) + "temp=" + str(args.temp_cl) + "layer_cl=" + str(args.layer_cl) + "lambda_cl=" + str(args.lambda_cl)+ "eps_XSimGCL=" + str(args.eps_XSimGCL)

    if args.modeltype == "SGL":
        saveID += "_SGL_" + "n_layers=" + str(args.n_layers) + "batch_size=" + str(args.batch_size) + "lr=" + str(args.lr) + "temp=" + str(args.temp_cl) + "lambda_cl=" + str(args.lambda_cl)+ "droprate=" + str(args.droprate)

    if args.modeltype == "NCL":
        saveID += "_NCL_" + "n_layers=" + str(args.n_layers) + "batch_size=" + str(args.batch_size) + "lr=" + str(args.lr) + "temp=" + str(args.temp_cl) + "lambda_cl=" + str(args.lambda_cl)+ "proto_reg=" + str(args.proto_reg) + "n_clusters=" + str(args.num_clusters) + "alpha=" + str(args.ncl_alpha)

    if args.n_layers > 0 and args.modeltype != "LGN":
        base_path = './weights/{}/{}-LGN/{}'.format(args.dataset, args.modeltype, saveID)
    else:
        base_path = './weights/{}/{}/{}'.format(args.dataset, args.modeltype, saveID)

    if args.modeltype == 'LGN':
        saveID += "n_layers=" + str(args.n_layers) + "lr=" + str(args.lr) + "batch_size=" + str(args.batch_size)
        base_path = './weights/{}/{}/{}'.format(args.dataset, args.modeltype, saveID)

    checkpoint_buffer=[]
    #@ epochs for training the margin for BC Loss
    freeze_epoch=args.freeze_epoch if (args.modeltype=="BC_LOSS") else 0
    ensureDir(base_path)

    perf_str = str(args)
    with open(base_path + 'stats_{}.txt'.format(args.saveID),'a') as f:
            f.write(perf_str+"\n")

    #@ calculate the popularity (not used in the paper)
    p_item = np.array([len(data.train_item_list[u]) if u in data.train_item_list else 0 for u in range(data.n_items)])
    p_user = np.array([len(data.train_user_list[u]) if u in data.train_user_list else 0 for u in range(data.n_users)])
    m_user=np.argmax(p_user)

    pop_sorted=np.sort(p_item)
    n_groups=3
    grp_view=[]
    for grp in range(n_groups):
        split=int((data.n_items-1)*(grp+1)/n_groups)
        grp_view.append(pop_sorted[split])
    print("group_view:",grp_view)
    idx=np.searchsorted(grp_view,p_item)

    if(args.dataset != "tencent_synthetic" and args.dataset != "kuairec_ood"):
        eval_test_ood_split=split_grp_view(data.test_ood_user_list,idx)
        eval_test_id_split=split_grp_view(data.test_id_user_list,idx)

    grp_view=[0]+grp_view

    pop_dict={}
    for user,items in data.train_user_list.items():
        for item in items:
            if item not in pop_dict:
                pop_dict[item]=0
            pop_dict[item]+=1
    # The key of pop_dict is item, and the value is the popularity of item

    sort_pop=sorted(pop_dict.items(), key=lambda item: item[1],reverse=True)

    # sort_pop is a list, the elements in the list are tuples, the first element of the tuple is item, and the second element is the popularity of the item
    pop_mask=[item[0] for item in sort_pop[:20]]
    print(pop_mask)

    if "kuairec" in args.dataset:
        with open("data/" + args.dataset + '/not_candidate.txt', 'r') as f:
            not_candidate = f.readlines()
            not_candidate = [int(item.strip()) for item in not_candidate]
            not_candidate_dict = {u:not_candidate for u in data.users}

    if not args.pop_test:

        if(args.dataset == "tencent_synthetic"):
            eval_valid = ProxyEvaluator(data,data.train_user_list,data.valid_user_list,top_k=[20])
            eval_test_ood_1 = ProxyEvaluator(data,data.train_user_list,data.test_ood_user_list_1,top_k=[20],\
                                dump_dict=merge_user_list([data.train_user_list,data.valid_user_list,data.test_ood_user_list_2,data.test_ood_user_list_3]))
            eval_test_ood_2 = ProxyEvaluator(data,data.train_user_list,data.test_ood_user_list_2,top_k=[20],\
                                dump_dict=merge_user_list([data.train_user_list,data.valid_user_list,data.test_ood_user_list_1,data.test_ood_user_list_3]))
            eval_test_ood_3 = ProxyEvaluator(data,data.train_user_list,data.test_ood_user_list_3,top_k=[20],\
                                dump_dict=merge_user_list([data.train_user_list,data.valid_user_list,data.test_ood_user_list_1,data.test_ood_user_list_2]))
        elif(args.dataset == "kuairec_ood"):
            eval_valid = ProxyEvaluator(data,data.train_user_list,data.valid_user_list,top_k=[20],dump_dict=merge_user_list([data.train_user_list,not_candidate_dict]))
            eval_test_ood_1 = ProxyEvaluator(data,data.train_user_list,data.test_ood_user_list_1,top_k=[20],\
                                dump_dict=merge_user_list([data.train_user_list,data.valid_user_list,not_candidate_dict]))
            eval_test_ood_2 = ProxyEvaluator(data,data.train_user_list,data.test_ood_user_list_2,top_k=[20],\
                                dump_dict=merge_user_list([data.train_user_list,data.valid_user_list,not_candidate_dict, data.test_ood_user_list_1]))
            eval_test_ood_3 = ProxyEvaluator(data,data.train_user_list,data.test_ood_user_list_3,top_k=[20],\
                                dump_dict=merge_user_list([data.train_user_list,data.valid_user_list,not_candidate_dict, data.test_ood_user_list_1, data.test_ood_user_list_2]))
        else:
            if "kuairec" in args.dataset:
                eval_valid = ProxyEvaluator(data,data.train_user_list,data.valid_user_list,top_k=[20],dump_dict=merge_user_list([data.train_user_list,not_candidate_dict]))
                eval_test_ood = ProxyEvaluator(data,data.train_user_list,data.test_ood_user_list,top_k=[20],dump_dict=merge_user_list([data.train_user_list,data.valid_user_list,data.test_id_user_list,not_candidate_dict]))
                eval_test_id = ProxyEvaluator(data,data.train_user_list,data.test_id_user_list,top_k=[20],dump_dict=merge_user_list([data.train_user_list,data.valid_user_list,data.test_ood_user_list,not_candidate_dict]))
            else:
                eval_valid = ProxyEvaluator(data,data.train_user_list,data.valid_user_list,top_k=[20])
                eval_test_ood = ProxyEvaluator(data,data.train_user_list,data.test_ood_user_list,top_k=[20],dump_dict=merge_user_list([data.train_user_list,data.valid_user_list,data.test_id_user_list]))
                eval_test_id = ProxyEvaluator(data,data.train_user_list,data.test_id_user_list,top_k=[20],dump_dict=merge_user_list([data.train_user_list,data.valid_user_list,data.test_ood_user_list]))

    else:
        if(args.dataset == "tencent_synthetic"):
            eval_valid = ProxyEvaluator(data,data.train_user_list,data.valid_user_list,top_k=[20],pop_mask=pop_mask)
            eval_test_ood_1 = ProxyEvaluator(data,data.train_user_list,data.test_ood_user_list_1,top_k=[20],\
                            dump_dict=merge_user_list([data.train_user_list,data.valid_user_list,data.test_ood_user_list_2,data.test_ood_user_list_3]),pop_mask=pop_mask)
            eval_test_ood_2 = ProxyEvaluator(data,data.train_user_list,data.test_ood_user_list_2,top_k=[20],\
                                dump_dict=merge_user_list([data.train_user_list,data.valid_user_list,data.test_ood_user_list_1,data.test_ood_user_list_3]),pop_mask=pop_mask)
            eval_test_ood_3 = ProxyEvaluator(data,data.train_user_list,data.test_ood_user_list_3,top_k=[20],\
                                dump_dict=merge_user_list([data.train_user_list,data.valid_user_list,data.test_ood_user_list_1,data.test_ood_user_list_2]),pop_mask=pop_mask)
        elif(args.dataset == "kuairec_ood"):
            eval_valid = ProxyEvaluator(data,data.train_user_list,data.valid_user_list,top_k=[20],dump_dict=merge_user_list([data.train_user_list,not_candidate_dict]),pop_mask=pop_mask)
            eval_test_ood_1 = ProxyEvaluator(data,data.train_user_list,data.test_ood_user_list_1,top_k=[20],\
                                dump_dict=merge_user_list([data.train_user_list,data.valid_user_list,not_candidate_dict]),pop_mask=pop_mask)
            eval_test_ood_2 = ProxyEvaluator(data,data.train_user_list,data.test_ood_user_list_2,top_k=[20],\
                                dump_dict=merge_user_list([data.train_user_list,data.valid_user_list,not_candidate_dict,data.test_ood_user_list_1]),pop_mask=pop_mask)
            eval_test_ood_3 = ProxyEvaluator(data,data.train_user_list,data.test_ood_user_list_3,top_k=[20],\
                                dump_dict=merge_user_list([data.train_user_list,data.valid_user_list,not_candidate_dict,data.test_ood_user_list_1,data.test_ood_user_list_2]),pop_mask=pop_mask)
        else:
            if "kuairec" in args.dataset:
                eval_valid = ProxyEvaluator(data,data.train_user_list,data.valid_user_list,top_k=[20],dump_dict=merge_user_list([data.train_user_list,not_candidate_dict]),pop_mask=pop_mask)
                eval_test_ood = ProxyEvaluator(data,data.train_user_list,data.test_ood_user_list,top_k=[20],dump_dict=merge_user_list([data.train_user_list,data.valid_user_list,data.test_id_user_list,not_candidate_dict]),pop_mask=pop_mask)
                eval_test_id = ProxyEvaluator(data,data.train_user_list,data.test_id_user_list,top_k=[20],dump_dict=merge_user_list([data.train_user_list,data.valid_user_list,data.test_ood_user_list,not_candidate_dict]),pop_mask=pop_mask)
            else:
                eval_valid = ProxyEvaluator(data,data.train_user_list,data.valid_user_list,top_k=[20],pop_mask=pop_mask)
                eval_test_ood = ProxyEvaluator(data,data.train_user_list,data.test_ood_user_list,top_k=[20],dump_dict=merge_user_list([data.train_user_list,data.valid_user_list,data.test_id_user_list]),pop_mask=pop_mask)
                eval_test_id = ProxyEvaluator(data,data.train_user_list,data.test_id_user_list,top_k=[20],dump_dict=merge_user_list([data.train_user_list,data.valid_user_list,data.test_ood_user_list]),pop_mask=pop_mask)

    if(args.dataset == "tencent_synthetic" or args.dataset == "kuairec_ood"):
        evaluators=[eval_valid, eval_test_ood_1, eval_test_ood_2, eval_test_ood_3]
        eval_names=["valid","test_ood_1", "test_ood_2", "test_ood_3"]
    else:
        evaluators=[ eval_valid,eval_test_id, eval_test_ood]
        eval_names=["valid","test_id", "test_ood" ]

    # need change 3
    if args.modeltype == 'LGN':
        model = LGN(args, data)
    if args.modeltype == 'INFONCE':
        model = INFONCE(args, data)
    if args.modeltype == 'INFONCE_linear':
        model = INFONCE_linear(args, data)
    if args.modeltype == 'IPS':
        model = IPS(args, data)
    if args.modeltype == 'CausE':
        model = CausE(args, data)
    if args.modeltype == 'BC_LOSS':
        model = BC_LOSS(args, data)
    if args.modeltype == 'MACR':
        model = MACR(args, data)
    if args.modeltype == 'SAMREG':
        model = SAMREG(args, data)
    if args.modeltype == "SimpleX":
        model = SimpleX(args,data)
    if args.modeltype == "SGL":
        model = SGL(args,data)
    if args.modeltype == "NCL":
        model = NCL(args,data)
    if args.modeltype == "XSimGCL":
        model = XSimGCL(args,data)
    if args.modeltype == "PID":
        model = PID(args,data)
    if args.modeltype == "NBC_LOSS":
        model = NBC_LOSS(args, data)

#    b=args.sample_beta

    model.cuda(device)

    #@ load the model
    model, start_epoch = restore_checkpoint(model, base_path, device)

    if args.test_only:

        for i,evaluator in enumerate(evaluators):
            is_best, temp_flag = evaluation(args, data, model, start_epoch, base_path, evaluator,eval_names[i])

        exit()

    flag = False

    # for param in model.named_parameters():
        # print(param)
    if args.modeltype == "PID":
        optimizer = torch.optim.Adam(model.parameters(), lr=model.lr, weight_decay=5e-6)
    else:
        optimizer = torch.optim.Adam([param for param in model.parameters() if param.requires_grad == True], lr=model.lr)

    #item_pop_idx = torch.tensor(data.item_pop_idx).cuda(device)

    for epoch in range(start_epoch, args.epoch):
        print(f"current epoch: {epoch}/{args.epoch}")

        # If the early stopping has been reached, restore to the best performance model
        if flag:
            break

        # need change 4
        # All models
        running_loss, running_mf_loss, running_reg_loss, num_batches = 0, 0, 0, 0
        # CausE
        running_cf_loss = 0
        # BC_LOSS
        running_loss1, running_loss2 = 0, 0
        # XSimGCL
        running_cl_loss = 0
        # NCL
        running_proto_loss = 0
        # PID
        running_pid_loss = 0

        if args.modeltype == "SGL":
            print("get enhanced graph")
            dropped_adj1 = model.get_enhanced_adj(data.ui_mat, args.droprate)
            dropped_adj2 = model.get_enhanced_adj(data.ui_mat, args.droprate)

        if args.modeltype == "NCL" and epoch >= args.ncl_start_epoch:
            model.e_step()

        t1=time.time()

        pbar = tqdm(enumerate(data.train_loader), mininterval=2, total = len(data.train_loader))

        for batch_i, batch in pbar:
            # print(batch_i, batch)

            batch = [x.cuda(device) for x in batch]

            users = batch[0]
            pos_items = batch[1]

            if args.modeltype != 'CausE':
                users_pop = batch[2]
                pos_items_pop = batch[3]
                pos_weights = batch[4]
                if args.infonce == 0 or args.neg_sample != -1:
                    neg_items = batch[5]
                    neg_items_pop = batch[6]

            model.train()

            # need change 5

            if args.modeltype == 'INFONCE':

                mf_loss, reg_loss = model(users, pos_items, neg_items)
                loss = mf_loss + reg_loss

            elif args.modeltype == 'INFONCE_linear':

                mf_loss, reg_loss = model(users, pos_items, neg_items, neg_items_pop)
                loss = mf_loss + reg_loss

            elif args.modeltype == 'BC_LOSS':
                loss1, loss2, reg_loss, reg_loss_freeze, reg_loss_norm  = model(users, pos_items, neg_items, \
                                                                                users_pop, pos_items_pop, neg_items_pop)

                if epoch < args.freeze_epoch:
                    loss =  loss2 + reg_loss_freeze #@ loss2是poplarity bias loss
                else:
                    model.freeze_pop()
                    loss = loss1 + loss2 + reg_loss

            elif args.modeltype == 'NBC_LOSS':
                loss1, loss2, reg_loss, reg_loss_freeze, reg_loss_norm  = model(users, pos_items, neg_items, \
                                                                                users_pop, pos_items_pop, neg_items_pop)

                if epoch < args.freeze_epoch:
                    loss = loss1 + loss2 + reg_loss_freeze
                else:
                    model.freeze_pop()
                    loss = loss1 + loss2 + reg_loss

            elif args.modeltype == 'IPS' or args.modeltype =='SAMREG':

                mf_loss, reg_loss = model(users, pos_items, neg_items, pos_weights)
                loss = mf_loss + reg_loss

            elif args.modeltype == 'CausE':
                neg_items = batch[2]
                all_reg = torch.squeeze(batch[3].T.reshape([1, -1]))
                all_ctrl = torch.squeeze(batch[4].T.reshape([1, -1]))
                mf_loss, reg_loss, cf_loss = model(users, pos_items, neg_items, all_reg, all_ctrl)
                loss = mf_loss + reg_loss + cf_loss

            elif args.modeltype == "SimpleX":
                mf_loss, reg_loss = model(users, pos_items, neg_items)
                loss = mf_loss + reg_loss

            elif args.modeltype == "XSimGCL":
                mf_loss, cl_loss, reg_loss = model(users, pos_items, neg_items)
                loss = mf_loss + cl_loss + reg_loss

            elif args.modeltype == "SGL":
                mf_loss, cl_loss, reg_loss = model(users, pos_items, neg_items, dropped_adj1, dropped_adj2)
                loss = mf_loss + cl_loss + reg_loss

            elif args.modeltype == "NCL":
                mf_loss, cl_loss, proto_loss, reg_loss = model(users, pos_items, neg_items, epoch)
                if(epoch >= args.ncl_start_epoch):
                    loss = mf_loss + cl_loss + proto_loss + reg_loss
                else:
                    loss = mf_loss + cl_loss + reg_loss

            elif args.modeltype == "PID":
                mf_loss, pid_loss, reg_loss = model(users, pos_items, neg_items)
                loss = mf_loss + pid_loss + reg_loss
                # loss = mf_loss + reg_loss

            else:
                mf_loss, reg_loss = model(users, pos_items, neg_items)
                loss = mf_loss + reg_loss

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            running_loss += loss.detach().item()
            running_reg_loss += reg_loss.detach().item()

            # need change 6
            if args.modeltype != 'BC_LOSS' and args.modeltype != 'NBC_LOSS':
                running_mf_loss += mf_loss.detach().item()

            if args.modeltype == 'CausE':
                running_cf_loss += cf_loss.detach().item()

            if args.modeltype == 'BC_LOSS' or args.modeltype == 'NBC_LOSS':
                running_loss1 += loss1.detach().item()
                running_loss2 += loss2.detach().item()

            if args.modeltype == 'XSimGCL' or args.modeltype == 'SGL':
                running_cl_loss += cl_loss.detach().item()

            if args.modeltype == 'NCL':
                running_cl_loss += cl_loss.detach().item()
                running_proto_loss += proto_loss.detach().item()

            if args.modeltype == 'PID':
                running_pid_loss += pid_loss.detach().item()

            num_batches += 1

        t2=time.time()

        # need change 7
        # Training data for one epoch
        if args.modeltype == "CausE":
            perf_str = 'Epoch %d [%.1fs]: train==[%.5f=%.5f + %.5f + %.5f]' % (
                epoch, t2 - t1, running_loss / num_batches,
                running_mf_loss / num_batches, running_reg_loss / num_batches, running_cf_loss / num_batches)

        elif args.modeltype=="BC_LOSS" or args.modeltype=="NBC_LOSS":
            perf_str = 'Epoch %d [%.1fs]: train==[%.5f=%.5f + %.5f + %.5f]' % (
                epoch, t2 - t1, running_loss / num_batches,
                running_loss1 / num_batches, running_loss2 / num_batches, running_reg_loss / num_batches)
                #@ 1 for main branch 2 for popularity branch reg for reg branch

        elif args.modeltype == "XSimGCL" or args.modeltype == "SGL":
            perf_str = 'Epoch %d [%.1fs]: train==[%.5f=%.5f + %.5f + %.5f]' % (
                epoch, t2 - t1, running_loss / num_batches,
                running_mf_loss / num_batches, running_cl_loss / num_batches, running_reg_loss / num_batches)

        elif args.modeltype == "NCL":
            perf_str = 'Epoch %d [%.1fs]: train==[%.5f=%.5f + %.5f + %.5f + %.5f]' % (
                epoch, t2 - t1, running_loss / num_batches,
                running_mf_loss / num_batches, running_cl_loss / num_batches, running_proto_loss / num_batches, running_reg_loss / num_batches)

        elif args.modeltype == "PID":
            perf_str = 'Epoch %d [%.1fs]: train==[%.5f=%.5f + %.5f + %.5f]' % (
                epoch, t2 - t1, running_loss / num_batches,
                running_mf_loss / num_batches, running_pid_loss / num_batches, running_reg_loss / num_batches)

        else:
            perf_str = 'Epoch %d [%.1fs]: train==[%.5f=%.5f + %.5f]' % (
                epoch, t2 - t1, running_loss / num_batches,
                running_mf_loss / num_batches, running_reg_loss / num_batches)

        #@ 表现写入txt文件
        with open(base_path + 'stats_{}.txt'.format(args.saveID),'a') as f:
            f.write(perf_str+"\n")

        # Evaluate the trained model
        if (epoch + 1) % args.verbose == 0 and epoch >= freeze_epoch:
            model.eval()

            for i,evaluator in enumerate(evaluators):
                is_best, temp_flag = evaluation(args, data, model, epoch, base_path, evaluator,eval_names[i])

                if is_best:
                    checkpoint_buffer=save_checkpoint(model, epoch, base_path, checkpoint_buffer, args.max2keep)

                #@ early stop?
                if temp_flag:
                    flag = True

            all_users, all_items = model.compute()
            users_emb = all_users[users]
            pos_emb = all_items[pos_items]
            users_emb = F.normalize(users_emb, dim = -1)
            pos_emb = F.normalize(pos_emb, dim = -1)
            aloss = align_loss(users_emb, pos_emb)
            uloss = uniform_loss(users_emb)
            uloss += uniform_loss(pos_emb)

            with open(base_path + '/2loss.txt','a') as f:
                f.write(str(aloss.item())+"\t"+str(uloss.item())+"\n")

            model.train()

    # Get result
    model = restore_best_checkpoint(data.best_valid_epoch, model, base_path, device)
    print_str = "The best epoch is % d" % data.best_valid_epoch
    with open(base_path +'stats_{}.txt'.format(args.saveID), 'a') as f:
        f.write(print_str + "\n")

    for i,evaluator in enumerate(evaluators[:]):
        evaluation(args, data, model, epoch, base_path, evaluator, eval_names[i])
    with open(base_path +'stats_{}.txt'.format(args.saveID), 'a') as f:
        f.write(print_str + "\n")

    read_log(base_path +'stats_{}.txt'.format(args.saveID))


from django.urls import path, include
from django.contrib import admin
from main import views
from django.contrib.auth.views import LogoutView,LoginView

urlpatterns = [

    path('admin/', admin.site.urls),
    path('partner/', include('partner.urls')),
    path('user/', include('user.urls')),

    path('', views.home_view, name=''),
    path('logout', LogoutView.as_view(template_name='main/logout.html'), name='logout'),
    path('aboutus', views.aboutus_view),
    path('afterlogin', views.afterlogin_view, name='afterlogin'),

    path('adminclick', views.adminclick_view),
    path('adminlogin', LoginView.as_view(template_name='main/adminlogin.html'), name='adminlogin'),
    path('admin-dashboard', views.admin_dashboard_view, name='admin-dashboard'),
    path('admin-partner', views.admin_partner_view, name='admin-partner'),
    path('admin-view-partner', views.admin_view_partner_view, name='admin-view-partner'),
    path('update-partner/<int:pk>', views.update_partner_view, name='update-partner'),
    path('delete-partner/<int:pk>', views.delete_partner_view, name='delete-partner'),
    path('admin-view-pending-partner', views.admin_view_pending_partner_view, name='admin-view-pending-partner'),
    path('approve-partner/<int:pk>', views.approve_partner_view, name='approve-partner'),
    path('reject-partner/<int:pk>', views.reject_partner_view, name='reject-partner'),

    path('admin-user', views.admin_user_view, name='admin-user'),
    path('admin-view-user', views.admin_view_user_view, name='admin-view-user'),
    path('update-user/<int:pk>', views.update_user_view, name='update-user'),
    path('delete-user/<int:pk>', views.delete_user_view, name='delete-user'),

    path('admin-course', views.admin_course_view, name='admin-course'),
    path('admin-add-course', views.admin_add_course_view, name='admin-add-course'),
    path('admin-view-course', views.admin_view_course_view, name='admin-view-course'),
    path('delete-course/<int:pk>', views.delete_course_view, name='delete-course'),

    ]


#Problem: https://www.hackerrank.com/challenges/between-two-sets/problem

def getTotalX(a, b):
    numbers1 = []
    numbers2 = []
    for i in range(max(a), min(b)+1):
        if all([i%j==0 for j in a]):
            numbers1.append(i)
    for k in numbers1:
        if all([j%k==0 for j in b]):
            numbers2.append(k)
    return len(numbers2)


import os
import time
import shutil

class CleanDesktop:
    '''
    Cleaning up the desktop by moving directories and files to directory
    named "CLEAN DESKTOP"
    '''

    def __init__(self):
        self.exclude_extensions = ['ini', 'lnk']
        self.desktop_location = os.path.join(os.environ['USERPROFILE'], 'Desktop')
        self.exclude_directories = ['CLEANUP DESKTOP', 'My Projects', 'Tor Browser']

    def create_today_directory(self):
        '''
        Create directory as per today's date inside "CLEAN DESKTOP" directory
        '''

        today_date = time.strftime('%a %b %m %Y')
        self.cleanup_location = os.path.join(self.desktop_location, 'CLEANUP DESKTOP', today_date)

        if not os.path.exists(self.cleanup_location):
            os.makedirs(self.cleanup_location)

    def is_valid_file(self, file):
        '''
        Checks if the file type is not within the exclude_extensions list
        '''

        file = os.path.basename(file)

        if file in self.exclude_directories:   # Not including the cleanup folder itself
            return False

        for extension in self.exclude_extensions:
            if file.endswith(extension):
                return False

        return True

    def get_common_files_or_directories(self):
        '''
        Get common files that are in desktop and the cleanup locations
        '''

        desktop_files = set(os.listdir(self.desktop_location))
        cleanup_files = set(os.listdir(self.cleanup_location))
        common_files = [os.path.join(self.desktop_location, f) for f in desktop_files.intersection(cleanup_files)]

        return [f for f in common_files if self.is_valid_file(f)]

    def rename_common_files(self):
        '''
        Rename files or directories if they exists in desktop and cleanup directories
        '''

        common_files = self.get_common_files_or_directories()

        for file in common_files:
            if self.is_valid_file(file):
                new_name = self.extract_file_number(file)
                os.rename(file, new_name)

    def extract_file_number(self, file):
        '''
        Extract file_number having same name of **file** parameter's basename
        and add 1 to it
        '''

        _dir, basename = os.path.split(file)
        file_name = basename.split('.')[0]

        if os.path.isfile(file):
            extension = '.' + basename.split('.')[0]
            file_with_same_names = [f for f in os.listdir(self.cleanup_location) if f.startswith(file_name) and os.path.isfile(os.path.join(self.cleanup_location, f))]

        else:
            extension = ''
            file_with_same_names = [f for f in os.listdir(self.cleanup_location) if f.startswith(basename) and os.path.isdir(os.path.join(self.cleanup_location, f))]

        file_with_same_names.sort(key=len)
        last_file_name = file_with_same_names[-1].split('.')[0]  # Removing file extension
        file_number = last_file_name.replace(file_name, '')[1:-1]

        if file_number.isdigit():
            file_number = str(int(file_number) + 1)

            split = last_file_name.split('(')
            split[-1] = file_number
            new_name = f'{"(".join(split)}){extension}'

        else:
            file_number = '1'
            new_name = f'{last_file_name}({file_number}){extension}'

        return os.path.join(_dir, new_name)

    def main(self):
        '''
        Cleaning desktop
        '''

        self.create_today_directory()
        self.rename_common_files()

        desktop_files = [os.path.join(self.desktop_location, f) for f in os.listdir(self.desktop_location)]

        for file in desktop_files:
            if self.is_valid_file(file):
                move_to = os.path.join(self.cleanup_location, os.path.basename(file))
                shutil.move(file, move_to)

        os.startfile(self.cleanup_location)

if __name__ == '__main__':
    cleanup = CleanDesktop()
    cleanup.main()