import random
import sys
from abc import ABC

import numpy as np

# Maze parameters
class MazeParameter(ABC):
    def sample(self):
        pass

    def __str__(self):
        print(f"<{self.__class__} object with value {self.value}>")

    def __repr__(self):
        return f"<{self.__class__} object with value {self.value}>"

class BooleanParameter(MazeParameter):
    def __init__(self, value):
        assert isinstance(value, bool) or isinstance(value, float) or isinstance(value,
                                                                                int), "value must be either a boolean or a float"

        self.value = float(value)

    @property
    def sampler(self):
        return lambda: bool(np.random.binomial(1, self.value))

    def sample(self):
        return self.sampler()

    def __str__(self):
        print(f"<{self.__class__} object with value {self.value}>")

class FloatParameter(MazeParameter):
    def __init__(self, value):
        assert isinstance(value, float) or isinstance(value,
                                                      list) or isinstance(value,
                                                                          int), "value must be either a float or a list representing an interval"
        if isinstance(value, list):
            assert len(value) == 2, f"length of value must be 2 to represent an interval, currently length {len(value)}"
            a, b = value
            assert a <= b, "Your interval [a,b] is invalid because a>b. You must have a<=b for it to be a valid interval"
            assert a >= 0 and b <= 1, "Your interval must be contained in [0,1]"

        if isinstance(value, float) or isinstance(value, int):
            self.value = [value, value]

        else:
            self.value = value

    @property
    def sampler(self):
        a, b = self.value
        return lambda: random.uniform(a, b)

    def sample(self):
        return self.sampler()

class IntParameter(MazeParameter):
    def __init__(self, value):
        assert isinstance(value, int) or isinstance(value,
                                                    list), "value must be either an int or a list representing an interval"
        if isinstance(value, list):
            assert len(value) == 2, f"length of value must be 2 to represent an interval, currently length {len(value)}"
            a, b = value
            a, b = int(a), int(b)
            value = [a, b]
            assert a <= b, "Your interval [a,b] is invalid because a>b. You must have a<=b for it to be a valid interval"

        if isinstance(value, float) or isinstance(value, int):
            self.value = [int(value), int(value)]
        else:
            self.value = value

    @property
    def sampler(self):
        a, b = self.value
        return lambda: random.randint(a, b)

    def sample(self):
        return self.sampler()

class PyratState:
    # TODO : Everything in numpy for speed
    """
    Class representing a game state.
    Mud support is to be tested
    """
    # Dictionary for converting a direction to an offset and vice-versa
    DIRECTION_DICT = {
        "L": (-1, 0),
        "U": (0, 1),
        "R": (1, 0),
        "D": (0, -1),
        (-1, 0): "L",
        (0, 1): "U",
        (1, 0): "R",
        (0, -1): "D",
        0: (-1, 0),
        1: (0, 1),
        2: (1, 0),
        3: (0, -1),
    }

    def __init__(self, width=21, height=15, max_turns=2000, turns=0, maze={},
                 pieces=[], player1_location=None, player1_score=0, player1_moves=0, player1_misses=0, player1_mud=0,
                 player2_location=None, player2_score=0, player2_moves=0, player2_misses=0, player2_mud=0):
        # Maze configs
        self.width = width
        self.height = height

        # Global variables
        self.turns = turns
        self.max_turns = max_turns

        self.maze = maze
        self.pieces_of_cheese = pieces

        # Player 1 variables
        self.player1_location = player1_location
        self.player1_score = player1_score
        self.player1_moves = player1_moves
        self.player1_misses = player1_misses
        self.player1_mud = player1_mud

        # Player 2 variables
        self.player2_location = player2_location
        self.player2_score = player2_score
        self.player2_moves = player2_moves
        self.player2_misses = player2_misses
        self.player2_mud = player2_mud

        # save maze matrix
        self.maze_matrix_L, self.maze_matrix_U, self.maze_matrix_R, self.maze_matrix_D = self._calculate_maze_matrices()
        self.cheese_matrix = self._calculate_cheese_matrix()

    @property
    def original_nb_cheeses(self):
        return self.player1_score + self.player2_score + len(self.pieces_of_cheese)

    def step(self, decisions):
        """Makes a step for both players at the same time
        """
        # Move the players
        self._move(decisions)

        # Check for cheese captures and calculate the new scores
        reward = self._capture_cheeses_and_update_scores()

        done = self._check_done()

        return reward, done

    def set_state(self, state: dict):
        assert self._valid_state_dict(state), f"Invalid state : state needs to contain keys {self.__dict__.keys()}"
        for k in self.__dict__.keys():
            self.__setattr__(k, state[k])

    def get_obs(self):
        return self.__dict__.copy()

    def _move(self, decisions):
        p1_action, p2_action = decisions
        next_cell_p1, next_cell_p2 = np.add(self.player1_location, self.DIRECTION_DICT[p1_action]), np.add(
            self.player2_location, self.DIRECTION_DICT[p2_action])

        next_cell_p1 = tuple(next_cell_p1)
        next_cell_p2 = tuple(next_cell_p2)

        # Player 1
        # If the player is stuck in mud
        if self.player1_mud > 0:
            # One turn towards being able to move
            self.player1_mud -= 1
        # Elif next move is valid
        elif next_cell_p1 in self.maze[self.player1_location]:
            # Number of turns after this one the player is going to be stuck
            self.player1_mud = self.maze[self.player1_location][next_cell_p1] - 1
            self.player1_location = next_cell_p1
            self.player1_moves += 1
        else:
            self.player1_misses += 1

        # Player 2
        # If the player is stuck in mud
        if self.player2_mud > 0:
            # One turn towards being able to move
            self.player2_mud -= 1
        # Elif next move is valid
        elif next_cell_p2 in self.maze[self.player2_location]:
            # Number of turns after this one the player is going to be stuck
            self.player2_mud = self.maze[self.player2_location][next_cell_p2] - 1
            self.player2_location = next_cell_p2
            self.player2_moves += 1
        else:
            self.player2_misses += 1

        self.turns += 1

    def _capture_cheeses_and_update_scores(self):
        reward = 0
        if self.player1_location in self.pieces_of_cheese and self.player1_mud <= 0:
            self.pieces_of_cheese.remove(self.player1_location)
            self.cheese_matrix[self.player1_location] = 0
            if self.player2_location == self.player1_location and self.player2_mud <= 0:
                self.player1_score += 0.5
                self.player2_score += 0.5
            else:
                self.player1_score += 1
                reward = 1
        if self.player2_location in self.pieces_of_cheese and self.player2_mud <= 0:
            self.pieces_of_cheese.remove(self.player2_location)
            self.cheese_matrix[self.player2_location] = 0
            self.player2_score += 1
            reward = -1

        return reward

    def _check_done(self):
        return (self.turns >= self.max_turns) or (
                self.player1_score > self.original_nb_cheeses / 2) or (
                       self.player2_score > self.original_nb_cheeses / 2)

    # State conversion methods
    def _calculate_maze_matrices(self):
        """
        Generates the maze matrix.
        Returns 4 matrices of size self.width * self.height, which respectively indicate in cell [ij] whether you can go left, up, right, or down
        on cell [ij]
        :return : maze_matrix_L, maze_matrix_U, maze_matrix_R, maze_matrix_D
        """
        maze_dict = self.maze
        maze_matrix_U = np.zeros((self.width, self.height), dtype=np.int8)
        maze_matrix_D = np.zeros((self.width, self.height), dtype=np.int8)
        maze_matrix_L = np.zeros((self.width, self.height), dtype=np.int8)
        maze_matrix_R = np.zeros((self.width, self.height), dtype=np.int8)
        for position in maze_dict:
            for destination in maze_dict[position]:
                delta = np.subtract(destination, position)
                direction = self.DIRECTION_DICT[tuple(delta)]
                if direction == 'U':
                    maze_matrix_U[position[0], position[1]] = maze_dict[position][destination]
                elif direction == 'D':
                    maze_matrix_D[position[0], position[1]] = maze_dict[position][destination]
                elif direction == 'R':
                    maze_matrix_R[position[0], position[1]] = maze_dict[position][destination]
                elif direction == 'L':
                    maze_matrix_L[position[0], position[1]] = maze_dict[position][destination]

        return maze_matrix_L, maze_matrix_U, maze_matrix_R, maze_matrix_D

    def get_maze_matrix(self):
        return self.maze_matrix_L, self.maze_matrix_U, self.maze_matrix_R, self.maze_matrix_D

    def _calculate_cheese_matrix(self):
        cheese_matrix = np.zeros((self.width, self.height), dtype=np.int8)
        for cheese in self.pieces_of_cheese:
            cheese_matrix[cheese] = 1

        return cheese_matrix

    def get_cheese_matrix(self) -> np.ndarray:
        """Returns the matrix of size self.width * self.height, which indicates whether cell [ij] contains a cheese
        :return cheese_matrix: matrix such as cheese_matrix[ij] = 1 if [ij] contains a cheese
                                                                  0 else"""
        return self.cheese_matrix

    # Utility methods
    def _calculate_direction(self, source, destination):
        direction = None
        delta = (destination[0] - source[0], destination[1] - source[1])
        if delta == (0, 1):
            direction = 'U'
        elif delta == (0, -1):
            direction = 'D'
        elif delta == (1, 0):
            direction = 'R'
        elif delta == (-1, 0):
            direction = 'L'
        return direction

    def _valid_state_dict(self, state: dict):
        valid = True
        for key in self.__dict__.keys():
            valid = valid and (key in state.keys())

        return valid

    def __str__(self):
        return f"""Game State : {self.turns}/{self.max_turns} turns
- Maze : {self.maze}
- Pieces of cheese : ({len(self.pieces_of_cheese)} remaining)
                     {self.pieces_of_cheese}
---------------- Player 1 ----------------
- Score : {self.player1_score}
- Location : {self.player1_location}
- Moves : {self.player1_moves}
- Misses : {self.player1_misses}
- Turns remaining stuck in mud : {self.player1_mud}

---------------- Player 2 ----------------
- Score : {self.player2_score}
- Location : {self.player2_location}
- Moves : {self.player2_moves}
- Misses : {self.player2_misses}
- Turns remaining stuck in mud : {self.player2_mud}\n"""

    def __eq__(self, other):
        return self.__dict__ == other.__dict__

class MazeGenerator:
    """Class responsible for returning a maze.
    Maze generation returns the width, the height, the list of pieces of cheeses and the maze dictionary.
    Most of the code is directly ported from https://github.com/vgripon/PyRat/blob/master/imports/maze.py"""

    def __init__(self):
        self.width = None
        self.height = None
        self.maze = {}
        self.pieces_of_cheese = []

    def connected_region(self, cell, connected, possible_border):
        for (i, j) in self.maze[cell]:
            if connected[i][j] == 0:
                connected[i][j] = 1
                possible_border.append((i, j))
                self.connected_region((i, j), connected, possible_border)

    def gen_mud(self, mud_density, mud_range):
        """Generates a value for the distance between 2 cells"""
        if random.uniform(0, 1) < mud_density:
            return random.randrange(2, mud_range + 1)
        else:
            return 1

    def generate_pieces_of_cheese(self, nb_pieces, symmetry, start_random):
        """Generates the list of pieces of cheese, given maze configs

        ----
        Returns:
            (pieces_of_cheese : list, player1_location : (x,y) coordinates, player2_location :(x,y) coordinates)
            """
        player1_location = (-1, -1)
        player2_location = (-1, -1)
        if start_random:
            remaining = nb_pieces + 2
        else:
            remaining = nb_pieces
            player1_location = (0, 0)
            player2_location = (self.width - 1, self.height - 1)
        pieces = []
        candidates = []
        considered = []
        if symmetry:
            if nb_pieces % 2 == 1 and (self.width % 2 == 0 or self.height % 2 == 0):
                sys.exit(
                    "The maze has even width or even height and thus cannot contain an odd number of pieces of cheese if symmetric.")
            if nb_pieces % 2 == 1:
                pieces.append((self.width // 2, self.height // 2))
                considered.append((self.width // 2, self.height // 2))
                remaining = remaining - 1
        for i in range(self.width):
            for j in range(self.height):
                if (not (symmetry) or not ((i, j) in considered)) and (i, j) != player1_location and (
                        i, j) != player2_location:
                    candidates.append((i, j))
                    if symmetry:
                        considered.append((i, j))
                        considered.append((self.width - 1 - i, self.height - 1 - j))
        while remaining > 0:
            if len(candidates) == 0:
                sys.exit("Too many pieces of cheese for that dimension of maze")
            chosen = candidates[random.randrange(len(candidates))]
            pieces.append(chosen)
            if symmetry:
                a, b = chosen
                pieces.append((self.width - a - 1, self.height - 1 - b))
                symmetric = (self.width - a - 1, self.height - 1 - b)
                candidates = [i for i in candidates if i != symmetric]
                remaining = remaining - 1
            candidates = [i for i in candidates if i != chosen]
            remaining = remaining - 1
        if not (start_random):
            pieces.append(player1_location)
            pieces.append(player2_location)
        return pieces[:-2], pieces[-2], pieces[-1]

    def generate_maze(self, target_density, connected, symmetry, mud_density, mud_range, maze_file, seed):
        """Generates the maze and
        ------
        Returns :
            width (int): the maze's width
            height (int): the maze's height
            pieces_of_cheese (list): the list of cheeses in the maze. Is [] if the maze was not loaded from a file.
            maze (dict) : the dictionary representing the maze"""
        width = self.width
        height = self.height
        if maze_file != "":
            with open(maze_file, 'r') as content_file:
                content = content_file.read()
            lines = content.split("\n")
            width = int(lines[0])
            height = int(lines[1])
            self.maze = {}
            for i in range(width):
                for j in range(height):
                    self.maze[(i, j)] = {}
                    line = lines[i + j * width + 2].split(" ")
                    if line[0] != "0":
                        self.maze[(i, j)][(i, j + 1)] = int(line[0])
                    if line[1] != "0":
                        self.maze[(i, j)][(i, j - 1)] = int(line[1])
                    if line[2] != "0":
                        self.maze[(i, j)][(i - 1, j)] = int(line[2])
                    if line[3] != "0":
                        self.maze[(i, j)][(i + 1, j)] = int(line[3])
            line = lines[height * width + 2].split(" ")
            pieces_of_cheese = []
            for i in range(len(line)):
                l = int(line[i])
                pieces_of_cheese.append((l % width, l // width))
        else:
            random.seed(seed)
            # Start with purely random maze
            self.maze = {};
            not_considered = {};
            for i in range(width):
                for j in range(height):
                    self.maze[(i, j)] = {}
                    not_considered[(i, j)] = True
            for i in range(width):
                for j in range(height):
                    if not (symmetry) or not_considered[(i, j)]:
                        if random.uniform(0, 1) > target_density and i + 1 < width:
                            m = self.gen_mud(mud_density, mud_range)
                            self.maze[(i, j)][(i + 1, j)] = m
                            self.maze[(i + 1, j)][(i, j)] = m
                            if symmetry:
                                self.maze[(width - 1 - i, height - 1 - j)][(width - 2 - i, height - 1 - j)] = m
                                self.maze[(width - 2 - i, height - 1 - j)][(width - 1 - i, height - 1 - j)] = m
                        if random.uniform(0, 1) > target_density and j + 1 < height:
                            m = self.gen_mud(mud_density, mud_range)
                            self.maze[(i, j)][(i, j + 1)] = m
                            self.maze[(i, j + 1)][(i, j)] = m
                            if symmetry:
                                self.maze[(width - 1 - i, height - 2 - j)][(width - 1 - i, height - 1 - j)] = m
                                self.maze[(width - 1 - i, height - 1 - j)][(width - 1 - i, height - 2 - j)] = m
                        if symmetry:
                            not_considered[(i, j)] = False
                            not_considered[(width - 1 - i, height - 1 - j)] = False
            for i in range(width):
                for j in range(height):
                    if len(self.maze[(i, j)]) == 0 and (i == 0 or j == 0 or i == width - 1 or j == height - 1):
                        m = self.gen_mud(mud_density, mud_range)
                        possibilities = []
                        if i + 1 < width:
                            possibilities.append((i + 1, j))
                        if j + 1 < height:
                            possibilities.append((i, j + 1))
                        if i - 1 >= 0:
                            possibilities.append((i - 1, j))
                        if j - 1 >= 0:
                            possibilities.append((i, j - 1))
                        chosen = possibilities[random.randrange(len(possibilities))]
                        self.maze[(i, j)][chosen] = m
                        self.maze[chosen][(i, j)] = m
                        if symmetry:
                            ii, jj = chosen
                            self.maze[(width - 1 - i, height - 1 - j)][(width - 1 - ii, height - 1 - jj)] = m
                            self.maze[(width - 1 - ii, height - 1 - jj)][(width - 1 - i, height - 1 - j)] = m

            # Then connect it
            if connected:
                connected = [[0 for x in range(height)] for y in range(width)]
                possible_border = [(0, height - 1)]
                connected[0][height - 1] = 1
                self.connected_region((0, height - 1), connected, possible_border)
                while 1:
                    border = []
                    new_possible_border = []
                    for (i, j) in possible_border:
                        is_candidate = False
                        if not ((i + 1, j) in self.maze[(i, j)]) and i + 1 < width:
                            if connected[i + 1][j] == 0:
                                border.append(((i, j), (i + 1, j)))
                                is_candidate = True
                        if not ((i - 1, j) in self.maze[(i, j)]) and i > 0:
                            if connected[i - 1][j] == 0:
                                border.append(((i, j), (i - 1, j)))
                                is_candidate = True
                        if not ((i, j + 1) in self.maze[(i, j)]) and j + 1 < height:
                            if connected[i][j + 1] == 0:
                                border.append(((i, j), (i, j + 1)))
                                is_candidate = True
                        if not ((i, j - 1) in self.maze[(i, j)]) and j > 0:
                            if connected[i][j - 1] == 0:
                                border.append(((i, j), (i, j - 1)))
                                is_candidate = True
                        if is_candidate:
                            new_possible_border.append((i, j))
                    possible_border = new_possible_border
                    if not border:
                        break
                    a, b = border[random.randrange(len(border))]
                    m = self.gen_mud(mud_density, mud_range)
                    self.maze[a][b] = m
                    self.maze[b][a] = m
                    ai, aj = a
                    bi, bj = b
                    if symmetry:
                        bsym = (width - 1 - bi, height - 1 - bj)
                        asym = (width - 1 - ai, height - 1 - aj)
                        self.maze[asym][bsym] = m
                        self.maze[bsym][asym] = m
                    connected[bi][bj] = 1
                    self.connected_region(b, connected, possible_border)
                    possible_border.append(b)
                    if symmetry:
                        if connected[width - 1 - bi][height - 1 - bj] == 0 and connected[width - 1 - ai][
                            height - 1 - aj] == 1:
                            connected[width - 1 - bi][height - 1 - bj] = 1
                            self.connected_region(bsym, connected, possible_border)
                            possible_border.append(bsym)
            pieces_of_cheese = []

        return width, height, pieces_of_cheese, self.maze

    def __call__(self, width, height, nb_pieces_of_cheese, target_density, mud_density, connected, symmetry, mud_range,
                 start_random, maze_file, seed):
        self.width = width
        self.height = height
        self.maze = {}
        self.pieces_of_cheese = []

        self.width, self.height, self.pieces_of_cheese, self.maze = self.generate_maze(target_density, connected,
                                                                                symmetry, mud_density, mud_range,
                                                                                maze_file, seed)

        if self.pieces_of_cheese == []:
            self.pieces_of_cheese, p1_location, p2_location = self.generate_pieces_of_cheese(nb_pieces_of_cheese,
                                                                                symmetry, start_random)

        return self.width, self.height, self.maze, self.pieces_of_cheese, p1_location, p2_location

class GameGenerator:
    """Maze Generator class
    Handles Pyrat Maze generation.
    Most functions are ported from https://github.com/vgripon/PyRat/blob/master/imports/maze.py

    To get a maze :
        >> game_generator = GameGenerator()
        >> state = game_generator()"""

    def __init__(self, width=21, height=15, nb_pieces_of_cheese=41, target_density=0.7, mud_density=0, connected=True,
                 symmetry=True, mud_range=10, start_random=False, maze_file="", seed="random", max_turns=2000):
        """
        Initializes the maze generator
        Most arguments can take either a fixed value or can be sampled uniformly at random from a distribution.
        Only height and width
        For INT parameters :
            - scalar value : parameter is always that value
            - list [a,b] : integer sampled uniformly at random from integer interval  [a,b]
        For FLOAT parameters :
            - scalar value : parameter is always that value
            - list [a,b] : float sampled uniformly at random from integer interval  [a,b]
        For BOOL parameters :
            - True : parameter is always True
            - False : parameter is always False
            - scalar value p within [0,1]: parameter follows a Bernoulli Law of parameter p, i.e. has probability p to be True

        ---------------------
        Arguments:
            width : int
                NOT RANDOMIZABLE
                Width of the maze. Default : 21
            height : int
                NOT RANDOMIZABLE
                Height of the maze. Default : 15
            nb_pieces_of_cheese : int
                RANDOMIZABLE
                Number of initial pieces of cheese put into the maze. Default : 41
            target_density : float
                RANDOMIZABLE
                Probability of having a wall between two given maze cells. Needs to be in [0,1]. Default : 0.7
            mud_density : float
                RANDOMIZABLE
                Probability of having mud between two given maze cells. Needs to be in [0,1]. Default : 0
            connected : boolean
                RANDOMIZABLE
                Whether the maze should be in one connected component or not. Default : True
            symmetry : boolean
                RANDOMIZABLE
                Whether the maze should be symmetric. Default : True
            mud_range : int
                RANDOMIZABLE
                Upper bound on the mud values i.e. the number of turns it can take to cross from one cell to another.
            start_random : boolean
                RANDOMIZABLE
                Whether the players should be started in fixed positions (upper right corner and lower left) or if they can be started from anywhere in the maze.
                Default : False
            maze_file : string
                Path to a maze file, if you want the maze that's generated to always be the same
            seed : int or "random"
                Seed for torch, numpy and random.
                Can be either chose randomly if "random" or be a fixed int
        ---------------------
        Use examples :
        >> GameGenerator() : Default values.
                15x21 maze, 41 cheeses, 0.7 wall density, no mud, one connected component, a symmetric maze and players starting in the lower left and upper right corners respectively.
        >> GameGenerator(nb_pieces_of_cheese = [21,41]) :
                Same as previously, but with a number of pieces of cheese chosen uniformly in [21,41].
        >> GameGenerator(target_density = [0.2,0.7]) :
                Same as the first one, but with a wall density chosen uniformly in [0.2,0.7].
        >> GameGenerator(symmetry = 0.8)
                Same as the first example, but with a maze having 0.8 probability to be symmetric.
        """
        # TODO : Implement seed everywhere
        self.width = width
        self.height = height
        self.max_turns = max_turns
        self.nb_pieces_of_cheese = IntParameter(nb_pieces_of_cheese)
        self.target_density = FloatParameter(target_density)
        self.mud_density = FloatParameter(mud_density)
        self.connected = BooleanParameter(connected)
        self.symmetry = BooleanParameter(symmetry)
        self.mud_range = IntParameter(mud_range)
        self.start_random = BooleanParameter(start_random)
        self.maze_file = maze_file

        assert isinstance(seed, str) or isinstance(seed,
                                                   int), "seed parameter needs to be either and integer or 'random'"
        if isinstance(seed, str):
            assert seed == "random", f"unrecognized value for seed '{seed}'. Needs to be 'random'"
        self.seed = seed

        self.maze_gen = MazeGenerator()

    def _sample_parameters(self):
        (width,
         height,
         nb_pieces_of_cheese,
         target_density,
         mud_density,
         connected,
         symmetry,
         mud_range,
         start_random,
         maze_file) = (self.width, self.height,
                       self.nb_pieces_of_cheese.sample(),
                       self.target_density.sample(),
                       self.mud_density.sample(),
                       self.connected.sample(),
                       self.symmetry.sample(),
                       self.mud_range.sample(),
                       self.start_random.sample(),
                       self.maze_file)
        if self.seed == "random":
            seed = random.randint(0, sys.maxsize)
        else:
            seed = self.seed

        return width, height, nb_pieces_of_cheese, target_density, mud_density, connected, symmetry, mud_range, start_random, maze_file, seed

    def __call__(self):
        width, height, nb_pieces_of_cheese, target_density, mud_density, connected, symmetry, mud_range, start_random, maze_file, seed = self._sample_parameters()
        width, height, maze, pieces_of_cheese, player1_location, player2_location = self.maze_gen(width, height,
                                                                                nb_pieces_of_cheese,
                                                                                target_density,
                                                                                mud_density,
                                                                                connected, symmetry,
                                                                                mud_range,
                                                                                start_random,
                                                                                maze_file, seed)
        game_state = PyratState(width=width, height=height, maze=maze, pieces=pieces_of_cheese,
                                player1_location=player1_location, player2_location=player2_location)
        return game_state


import os
from sendgrid import SendGridAPIClient
from django.conf import settings
from sendgrid.helpers.mail import Mail
from django.template.loader import render_to_string

def send_welcome_email(name, email):
    message = Mail(
        from_email="contact@victormaina.com",
        to_emails=email,
        subject="Test Subject",
        html_content = "
This is a test email
")

    try:
        sg = SendGridAPIClient(api_key=settings.SENDGRID_API_KEY)
        response = sg.send(message)
        print(response.status_code)
    except Exception as e:
        print("\nSENDGRID ERROR: ", e, "\n")

# from django.core.mail import EmailMultiAlternatives
# from django.template.loader import render_to_string

# def send_welcome_email(name, receiver):
#     subject = "Welcome to the MoringaTribune Newsletter"
#     sender = "james@moringaschool.com"

#     text_content = render_to_string("email/newsemail.txt", {"name":name})
#     html_content = render_to_string("email/newsemail.html", {"name":name})

#     msg = EmailMultiAlternatives(subject, text_content, sender, [receiver])
#     msg.attach_alternative(html_content, "text/html")
#     msg.send


#!/usr/bin/env python3
import psutil

import rospy
from std_msgs.msg import Bool

RATE = 5

def main():
    rospy.init_node('PsNode')
    proc_name = './test.sh'
    pub       = rospy.Publisher('/system/ps/psychopy', Bool, queue_size = 1)
    ros_rate  = rospy.Rate(RATE)

    rospy.loginfo("PsNode started")

    # keep going till shutdown
    while not rospy.is_shutdown():

        for pid in psutil.pids():
            try:
                if psutil.Process(pid).cmdline()[1] == proc_name:
                    pub.publish(Bool(True))
                    break
            except:
                pass
        ros_rate.sleep()

if __name__ == '__main__':
    main()


import talib as ta
import numpy as np
import pandas as pd

"""
将kdj策略需要用到的信号生成器抽离出来
"""

class barPowerSignal():

    def __init__(self):
        self.author = 'Haolin Pan'

    def maEnvironment(self, am, paraDict):
        envFastPeriod = paraDict["envFastPeriod"]
        envSlowPeriod = paraDict["envSlowPeriod"]

        sma = ta.MA(am.close, envFastPeriod)
        lma = ta.MA(am.close, envSlowPeriod)

        envSignal = 0
        if sma[-1]>lma[-1]:
            envSignal = 1
        else:
            envSignal = -1
        return envSignal, sma, lma

    def maCross(self,am,paraDict):
        fastPeriod = paraDict["fastPeriod"]
        slowPeriod = paraDict["slowPeriod"]

        sma = ta.MA(am.close, fastPeriod)
        lma = ta.MA(am.close, slowPeriod)
        goldenCross = sma[-1]>lma[-1] and sma[-2]<=lma[-2]
        deathCross = sma[-1]<lma[-1] and sma[-2]>=lma[-2]

        maCrossSignal = 0
        if goldenCross:
            maCrossSignal = 1
        elif deathCross:
            maCrossSignal = -1
        else:
            maCrossSignal = 0
        return maCrossSignal, sma, lma

    def maStatus(self,am,paraDict):
        fastPeriod = paraDict["fastPeriod"]
        slowPeriod = paraDict["slowPeriod"]

        sma = ta.MA(am.close, fastPeriod)
        lma = ta.MA(am.close, slowPeriod)
        goldenCross = sma[-1]>lma[-1] and sma[-2]<=lma[-2]
        deathCross = sma[-1]<lma[-1] and sma[-2]>=lma[-2]

        maCrossSignal = 0
        if sma[-1]>lma[-1]:
            maCrossSignal = 1
        else:
            maCrossSignal = -1
        return maCrossSignal, sma, lma

    def atrCross(self,am,paraDict):
        atrFastPeriod = paraDict["atrFastPeriod"]
        atrSlowPeriod = paraDict["atrSlowPeriod"]

        fa = ta.ATR(am.high, am.low, am.close, atrFastPeriod)
        sa = ta.ATR(am.high, am.low, am.close, atrSlowPeriod)
        atrDif = fa-sa*1.2

        atrCrossSignal = 0
        if atrDif[-1]>0:
            atrCrossSignal = 1
        else:
            atrCrossSignal = 0
        return atrCrossSignal, atrDif

    def rsiStatus(self,am,paraDict):
        rsiPeriod = paraDict["rsiPeriod"]

        rsi = ta.RSI(am.close, rsiPeriod)
        mrsi = ta.MA(rsi, 5)

        overSell = rsi[-1]>30 and rsi[-2]<30
        overBuy = rsi[-1]<70 and rsi[-2]>70

        rsiSignal = 0
        if overSell:
            rsiSignal = 1
        elif overBuy:
            rsiSignal = -1

        return rsiSignal, rsi

    def rsiExit(self,am,paraDict):
        rsiPeriod = paraDict["rsiPeriod"]

        rsi = ta.RSI(am.close, rsiPeriod)

        overSell = rsi[-1]<20
        overBuy = rsi[-1]>80

        rsiSignal = 0
        if overSell:
            rsiSignal = 1
        elif overBuy:
            rsiSignal = -1
        else:
            rsiSignal = 0
        return rsiSignal, rsi

    def priceStatus(self,am,paraDict):
        price = am.close
        priceSignal = 0
        if price[-1]>price[-3]:
            priceSignal = 1
        else:
            priceSignal = -1
        return priceSignal, price

    def barStatus(self,am,paraDict):
        barSignal = 0
        if am.close[-1]>am.open[-1]:
            barSignal = 1
        else:
            barSignal = -1
        return barSignal

    def cciCross(self,am,paraDict):
        cciPeriod = paraDict["cciPeriod"]
        cci = ta.CCI(am.high, am.low, am.close, cciPeriod)

        cciSignal = 0
        if cci[-1]>100:
            cciSignal = 1
        elif cci[-1]<-100:
            cciSignal = -1
        else:
            cciSignal = 0
        return cciSignal, cci

    def cciExit(self,am,paraDict):
        cciPeriod = paraDict["cciPeriod"]
        cci = ta.CCI(am.high, am.low, am.close, cciPeriod)

        cciExitSignal = 0
        if cci[-1]>-100 and cci[-2]<-100:
            cciExitSignal = 1
        elif cci[-1]<100 and cci[-2]>100:
            cciExitSignal = -1
        else:
            cciExitSignal = 0
        return cciExitSignal, cci

    def bollStatus(self,am,paraDict):
        bollPeriod = paraDict["bollPeriod"]
        up, mid, low = ta.BBANDS(am.close, bollPeriod)

        bollExitSignal = 0
        if am.close[-1]<low[-1]:
            bollExitSignal = 1
        elif am.close[-1]>up[-1]:
            bollExitSignal = -1
        else:
            bollExitSignal = 0
        return bollExitSignal, up, mid, low

    def barPower(self,am,paraDict):
        bpPeriod = paraDict["bpPeriod"]
        o = am.open
        h = am.high
        l = am.low
        c = am.close
        v = am.volume

        spread = c-o
        rv = v/ta.SMA(v,bpPeriod)
        rv[spread==0]=0
        rv[np.isnan(rv)]=0
        bp = ta.EMA(spread,bpPeriod)
        mbp = ta.EMA(bp,3)

        mbpSignal = 0
        if bp[-1]>mbp[-1]:
            mbpSignal = 1
        else:
            mbpSignal = -1
        return mbpSignal, bp


from leetcode.data_structures.list_node import ListNode
from leetcode.data_structures.tree_node import TreeNode
from typing import List, Optional
from collections import deque

def create_list_node(nodes: List[int]) -> Optional[ListNode]:
    """Create a linked-list with a list"""
    head = ListNode(0)
    dummy = head
    for node in nodes:
        dummy.next = ListNode(node)
        dummy = dummy.next
    return head.next

def print_list_node(head: Optional[ListNode]) -> str:
    """Convert a linked-list to string"""
    res = ""
    while head is not None:
        res += str(head.val)
        if head is not None:
            res += "->"
    return res

def list_node_to_list(head: Optional[ListNode]) -> List[int]:
    """Convert a linked-list to a list"""
    res = []
    while head:
        res.append(head.val)
        head = head.next
    return res

def create_treenode(nums: List[int | None]) -> Optional[TreeNode]:
    """Create Binary Tree Node"""
    n = len(nums)
    if n == 0 or nums[0] is None:
        return None
    root = TreeNode(nums[0])
    q = deque()
    q.appendleft(root)
    cursor = 1
    while cursor < n:
        node = q.pop()
        if cursor > n - 1 or nums[cursor] is None:
            node.left = None
        else:
            left_node = TreeNode(nums[cursor])
            if left_node:
                node.left = left_node
            q.appendleft(left_node)
        if cursor + 1 > n - 1 or nums[cursor + 1] is None:
            node.right = None
        else:
            right_node = TreeNode(nums[cursor + 1])
            if right_node:
                node.right = right_node
            q.appendleft(right_node)
        cursor += 2
    return root

def preorder_traversal(root: TreeNode) -> List[int]:
    """Traversal a binary tree node with preoder"""
    res = list()
    if not root:
        return res
    stack = []
    node = root
    while stack or node:
        while node:
            res.append(node.val)
            stack.append(node)
            node = node.left
        node = stack.pop()
        node = node.right
    return res

# def is_same_tree(root1: Optional[TreeNode], root2: Optional[TreeNode]) -> bool:
#     """Judges if two binary tree nodes are equal"""
#     if not root1 and not root2:
#         return True
#     if not root1 or not root2:
#         return False
#
#     queue1 = deque([root1])
#     queue2 = deque([root2])
#
#     while queue1 and queue2:
#         node1 = queue1.popleft()
#         node2 = queue2.popleft()
#         if node1.val != node2.val:
#             return False
#         left1, right1 = node1.left, node1.right
#         left2, right2 = node2.left, node2.right
#         if (not left1) ^ (not left2):
#             return False
#         if (not right1) ^ (not right2):
#             return False
#         if left1:
#             queue1.append(left1)
#         if right1:
#             queue1.append(right1)
#         if left2:
#             queue2.append(left2)
#         if right2:
#             queue2.append(right2)
#
#     return not queue1 and not queue2


magicnumber = [[2,7,6],
               [9,5,1],
               [4,3,8]]
print('Input Matrix: ')
#printing the matrix input
for row in magicnumber:
    print(row)

#ROW SUM
rowsum = [0 for x in range(len(magicnumber))]
for rowIndex , row in enumerate(magicnumber):
    for num in row:
        rowsum[rowIndex] += num

#COLUM SUM
columsum = [0 for x in range(len(magicnumber))]
for row in magicnumber:
    for index in range(len(magicnumber)):
        columsum[index] += row[index]

#digonal sum (left to right digonal sum)
digonalsum = 0
for i in range(len(magicnumber)):
    digonalsum += magicnumber[i][i]

#another digonal sum(right to left digonal sum)
for x in range(len(magicnumber)):
    for y in range(len(magicnumber)):
        if x==0 and y==2:
            r1=magicnumber[x][y]
        if x==2 and y==0:
            r2=magicnumber[x][y]
        if x==1 and y==1:
            r=magicnumber[x][y]
anotherdigonalsum = r+r1+r2 #another digonal sum (from right to left)

print(digonalsum)
print(anotherdigonalsum)
#checking that row sum, colum sum, digonal sum, another digonal are same
if rowsum[0]==columsum[0]==digonalsum==anotherdigonalsum:
    print("This is a Lo Shu Magic Square.")
    n=len(magicnumber)
    print("And, The magic number is ",n*(n**2+1)/2)


from ..pipeline import Flush, Writer
from .database_connector import DatabaseConnector
from .debounced_ingest_strategy import DebouncedIngestStrategy
from .ingest_strategy import INGESTION_STRATEGY_REGISTRY, IngestionStrategy

class GraphDatabaseWriter(Writer):
    @classmethod
    def from_file_data(
        cls,
        batch_size: int,
        database: str,
        ingest_strategy_name: str = INGESTION_STRATEGY_REGISTRY.name_for(
            DebouncedIngestStrategy
        ),
        collect_stats: bool = True,
        **database_args
    ):
        # Import all query executors so that they can register themselves
        DatabaseConnector.import_all()
        connector = DatabaseConnector.from_database_args(
            database=database, **database_args
        )
        executor = connector.get_query_executor(collect_stats=collect_stats)
        ingest_strategy_cls = INGESTION_STRATEGY_REGISTRY.get(ingest_strategy_name)
        ingest_strategy = ingest_strategy_cls(executor)

        return cls(
            batch_size=batch_size,
            ingest_strategy=ingest_strategy,
        )

    def __init__(self, batch_size: int, ingest_strategy: IngestionStrategy) -> None:
        self.batch_size = batch_size
        self.ingest_strategy = ingest_strategy
        self.pending_records = 0

    async def flush(self):
        await self.ingest_strategy.flush()
        self.pending_records = 0

    async def write_record(self, ingestible):
        if ingestible is Flush:
            await self.flush()
            return

        await ingestible.ingest(self.ingest_strategy)
        self.pending_records += 1
        if self.pending_records >= self.batch_size:
            await self.flush()

    async def finish(self):
        await self.flush()


#!/bin/python

# This interpolation script can be used on a HeFESTo table in its original format,
# subsamples it, reorders the columns, and adds a header so that it can be used
# by the AsciiDataLookup class inside ASPECT.

import numpy as np

# Change this to the original input file
input_filename = "opxtable2_s.aspect"
output_filename = "material_table.txt"

data = np.genfromtxt(input_filename)

n_pressure_points_original = 1401
n_entropy_points_original = 241

table = data.reshape(n_pressure_points_original,n_entropy_points_original,9)

pressure_subsampling = 10
entropy_subsampling = 5

modified = table[::pressure_subsampling,:,:]
modified = modified[:,::entropy_subsampling,:]
modified = modified.reshape(modified.shape[0]*modified.shape[1],9)

# Change column order by providing permutation from old column index to new
# column index.
# Original order: T(K) P(bar) rho,kg/m3 alpha,1/K cp,J/K/kg vp,km/s vs,km/s h,J/kg s,J/kg/K
# Required order: s,J/kg/K P(bar) T(K) rho,kg/m3 alpha,1/K cp,J/K/kg vp,km/s vs,km/s h,J/kg
permutation = [2, 1, 3, 4, 5, 6, 7, 8, 0]
idx = np.empty_like(permutation)
idx[permutation] = np.arange(len(permutation))
modified = modified[:, idx]

headerfile = open("header_table.txt","r")
headertext = headerfile.read()
headertext = headertext[:-1]

np.savetxt(output_filename, modified,fmt='%.6g', header=headertext, comments='')


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import itertools
import logging
from abc import ABCMeta
from abc import abstractmethod

import numpy as np
from scipy.special import gammaln
from six import add_metaclass

from six.moves import range
from treecat.structure import OP_IN
from treecat.structure import OP_OUT
from treecat.structure import OP_ROOT
from treecat.structure import OP_UP
from treecat.structure import TreeStructure
from treecat.structure import estimate_tree
from treecat.structure import make_propagation_program
from treecat.structure import sample_tree
from treecat.tables import TY_MULTINOMIAL
from treecat.tables import Table
from treecat.util import SERIES
from treecat.util import TODO
from treecat.util import jit
from treecat.util import parallel_map
from treecat.util import prange
from treecat.util import profile
from treecat.util import sample_from_probs
from treecat.util import set_random_seed

logger = logging.getLogger(__name__)

def count_pairs(assignments, v1, v2, M):
    """Construct sufficient statistics for (v1, v2) pairs.

    Args:
        assignments: An _ x V assignment matrix with values in range(M).
        v1, v2: Column ids of the assignments matrix.
        M: The number of possible assignment bins.

    Returns:
        An M x M array of counts.
    """
    assert v1 != v2
    pairs = assignments[:, v1].astype(np.int32) * M + assignments[:, v2]
    return np.bincount(pairs, minlength=M * M).reshape((M, M))

def logprob_dc(counts, prior, axis=None):
    """Non-normalized log probability of a Dirichlet-Categorical distribution.

    See https://en.wikipedia.org/wiki/Dirichlet-multinomial_distribution
    """
    # Note that this excludes the factorial(counts) term, since we explicitly
    # track permutations in assignments.
    return gammaln(np.add(counts, prior, dtype=np.float32)).sum(axis)

def make_annealing_schedule(num_rows, epochs, sample_tree_rate):
    """Iterator for subsample annealing, yielding (action, arg) pairs.

    This generates a subsample annealing schedule starting from an empty
    assignment state (no rows are assigned). It then interleaves 'add_row' and
    'remove_row' actions so as to gradually increase the number of assigned
    rows. The increase rate is linear.

    Args:
        num_rows (int): Number of rows in dataset.
            The annealing schedule terminates when all rows are assigned.
        epochs (float): Number of epochs in the schedule (i.e. the number of
            times each datapoint is assigned). The fastest schedule is
            `epochs=1` which simply sequentially assigns all datapoints. More
            epochs takes more time.
        sample_tree_rate (float): The rate at which 'sample_tree' actions are
            generated. At `sample_tree_rate=1`, trees are sampled after each
            complete flushing of the subsample.

    Yields: (action, arg) pairs.
        Actions are one of: 'add_row', 'remove_row', or 'sample_tree'.
        When `action` is 'add_row' or 'remove_row', `arg` is a row_id in
        `range(num_rows)`. When `action` is 'sample_tree' arg is undefined.
    """
    assert epochs >= 1.0
    assert sample_tree_rate >= 1.0
    # Randomly shuffle rows.
    row_ids = list(range(num_rows))
    np.random.shuffle(row_ids)
    row_to_add = itertools.cycle(row_ids)
    row_to_remove = itertools.cycle(row_ids)

    # Use a linear annealing schedule.
    epochs = float(epochs)
    add_rate = epochs
    remove_rate = epochs - 1.0
    state = 2.0 * epochs

    # Sample the tree sample_tree_rate times per batch.
    num_assigned = 0
    next_batch = 0
    while num_assigned < num_rows:
        if state >= 0.0:
            yield 'add_row', next(row_to_add)
            state -= remove_rate
            num_assigned += 1
            next_batch -= sample_tree_rate
        else:
            yield 'remove_row', next(row_to_remove)
            state += add_rate
            num_assigned -= 1
        if num_assigned > 0 and next_batch <= 0:
            yield 'sample_tree', None
            next_batch = num_assigned

@add_metaclass(ABCMeta)
class TreeTrainer(object):
    """Abstract base class for training a tree model various latent state.

    Derived classes must implement:
    - `add_row(row_id)`
    - `remove_row(row_id)`
    - `compute_edge_logits()`
    - `logprob()`
    """

    def __init__(self, N, V, tree_prior, config):
        """Initialize a model with an empty subsample.

        Args:
            N (int): Number of rows in the dataset.
            V (int): Number of columns (features) in the dataset.
            tree_prior: A [K]-shaped numpy array of prior edge log odds, where
                K is the number of edges in the complete graph on V vertices.
            config: A global config dict.
        """
        assert isinstance(N, int)
        assert isinstance(V, int)
        assert isinstance(tree_prior, np.ndarray)
        assert isinstance(config, dict)
        K = V * (V - 1) // 2  # Number of edges in complete graph.
        assert V <= 32768, 'Invalid # features > 32768: {}'.format(V)
        assert tree_prior.shape == (K, )
        assert tree_prior.dtype == np.float32
        self._config = config.copy()
        self._num_rows = N
        self._tree_prior = tree_prior
        self._tree = TreeStructure(V)
        assert self._tree.num_vertices == V
        self._program = make_propagation_program(self._tree.tree_grid)
        self._added_rows = set()

    @abstractmethod
    def add_row(self, row_id):
        """Add a given row to the current subsample."""

    @abstractmethod
    def remove_row(self, row_id):
        """Remove a given row from the current subsample."""

    @abstractmethod
    def compute_edge_logits(self):
        """Compute edge log probabilities on the complete graph."""

    @abstractmethod
    def logprob(self):
        """Compute non-normalized log probability of data and latent state.

        This is used for testing goodness of fit of the latent state kernel.
        This should only be called after training, i.e. after all rows have
        been added.
        """
        assert len(self._added_rows) == self._num_rows

    def get_edges(self):
        """Get a list of the edges in the current tree.

        Returns:
            An E-long list of (vertex,vertex) pairs.
        """
        return [tuple(edge) for edge in self._tree.tree_grid[1:3, :].T]

    def set_edges(self, edges):
        """Set edges of the latent structure and update statistics.

        Args:
            edges: An E-long list of (vertex,vertex) pairs.
        """
        self._tree.set_edges(edges)
        self._program = make_propagation_program(self._tree.tree_grid)

    @profile
    def sample_tree(self):
        """Samples a random tree.

        Returns:
            A pair (edges, edge_logits), where:
                edges: A list of (vertex, vertex) pairs.
                edge_logits: A [K]-shaped numpy array of edge logits.
        """
        logger.info('TreeCatTrainer.sample_tree given %d rows',
                    len(self._added_rows))
        SERIES.sample_tree_num_rows.append(len(self._added_rows))
        complete_grid = self._tree.complete_grid
        edge_logits = self.compute_edge_logits()
        assert edge_logits.shape[0] == complete_grid.shape[1]
        assert edge_logits.dtype == np.float32
        edges = self.get_edges()
        edges = sample_tree(complete_grid, edge_logits, edges)
        return edges, edge_logits

    def estimate_tree(self):
        """Compute a maximum likelihood tree.

        Returns:
            A pair (edges, edge_logits), where:
                edges: A list of (vertex, vertex) pairs.
                edge_logits: A [K]-shaped numpy array of edge logits.
        """
        logger.info('TreeCatTrainer.estimate_tree given %d rows',
                    len(self._added_rows))
        complete_grid = self._tree.complete_grid
        edge_logits = self.compute_edge_logits()
        edges = estimate_tree(complete_grid, edge_logits)
        return edges, edge_logits

    @profile
    def train(self):
        """Train a model using subsample-annealed MCMC.

        Returns:
            A trained model as a dictionary with keys:
                config: A global config dict.
                tree: A TreeStructure instance with the learned latent
                    structure.
                edge_logits: A [K]-shaped array of all edge logits.
        """
        logger.info('TreeTrainer.train')
        set_random_seed(self._config['seed'])
        init_epochs = self._config['learning_init_epochs']
        full_epochs = self._config['learning_full_epochs']
        sample_tree_rate = self._config['learning_sample_tree_rate']
        num_rows = self._num_rows

        # Initialize using subsample annealing.
        assert len(self._added_rows) == 0
        schedule = make_annealing_schedule(num_rows, init_epochs,
                                           sample_tree_rate)
        for action, row_id in schedule:
            if action == 'add_row':
                self.add_row(row_id)
            elif action == 'remove_row':
                self.remove_row(row_id)
            elif action == 'sample_tree':
                edges, edge_logits = self.sample_tree()
                self.set_edges(edges)
            else:
                raise ValueError(action)

        # Run full gibbs scans.
        assert len(self._added_rows) == num_rows
        for step in range(full_epochs):
            edges, edge_logits = self.sample_tree()
            self.set_edges(edges)
            for row_id in range(num_rows):
                self.remove_row(row_id)
                self.add_row(row_id)

        # Compute optimal tree.
        assert len(self._added_rows) == num_rows
        edges, edge_logits = self.estimate_tree()
        if self._config['learning_estimate_tree']:
            self.set_edges(edges)

        self._tree.gc()

        return {
            'config': self._config,
            'tree': self._tree,
            'edge_logits': edge_logits,
        }

@jit(nopython=True, cache=True)
def treecat_add_cell(
        feature_type,
        ragged_index,
        data_row,
        message,
        feat_probs,
        meas_probs,
        v, ):
    if feature_type == TY_MULTINOMIAL:
        beg, end = ragged_index[v:v + 2]
        feat_block = feat_probs[beg:end, :]
        meas_block = meas_probs[v, :]
        for c, count in enumerate(data_row[beg:end]):
            for _ in range(count):
                message *= feat_block[c, :]
                message /= meas_block
                feat_block[c, :] += np.float32(1)
                meas_block += np.float32(1)
    else:
        raise NotImplementedError

@profile
@jit(nopython=True, cache=True)
def treecat_add_row(
        feature_types,
        ragged_index,
        data_row,
        tree_grid,
        program,
        assignments,
        vert_ss,
        edge_ss,
        feat_ss,
        meas_ss,
        vert_probs,
        edge_probs,
        feat_probs,
        meas_probs, ):
    # Sample latent assignments using dynamic programming.
    messages = vert_probs.copy()
    for i in range(len(program)):
        op, v, v2, e = program[i]
        message = messages[v, :]
        if op == OP_UP:
            # Propagate upward from observed to latent.
            treecat_add_cell(
                feature_types[v],
                ragged_index,
                data_row,
                message,
                feat_probs,
                meas_probs,
                v, )
        elif op == OP_IN:
            # Propagate latent state inward from children to v.
            trans = edge_probs[e, :, :]
            if v > v2:
                trans = trans.T
            message *= np.dot(trans, messages[v2, :] / vert_probs[v2, :])
            message /= vert_probs[v, :]
            message /= message.max()  # Scale for numerical stability.
        elif op == OP_ROOT:
            # Process root node.
            assignments[v] = sample_from_probs(message)
        elif op == OP_OUT:
            # Propagate latent state outward from parent to v.
            trans = edge_probs[e, :, :]
            if v2 > v:
                trans = trans.T
            message *= trans[assignments[v2], :]
            message /= vert_probs[v, :]
            assignments[v] = sample_from_probs(message)

    # Update sufficient statistics.
    for v, m in enumerate(assignments):
        vert_ss[v, m] += 1
    for e in range(tree_grid.shape[1]):
        m1 = assignments[tree_grid[1, e]]
        m2 = assignments[tree_grid[2, e]]
        edge_ss[e, m1, m2] += 1
        edge_probs[e, m1, m2] += 1
    for v, m in enumerate(assignments):
        feature_type = feature_types[v]
        if feature_type == TY_MULTINOMIAL:
            beg, end = ragged_index[v:v + 2]
            feat_ss[beg:end, m] += data_row[beg:end]
            meas_ss[v, m] += data_row[beg:end].sum()
        else:
            raise NotImplementedError

@profile
@jit(nopython=True, cache=True)
def treecat_remove_row(
        feature_types,
        ragged_index,
        data_row,
        tree_grid,
        assignments,
        vert_ss,
        edge_ss,
        feat_ss,
        meas_ss,
        edge_probs, ):
    # Update sufficient statistics.
    for v, m in enumerate(assignments):
        vert_ss[v, m] -= 1
    for e in range(tree_grid.shape[1]):
        m1 = assignments[tree_grid[1, e]]
        m2 = assignments[tree_grid[2, e]]
        edge_ss[e, m1, m2] -= 1
        edge_probs[e, m1, m2] -= 1
    for v, m in enumerate(assignments):
        feature_type = feature_types[v]
        if feature_type == TY_MULTINOMIAL:
            beg, end = ragged_index[v:v + 2]
            feat_ss[beg:end, m] -= data_row[beg:end]
            meas_ss[v, m] -= data_row[beg:end].sum()
        else:
            raise NotImplementedError

@jit(nopython=True, cache=True)
def treecat_compute_edge_logit(M, gammaln_table, assign1, assign2):
    counts = np.zeros((M, M), np.int32)
    for n in range(assign1.shape[0]):
        counts[assign1[n], assign2[n]] += 1
    result = np.float32(0)
    for m1 in range(M):
        for m2 in range(M):
            result += gammaln_table[counts[m1, m2]]
    return result

@jit(nopython=True, parallel=True)
def treecat_compute_edge_logits_par(M, grid, gammaln_table, assignments,
                                    vert_logits):
    K = grid.shape[1]
    N, V = assignments.shape
    edge_logits = np.zeros(K, np.float32)
    for k in prange(K):
        v1, v2 = grid[1:3, k]
        assign1 = assignments[:, v1]
        assign2 = assignments[:, v2]
        edge_logit = treecat_compute_edge_logit(M, gammaln_table, assign1,
                                                assign2)
        edge_logits[k] = edge_logit - vert_logits[v1] - vert_logits[v2]
    return edge_logits

@jit(nopython=True, cache=True)
def treecat_compute_edge_logits_seq(M, grid, gammaln_table, assignments,
                                    vert_logits):
    K = grid.shape[1]
    N, V = assignments.shape
    edge_logits = np.zeros(K, np.float32)
    for k in range(K):
        v1, v2 = grid[1:3, k]
        assign1 = assignments[:, v1]
        assign2 = assignments[:, v2]
        edge_logit = treecat_compute_edge_logit(M, gammaln_table, assign1,
                                                assign2)
        edge_logits[k] = edge_logit - vert_logits[v1] - vert_logits[v2]
    return edge_logits

@profile
def treecat_compute_edge_logits(M, grid, gammaln_table, assignments,
                                vert_logits, parallel):
    if parallel:
        return treecat_compute_edge_logits_par(M, grid, gammaln_table,
                                               assignments, vert_logits)
    else:
        return treecat_compute_edge_logits_seq(M, grid, gammaln_table,
                                               assignments, vert_logits)

class TreeCatTrainer(TreeTrainer):
    """Class for training a TreeCat model."""

    def __init__(self, table, tree_prior, config):
        """Initialize a model with an empty subsample.

        Args:
            table: A Table instance holding N rows of V features of data.
            tree_prior: A [K]-shaped numpy array of prior edge log odds, where
                K is the number of edges in the complete graph on V vertices.
            config: A global config dict.
        """
        logger.info('TreeCatTrainer of %d x %d data', table.num_rows,
                    table.num_cols)
        assert isinstance(table, Table)
        N = table.num_rows  # Number of rows.
        V = table.num_cols  # Number of features, i.e. vertices.
        TreeTrainer.__init__(self, N, V, tree_prior, config)
        assert self._num_rows == N
        assert len(self._added_rows) == 0
        self._table = table
        self._assignments = np.zeros([N, V], dtype=np.int8)

        # These are useful dimensions to import into locals().
        E = V - 1  # Number of edges in the tree.
        K = V * (V - 1) // 2  # Number of edges in the complete graph.
        M = self._config['model_num_clusters']  # Clusters per latent.
        assert M <= 128, 'Invalid model_num_clusters > 128: {}'.format(M)
        self._VEKM = (V, E, K, M)

        # Use Jeffreys priors.
        self._vert_prior = 0.5
        self._edge_prior = 0.5 / M
        self._feat_prior = 0.5 / M
        self._meas_prior = self._feat_prior * np.array(
            [(table.ragged_index[v + 1] - table.ragged_index[v])
             for v in range(V)],
            dtype=np.float32).reshape((V, 1))
        self._gammaln_table = gammaln(
            np.arange(1 + N, dtype=np.float32) + self._edge_prior)
        assert self._gammaln_table.dtype == np.float32

        # Sufficient statistics are maintained always.
        self._vert_ss = np.zeros([V, M], np.int32)
        self._edge_ss = np.zeros([E, M, M], np.int32)
        self._feat_ss = np.zeros([table.ragged_index[-1], M], np.int32)
        self._meas_ss = np.zeros([V, M], np.int32)

        # Temporaries.
        self._vert_probs = np.empty(self._vert_ss.shape, np.float32)
        self._edge_probs = np.empty(self._edge_ss.shape, np.float32)
        self._feat_probs = np.empty(self._feat_ss.shape, np.float32)
        self._meas_probs = np.empty(self._meas_ss.shape, np.float32)

        # Maintain edge_probs.
        np.add(self._edge_ss, self._edge_prior, out=self._edge_probs)

    @profile
    def add_row(self, row_id):
        logger.debug('TreeCatTrainer.add_row %d', row_id)
        assert row_id not in self._added_rows, row_id
        self._added_rows.add(row_id)

        # These are used for scratch work, so we create them each step.
        np.add(self._vert_ss, self._vert_prior, out=self._vert_probs)
        np.add(self._feat_ss, self._feat_prior, out=self._feat_probs)
        np.add(self._meas_ss, self._meas_prior, out=self._meas_probs)

        treecat_add_row(
            self._table.feature_types,
            self._table.ragged_index,
            self._table.data[row_id, :],
            self._tree.tree_grid,
            self._program,
            self._assignments[row_id, :],
            self._vert_ss,
            self._edge_ss,
            self._feat_ss,
            self._meas_ss,
            self._vert_probs,
            self._edge_probs,
            self._feat_probs,
            self._meas_probs, )

    @profile
    def remove_row(self, row_id):
        logger.debug('TreeCatTrainer.remove_row %d', row_id)
        assert row_id in self._added_rows, row_id
        self._added_rows.remove(row_id)

        treecat_remove_row(
            self._table.feature_types,
            self._table.ragged_index,
            self._table.data[row_id, :],
            self._tree.tree_grid,
            self._assignments[row_id, :],
            self._vert_ss,
            self._edge_ss,
            self._feat_ss,
            self._meas_ss,
            self._edge_probs, )

    def set_edges(self, edges):
        TreeTrainer.set_edges(self, edges)
        V, E, K, M = self._VEKM
        assignments = self._assignments[sorted(self._added_rows), :]
        for e, v1, v2 in self._tree.tree_grid.T:
            self._edge_ss[e, :, :] = count_pairs(assignments, v1, v2, M)
        np.add(self._edge_ss, self._edge_prior, out=self._edge_probs)

    @profile
    def compute_edge_logits(self):
        """Compute non-normalized logprob of all V(V-1)/2 candidate edges.

        This is used for sampling and estimating the latent tree.
        """
        V, E, K, M = self._VEKM
        vert_logits = logprob_dc(self._vert_ss, self._vert_prior, axis=1)
        if len(self._added_rows) == V:
            assignments = self._assignments
        else:
            assignments = self._assignments[sorted(self._added_rows), :]
        assignments = np.array(assignments, order='F')
        parallel = self._config['learning_parallel']
        result = treecat_compute_edge_logits(M, self._tree.complete_grid,
                                             self._gammaln_table, assignments,
                                             vert_logits, parallel)
        result += self._tree_prior
        return result

    def logprob(self):
        """Compute non-normalized log probability of data and latent state.

        This is used for testing goodness of fit of the latent state kernel.
        """
        assert len(self._added_rows) == self._num_rows
        V, E, K, M = self._VEKM
        vertex_logits = logprob_dc(self._vert_ss, self._vert_prior, axis=1)
        logprob = vertex_logits.sum()
        for e, v1, v2 in self._tree.tree_grid.T:
            logprob += (logprob_dc(self._edge_ss[e, :, :], self._edge_prior) -
                        vertex_logits[v1] - vertex_logits[v2])
        for v in range(V):
            beg, end = self._table.ragged_index[v:v + 2]
            logprob += logprob_dc(self._feat_ss[beg:end, :], self._feat_prior)
            logprob -= logprob_dc(self._meas_ss[v, :], self._meas_prior[v])
        return logprob

    def train(self):
        """Train a TreeCat model using subsample-annealed MCMC.

        Returns:
            A trained model as a dictionary with keys:
                config: A global config dict.
                tree: A TreeStructure instance with the learned latent
                    structure.
                edge_logits: A [K]-shaped array of all edge logits.
                suffstats: Sufficient statistics of features, vertices, and
                    edges and a ragged_index for the features array.
                assignments: An [N, V]-shaped numpy array of latent cluster
                    ids for each cell in the dataset, where N be the number of
                    data rows and V is the number of features.
        """
        model = TreeTrainer.train(self)
        model['assignments'] = self._assignments
        model['suffstats'] = {
            'ragged_index': self._table.ragged_index,
            'vert_ss': self._vert_ss,
            'edge_ss': self._edge_ss,
            'feat_ss': self._feat_ss,
            'meas_ss': self._meas_ss,
        }
        return model

def treegauss_add_row(
        data_row,
        tree_grid,
        program,
        latent_row,
        vert_ss,
        edge_ss,
        feat_ss, ):
    # Sample latent state using dynamic programming.
    TODO('https://github.com/posterior/treecat/issues/26')

    # Update sufficient statistics.
    for v in range(latent_row.shape[0]):
        z = latent_row[v, :]
        vert_ss[v, :, :] += np.outer(z, z)
    for e in range(tree_grid.shape[1]):
        z1 = latent_row[tree_grid[1, e], :]
        z2 = latent_row[tree_grid[2, e], :]
        edge_ss[e, :, :] += np.outer(z1, z2)
    for v, x in enumerate(data_row):
        if np.isnan(x):
            continue
        z = latent_row[v, :]
        feat_ss[v] += 1
        feat_ss[v, 1] += x
        feat_ss[v, 2:] += x * z  # TODO Use central covariance.

def treegauss_remove_row(
        data_row,
        tree_grid,
        latent_row,
        vert_ss,
        edge_ss,
        feat_ss, ):
    # Update sufficient statistics.
    for v in range(latent_row.shape[0]):
        z = latent_row[v, :]
        vert_ss[v, :, :] -= np.outer(z, z)
    for e in range(tree_grid.shape[1]):
        z1 = latent_row[tree_grid[1, e], :]
        z2 = latent_row[tree_grid[2, e], :]
        edge_ss[e, :, :] -= np.outer(z1, z2)
    for v, x in enumerate(data_row):
        if np.isnan(x):
            continue
        z = latent_row[v, :]
        feat_ss[v] -= 1
        feat_ss[v, 1] -= x
        feat_ss[v, 2:] -= x * z  # TODO Use central covariance.

class TreeGaussTrainer(TreeTrainer):
    """Class for training a TreeGauss model."""

    def __init__(self, data, tree_prior, config):
        """Initialize a model with an empty subsample.

        Args:
            data: An [N, V]-shaped numpy array of real-valued data.
            tree_prior: A [K]-shaped numpy array of prior edge log odds, where
                K is the number of edges in the complete graph on V vertices.
            config: A global config dict.
        """
        assert isinstance(data, np.ndarray)
        data = np.asarray(data, np.float32)
        assert len(data.shape) == 2
        N, V = data.shape
        D = config['model_latent_dim']
        E = V - 1  # Number of edges in the tree.
        TreeTrainer.__init__(self, N, V, tree_prior, config)
        self._data = data
        self._latent = np.zeros([N, V, D], np.float32)

        # This is symmetric positive definite.
        self._vert_ss = np.zeros([V, D, D], np.float32)
        # This is arbitrary (not necessarily symmetric).
        self._edge_ss = np.zeros([E, D, D], np.float32)
        # This represents (count, mean, covariance).
        self._feat_ss = np.zeros([V, D, 1 + 1 + D], np.float32)

    def add_row(self, row_id):
        logger.debug('TreeGaussTrainer.add_row %d', row_id)
        assert row_id not in self._added_rows, row_id
        self._added_rows.add(row_id)

        treegauss_add_row(
            self._data[row_id, :],
            self._tree.tree_grid,
            self._program,
            self._latent[row_id, :, :],
            self._vert_ss,
            self._edge_ss,
            self._feat_ss, )

    def remove_row(self, row_id):
        logger.debug('TreeGaussTrainer.remove_row %d', row_id)
        assert row_id in self._added_rows, row_id
        self._added_rows.remove(row_id)

        treecat_remove_row(
            self._data[row_id, :],
            self._tree.tree_grid,
            self._latent[row_id, :, :],
            self._vert_ss,
            self._edge_ss,
            self._feat_s, )

    def set_edges(self, edges):
        TreeTrainer.set_edges(self, edges)
        latent = self._latent[sorted(self._added_rows), :, :]
        for e, v1, v2 in self._tree.tree_grid.T:
            self._edge_ss[e, :, :] = np.dot(latent[:, v1, :].T,
                                            latent[:, v2, :])

    def compute_edge_logits(self):
        """Compute non-normalized logprob of all V(V-1)/2 candidate edges.

        This is used for sampling and estimating the latent tree.
        """
        TODO('https://github.com/posterior/treecat/issues/26')

    def logprob(self):
        """Compute non-normalized log probability of data and latent state.

        This is used for testing goodness of fit of the latent state kernel.
        """
        assert len(self._added_rows) == self._num_rows
        TODO('https://github.com/posterior/treecat/issues/26')

    def train(self):
        """Train a TreeGauss model using subsample-annealed MCMC.

        Returns:
            A trained model as a dictionary with keys:
                config: A global config dict.
                tree: A TreeStructure instance with the learned latent
                    structure.
                edge_logits: A [K]-shaped array of all edge logits.
                suffstats: Sufficient statistics of features and vertices.
                latent: An [N, V, M]-shaped numpy array of latent states, where
                    N is the number of data rows, V is the number of features,
                    and M is the dimension of each latent variable.
        """
        model = TreeTrainer.train(self)
        model['latent'] = self._latent
        model['suffstats'] = {
            'vert_ss': self._vert_ss,
            'edge_ss': self._edge_ss,
            'feat_ss': self._feat_ss,
        }
        return model

class TreeMogTrainer(TreeTrainer):
    """Class for training a tree mixture-of-Gaussians model."""

    def __init__(self, data, tree_prior, config):
        TODO('https://github.com/posterior/treecat/issues/27')

    def add_row(self, row_id):
        """Add a given row to the current subsample."""
        TODO('https://github.com/posterior/treecat/issues/27')

    def remove_row(self, row_id):
        """Remove a given row from the current subsample."""
        TODO('https://github.com/posterior/treecat/issues/27')

    def compute_edge_logits(self):
        """Compute edge log probabilities on the complete graph."""
        TODO('https://github.com/posterior/treecat/issues/27')

    def logprob(self):
        """Compute non-normalized log probability of data and latent state."""
        assert len(self._added_rows) == self._num_rows
        TODO('https://github.com/posterior/treecat/issues/27')

def train_model(table, tree_prior, config):
    """Train a TreeCat model using subsample-annealed MCMC.

    Let N be the number of data rows and V be the number of features.

    Args:
        table: A Table instance holding N rows of V features of data.
        tree_prior: A [K]-shaped numpy array of prior edge log odds, where
            K is the number of edges in the complete graph on V vertices.
        config: A global config dict.

    Returns:
        A trained model as a dictionary with keys:
            config: A global config dict.
            tree: A TreeStructure instance with the learned latent structure.
            edge_logits: A [K]-shaped array of all edge logits.
            suffstats: Sufficient statistics of features, vertices, and edges.
            assignments: An [N, V] numpy array of latent cluster ids for each
                cell in the dataset.
    """
    assert isinstance(table, Table)
    M = config['model_num_clusters']
    D = config['model_latent_dim']
    assert M >= 1
    assert D >= 0
    if D == 0:
        Trainer = TreeCatTrainer
    elif M == 1:
        Trainer = TreeGaussTrainer
    else:
        Trainer = TreeMogTrainer
    return Trainer(table, tree_prior, config).train()

def _train_model(task):
    table, tree_prior, config = task
    return train_model(table, tree_prior, config)

def train_ensemble(table, tree_prior, config):
    """Train a TreeCat ensemble model using subsample-annealed MCMC.

    The ensemble size is controlled by config['model_ensemble_size'].
    Let N be the number of data rows and V be the number of features.

    Args:
        table: A Table instance holding N rows of V features of data.
        tree_prior: A [K]-shaped numpy array of prior edge log odds, where
            K is the number of edges in the complete graph on V vertices.
        config: A global config dict.

    Returns:
        A trained model as a dictionary with keys:
            tree: A TreeStructure instance with the learned latent structure.
            suffstats: Sufficient statistics of features, vertices, and edges.
            assignments: An [N, V] numpy array of latent cluster ids for each
                cell in the dataset.
    """
    tasks = []
    for sub_seed in range(config['model_ensemble_size']):
        sub_config = config.copy()
        sub_config['seed'] += sub_seed
        tasks.append((table, tree_prior, sub_config))
    return parallel_map(_train_model, tasks)


from __future__ import absolute_import
from __future__ import print_function

import logging
import multiprocessing
from rich import traceback

traceback.install()

def get_detail_logger(name: str, multi: bool = False) -> logging.Logger:
    if multi:
        logger = multiprocessing.get_logger()
        logger.setLevel(logging.INFO)
    else:
        logger = logging.getLogger(name)
        logger.setLevel(logging.INFO)

    logging_format = logging.Formatter(
        "%(asctime)-25s"
        "%(levelname)-10s"
        "%(filename)s:%(lineno)-20d"
        "%(process)d:%(thread)-20d"
        "%(message)s"
    )
    stream_handler = logging.StreamHandler()
    stream_handler.setFormatter(logging_format)
    logger.addHandler(stream_handler)
    return logger

class DetailLogger(object):
    """Preset logging module's logger.

    Attributes:
        _multi: bool, if true, multi processing compatible logger kicks in.
        _nib: logging.Logger, the single or multi processing logger.
    """

    def __init__(self,
                 name: str,
                 multi: bool = False) -> None:
        self._multi = multi

        if self._multi:
            self._nib = multiprocessing.get_logger()
            self._nib.setLevel(logging.INFO)
        else:
            self._nib = logging.getLogger(name)
            self._nib.setLevel(logging.INFO)

        if not self._nib.handlers:
            logging_format = logging.Formatter(
                "%(asctime)-25s"
                "%(levelname)-10s"
                "%(filename)s:%(lineno)-20d"
                "%(process)d:%(thread)-20d"
                "%(message)s"
            )
            stream_handler = logging.StreamHandler()
            stream_handler.setFormatter(logging_format)
            self._nib.addHandler(stream_handler)

    def debug(self, message: str) -> None:
        self._nib.debug(message)

    def info(self, message: str) -> None:
        self._nib.info(message)

    def warning(self, message: str) -> None:
        self._nib.warning(message)

    def error(self, message: str) -> None:
        self._nib.error(message)

    def critical(self, message: str) -> None:
        self._nib.critical(message)

    def progress(self,
                current_iter,
                total_iter,
                prefix="Progress",
                suffix="Complete",
                decimals=2,
                total_length=50):
        """ Terminal progress bar to stdout.

        Args:
            current_iter: int, current iteration count.
            total_iter: int, total iteration count.
            prefix: str, prefix string.
            suffix: str, suffix string.
            decimals: int, positive number of decimals in percent.
            total_length: int, leng of bar.
        """

        percent = ("{0:." + str(decimals) + "f}").format(
            100 * (current_iter / float(total_iter)))
        filled_length = int(total_length * current_iter // total_iter)
        bars = ("█" * filled_length) + ("-" * (total_length - filled_length))
        self._nib.info(f"\r{prefix} |{bars}| {percent}% {suffix}")
        if current_iter == total_iter:
            self._nib.info("\n")


import copy
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '3'
import random
import time

import numpy as np
from torch import optim
import torch
import torch.nn as nn
import torch.utils.data
import torch.nn.utils.rnn as rnn_utils
import pickle
import pandas as pd
import torch.nn.functional as F

class EmbeddingModel(nn.Module):
    def __init__(self, vocab_size, embed_size):
        super(EmbeddingModel, self).__init__()

        self.vocab_size = vocab_size
        self.embed_size = embed_size

        self.in_embed = nn.Embedding(self.vocab_size, self.embed_size)
        self.out_embed = nn.Embedding(self.vocab_size, self.embed_size)

    def forward(self, input_labels, pos_labels, neg_labels):
        ''' input_labels: center words, [batch_size]
            pos_labels: positive words, [batch_size, (window_size * 2)]
            neg_labels：negative words, [batch_size, (window_size * 2 * K)]

            return: loss, [batch_size]
        '''
        input_embedding = self.in_embed(input_labels)  # [batch_size, embed_size]
        pos_embedding = self.out_embed(pos_labels)  # [batch_size, (window * 2), embed_size]
        neg_embedding = self.out_embed(neg_labels)  # [batch_size, (window * 2 * K), embed_size]

        input_embedding = input_embedding.unsqueeze(2)  # [batch_size, embed_size, 1]

        pos_dot = torch.bmm(pos_embedding, input_embedding)  # [batch_size, (window * 2), 1]
        pos_dot = pos_dot.squeeze(2)  # [batch_size, (window * 2)]

        neg_dot = torch.bmm(neg_embedding, -input_embedding)  # [batch_size, (window * 2 * K), 1]
        neg_dot = neg_dot.squeeze(2)  # batch_size, (window * 2 * K)]

        log_pos = F.logsigmoid(pos_dot).sum(1)  # .sum()结果只为一个数，.sum(1)结果是一维的张量
        log_neg = F.logsigmoid(neg_dot).sum(1)

        loss = log_pos + log_neg

        return -loss

    def input_embedding(self):
        return self.in_embed.weight.detach().numpy()

class pro_DataSet(torch.utils.data.Dataset):
    def __init__(self, filename, pro_num):
        # with open(r'E:\pycharmworkSpace\StatementContrastive\identifier\token_vec\idx2word.pkl', 'rb') as f:
        #     self.idx2word = pickle.load(f)
        # with open(r'E:\pycharmworkSpace\StatementContrastive\identifier\token_vec\word2idx.pkl', 'rb') as f1:
        #     self.word2idx = pickle.load(f1)
        # MAX_VOCAB_SIZE = 10000
        # EMBEDDING_SIZE = 512
        # model = EmbeddingModel(MAX_VOCAB_SIZE, EMBEDDING_SIZE)
        # model.load_state_dict(torch.load(r'E:\pycharmworkSpace\StatementContrastive\identifier\token_vec\embedding-512.th'))
        # self.embedding_weights = model.input_embedding()
        de = pd.read_json(filename, lines=True)
        # 给每个token嵌入词向量
        data_list = list()

        # for i in range(4):
        for i in range(len(de)):

            if i % 2 == 1:
                p_num = None
                p_name = None
                codeSquence = None
                covers = None
                mutants = None
                feature_labels = None
                sen_labels = None
                for a in de.loc[i, ["pro_num"]].values:
                    p_num = a
                for a in de.loc[i, ["pro_name"]].values:
                    p_name = a
                for a in de.loc[i, ["codeSquence"]].values:
                    codeSquence = a
                for a in de.loc[i, ["Mutant_cover1"]].values:
                    mutants = a
                for a in de.loc[i, ["codecover"]].values:
                    covers = a
                # for a in de.loc[i, ["seq_label"]].values:
                #     feature_labels = a
                for a in de.loc[i, ["sen_label"]].values:
                    sen_labels = a
                this_name = str(p_name) + str(p_num)
                if this_name in pro_num:
                    print(i)
                    continue

                for a in range(len(codeSquence)):
                    data_dict = {
                        "token_seq": None,
                        # "tokenseq_len": None,
                        "cover": None,
                        "mutant": None,
                        "sen_labels": None
                    }
                    token_seq = list()
                    for b in range(8):
                        try:
                            token_seq.append(codeSquence[a][b])
                        except:
                            token_seq.append("<pad>")

                    tokenseq_len = len(token_seq)
                    #                     print(covers[a])
                    cc = covers[a].split(",")

                    mutant = list()

                    if len(mutants) < 32:
                        for b in range(len(mutants)):
                            lst = list()

                            for c in range(512):
                                try:
                                    lst.append(float(mutants[b][a][c]))
                                except:
                                    lst.append(0)
                            mutant.append(lst)
                        for b in range(32 - len(mutants)):
                            lst = list()
                            for c in range(512):
                                try:
                                    lst.append(float(mutants[b][a][c]))
                                except:
                                    lst.append(0)
                            mutant.append(lst)
                    else:
                        for b in range(32):
                            lst = list()

                            for c in range(512):
                                try:
                                    lst.append(float(mutants[b][a][c]))
                                except:
                                    lst.append(0)
                            mutant.append(lst)

                    cover = list()
                    for b in range(512):
                        try:
                            if cc[b] == "1":
                                cover.append(1)
                            elif cc[b] == "-1":
                                cover.append(-1)
                            else:
                                cover.append(0)
                        except:
                            cover.append(0)

                    sen_label = np.array(sen_labels[a])
                    data_dict["token_seq"] = token_seq
                    # data_dict["tokenseq_len"] =tokenseq_len
                    data_dict["cover"] = cover
                    data_dict["mutant"] = np.array(mutant)
                    data_dict["sen_labels"] = sen_label
                    data_list.append(data_dict)
        self.data = data_list

    def __getitem__(self, idx):  # if the index is idx, what will be the data?
        return self.data[idx]["token_seq"], self.data[idx]["cover"], self.data[idx]["mutant"], self.data[idx][
            "sen_labels"]

    def __len__(self):  # What is the length of the dataset
        return len(self.data)

class CNN2(nn.Module):
    def __init__(self, in_channel=1, feature=15, num_class=32):
        super(CNN2, self).__init__()
        self.Conv1 = nn.Conv2d(in_channel, feature, kernel_size=3, stride=1, padding=1)
        self.Relu1 = nn.ReLU()
        self.Pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Fc = nn.Linear(61440, num_class)

    def forward(self, x):
        x = self.Conv1(x)  # [B,1,32,512] --> [B,15,32,512]
        x = self.Relu1(x)
        x = self.Pool1(x)  # [B,15,32,512] --> [B,15,16,256]
        x = x.reshape(x.shape[0], -1)
        x = self.Fc(x)
        return x

class CNN1(nn.Module):
    def __init__(self, in_channel=1, feature=15, num_class=32):
        super(CNN1, self).__init__()
        self.Conv1 = nn.Conv1d(in_channel, feature, kernel_size=3, stride=1, padding=1)
        self.Relu1 = nn.ReLU()
        self.Pool1 = nn.MaxPool1d(kernel_size=2, stride=2)
        self.Fc = nn.Linear(3840, num_class)

    def forward(self, x):
        x = self.Conv1(x)
        x = self.Relu1(x)
        x = self.Pool1(x)
        x = x.reshape(x.shape[0], -1)
        x = self.Fc(x)
        return x

class MLP(nn.Module):
    def __init__(self, num_input, num_hidden, num_class=32):
        super(MLP, self).__init__()

        self.linear1 = nn.Linear(num_input, num_hidden)
        self.linear2 = nn.Linear(num_hidden, num_class)
        self.Relu1 = nn.ReLU()
        self.Relu2 = nn.ReLU()

    def forward(self, x):
        x = self.linear1(x)
        x = self.Relu1(x)
        x = self.linear2(x)
        x = self.Relu2(x)
        return x

class FullModel(nn.Module):
    def __init__(self, CNN1, CNN2, MLP, in_channel=1, feature=15, num_class=1):
        super().__init__()
        self.CNN1 = CNN1
        self.CNN2 = CNN2
        self.MLP = MLP

        self.Conv1 = nn.Conv3d(in_channel, feature, kernel_size=3, stride=1, padding=1)
        self.Relu1 = nn.ReLU()
        self.Pool1 = nn.MaxPool3d(kernel_size=2, stride=2)
        self.Fc = nn.Linear(61440, 1)

    def forward(self, cover, mutant, word):
        cov = self.CNN1.forward(cover)
        mut = self.CNN2.forward(mutant)
        wo = self.MLP.forward(word)
        cov = cov.unsqueeze(1).unsqueeze(1)
        mut = mut.unsqueeze(1).unsqueeze(1)
        wo = wo.unsqueeze(1).unsqueeze(1)
        cov = torch.reshape(cov, (cov.shape[0], cov.shape[1], cov.shape[3], cov.shape[2]))
        mut = torch.reshape(mut, (mut.shape[0], mut.shape[3], mut.shape[1], mut.shape[2]))

        c = torch.mul(cov, mut)

        b = torch.mul(c, wo)  # [15, 1, 3, 3, 3]
        b = b.unsqueeze(1)

        x = self.Conv1(b)

        x = self.Relu1(x)
        x = self.Pool1(x)
        x = x.reshape(x.shape[0], -1)
        x = self.Fc(x)
        x = torch.sigmoid(x)
        return x

# self.data[idx]["token_seq"],self.data[idx]["cover"],self.data[idx]["mutant"],self.data[idx]["sen_labels"]

if __name__ == '__main__':
    device = torch.device('cuda:0')

    def init_weights(m):
        for name, param in m.named_parameters():
            nn.init.uniform_(param.data, -0.08, 0.08)

    def epoch_time(start_time, end_time):
        elapsed_time = end_time - start_time
        elapsed_mins = int(elapsed_time / 60)
        elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
        return elapsed_mins, elapsed_secs

    with open(r'data/token_vec/idx2word.pkl', 'rb') as f2:
        Tidx2word = pickle.load(f2)
    with open(r'data/token_vec/word2idx.pkl', 'rb') as f3:
        Tword2idx = pickle.load(f3)
    MAX_VOCAB_SIZE = 10000
    EMBEDDING_SIZE = 512
    model1 = EmbeddingModel(MAX_VOCAB_SIZE, EMBEDDING_SIZE)
    model1.load_state_dict(torch.load(r'data/token_vec/embedding-512.th'))
    Tembedding_weights = model1.input_embedding()

    def embedding_Tword(word):
        if word == "<pad>":
            embedding = torch.zeros(512)
        else:
            try:
                index = Tword2idx[word]
            except:
                index = Tword2idx['<UNK>']
            embedding = torch.tensor(Tembedding_weights[index])
        return embedding

    def train(model, iterator, optimizer, criterion, clip, batch_s):

        # model.train()
        epoch_loss = 0
        for i, (token_seq, cover, mutant, sen_label) in enumerate(iterator):
            cover = torch.tensor([item.cpu().detach().numpy() for item in cover], dtype=torch.float)
            cover = cover.unsqueeze(1)
            cover = torch.reshape(cover, (cover.shape[2], cover.shape[1], cover.shape[0])).to(device)
            mutant = mutant.unsqueeze(1)
            mutant = mutant.float().to(device)

            #             mutant = torch.Tensor(mutant)
            #             sen_label = torch.Tensor(sen_label)
            #             sen_label = torch.from_numpy(sen_label)
            tss = list()

            for b in range(batch_s):
                ts = None
                for seq in range(len(token_seq)):
                    if seq == 0:
                        ts = embedding_Tword(token_seq[seq][b])
                    else:
                        ts = torch.cat((ts, embedding_Tword(token_seq[seq][b])), dim=-1)
                tss.append(ts)
            token_seq = torch.tensor([item.cpu().detach().numpy() for item in tss], dtype=torch.float).to(device)
            optimizer.zero_grad()

            output = model.forward(cover, mutant, token_seq)

            sen_label = sen_label.unsqueeze(1).float().to(device)

            loss = criterion.forward(output, sen_label)

            loss.backward()

            # gradient clipping 防止梯度爆炸问题
            torch.nn.utils.clip_grad_norm_(model.parameters(), clip)

            optimizer.step()

            epoch_loss += loss.item()
        return epoch_loss / len(iterator)

    SEED = 1234
    random.seed(SEED)
    torch.manual_seed(SEED)
    torch.backends.cudnn.deterministic = True
    BATCH_SIZE = 2
    c1 = CNN1(in_channel=1, feature=15, num_class=32).to(device)
    c2 = CNN2(in_channel=1, feature=15, num_class=32).to(device)
    mlp = MLP(4096, 512, 32).to(device)
    fmodel = FullModel(c1, c2, mlp, in_channel=1, feature=15, num_class=2).to(device)
    # b = np.load('a3.npy')
    # b = b.tolist()
    # dataset1 = pro_DataSet(r"data/data_new1.json",pro_num=b)  # create the dataset
    dataset1 = pro_DataSet(r"data/data_new3.json", pro_num=[])  # create the dataset
    dataloader = torch.utils.data.DataLoader(dataset=dataset1, shuffle=True, batch_size=BATCH_SIZE, drop_last=True)
    # fmodel.apply(init_weights)
    optimizer = optim.Adam(fmodel.parameters())
    N_EPOCHS = 10
    CLIP = 1
    criterion = nn.BCELoss()
    for epoch in range(N_EPOCHS):
        start_time = time.time()

        train_loss = train(fmodel, dataloader, optimizer, criterion, CLIP, BATCH_SIZE)
        # valid_loss = evaluate(model, dataloader1, criterion)

        end_time = time.time()

        epoch_mins, epoch_secs = epoch_time(start_time, end_time)
        torch.save(fmodel, 'tut3-model.pt')

        print(f'Epoch: {epoch + 1:02} | Time: {epoch_mins}m {epoch_secs}s')
        print(f'\tTrain Loss: {train_loss:.6f}')


from enum import Enum
from logging import Logger
from typing import Any

from lambda_pipeline.types import FrozenDict, LambdaContext, PipelineData

from nrlf.core.common_search_steps import get_paginated_document_references
from nrlf.core.common_steps import make_common_log_action, parse_headers
from nrlf.core.event_parsing import fetch_body_from_event
from nrlf.core.model import APIGatewayProxyEventModel, ConsumerRequestParams
from nrlf.core.repository import Repository
from nrlf.core.transform import create_bundle_from_paginated_response

log_action = make_common_log_action()

class LogReference(Enum):
    SEARCHPOST001 = "Searching for document references"

@log_action(log_reference=LogReference.SEARCHPOST001)
def search_document_references(
    data: PipelineData,
    context: LambdaContext,
    event: APIGatewayProxyEventModel,
    dependencies: FrozenDict[str, Any],
    logger: Logger,
) -> PipelineData:

    body = fetch_body_from_event(event)
    requestParams = ConsumerRequestParams(**body)
    repo: Repository = dependencies["repository"]

    response = get_paginated_document_references(
        request_params=requestParams,
        query_string_params=body,
        repository=repo,
        type_identifier=requestParams.type,
        raw_pointer_types=data["pointer_types"],
        nhs_number=requestParams.nhs_number,
    )
    bundle = create_bundle_from_paginated_response(response)
    return PipelineData(bundle)

steps = [
    parse_headers,
    search_document_references,
]


import networkx as nx
from fa2 import ForceAtlas2
import matplotlib.pyplot as plt

G = nx.random_geometric_graph(400, 0.2)

forceatlas2 = ForceAtlas2(
                        # Behavior alternatives
                        outboundAttractionDistribution=True,  # Dissuade hubs
                        linLogMode=False,  # NOT IMPLEMENTED
                        adjustSizes=False,  # Prevent overlap (NOT IMPLEMENTED)
                        edgeWeightInfluence=1.0,

                        # Performance
                        jitterTolerance=1.0,  # Tolerance
                        barnesHutOptimize=True,
                        barnesHutTheta=1.2,
                        multiThreaded=False,  # NOT IMPLEMENTED

                        # Tuning
                        scalingRatio=2.0,
                        strongGravityMode=False,
                        gravity=1.0,

                        # Log
                        verbose=True)

positions = forceatlas2.forceatlas2_networkx_layout(G, pos=None, iterations=2000)
nx.draw_networkx_nodes(G, positions, node_size=20, label=False, node_color="blue", alpha=0.4)
nx.draw_networkx_edges(G, positions, edge_color="green", alpha=0.05)
plt.axis('off')
plt.show()