"""Dynamics model environment base class.

Use this class if you want to build an environment
based on a model of class `BaseDynamicsModel`.
"""
import os
from abc import ABC, abstractmethod
from typing import Dict, Union, List, Tuple, Optional, Callable, Any

import gym
import numpy as np

from agents import base_agent
from agents.base_agent import RL_MODEL_DIR
from dynamics.base_model import BaseDynamicsModel
from util.numerics import npf32
from util.util import Arr, create_dir, make_param_ext, str_to_np_dt, Num, day_offset_ts, ts_per_day, ProgWrap, \
    load_pickle, save_pickle
from util.visualize import rl_plot_path, plot_env_evaluation, plot_reward_details, OVERLEAF_IMG_DIR, MergeListT, \
    eval_env_evaluation

Agents = Union[List[base_agent.AgentBase], base_agent.AgentBase]

DEFAULT_ENV_SAMPLE_DATA: str = "train"

RejSamplerFun = Callable[[np.ndarray], bool]

class RejSampler:
    """Rejection sampling class.

    Can be used for unbalanced environment resets by
    rejecting unwanted states. Will impact performance if
    a lot are rejected.
    """

    name: str = "Rej"
    max_rej: int = 1

    def __init__(self, name: str, fun: RejSamplerFun = lambda x: True,
                 max_rej: int = 1):
        self.name = f"{name}{max_rej}"
        self.fun = fun
        self.max_rej = max_rej

    def __call__(self, hist: np.ndarray):
        return self.fun(hist)

class DynEnv(ABC, gym.Env):
    """The environment wrapper class for `BaseDynamicsModel`.

    Takes an instance of `BaseDynamicsModel` and adds
    all the functionality needed to turn it into an environment
    to be used for reinforcement learning.
    """
    # Fix data
    m: BaseDynamicsModel  #: Prediction model.
    act_dim: int  #: The dimension of the action space.
    state_dim: int  #: The dimension of the state space.
    n_ts_per_eps: int  #: The maximum number of timesteps per episode.
    n_ts_per_day: int  #: Number of time steps in a day.
    t_init_n: int  #: The index of the initial timestep in the dataset.

    short_name: str = None

    # State data, might change if `step` is called.
    n_ts: int = 0  #: The current number of timesteps.
    hist: np.ndarray  #: 2D array with current state.

    # The current data to sample initial conditions from.
    train_data: np.ndarray  #: The training data.
    train_indices: np.ndarray  #: The indices corresponding to `train_data`.
    n_start_data: int  #: The number of possible initializations using the training data.

    # Info about the current episode.
    use_noise: bool = False  #: Whether to add noise when simulating.
    curr_ind: int = None
    curr_n: int = None

    orig_actions: np.ndarray  #: Array with fallback actions

    # The one day data
    n_train_days: int  #: Number of training days
    train_days: List  #: The data of all days, where all data is available.
    day_inds: np.ndarray  #: Index vector storing the timestep offsets to the days
    day_ind: int = 0  #: The index of the day in `train_days`.

    # Scaling info
    info: Dict = None  #: A dict with info about the current agent.
    do_scaling: bool = False
    a_scaling_pars: Tuple[np.ndarray, np.ndarray] = None  #: Extra scaling parameters.

    # Plotting default values
    default_state_mask: np.ndarray = None
    default_series_merging: MergeListT = None

    _dummy_use: bool
    _plot_path: str
    _disturbance_is_modelled: bool = False

    def __init__(self, m: BaseDynamicsModel, name: str = None, max_eps: int = None,
                 disturb_fac: float = 1.0,
                 init_res: bool = True,
                 dummy_use: bool = False,
                 sample_from: str = DEFAULT_ENV_SAMPLE_DATA,
                 rejection_sampler: RejSampler = None,
                 verbose: int = 0):
        """Initialize the environment.

        Args:
            m: Full model predicting all the non-control features.
            max_eps: Number of continuous predictions in an episode.
        """
        self._dummy_use = dummy_use
        self.verbose = verbose
        if not self._dummy_use:
            # m.model_disturbance()
            pass
        self.m = m
        self.rej_sampler = rejection_sampler

        # Define name extensions
        sam_ext = f"_SAM_{sample_from}" if sample_from != DEFAULT_ENV_SAMPLE_DATA else ""
        h_sam_ext = f"_RejS_{rejection_sampler.name}" if rejection_sampler is not None else ""

        if name is not None:
            dist_ex = make_param_ext([("DF", disturb_fac)])
            self.name = f"{name}{dist_ex}{sam_ext}{h_sam_ext}_DATA_{m.name}"
        else:
            self.name = f"RLEnv_{m.name}"
        # self.plot_path = os.path.join(rl_plot_path, self.name)
        self._plot_path = os.path.join(rl_plot_path, self.name)

        # Set attributes.
        dat = m.data
        self.disturb_fac = disturb_fac
        self.act_dim = dat.n_c
        self.state_dim = dat.d

        # Time indices
        self.n_ts_per_eps = 100 if max_eps is None else max_eps
        self.t_init_n = day_offset_ts(dat.t_init, dat.dt,
                                      remaining=False)
        self.n_ts_per_day = ts_per_day(dat.dt)

        # Set data and initialize env.
        self._set_data(sample_from)
        if init_res:
            self.reset(use_noise=False)

    def __str__(self):
        """Generic string conversion."""
        n = self.short_name if self.short_name is not None else self.name
        return f"RL env of class {self.__class__.__name__} with name {n}"

    @property
    def plot_path(self) -> str:
        """Returns the path of the plot folder."""
        if self.short_name is not None:
            return os.path.join(rl_plot_path, self.short_name)
        return self._plot_path

    @property
    def model_path(self) -> str:
        """Returns the path of the model folder."""
        if self.short_name is not None:
            return os.path.join(RL_MODEL_DIR, self.short_name)
        else:
            raise ValueError("No short name specified!!!")

    def set_agent(self, a: base_agent.AgentBase) -> None:
        """Sets the given agent to the environment.

        Handles the extra scaling if it is required by the
        agent, the environment does not actually get access to the
        agent.

        Args:
            a: The agent to be used with the env.
        """
        # Get info about agent and check whether scaling is needed
        self.info = a.get_info()
        scaling = self.info.get('action_scaled_01')
        self.do_scaling = scaling is not None

        # Handle scaling
        if self.do_scaling:
            p1 = np.array([i[0] for i in scaling], dtype=np.float32)
            p2 = np.array([i[1] - i[0] for i in scaling], dtype=np.float32)
            self.a_scaling_pars = p1, p2

    def _set_data(self, part_str: str = "train") -> None:
        """Sets the data to use in the env.

        Args:
            part_str: The string specifying the part of the data.
        """
        self.m.data.check_part(part_str)
        self.train_data, _, self.train_indices = self.m.data.get_split(part_str)
        self.n_start_data = len(self.train_data)

    def get_plt_path(self, name: str) -> str:
        """Specifies the path of the plot with name 'name' where it should be saved.

        If there is not a directory
        for the current model, it is created.

        Args:
            name: Name of the plot.

        Returns:
            Full path of the plot file.
        """
        dir_name = self.plot_path
        create_dir(dir_name)
        return os.path.join(dir_name, name)

    def get_model_path(self, mod_name: str) -> str:
        """Returns the model path."""
        dir_name = self.model_path
        create_dir(dir_name)
        return os.path.join(dir_name, mod_name)

    reward_descs: List = []  #: The description of the detailed reward.

    @abstractmethod
    def detailed_reward(self, curr_pred: np.ndarray, action: Arr) -> np.ndarray:
        """Computes the different components of the reward and returns them all.

        Args:
            curr_pred: The current predictions.
            action: The most recent action taken.

        Returns:
            The reward components in an 1d array.
        """
        pass

    def compute_reward(self, curr_pred: np.ndarray, action: Arr) -> float:
        """Computes the reward to be maximized during the RL training.

        Args:
            curr_pred: The current predictions.
            action: The most recent action taken.

        Returns:
            Value of the reward of that outcome.
        """
        # The base implementation just sums up the different rewards.
        return np.sum(self.detailed_reward(curr_pred, action)).item()

    @abstractmethod
    def episode_over(self, curr_pred: np.ndarray) -> bool:
        """Defines criterion for episode to be over.

        Args:
            curr_pred: The next predicted state.

        Returns:
            True if episode is over, else false
        """
        return False

    def scale_action_for_step(self, action: Arr) -> Arr:
        """Scales the action to be used by the underlying ML model.

        This does nothing, override if needed.

        Args:
            action: The action to be scaled.
        Returns:
            The scaled action, same shape as the input.
        """
        return action

    @property
    def dummy_use(self):
        return self._dummy_use

    def step(self, action: Arr) -> Tuple[np.ndarray, float, bool, Dict]:
        """Evolve the model with the given control input `action`.

        This function should not be overridden, instead override
        `scale_action_for_step`! If it is, it should be called at some
        point in the overriding method.

        Args:
            action: The control input (action), will be scaled for the underlying model
                of the environment.

        Returns:
            The next state, the reward of having chosen that action and a bool
            determining if the episode is over. (And an empty dict)
        """
        if self._dummy_use:
            raise ValueError("Stepping not supported for dummy usage!")

        # Scale the action
        action = self.scale_action_for_step(action)

        # Add the control to the history
        self.hist[-1, -self.act_dim:] = action

        # Predict using the model
        hist_res = np.copy(self.hist).reshape((1, -1, self.state_dim))
        curr_pred = self.m.predict(hist_res)[0]

        # Add noise
        if self.use_noise:
            assert self._disturbance_is_modelled, "Need to model disturbance first!"
            curr_pred += self.disturb_fac * self.m.disturb()

        # Save the chosen action
        self.orig_actions[self.n_ts, :] = action

        # Update history
        self.hist[:-1, :] = self.hist[1:, :]
        self.hist[-1, :-self.act_dim] = curr_pred
        self.n_ts += 1

        # Compute reward and episode termination
        r = np.array(self.compute_reward(curr_pred, action)).item()
        ep_over = self.n_ts == self.n_ts_per_eps or self.episode_over(curr_pred)
        return curr_pred, r, ep_over, {}

    @property
    def curr_state(self) -> np.ndarray:
        """Returns the current state.

        Returns:
            Current state of environment, 1d array.
        """
        return self.hist[-1, :-self.act_dim]

    def scale_state(self, state: np.ndarray,
                    remove_mean: bool = False) -> np.ndarray:
        """Scales the state using the Dataset."""
        return self.m.data.scale(state, remove_mean=remove_mean, state_only=True)

    def get_curr_day_n(self) -> int:
        """Returns the current timestep index."""
        return (self.curr_n + self.n_ts) % self.n_ts_per_day

    def _model_disturbance(self):
        """Models the disturbance of the underlying model."""
        with ProgWrap("Modeling disturbance...", self.verbose > 0):
            self.m.model_disturbance()
            self._disturbance_is_modelled = True

    def reset(self, start_ind: int = None, use_noise: bool = True) -> np.ndarray:
        """Resets the environment.

        Always needs to be called if the episode is over.
        Initializes the environment with a new initial state.
        Uses rejection sampling to favor some initial conditions over
        others, as defined by `self.rej_sampler`.

        Args:
            start_ind: The index specifying the initial condition.
            use_noise: Whether noise should be added when simulating.

        Returns:
            A new initial state.
        """
        # Reset time step and disturbance
        self.use_noise = use_noise
        self.n_ts = 0
        self.m.reset_disturbance()

        # Model disturbance if not yet done
        if use_noise and not self._disturbance_is_modelled:
            if not self._dummy_use:
                self._model_disturbance()

        # Reset original actions (necessary?)
        self.orig_actions = np.empty((self.n_ts_per_eps, self.act_dim), dtype=np.float32)

        accepted, ret_val = False, None
        ct = 0
        while not accepted:
            # Select new start data
            if start_ind is None or self.rej_sampler is not None:
                start_ind = np.random.randint(self.n_start_data)
            else:
                if self.n_start_data <= start_ind:
                    raise ValueError("start_ind is too fucking large!")

            self.hist = np.copy(self.train_data[start_ind])
            self.curr_ind = self.train_indices[start_ind] + self.m.data.seq_len - 1
            self.curr_n = (self.t_init_n + self.curr_ind) % self.n_ts_per_day
            ret_val = np.copy(self.hist[-1, :-self.act_dim])

            if self.rej_sampler is not None:
                accepted = self.rej_sampler(ret_val)
                if ct >= self.rej_sampler.max_rej:
                    accepted = True
            else:
                accepted = True

            if not accepted:
                ct += 1

        # print(f"Chose start ind: {start_ind}, rejected: {ct}")
        return ret_val

    def get_scaled_init_state(self, init_ind: int, heat_inds) -> np.ndarray:
        """Returns the initial states when reset with `init_ind`.

        Args:
            init_ind: The reset index.
            heat_inds: The indices of the series to return.

        Returns:
            The scaled initial values of the series.
        """
        # Find the current heating water temperatures
        self.reset(init_ind, use_noise=False)
        return self.scale_state(self.curr_state, remove_mean=False)[heat_inds]

    def render(self, mode='human'):
        """Render method for compatibility with OpenAI gym."""
        print("Rendering not implemented!")

    def _to_scaled(self, action: Arr, to_original: bool = False,
                   extra_scaling: bool = False) -> np.ndarray:
        """Converts actions to the right range.

        Needs to be overridden by subclasses.
        """
        raise NotImplementedError("Implement this shit!")

    def analyze_agents_visually(self, agents: Agents,
                                fitted: bool = True,
                                use_noise: bool = False,
                                start_ind: int = None,
                                max_steps: int = None,
                                state_mask: np.ndarray = None,
                                plot_constrain_actions: bool = True,
                                show_rewards: bool = True,
                                title_ext: str = "",
                                put_on_ol: bool = False,
                                plot_rewards: bool = False,
                                bounds: List[Tuple[int, Tuple[Num, Num]]] = None,
                                series_merging_list: MergeListT = None,
                                overwrite: bool = False,
                                plot_all_rewards: bool = True) -> None:
        """Analyzes and compares a set of agents / control strategies.

        Uses the same initial condition of the environment and evaluates
        all the given agents in `agents` starting from these initial conditions.
        Then makes a plot comparing how the environment behaves under all the
        agents.

        Args:
            agents: A list of agents or a single agent.
            fitted: Whether the agents are already fitted.
            use_noise: Whether to use noise in the predictions.
            start_ind: Index of initial configuration, random if None.
            max_steps: The maximum number of steps of an episode.
            state_mask: Mask defining which series to plot.
            plot_constrain_actions: Whether to plot the actions constrained by the env.
            show_rewards: Use default!
            title_ext: The title to put in the plot.
            put_on_ol: Whether to save in Overleaf plot dir.
            plot_rewards: Whether to plot the reward bar plot for this scenario.
            bounds: Optional state bounds to put in plot.
            series_merging_list: Defines series that will be merged. Should be
                independent of agent's actions.
            overwrite: Whether to overwrite existing plot files.
            plot_all_rewards:

        Raises:
            ValueError: If `start_ind` is too large or if an agent is not suited
                for this environment.
        """
        print("This is deprecated!")

        # Make function compatible for single agent input
        if not isinstance(agents, list):
            agents = [agents]

        if max_steps is None:
            max_steps = 100000

        if not show_rewards:
            raise NotImplementedError("Do not do this!")

        # Use (possibly overridden) defaults
        if state_mask is None:
            state_mask = self.default_state_mask
        if series_merging_list is None:
            series_merging_list = self.default_series_merging

        # Choose same random start for all agents
        if start_ind is None:
            start_ind = np.random.randint(self.n_start_data)
        elif start_ind >= self.n_start_data:
            raise ValueError(f"start_ind: {start_ind} cannot be larger than: {self.n_start_data}!")

        # Check if file already exists
        analysis_plot_path = self._construct_plot_name("AgentAnalysis", start_ind, agents, put_on_ol)
        if os.path.isfile(analysis_plot_path + ".pdf") and not overwrite:
            return

        # Define arrays to save trajectories
        n_non_c_states = self.state_dim - self.act_dim
        n_agents = len(agents)
        action_sequences = npf32((n_agents, self.n_ts_per_eps, self.act_dim), fill=np.nan)
        clipped_action_sequences = npf32((n_agents, self.n_ts_per_eps, self.act_dim), fill=np.nan)
        trajectories = npf32((n_agents, self.n_ts_per_eps, n_non_c_states), fill=np.nan)
        rewards = npf32((n_agents, self.n_ts_per_eps), fill=np.nan)
        n_tot_rewards = len(self.reward_descs) + 1
        all_rewards = npf32((n_agents, self.n_ts_per_eps, n_tot_rewards), fill=np.nan)

        for a_id, a in enumerate(agents):
            # Check that agent references this environment
            if not a.env == self:
                raise ValueError(f"Agent {a_id} was not assigned to this env!")

            # Set agent, required for the ones that need special scaling
            self.set_agent(a)

            # Fit agent if not already fitted
            if not fitted:
                a.fit()

            # Reset env
            curr_state = self.reset(start_ind=start_ind, use_noise=use_noise)
            episode_over = False
            count = 0

            # Evaluate agent and save states, actions and reward
            while not episode_over and count < max_steps:
                curr_action = a.get_action(curr_state)
                scaled_a = self.scale_action_for_step(curr_action)

                curr_state, rew, episode_over, extra = self.step(curr_action)
                det_rew = self.detailed_reward(curr_state, scaled_a)
                action_sequences[a_id, count, :] = curr_action

                # Save trajectories and reward
                trajectories[a_id, count, :] = np.copy(curr_state)
                rewards[a_id, count] = rew
                all_rewards[a_id, count, 0] = rew
                all_rewards[a_id, count, 1:] = det_rew
                count += 1

            # Get original actions
            clipped_action_sequences[a_id, :self.n_ts, :] = self.orig_actions[:self.n_ts, :]

            # Scale actions
            if self.do_scaling:
                action_sequences[a_id] = self.a_scaling_pars[0] + action_sequences[a_id] * self.a_scaling_pars[1]

        # Scale the data to the right values
        trajectories = self.m.rescale_output(trajectories, out_put=True)
        s_ac = clipped_action_sequences.shape
        for k in range(s_ac[0]):
            for i in range(s_ac[1]):
                clipped_action_sequences[k, i] = self._to_scaled(clipped_action_sequences[k, i],
                                                                 to_original=True)
        if not plot_constrain_actions or np.allclose(clipped_action_sequences, action_sequences):
            clipped_action_sequences = None

        # Time stuff
        shifted_t_init = self.m.data.get_shifted_t_init(self.curr_ind)
        np_st_init = str_to_np_dt(shifted_t_init)

        # Plot all the things
        name_list = [a.get_short_name() for a in agents]
        add_pth = self._construct_plot_name("AgentAnalysisReward", start_ind, agents, put_on_ol)
        add_pth = add_pth if plot_rewards else None
        rewards_for_eval = all_rewards if plot_all_rewards else rewards
        plot_env_evaluation(action_sequences, trajectories, rewards_for_eval, self.m.data,
                            name_list, analysis_plot_path, clipped_action_sequences,
                            state_mask, show_rewards=show_rewards, title_ext=title_ext,
                            np_dt_init=np_st_init, rew_save_path=add_pth, bounds=bounds,
                            series_merging_list=series_merging_list,
                            reward_descs=self.reward_descs)

        # Plot the rewards for this episode
        add_pth = self._construct_plot_name("AgentAnalysisRewardDetailed", start_ind, agents, put_on_ol)
        names = [a.get_short_name() for a in agents]
        plot_reward_details(names, all_rewards, add_pth,
                            self.reward_descs, dt=self.m.data.dt, n_eval_steps=self.n_ts_per_eps,
                            title_ext=title_ext)

    def _construct_plot_name(self, base_name: str, start_ind: Optional[int], agent_list: List,
                             put_on_ol: bool = False):
        name_list = [a.get_short_name() for a in agent_list]
        agent_names = '_'.join(name_list)
        s_ext = "" if start_ind is None else f"_{start_ind}"
        base = f"{base_name}{s_ext}_{agent_names}"
        orig_name = self.get_plt_path(base)
        if put_on_ol:
            print(f"Original: {orig_name}")
            return os.path.join(OVERLEAF_IMG_DIR, base)
        return orig_name

    def _get_detail_eval_title_ext(self):
        # For the water temperature info in the title, see overriding method.
        # This is so fucking ugly!
        return None

    def detailed_eval_agents(self, agent_list: Agents,
                             n_steps: int = 100,
                             use_noise: bool = False,
                             put_on_ol: bool = False,
                             overwrite: bool = False,
                             verbose: int = 0,
                             plt_fun: Callable = plot_reward_details,
                             episode_marker: Callable = None,
                             visual_eval: bool = False,
                             bounds: List[Tuple[int, Tuple[Num, Num]]] = None,
                             agent_filter_ind: int = None,
                             filter_good_cases: bool = True,
                             plot_tot_reward_cases: bool = True,
                             plot_tot_eval: bool = True,
                             plot_constrained_actions: bool = False,
                             max_visual_evals: int = 4,
                             heating_title_ext: bool = False,
                             indicate_bad_case: bool = False,
                             disconnect_data: Any = None,
                             eval_quality: bool = False
                             ) -> Optional[np.ndarray]:
        """Evaluates the given agents for this environment.

        Let's the agents act in the environment and observes the
        rewards. These are then averaged, plotted and returned.
        If `overwrite` is False and the plot already exists, nothing
        is returned.

        Args:
            agent_list: List of agents that are based on this env.
            n_steps: Number of evaluation steps.
            use_noise: Whether to use noise in the evaluation.
            put_on_ol: Whether to put the resulting plot into the Overleaf plot folder.
            overwrite: Whether to overwrite pre-existing plot files.
            verbose: Verbosity.
            plt_fun: The plotting function.
            episode_marker: Marker for the case if `plt_fun`
                is :func:`util.visualize.plot_heat_cool_rew_det`
            visual_eval: Whether to do visual analysis with the evaluated data.
            bounds: The temperature bounds for the visual analysis plot.
                Ignored if `visual_eval` is False.
            agent_filter_ind: Index of agent whose rewards need to be
                among the best to be plotted. Ignored if `visual_eval` is False.
            filter_good_cases: Whether to filter out the bad cases and only plot
                the good ones in the time series analysis.
            plot_tot_reward_cases: Whether to plot the rewards corresponding to
                 the time series cases.
            plot_tot_eval: Whether to plot the total rewards of the whole evaluation.
            plot_constrained_actions: Whether to plot the constrained actions in the time series
                plot.
            max_visual_evals: Max. number of visual evaluations.
            heating_title_ext:
            indicate_bad_case:
            disconnect_data:
            eval_quality:

        Returns:
            The rewards seen by all the agents, or None if `overwrite` is False
            and the plot already exists.
        """

        if agent_filter_ind is None:
            agent_filter_ind = 0

        # Make function compatible for single agent input.
        if not isinstance(agent_list, list):
            agent_list = [agent_list]

        # Handle overwriting
        p_name = self._construct_plot_name("DetailAnalysis", n_steps, agent_list, put_on_ol)
        if os.path.isfile(p_name + ".pdf") and not overwrite:
            if verbose > 0:
                print("Agent evaluation plot already exists!")
            return
        p_1 = self._construct_plot_name(f"AgentAnalysis_E0", None,
                                        agent_list, put_on_ol)
        p_2 = self._construct_plot_name(f"AgentAnalysis_E0_n", None,
                                        agent_list, put_on_ol)
        if (os.path.isfile(p_1 + ".pdf") or os.path.isfile(p_2 + ".pdf")) and not overwrite:
            if visual_eval and not plot_tot_eval:
                print("Agent time series plot already exists!")
                return

        pickle_data_path = self._construct_plot_name(f"pickled_data_{n_steps}", None,
                                                     agent_list, False) + ".pkl"

        # Init scores.
        n_agents = len(agent_list)
        n_extra_rewards = len(self.reward_descs)
        n_tot_rewards = n_extra_rewards + 1

        # Load data from pickle if exists
        if os.path.isfile(pickle_data_path):
            if verbose > 0:
                print("Loading from previous evaluation...")
            all_rewards, all_rewards, all_states, all_marks, \
                all_actions, all_scaled_actions, ret_inds = load_pickle(pickle_data_path)
        else:
            all_rewards = npf32((n_agents, n_steps, n_tot_rewards))
            all_actions = npf32((n_agents, n_steps, self.act_dim))
            all_scaled_actions = npf32((n_agents, n_steps, self.act_dim))
            all_states = npf32((n_agents, n_steps, self.state_dim - self.act_dim))
            all_marks = np.empty((n_agents, n_steps), dtype=np.int32)
            ret_inds = None

            do_scaling = self.m.data.fully_scaled

            if verbose > 0:
                print("Evaluating agents...")
            for a_id, a in enumerate(agent_list):

                # Set agent
                self.set_agent(a)

                # Fit agent if not already fitted
                if verbose:
                    print(f"Fitting agent: {a}")
                a.fit(verbose=verbose, train_data=a.fit_data)

                # Check that agent references this environment
                if not a.env == self:
                    raise ValueError(f"Agent {a_id} was not assigned to this env!")

                # Evaluate agent.
                if verbose:
                    print(f"Evaluating agent: {a}")
                eval_res = a.eval(n_steps, reset_seed=True, detailed=True,
                                  use_noise=use_noise, scale_states=do_scaling,
                                  episode_marker=episode_marker,
                                  return_inds=True,
                                  verbose=verbose)
                rew, ex_rew, states, ep_marks, acs, s_acs, ret_inds = eval_res

                all_rewards[a_id, :, 0] = rew
                all_rewards[a_id, :, 1:] = ex_rew
                all_states[a_id] = states
                all_marks[a_id] = ep_marks
                all_actions[a_id] = acs
                all_scaled_actions[a_id] = s_acs

                if verbose:
                    print(f"Saving evaluation data...")
                all_data = (all_rewards, all_rewards, all_states, all_marks,
                            all_actions, all_scaled_actions, ret_inds)
                save_pickle(pickle_data_path, all_data)

        # Plot total rewards for all steps
        title_ext = ""  # self._get_detail_eval_title_ext()
        if plot_tot_eval:
            if verbose > 0:
                print("Plotting evaluation...")
            names = [a.get_plot_name() for a in agent_list]
            plt_fun(names, all_rewards, p_name,
                    self.reward_descs, dt=self.m.data.dt, n_eval_steps=n_steps,
                    title_ext=title_ext, all_states=all_states,
                    verbose=0, ep_marks=all_marks,
                    add_base_title=False)

        if eval_quality:
            quality_plot_path = self._construct_plot_name(f"Env_Quality_{n_steps}", None,
                                                          agent_list, put_on_ol)
            eval_env_evaluation(all_rewards, all_states,
                                all_marks, self.n_ts_per_eps,
                                plt_base_name=quality_plot_path)

        print(f"Mean rewards: {np.mean(all_rewards[:, :, 0], axis=1)}")

        # Plot time series
        if visual_eval:
            ct = 0
            n_vis_eval = n_steps // self.n_ts_per_eps
            for k in range(n_vis_eval):
                k0 = k * self.n_ts_per_eps
                k1 = (k + 1) * self.n_ts_per_eps

                # Extract stuff
                action_sequences = all_actions[:, k0:k1]
                clipped_action_sequences = all_scaled_actions[:, k0:k1]
                curr_states = all_states[:, k0:k1]
                curr_rew = all_rewards[:, k0:k1]
                ep_ind = ret_inds[k]

                mean_rew = np.mean(curr_rew[:, :, 0], axis=1)
                winner = np.argwhere(mean_rew == np.amax(mean_rew))

                ext_ = "" if [agent_filter_ind] not in winner else "_n"
                if filter_good_cases:
                    if [agent_filter_ind] not in winner or ct > max_visual_evals:
                        # if [agent_filter_ind] in winner:
                        #     print(f"Len: {len(winner)}")
                        # else:
                        #     print(f"Bad case!")
                        continue
                    else:
                        # print(f"Found time series!!")
                        ct += 1
                else:
                    ct += 1
                if ct > max_visual_evals:
                    continue

                # Time stuff
                shifted_t_init = self.m.data.get_shifted_t_init(ep_ind)
                np_st_init = str_to_np_dt(shifted_t_init)

                state_mask = self.default_state_mask
                name_list = [a.get_plot_name() for a in agent_list]
                add_pth = None
                series_merging_list = self.default_series_merging

                if heating_title_ext:
                    unscaled_state = curr_states[0, 0]
                    w_in = unscaled_state[2]
                    heating = w_in > unscaled_state[4]
                    ext = f"{'Heating' if heating else 'Cooling'}, inflow water temp.: {w_in:.1f} °C"
                    curr_title_ext = f"{title_ext} {ext}"
                    ext_ += "_h" if heating else "_c"
                    if indicate_bad_case:
                        assert n_agents == 4, "Fuck you!"
                        r_temp_closed = curr_states[1, -1, 4]
                        r_temp_open = curr_states[0, -1, 4]
                        h_case = r_temp_open > r_temp_closed
                        if h_case != heating:
                            print(f"Fuck! k = {k}, {heating}")
                            ext_ += "_b"
                        else:
                            print(f"Nice! k = {k}, {heating}")
                else:
                    curr_title_ext = title_ext

                analysis_plot_path = self._construct_plot_name(f"AgentAnalysis_E{k}{ext_}", None,
                                                               agent_list, put_on_ol)

                # Plot time series for this episode
                ex_actions = clipped_action_sequences if plot_constrained_actions else None
                plot_env_evaluation(action_sequences, curr_states, curr_rew, self.m.data,
                                    name_list, analysis_plot_path, extra_actions=ex_actions,
                                    series_mask=state_mask, show_rewards=True, title_ext=curr_title_ext,
                                    np_dt_init=np_st_init, rew_save_path=add_pth,
                                    bounds=bounds,
                                    series_merging_list=series_merging_list,
                                    reward_descs=self.reward_descs,
                                    disconnect_data=disconnect_data,
                                    tot_reward_only=True)

                # Plot the rewards for this episode
                if not plot_tot_reward_cases:
                    continue
                add_pth = self._construct_plot_name(f"AgentAnalysis_E{k}_RewardDetailed", None,
                                                    agent_list, put_on_ol)
                plot_reward_details(name_list, curr_rew, add_pth,
                                    self.reward_descs, dt=self.m.data.dt,
                                    n_eval_steps=self.n_ts_per_eps,
                                    title_ext=title_ext,
                                    add_base_title=False)

        return all_rewards


from django.contrib import admin
from .models import (
    User, Subject, Quiz, Question, Answer, Student, TakenQuiz, StudentAnswer
)

@admin.register(User)
class UserAdmin(admin.ModelAdmin):
    list_display = ['username', 'is_student', 'is_teacher']
    list_filter = ['is_student', 'is_teacher']
    search_fields = ['username']
    ordering = ['username']
    fieldsets = [
        ['User Information', {'fields': ['username', 'password']}],
        ['Permissions', {'fields': ['is_student', 'is_teacher']}],
    ]

@admin.register(Subject)
class SubjectAdmin(admin.ModelAdmin):
    list_display = ['name', 'color']
    list_filter = ['name']
    search_fields = ['name']
    ordering = ['name']
    fieldsets = [
        ['Subject Information', {'fields': ['name', 'color']}],
    ]

@admin.register(Quiz)
class QuizAdmin(admin.ModelAdmin):
    list_display = ['name', 'subject', 'owner']
    list_filter = ['subject', 'owner']
    search_fields = ['name']
    ordering = ['name']
    fieldsets = [['Quiz Information', {'fields': ['name', 'subject', 'owner']}],
                 ]

@admin.register(Question)
class QuestionAdmin(admin.ModelAdmin):
    list_display = ['text', 'quiz']
    list_filter = ['quiz']
    search_fields = ['text']
    ordering = ['text']
    fieldsets = [['Question Information', {'fields': ['text', 'quiz']}],
                 ]

@admin.register(Answer)
class AnswerAdmin(admin.ModelAdmin):
    list_display = ['text', 'question', 'is_correct']
    list_filter = ['is_correct', 'question']
    search_fields = ['text']
    ordering = ['text']
    fieldsets = [['Answer Information', {'fields': ['text', 'question', 'is_correct']}],
                 ]

@admin.register(Student)
class StudentAdmin(admin.ModelAdmin):
    list_display = ['user', 'score']
    list_filter = ['score']
    search_fields = ['user__username']
    ordering = ['user__username']
    fieldsets = [['Student Information', {'fields': ['user', 'score']}],
                 ]

@admin.register(TakenQuiz)
class TakenQuizAdmin(admin.ModelAdmin):
    list_display = ['student', 'quiz', 'score', 'percentage', 'date']
    list_filter = ['quiz', 'score', 'percentage', 'date']
    search_fields = ['student__user__username']
    ordering = ['student__user__username']
    fieldsets = [['Taken Quiz Information', {'fields': ['student', 'quiz', 'score', 'percentage', 'date']}], ]

@admin.register(StudentAnswer)
class StudentAnswerAdmin(admin.ModelAdmin):
    list_display = ['student', 'answer']
    list_filter = ['answer']
    search_fields = ['student__user__username', 'answer__text']
    ordering = ['student__user__username']
    fieldsets = [
        ['Student Answer Information', {'fields': ['student', 'answer']}],
    ]


###.1

# import sys
#
# from collections import deque
#
# N,M = map(int,sys.stdin.readline().split())
#
# adjacency_list = [[] for i in range(N+1)]
#
# for i in range(M):
#
#     a,b = map(int,sys.stdin.readline().split())
#
#     adjacency_list[b].append(a)
#
#
# def dfs(start):
#
#     visited = [0]*(N+1)
#
#     count = 0
#
#     stack = deque([start])
#
#     while stack:
#
#         current = stack.pop()
#
#         visited[current] = 1
#
#         for i in adjacency_list[current]:
#
#             if visited[i] == 0:
#
#                 stack.append(i)
#
#                 count = count + 1
#
#     return count
#
# result = []
#
# for j in range(1,N+1):
#
#     output = dfs(j)
#
#     result.append(output)
#
# max_count = 0
#
# box = []
#
# for i in range(N):
#
#     if max_count <= result[i]:
#
#         max_count = result[i]
#
#         box.append(i+1)
#
#     else:
#
#         continue
#
# for m in sorted(box):
#
#     print(m, end=' ')

###.2

import sys

from collections import deque

N,M = map(int,sys.stdin.readline().split())

adjacency_list = [[] for i in range(N+1)]

for i in range(M):

    a,b = map(int,sys.stdin.readline().split())

    adjacency_list[b].append(a)

def bfs(start):

    queue = deque([start])

    visited = [0] * (N + 1)

    visited[start] = 1

    count = 0

    while queue:

        current = queue.popleft()

        count = count + 1

        for i in adjacency_list[current]:

            if visited[i] == 0:

                visited[i] = 1

                queue.append(i)

            else:

                continue

    return count

result = []

for m in range(1,N+1):

    output = bfs(m)

    result.append(output)

for index,value in enumerate(result):

    if value == max(result):

        print(index+1,end=' ')

# 방문처리를 하는 코드는 반드시 큐에 방문할 노드를 집어넣기 전해줘야 한다. 이렇게 하지 않을 경우 간선관계가 많은 노드, 즉 인기가 많은 노드가 중복적으로 큐에 담겨서 중복 방문하는 것처럼 표시되는 잘못된 경우가 발생할 수 있다.
# 반례: 1->2,1->3,1->4,2->5,3->5,4->5,5->6,5->7 과 같이 5노드가 다른 노드와의 간선 관계까 많이 형성되어 있을 때, 5노드가 큐에 담기기전에 방문처리를 하지 않게 되면, 5가 큐에 담기고 나서도 방문처리가 아직 안된상황이 발생하여 5가 중복하여 담겨질 수 있다.


import os
import json
import argparse
from tqdm import tqdm
from transformers import BertTokenizer, BertModel
from transformers import XLNetTokenizer

def init_arg_parser():
    arg_parser = argparse.ArgumentParser()

    arg_parser.add_argument('--data_path', default='orig_data')
    arg_parser.add_argument('--temp_data', default='temp_data')
    arg_parser.add_argument('--bert_path', default='/home/lishuan/pretrain_model/chinese_xlnet_base/')
    arg_parser.add_argument('--label_info', default={
        '解除质押': 0, '股份回购': 1, '股东减持': 2, '亏损': 3, '中标': 4, '高管变动': 5, '企业破产': 6,
        '股东增持': 7, '被约谈': 8, '企业收购': 9, '公司上市': 10, '企业融资': 11, '质押': 12
    })
    args = arg_parser.parse_args()
    return args

def data_pro(args):
    train_path = os.path.join(args.data_path, 'duee_fin_train.json')
    dev_path = os.path.join(args.data_path, 'duee_fin_dev.json')

    train_data = data_read(train_path)
    dev_data = data_read(dev_path)
    p(train_data)
    print(len(dev_data))

    train_ids, train_data_info = convert_text_to_id(train_data, args)
    dev_ids, dev_data_info = convert_text_to_id(dev_data, args)

    print(train_data_info, '\n', dev_data_info)

    train_ids_path = os.path.join(args.temp_data, 'xlnet_train_event_ids.json')
    dev_ids_path = os.path.join(args.temp_data, 'xlnet_dev_event_ids.json')

    data_save(train_ids, train_ids_path)
    data_save(dev_ids, dev_ids_path)

def data_read(path):
    text_data = []
    with open(path, mode="r", encoding="utf-8") as f:
        for line in f:
            text_data.append(json.loads(line))
    return text_data

def convert_text_to_id(text_data, args):
    tokenizer = XLNetTokenizer.from_pretrained(args.bert_path)
    data_ids = []
    data_info = {label_str: 0 for label_str in args.label_info.keys()}
    for sample in tqdm(text_data):
        if 'event_list' not in sample.keys() or 'event_type' not in sample['event_list'][0].keys():
            continue
        input_info = tokenizer.encode_plus(sample['text'], add_special_tokens=True)
        data_ids.append(
            [input_info['input_ids'], input_info['token_type_ids'],
             input_info['attention_mask'], args.label_info[sample['event_list'][0]['event_type']]]
        )
        data_info[sample['event_list'][0]['event_type']] += 1
    return data_ids, data_info

def data_save(data, path):
    with open(path, mode="a", encoding="utf-8") as f:
        json.dump(data, f)
    f.close()

def p(data, num=5):
    for i in range(num):
        print(data[i])

if __name__ == '__main__':
    args = init_arg_parser()
    data_pro(args)


#!/usr/bin/python

import re
import glob
import time
import xml.etree.ElementTree as ET
from py2neo import authenticate, Graph
from pandas import DataFrame

def main():
    # Connect to neo4j
    authenticate("localhost:7474", "neo4j", "neo4jpw")
    graph = Graph("localhost:7474/db/data/")

    pathway_list = read_pathway_list("C:/Databases/KEGG/pathway.list")
    main_pathway = "Metabolism"
    use_pathways = [i[2] for i in pathway_list if i[0] in main_pathway]
    use_pathways.remove('01100')

    metabolic_reactions = read_kegg_xml("C:/Databases/KEGG/kgml/ko", use_pathways=use_pathways)

    # Read in files and create a new database
    reactions = kegg_reactions("C:/Databases/KEGG/reaction/reaction")
    enzymes = kegg_enzymes("C:/Databases/KEGG/enzyme/enzyme")
    rclass = kegg_rclass("C:/Databases/KEGG/rclass/rclass")
    compounds = kegg_compounds("C:/Databases/KEGG/compound/compound")
    triples = find_triples(rclass)

    # Create a new database using rclass triples
    # create_db_from_triples(triples, rclass, compounds, graph)

    # Create a new database using metabolic reactions parsed from xml files
    # create_db_from_xml(metabolic_reactions, graph, reactions, enzymes, compounds)

    # Create a new database using reactions
    create_db_from_reactions(reactions, graph, enzymes, compounds, rclass)

    # Test database
    test_database(graph)

    return 0

def read_pathway_list(filename):
    # Read in pathway list and return a hierarchy
    pathways = list()
    level_1_heading = ""
    level_2_heading = ""
    with open(filename) as f:
        for line in f:
            if line[0:1] == "#":
                if line[0:2] == "##":
                    level_2_heading = line[2:].strip()
                else:
                    level_1_heading = line[1:].strip()
            else:
                pathways.append([level_1_heading, level_2_heading, line[0:5], line[6:].strip()])
    f.close()
    return pathways

def read_kegg_xml(folder, use_pathways=None, ignore_pathways=None):
    # Read in and parse xml files
    xml_files = glob.glob(folder + '/*.xml')
    if type(use_pathways) == list:
        r = re.compile("|".join(use_pathways))
        l_keepme = map(lambda x: len(r.findall(x)) > 0, xml_files)
        keep_me = [i for i, x in enumerate(list(l_keepme)) if x]
        xml_files = [xml_files[i] for i in keep_me]
    elif type(ignore_pathways) == list:
        r = re.compile("|".join(ignore_pathways))
        l_keepme = map(lambda x: len(r.findall(x)) > 0, xml_files)
        keep_me = [i for i, x in enumerate(list(l_keepme)) if not x]
        xml_files = [xml_files[i] for i in keep_me]

    reactions = list()
    for filename in xml_files:
        tree = ET.parse(filename).getroot()
        pathway_number = tree.attrib['number']
        pathway_name = tree.attrib['title']
        reactions_xml = tree.findall("reaction")
        for r in reactions_xml:
            reaction = dict()
            reaction['id'] = r.attrib['id']
            reaction['name'] = [x[3:] for x in r.attrib['name'].split(" ")]
            reaction['type'] = r.attrib['type']
            reaction['pathway_id'] = pathway_number
            reaction['pathway_name'] = pathway_name
            for child in r:
                # Currently takes just a single substrate or product when multiple are present
                cpd = dict()
                cpd['id'] = child.attrib['id']
                if child.attrib['name'][0:2] == "g1":
                    cpd['name'] = child.attrib['name'][3:9]  # glycan
                else:
                    cpd['name'] = child.attrib['name'][4:10]  # limit attribute to 6 characters (only first cpd entry)
                reaction[child.tag] = cpd
            reactions.append(reaction)
    return reactions

def kegg_reactions(filename):
    # Read in reactions file
    file = open(filename, "r")
    print("Reading Reaction File")
    reactions = file.readlines()
    file.close()

    # Find end of each record
    end_list = []
    for line_i, line in enumerate(reactions, 1):
        if re.match("///", line):
            end_list.append(line_i)

    # Loop through each record extracting tokens
    print("Parsing Reaction Class Data")
    count2 = 0
    reaction_data = dict()
    for ref in range(0, len(end_list)-1):
        count1 = count2
        count2 = end_list[ref]
        r = reactions[count1:count2]
        token_list = ["ENTR", "NAME", "DEFI", "EQUA", "RCLA", "ENZY"]
        tokens = dict()

        for line_i, line in enumerate(r, 1):
            if line[0:4] in token_list:
                tokens[line[0:4]] = line_i - 1
        # Parse tokens, extracting information into a dictionary
        if "ENTR" in tokens.keys() and "RCLA" in tokens.keys():
            record = dict()
            # Entry token
            record["entry"] = r[tokens["ENTR"]][12:18]
            # RClass token
            found_end = False
            count = 0
            v_rclass = []
            while not found_end:
                current_line = r[tokens["RCLA"] + count]
                if current_line[0:4] in ["RCLA", "    "]:
                    v_rclass.append(re.findall("([RC0-9]+)", current_line[12:]))
                    count += 1
                else:
                    found_end = True
                record['rclass'] = v_rclass
            # Name token
            if "NAME" in tokens.keys():
                record['name'] = r[tokens["NAME"]][12:].strip()
            # Definition token
            if "DEFI" in tokens.keys():
                record['definition'] = r[tokens["DEFI"]][12:].strip()
            # Equation token
            if "EQUA" in tokens.keys():
                record['equation'] = r[tokens["EQUA"]][12:].strip()
            # Enzyme token
            if "ENZY" in tokens.keys():
                record['enzyme'] = r[tokens["ENZY"]][12:].strip()
            # Add record to list
            reaction_data[record['entry']] = record
    print(len(reaction_data), "reaction records created\n")
    return reaction_data

def kegg_enzymes(filename):
    # Read in enzyme file
    file = open(filename, "r")
    print("Reading Enzyme File")
    enzymes = file.readlines()
    file.close()

    # Find end of each record
    end_list = []
    for line_i, line in enumerate(enzymes, 1):
        if re.match("///", line):
            end_list.append(line_i)

    # Loop through each record extracting tokens
    print("Parsing Enzyme Data")
    count2 = 0
    enzyme_data = dict()
    for ref in range(0, len(end_list)-1):
        count1 = count2
        count2 = end_list[ref]
        e = enzymes[count1:count2]
        token_list = ["ENTR", "NAME", "PATH"]
        tokens = dict()

        for line_i, line in enumerate(e, 1):
            if line[0:4] in token_list:
                tokens[line[0:4]] = line_i - 1
        # Parse tokens, extracting information into a dictionary
        if "ENTR" in tokens.keys():
            record = dict()
            entry_id = re.search("EC \d+.\d+.\d+.\d+", e[tokens["ENTR"]])
            if entry_id:
                record["entry"] = entry_id.group()
            else:
                break
            # Name token
            if "NAME" in tokens.keys():
                name_text = e[tokens["NAME"]].rstrip()
                if name_text[len(name_text) - 1] == ";":
                    record['name'] = name_text[12:len(name_text) - 1]
                else:
                    record['name'] = name_text[12:]
            # Pathway token
            if "PATH" in tokens.keys():
                found_end = False
                count = 0
                v_pathway = []
                while not found_end:
                    current_line = e[tokens["PATH"] + count]
                    if current_line[0:4] in ["PATH", "    "]:
                        v_pathway.append(current_line[12:19])
                        count += 1
                    else:
                        found_end = True
                record['pathway'] = v_pathway
            # Add record to list
            enzyme_data[record['entry']] = record
    print(len(enzyme_data), "enzyme records created\n")
    return enzyme_data

def kegg_rclass(filename):
    # Read in rclass file
    file = open(filename, "r")
    print("Reading Reaction Class File")
    rclass = file.readlines()
    file.close()

    # Find end of each record
    end_list = []
    for line_i, line in enumerate(rclass, 1):
        if re.match("///", line):
            end_list.append(line_i)

    # Loop through each record extracting tokens
    print("Parsing Reaction Class Data")
    count2 = 0
    rclass_data = dict()
    for ref in range(0, len(end_list)-1):
        count1 = count2
        count2 = end_list[ref]
        r = rclass[count1:count2]
        token_list = ["ENTR", "DEFI", "RPAI", "PATH"]
        tokens = dict()
        for line_i, line in enumerate(r, 1):
            if line[0:4] in token_list:
                tokens[line[0:4]] = line_i - 1
        # Parse tokens, extracting information into a dictionary
        if "ENTR" in tokens.keys():
            record = dict()
            record["entry"] = r[tokens["ENTR"]][12:19]
            if "DEFI" in tokens.keys():
                found_end = False
                count = 0
                v_definition = []
                while not found_end:
                    current_line = r[tokens["DEFI"] + count]
                    if current_line[0:4] in ["DEFI", "    "]:
                        v_definition.append(current_line[12:].rstrip())
                        count += 1
                    else:
                        found_end = True
                record['definition'] = v_definition
            if "RPAI" in tokens.keys():
                found_end = False
                count = 0
                v_rpair = []
                while not found_end:
                    current_line = r[tokens["RPAI"] + count]
                    if current_line[0:4] in ["RPAI", "    "]:
                        v_rpair.extend(re.findall("([C0-9_]+)", current_line[12:]))
                        count += 1
                    else:
                        found_end = True
                record['rpairs'] = v_rpair
            if "PATH" in tokens.keys():
                found_end = False
                count = 0
                v_pathway = []
                while not found_end:
                    current_line = r[tokens["PATH"] + count]
                    if current_line[0:4] in ["PATH", "    "]:
                        v_pathway.append(current_line[12:20])
                        count += 1
                    else:
                        found_end = True
                record['pathway'] = v_pathway
            # Add record to list
            rclass_data[record['entry']] = record
    print(len(rclass_data), "reaction class records created\n")
    return rclass_data

def kegg_compounds(filename):
    # Read in compound file
    file = open(filename, "r")
    print("Reading Compound File")
    compound = file.readlines()
    file.close()

    # Find end of each record
    end_list = []
    for line_i, line in enumerate(compound, 1):
        if re.match("///", line):
            end_list.append(line_i)

    # Loop through each record extracting tokens
    print("Parsing Compound Data")
    count2 = 0
    compound_data = dict()
    for ref in range(0, len(end_list)-1):
        count1 = count2
        count2 = end_list[ref]
        c = compound[count1:count2]
        token_list = ["ENTR", "NAME", "FORM", "EXAC", "PATH"]
        tokens = dict()
        for line_i, line in enumerate(c, 1):
            if line[0:4] in token_list:
                tokens[line[0:4]] = line_i - 1
        # Parse tokens, extracting information into a dictionary
        if "ENTR" in tokens.keys():
            record = dict()
            record["entry"] = c[tokens["ENTR"]][12:18]
            # Name token
            if "NAME" in tokens.keys():
                name_text = c[tokens["NAME"]].rstrip()
                if name_text[len(name_text) - 1] == ";":
                    record['name'] = name_text[12:len(name_text) - 1]
                else:
                    record['name'] = name_text[12:]
            # Formula token
            if "FORM" in tokens.keys():
                formula_text = c[tokens["FORM"]].rstrip()
                record['formula'] = formula_text[12:]
            # Exact mass token
            if "EXAC" in tokens.keys():
                mass_text = re.search("[0-9.]+", c[tokens["EXAC"]])
                if mass_text:
                    record['mass'] = float(mass_text.group(0))
            # Pathways token
            if "PATH" in tokens.keys():
                found_end = False
                count = 0
                v_pathway = []
                while not found_end:
                    current_line = c[tokens["PATH"] + count]
                    if current_line[0:4] in ["PATH", "    "]:
                        v_pathway.append(current_line[12:20])
                        count += 1
                    else:
                        found_end = True
                record['pathway'] = v_pathway

            # Add record to list
            compound_data[record['entry']] = record
    print(len(compound_data), "compound records created\n")
    return compound_data

def find_triples(rclass):
    triples = []
    for r in rclass:
        for p in rclass[r]['rpairs']:
            rpair_triple = p.split("_")
            rpair_triple.append(rclass[r]["entry"])
            triples.append(rpair_triple)
    return triples

def flatten_dict(d):
    out = list()
    for k, v in d.items():
        if type(v) is str:
            out.append("{k}: \"{v}\"".format(k=k, v=v))
        else:
            out.append("{k}: {v}".format(k=k, v=v))
    return ", ".join(out)

# Create and populate a neo4j database using data from reactions
# Reaction data are cross-referenced with compounds, rclass and enzyme
# Two sets of relationships are generated - connections which specify a non-directional edge between two compound
# nodes and reactions which are directional and contain additional data
def create_db_from_reactions(reactions, graph, enzymes=None, compounds=None, rclass=None):
    # clear old data
    graph.delete_all()
    # begin Cypher transaction
    tx = graph.begin()
    cursors = []
    # iterate over each reaction and add to database
    # include optional data as properties if available
    print("Processing", len(reactions), "reactions")
    start_time = time.time()
    for index, reaction_ref in enumerate(reactions):
        react_data = dict()
        r = reactions[reaction_ref]
        react_data['entry'] = r['entry']
        if "definition" in r:
            react_data['definition'] = r['definition']
        if "equation" in r:
            react_data['equation'] = r['equation']
        # include enzyme data
        if "enzyme" in r:
            enz_id = "EC " + r['enzyme']
            if enz_id in enzymes.keys():
                if 'entry' in enzymes[enz_id]:
                    react_data['enzyme_entry'] = enzymes[enz_id]['entry']
                if 'name' in enzymes[enz_id]:
                    react_data['enzyme_name'] = enzymes[enz_id]['name']
                if 'pathway' in enzymes[enz_id]:
                    react_data['enzyme_pathway'] = enzymes[enz_id]['pathway']
        # include reaction class data
        if "rclass" in r:
            for rc in r['rclass']:
                # loop over reaction class and extract compound data
                if rc[0] in rclass.keys():
                    # RClass
                    if "entry" in rclass[rc[0]]:
                        react_data['rclass_entry'] = rclass[rc[0]]['entry']
                    if "definition" in rclass[rc[0]]:
                        react_data['rclass_definition'] = rclass[rc[0]]['definition']
                    if "pathway" in rclass[rc[0]]:
                        react_data['rclass_pathway'] = rclass[rc[0]]['pathway']
                    if "definition" in rclass[rc[0]]:
                        react_data['rclass_rpairs'] = rclass[rc[0]]['rpairs']
                    # Compound 1
                    if rc[1] in compounds.keys():
                        cpd1 = "MERGE (c1:Compound {{{id}}})".format(id=flatten_dict(compounds[rc[1]]))
                    else:
                        cpd1 = "MERGE (c1:Compound {{entry: \"{id}\"}})".format(id=rc[1])
                    # Compound 2
                    if rc[2] in compounds.keys():
                        cpd2 = "MERGE (c2:Compound {{{id}}})".format(id=flatten_dict(compounds[rc[2]]))
                    else:
                        cpd2 = "MERGE (c2:Compound {{entry: \"{id}\"}})".format(id=rc[2])
                    # Mass change
                    delta_mass = None
                    if all([rc[1] in compounds.keys(), rc[2] in compounds.keys()]):
                        if all(['mass' in compounds[rc[1]], 'mass' in compounds[rc[2]]]):
                            delta_mass = round(compounds[rc[1]]['mass'] - compounds[rc[2]]['mass'], 4)
                            react_data['delta_mass'] = delta_mass
                            react_data['abs_delta_mass'] = abs(delta_mass)
                    # generate MERGE for cypher transaction
                    react = "[:REACTION {{{r}}}]".format(r=flatten_dict(react_data))
                    merge_text = "{cpd1} {cpd2} MERGE (c1)-{react}-(c2)"\
                        .format(cpd1=cpd1, react=react, cpd2=cpd2)
                    tx.append(merge_text)
                    # Create simple connection between pairs of compounds
                    if delta_mass is not None:
                        merge_text = "{cp1} {cp2} MERGE (c1)-[:CONNECTION {{abs_delta_mass: {m1}}}]-(c2)" \
                            .format(cp1=cpd1, cp2=cpd2, m1=abs(delta_mass))
                    else:
                        merge_text = "{cp1} {cp2} MERGE (c1)-[:CONNECTION]-(c2)".format(cp1=cpd1, cp2=cpd2)
                    cursors.append(tx.run(merge_text))
    end_time = time.time()
    print("Time to create transaction =", int(end_time - start_time), "seconds")
    start_time = time.time()
    tx.commit()
    end_time = time.time()
    print("Time to upload transaction =", int(end_time - start_time), "seconds")
    return

# Create and populate a neo4j database using data read from kgml (xml) files
# Data are cross-referenced with compounds, reactions and enzyme
# The KEGG xml files are incomplete and not all network reactions are detailed using this approach
def create_db_from_xml(metabolic_reactions, graph, reactions=None, enzymes=None, compounds=None):
    # clear old data
    graph.delete_all()
    # begin Cypher transaction
    tx = graph.begin()
    cursors = []
    # iterate over each member of metabolic_reactions and add to database
    # include optional data as properties if available
    print("Processing", len(metabolic_reactions), "reactions")
    start_time = time.time()
    for index, m in enumerate(metabolic_reactions):
        # Compound 1
        if m["substrate"]["name"] in compounds.keys():
            cpd1 = "MERGE (c1:Compound {{{id}}})".format(id=flatten_dict(compounds[m["substrate"]["name"]]))
        else:
            cpd1 = "MERGE (c1:Compound {{entry: \"{id}\"}})".format(id=m["substrate"]["name"])
        # Compound 2
        if m["product"]["name"] in compounds.keys():
            cpd2 = "MERGE (c2:Compound {{{id}}})".format(id=flatten_dict(compounds[m["product"]["name"]]))
        else:
            cpd2 = "MERGE (c2:Compound {{entry: \"{id}\"}})".format(id=m["product"]["name"])
        # Mass change
        delta_mass = None
        if all([m["substrate"]["name"] in compounds.keys(), m["product"]["name"] in compounds.keys()]):
            if all(['mass' in compounds[m["substrate"]["name"]], 'mass' in compounds[m["product"]["name"]]]):
                delta_mass = round(compounds[m["product"]["name"]]['mass'] -
                                   compounds[m["substrate"]["name"]]['mass'], 4)
        # Reaction
        if m["type"] == "reversible":
            relation = "-"
        else:
            relation = "->"
        react_data = dict()
        react_data["type"] = m["type"]

        # loop through reactions
        for react_name in m["name"]:
            react_data["entry"] = react_name
            if delta_mass is not None:
                react_data["delta_mass"] = delta_mass
                react_data["abs_delta_mass"] = abs(delta_mass)
            if react_name in reactions.keys():
                if "definition" in reactions[react_name]:
                    react_data["definition"] = reactions[react_name]["definition"]
                if "equation" in reactions[react_name]:
                    react_data["equation"] = reactions[react_name]["equation"]
                if "name" in reactions[react_name]:
                    react_data["reaction_name"] = reactions[react_name]["name"]
                if "rclass" in reactions[react_name]:
                    react_data["reaction_class"] = [x[0] for x in reactions[react_name]["rclass"]]
                if "enzyme" in reactions[react_name]:
                    enz = reactions[react_name]["enzyme"]
                    react_data["enzyme"] = enz
                    if enz in enzymes.keys():
                        if "name" in enzymes[enz]:
                            react_data["enzyme_name"] = enzymes[enz]["name"]
                        if "pathway" in enzymes[enz]:
                            react_data["pathway"] = enzymes[enz]["pathway"]
            # generate MERGE for cypher transaction
            react = "[:REACTION {{{r}}}]".format(r=flatten_dict(react_data))
            merge_text = "{cpd1} {cpd2} MERGE (c1)-{react}{relation}(c2)"\
                .format(cpd1=cpd1, react=react, relation=relation, cpd2=cpd2)
            tx.append(merge_text)
        # Create simple connection between pairs of compounds
        if delta_mass is not None:
            merge_text = "{cp1} {cp2} MERGE (c1)-[:CONNECTION {{abs_delta_mass: \"{m1}\"}}]-(c2)"\
                .format(cp1=cpd1, cp2=cpd2, m1=abs(delta_mass))
        else:
            merge_text = "{cp1} {cp2} MERGE (c1)-[:CONNECTION]-(c2)" .format(cp1=cpd1, cp2=cpd2)
        cursors.append(tx.run(merge_text))

    end_time = time.time()
    print("Time to create transaction =", int(end_time - start_time), "seconds")
    start_time = time.time()
    tx.commit()
    end_time = time.time()
    print("Time to upload transaction =", int(end_time - start_time), "seconds")
    return

# Create and populate a neo4j database using rclass and compound data
# This is the simplest approach, however rclass data are non-directional
def create_db_from_triples(triples, rclass, compounds, graph):
    # clear old data
    graph.delete_all()
    # begin Cypher transaction
    tx = graph.begin()
    cursors = []
    # iterate over each triple and add to database
    print("Processing", len(triples), "relationships")
    start_time = time.time()
    for index, t in enumerate(triples):
        # Lookup each value in triple.  If does not exist then return a record with just the Entry id
        # Compound 1
        if t[0] in compounds.keys():
            cpd1 = "MERGE (c1:Compound {{{id}}})".format(id=flatten_dict(compounds[t[0]]))
        else:
            cpd1 = "MERGE (c1:Compound {{entry: \"{id}\"}})".format(id=t[0])
        # Compound 2
        if t[1] in compounds.keys():
            cpd2 = "MERGE (c2:Compound {{{id}}})".format(id=flatten_dict(compounds[t[1]]))
        else:
            cpd2 = "MERGE (c2:Compound {{entry: \"{id}\"}})".format(id=t[1])
        # Reaction
        if t[2] in rclass.keys():
            react_data = rclass[t[2]]
        else:
            react_data = {'entry': t[2]}
        # Mass change
        if all([t[0] in compounds.keys(), t[1] in compounds.keys()]):
            if all(['mass' in compounds[t[0]], 'mass' in compounds[t[1]]]):
                delta_mass = round(compounds[t[0]]['mass'] -
                                   compounds[t[1]]['mass'], 4)
                react_data['delta_mass'] = delta_mass
                react_data['abs_delta_mass'] = abs(delta_mass)
        react = "[:REACTION {{{r}}}]".format(r=flatten_dict(react_data))
        merge_text = "{cpd1} {cpd2} MERGE (c1)-{react}-(c2)" \
            .format(cpd1=cpd1, react=react, cpd2=cpd2)
        cursors.append(tx.run(merge_text))
    end_time = time.time()
    print("Time to create transaction =", int(end_time - start_time), "seconds")
    start_time = time.time()
    tx.commit()
    end_time = time.time()
    print("Time to upload transaction =", int(end_time - start_time), "seconds")
    return

# A few tests to make sure we can execute cypher queries
def test_database(graph):
    tests = list()
    ppm_limit = 100
    mass_search = 18.0105
    tests.append({"title": "Return first 10 compound IDs",
                  "query": "MATCH (c:Compound) RETURN c.entry LIMIT 10", "dataframe": False})
    for r in graph.relationship_types:
        tests.append({"title": "Return first 10 relationships, type = {r}".format(r=r),
                      "query": "MATCH p=()-[r:{r}]->() RETURN p LIMIT 10".format(r=r), "dataframe": False})
        tests.append({"title": "Using dataframes", "query": "MATCH p=(c1:Compound)-[r:{r}]->(c2:Compound) "
                                                            "RETURN c1.entry, r.entry, c2.entry LIMIT 25".format(r=r),
                      "dataframe": True})
        tests.append({"title": "Search for mass match - dehydration reactions, type = {r}".format(r=r),
                      "query": "MATCH p=(c1:Compound)-[r:{r}]->(c2:Compound) WHERE "
                               "1E6*abs({mass_search} - r.abs_delta_mass)/{mass_search} < {ppm_limit} RETURN "
                               "c1.entry, r.entry, c2.entry LIMIT 25"
                     .format(r=r, mass_search=mass_search, ppm_limit=ppm_limit), "dataframe": True})

    for t in tests:
        print("\n" + t['title'])
        print(t['query'])
        if t['dataframe']:
            out = DataFrame(graph.data(t['query']))
        else:
            out = graph.data(t['query'])
        print(out)
    return

if __name__ == "__main__":
    main()


import os

import numpy as np
from matplotlib import pyplot as plt
from matplotlib.collections import LineCollection
from tqdm import trange

import arena

def turn(heading, turn_rate_sd):
    """
    Turn a heading in a random direction
    :param heading: current heading array
    :param turn_rate_sd: s.d. of turn rate from normal distribution
    """
    return heading + np.random.normal(0, turn_rate_sd, size=heading.shape)

def get_speed(shape, speed_mean=1, speed_sd=0.5):
    """
    Get a random speed for each mosquito
    :param shape: output shape
    :param speed_mean: mean speed
    :param speed_sd: standard deviation
    :return: speed array
    """
    # Non-zero minimum speed avoids issues analysing heading later
    return np.fmax(0.01, np.random.normal(speed_mean, speed_sd, size=shape))

def move(heading, speed):
    """
    Convert speed and heading to vector move
    :param pos: current position (array)
    :param heading: current heading (array)
    :param speed: current speed (array)
    :return: new position (array)
    """
    return speed * np.array([np.cos(heading), np.sin(heading)]).T

def grad_move(pos, x_grad, y_grad, speed):
    """
    Move points along a gradient
    :param pos: current position (Nx2 array)
    :param x_grad: x gradient (HxW array)
    :param y_grad: y gradient (HxW array)
    :param speed: current speed (Nx array)
    :return: vector movement for each input pos (Nx2 array)
    """
    # Linear index of position
    idx = np.floor(pos[:, 1]).astype(int) * x_grad.shape[1] + np.floor(pos[:, 0]).astype(int)

    # Get gradient at current position
    x_move = x_grad.flat[idx] * speed.flatten()
    y_move = y_grad.flat[idx] * speed.flatten()

    return np.c_[x_move, y_move]

def is_at_site(pos, sites, site_density=None):
    """
    Check if a point is at a site
    :param pos: current position (Nx2 array)
    :param sites: site locations (HxW array)
    :return: boolean array of length N
    """
    # Linear index of position
    idx = np.floor(pos[:, 1]).astype(int) * sites.shape[1] + np.floor(pos[:, 0]).astype(int)

    # Check if at site
    if site_density is not None:
        passes = (sites.flat[idx] != 0) & (site_density.flat[idx] < np.random.uniform(1, 10, len(idx)))
        site_density.flat[idx[passes]] += 1
        return passes, site_density
    else:
        return sites.flat[idx] != 0

def isalive(pos, age, bounds):
    """
    Check if a point is alive (i.e. within the bounds)
    :param pos: current position (Nx2 array)
    :param bounds: bounds of the arena (2x2 array)
    :return: boolean array of length N
    """
    return np.all(pos > bounds[:, 0], axis=1) & np.all(pos < bounds[:, 1], axis=1) & (age < 100)

def plot(bounds, mosquitoes, trails, fed):
    """Plots mosquito trails and current position"""
    fig, ax = plt.subplots(figsize=(14.4, 10))
    # Setup axes
    ax.set_xlim(bounds[0, :])
    ax.set_ylim(bounds[1, :])
    ax.set_aspect('equal')
    ax.set_xticks([])
    ax.set_yticks([])

    # Show habitat + houses
    plt.imshow(np.ma.masked_where(arena.breed_distance > 10, arena.breed_distance), alpha=0.2)
    y, x = np.where(arena.feed_sites != 0)
    ax.scatter(x, y, color='brown', s=2, alpha=0.5)

    # Plot current position
    ax.scatter(mosquitoes[fed, 0], mosquitoes[fed, 1], color='r', s=2, alpha=0.5)
    ax.scatter(mosquitoes[~fed, 0], mosquitoes[~fed, 1], color='b', s=2, alpha=0.5)
    # Plot trails
    segs = trails.transpose((1, 0, 2))
    segs = np.ma.masked_where(segs == -999, segs)
    lines = LineCollection(segs, colors='k', alpha=0.2, linewidths=1)
    ax.add_collection(lines)

    fig.tight_layout()

    return fig, ax

if __name__ == '__main__':
    os.makedirs('../output', exist_ok=True)

    height, width = arena.feed_sites.shape
    bounds = np.array([[0, width], [0, height]])
    n = 100
    mosquitoes = np.random.uniform(bounds[:, 0], bounds[:, 1], size=(n, 2))
    headings = np.random.uniform(0, 2*np.pi, size=n)
    fed = np.zeros_like(headings).astype(bool)
    timer = np.zeros_like(headings)
    age = np.random.uniform(0, 50, size=fed.shape)

    site_density = np.zeros_like(arena.breed_sites, dtype=float)

    trails = np.copy([mosquitoes])
    trails_fed = np.copy([fed])

    grad_weight = 0.5

    for t in trange(500):
        alive = isalive(mosquitoes, age, bounds)
        timer[timer > 0] -= 1
        age[alive] += 1
        site_density[site_density > 0] -= 0.1
        site_density[site_density < 0] = 0
        n_mos = len(alive)

        # Turn & Move
        headings = turn(headings, np.pi/2)
        speed = get_speed((n_mos, 1), 1, 0.5)

        heading_move = np.zeros_like(mosquitoes)
        heading_move[alive] = move(headings[alive], speed[alive])

        # Follow gradient
        feed_idx = alive & ~fed & (timer == 0)
        move_to_food = np.zeros_like(mosquitoes)
        move_to_food[feed_idx] = grad_move(mosquitoes[feed_idx], arena.feed_xgrad, arena.feed_ygrad, grad_weight * speed[feed_idx])

        breed_idx = alive & fed & (timer == 0)
        move_to_breed = np.zeros_like(mosquitoes)
        move_to_breed[breed_idx] = grad_move(mosquitoes[breed_idx], arena.breed_xgrad, arena.breed_ygrad, grad_weight * speed[breed_idx])

        # Combine
        mosquitoes += heading_move + move_to_food + move_to_breed
        alive = isalive(mosquitoes, age, bounds)

        # Check if mosquitoes have reached food/breeding site
        feed_idx = alive & ~fed & (timer == 0)  # recalculate as may have moved out of bounds
        can_feed = np.zeros_like(fed)
        can_feed[feed_idx] = is_at_site(mosquitoes[feed_idx], arena.feed_sites)

        breed_idx = alive & fed & (timer == 0)
        can_breed = np.zeros_like(fed)
        can_breed[breed_idx], site_density = is_at_site(mosquitoes[breed_idx], arena.breed_sites, site_density)

        fed[feed_idx & can_feed] = True
        fed[breed_idx & can_breed] = False

        timer[feed_idx & can_feed] = 10
        timer[breed_idx & can_breed] = 5

        # Create new mosquitoes
        new_mosquitoes = mosquitoes[can_breed].copy()
        new_mosquitoes = np.repeat(new_mosquitoes, 3, axis=0)

        n_new = new_mosquitoes.shape[0]
        new_headings = np.random.uniform(0, 2*np.pi, size=n_new)

        mosquitoes = np.concatenate([mosquitoes, new_mosquitoes], axis=0)
        headings = np.concatenate([headings, new_headings], axis=0)

        fed = np.concatenate([fed, np.zeros_like(new_headings).astype(bool)], axis=0)
        timer = np.concatenate([timer, np.zeros_like(new_headings)], axis=0)
        age = np.concatenate([age, np.zeros_like(new_headings)], axis=0)

        trails = np.pad(trails, ((0, 0), (0, n_new), (0, 0)), mode='constant', constant_values=-999)
        trails_fed = np.pad(trails_fed, ((0, 0), (0, n_new)), mode='constant', constant_values=0)

        # Records movement trail
        trails = np.concatenate([trails, [mosquitoes]], axis=0)
        trails_fed = np.vstack([trails_fed, fed])

        alive = isalive(mosquitoes, age, bounds)

        if False:
            print(n_new)
        else:
            fig, _ = plot(bounds, mosquitoes[alive], trails[:, alive, :], fed[alive])
            fig.savefig(f'../output/{t:03d}.png', dpi=400)
            plt.close(fig)

    # plot(bounds, mosquitoes[alive], trails[:, alive, :], fed[alive])
    plot(bounds, mosquitoes, trails, fed)
    plt.show()
    print(len(fed))


import numpy as np
import pandas as pd
#import sys
#from matplotlib import pyplot as plt
#from matplotlib import style
#style.use("ggplot")
#from openpyxl import *
from sklearn import svm
import tkinter
from tkinter import filedialog
import openpyxl
#----------------机器学习:用向量机进行匹配---------------
def getsvm(trainSet,tmpltSet,ptnset):
    clf = svm.SVC(kernel='linear', C = 1.0)
    clf.fit(trainSet, tmpltSet)
    result = clf.predict(ptnset)
    return result
#------------------获取excel文件路径---------------
def getExcelFilePath():
    while True:
        question = input("Do you want to get a excel file?(yes(y)/no(n)")[0]
        if question in ["y", "Y"]:
            win = tkinter.Tk()
            win.withdraw()
            path = filedialog.askopenfilename().replace("/", "\\")
            if path[-4:] != "xlsx":
                print("The file is not a excel file, try again:")
                continue
            print("The excel %s input successfully!" %path)
            return path
        elif question in ["n", "N"]:
            print("Exiting......")
            break
        else:
            print("The choice is wrong, try again:")

#----- get trainSet, ptnSet and tmpltSet-----
def getSetDatas(data, selectItemName):
    #SourceData = pd.read_excel(path, sheet_name=sheetName, engine="openpyxl")
    popItem = data.drop(index=selectItemName)

    #Train set
    TrainSetArray = np.array(popItem)

    #TmpltSet
    TmpltSet = popItem.index
    TmpltSetArray = np.array(TmpltSet)

    #PtnSet
    PtnSet = data.loc[selectItemName]
    PtnSetArray = np.array([PtnSet])

    result = (TrainSetArray, TmpltSetArray, PtnSetArray)

    return result

#----------------write to Excel File--------------------
def writeToExcelFile(path, resultList):
    selectedHerbName = resultList[:, 0]
    attribute_selected = resultList[:, 1]
    matchedHerbName = resultList[:, 2]
    attribute_matched  = resultList[:, 3]
    dataResult = pd.DataFrame({"待测药":selectedHerbName, "待测药趋向性":attribute_selected, "匹配药":matchedHerbName, "匹配药趋向性":attribute_matched})
    wb = openpyxl.load_workbook(path)
    writer = pd.ExcelWriter(path, engine="openpyxl")
    writer.book = wb
    dataResult.to_excel(writer, sheet_name="result", columns=None)
    writer.save()
    writer.close()
    print("Write to the excel file successfully!")

#---------------handle source data--------------
def handleSourceData(path, sheetName):
    with pd.ExcelFile(path) as xlsbook:
        SourceData = pd.read_excel(xlsbook, sheet_name=sheetName)
        SourceData = SourceData.set_index("中药名称").loc[:, "趋向性":]
        SourceData = SourceData.dropna(axis=0) #除去含有缺失值的行
        #SourceData = SourceData.dropna(axis=1) #除去含有缺失值的列
    return SourceData

#---------------main function-------------------
def main():
    sheetName = 'GdaSet'
    path = getExcelFilePath()
    Data = handleSourceData(path, sheetName)
    resultList = []
    itemNames = list(Data.index)
    for item in itemNames:
        attribute_item = int(Data.loc[item, ["趋向性"]].values[0])
        dataSample = Data.loc[:, "C数":]
        setDatas = getSetDatas(dataSample, item)
        matchResult = getsvm(setDatas[0], setDatas[1], setDatas[2])[0]
        attribute_result = int(Data.loc[matchResult, ["趋向性"]].values[0])
        print("The herb:{0:}, attribute_item:{1:}, match result: {2:}, attribute_result:{3:}\n".format(item, attribute_item, matchResult, attribute_result))
        resultList.append([item, attribute_item, matchResult, attribute_result])
    writeToExcelFile(path, np.array(resultList))

if __name__ == "__main__":
    main()


# -*- coding:utf-8 -*-
'''
有时候我们需要的并不是最后的结果，而是中间某一层的结果，这个例子演示如何得到中间结果
'''
from keras.layers import Dense,Activation,Input
from keras.models import  Sequential,Model
import  numpy as np

x=np.random.standard_normal(size=(200,12))
y=np.random.randint(low=0,high=2,size=(200,1))
# this is your initial model
input=Input(shape=(12,))
hidden_1=Dense(20,activation='tanh')(input)
hidden_2=Dense(64,activation='tanh',name='hidden_2')(hidden_1)
hidden_3=Dense(1,activation='softmax')(hidden_2)
model=Model(input=input,output=hidden_3)

# we train it
model.compile(loss='mse',optimizer='sgd')
model.fit(x,y,epochs=20,batch_size=10)

# we build a new model with the activations of old model
# this model is truncted after the first layer
test_x=np.random.standard_normal(size=(1,12))

model2= Model(inputs=model.input, outputs=model.get_layer('hidden_2').output)
activations=model2.predict(test_x)
print(activations)


import hashlib
import logging
import os
import re
import shutil

import distro

from config import Config
from package_manager.apt import Apt
from package_manager.pacman import Pacman
from utils import file_exists, file_info

used_exclude_files = []
used_content_filters = []

def get_base_distribution():
    base_distribution = distro.like()
    if base_distribution == '':
        base_distribution = distro.id()
    if base_distribution == '':
        logging.fatal(
            'Unable to detect your distribution. Please file a bug with the following information included:\n%s',
            distro.linux_distribution(full_distribution_name=False))
        exit(1)
    return base_distribution

def get_system_package_manager():
    base_distribution = get_base_distribution()
    logging.debug('Detected \'%s\' as base distribution', base_distribution)
    distribution_package_manager_map = {
        'arch': Pacman(),
        'debian': Apt(),
    }
    package_manager = distribution_package_manager_map.get(base_distribution)
    logging.debug('Using \'%s\' as package manager', package_manager.name)
    return package_manager

def get_installed_not_listed_packages(listed_packages: [str], package_manager) -> [str]:
    installed_not_listed_packages: [str] = []
    installed_packages = package_manager.get_packages()
    for installed_package in installed_packages:
        found = False
        for listed_package in listed_packages:
            if installed_package == listed_package:
                found = True
        if not found and package_manager.is_package_explicitly_installed(installed_package):
            logging.debug(f'\'{installed_package}\' installed but not listed')
            installed_not_listed_packages.append(installed_package)
    return installed_not_listed_packages

def get_listed_not_installed_packages(listed_packages: [str], package_manager) -> [str]:
    listed_not_installed_packages: [str] = []
    installed_packages = package_manager.get_packages()
    for listed_package in listed_packages:
        found = False
        for installed_package in installed_packages:
            if listed_package == installed_package:
                found = True
        if not found:
            logging.debug(f'\'{listed_package}\' listed but not installed')
            listed_not_installed_packages.append(listed_package)
    return listed_not_installed_packages

def get_etc_files(exclude_files: [str]):
    all_files = []
    for root, sub_dirs, files in os.walk('/etc'):
        for file in files:
            file_path = os.path.join(root, file)
            if not check_is_excluded(file_path, exclude_files):
                all_files.append(file_path)
    return all_files

def get_modified_not_listed_files(config: Config, exclude_files: [str], registered_files: {}, hash_method: str) -> {}:
    modified_files = []
    new_files = get_etc_files(exclude_files).copy()

    for path in registered_files:
        if check_is_excluded(path, exclude_files):
            continue
        if path in new_files:
            new_files.remove(path)

        should_hash_sum = registered_files[path]

        # We got an empty hash that means the file is a symlink
        if should_hash_sum == '':
            continue

        if not verify_hash(path, should_hash_sum, hash_method):
            modified_files.append(path)

    for new_file in new_files:
        modified_files.append(new_file)

    files = {}
    for file in modified_files:
        if file not in config.files:
            files[file] = file_info(file)

    return files

def get_listed_not_modified_files(config: Config, exclude_files: [str], registered_files: {}, hash_method: str) -> {}:
    not_modified_files = []

    for path in config.files:
        if check_is_excluded(path, exclude_files):
            not_modified_files.append(path)
            continue

        if not file_exists(path, follow_symlinks=False):
            not_modified_files.append(path)

        if path in registered_files:
            should_hash_sum = registered_files[path]

            # We got an empty hash that means the file is a symlink
            if should_hash_sum == '':
                continue

            if verify_hash(path, should_hash_sum, hash_method):
                not_modified_files.append(path)

    return not_modified_files

def verify_hash(path, should_hash_sum, hash_method):
    file_hash = None
    if hash_method == 'sha256':
        file_hash = hashlib.sha256()
    elif hash_method == 'md5':
        file_hash = hashlib.md5()
    else:
        logging.fatal(f'Unsupported hash method: {hash_method}')
        exit(1)
    with open(path, 'rb') as file:
        fb = file.read(65536)
        while len(fb) > 0:
            file_hash.update(fb)
            fb = file.read(65536)
    is_hash_sum = file_hash.hexdigest()
    return is_hash_sum == should_hash_sum

def save_files(config_directory: str, preset: str, files: {}, content_filters: [], relative_path: str):
    files_path = os.path.join(config_directory, f'{preset}_files')
    if os.path.exists(files_path):
        shutil.rmtree(files_path)
    os.mkdir(files_path)
    stat_info = os.stat(os.path.join(config_directory, '..'))
    for path in files:
        logging.debug(f'Saving {path}')
        local_path = os.path.join(files_path, path.strip('/'))
        local_dir = os.path.dirname(local_path)
        if not os.path.exists(local_dir):
            os.makedirs(local_dir)
        absolute_path = os.path.join(relative_path, path)
        if file_exists(absolute_path, follow_symlinks=False):
            if os.path.islink(absolute_path):
                shutil.copy(absolute_path, local_path, follow_symlinks=False)
            else:
                content = read_filtered_content(absolute_path, content_filters, relative_path)
                if isinstance(content, str):
                    with open(local_path, 'w') as file:
                        file.write(content)
                else:
                    with open(local_path, 'wb') as file:
                        file.write(content)
                # So we are not annoying with the permissions for users
                os.chown(local_path, stat_info.st_uid, stat_info.st_gid)

def apply_files(config_directory: str, preset: str, files: {}, relative_path: str):
    files_path = os.path.join(config_directory, f'{preset}_files')
    if not os.path.exists(files_path):
        if len(files) > 0:
            logging.fatal(f'Could not find required {preset}_files directory')
            exit(1)
        return
    for path in files:
        logging.debug(f'Applying {path}')
        local_path = os.path.join(files_path, path.strip('/'))
        if file_exists(local_path, follow_symlinks=False):
            uid = int(files[path].split(':')[0])
            gid = int(files[path].split(':')[1])
            mode = int('100' + files[path].split(':')[2], 8)
            absolute_path = os.path.join(relative_path, path)
            if file_exists(absolute_path, follow_symlinks=False):
                os.remove(absolute_path)
            else:
                logging.debug(f'File {path} was not available on the system')
            if not file_exists(os.path.dirname(absolute_path), follow_symlinks=False):
                os.makedirs(os.path.dirname(absolute_path))
                current_path = os.path.dirname(absolute_path)
                while current_path != os.path.abspath(relative_path):
                    _set_permission(current_path, None, uid, gid)
                    current_path = os.path.dirname(current_path)
            shutil.copy(local_path, absolute_path, follow_symlinks=False)
            _set_permission(absolute_path, mode, uid, gid)
        else:
            logging.fatal(f'Could not find required {preset}_files{path} file')
            exit(1)

def _set_permission(path, mode, uid, gid):
    if not os.path.islink(path):
        if mode is not None:
            try:
                os.chmod(path, mode)
            except PermissionError:
                logging.warning(
                    f'Unable to chmod {path}. This could be because it is a symlink or because you have set the '
                    f'wrong permissions on the file')
        try:
            os.chown(path, uid, gid)
        except PermissionError:
            logging.warning(
                f'Unable to chown {path}. This could be because it is a symlink or because you have set the '
                f'wrong permissions on the file')

def get_available_not_listed_files(config: Config, exclude_files: [str]):
    not_listed_files = []
    home_dir = os.path.expanduser('~')
    all_files = []

    exclude_files_expanded = []
    for exclude_file in exclude_files:
        exclude_files_expanded.append(os.path.join(home_dir, exclude_file))

    for name in os.listdir(home_dir):
        if name.startswith('.'):
            if os.path.isfile(os.path.join(home_dir, name)):
                all_files.append(os.path.join(home_dir, name))
    for root, _, files in os.walk(os.path.join(home_dir, '.config')):
        for file in files:
            all_files.append(os.path.join(root, file))

    for path in all_files:
        if check_is_excluded(path, exclude_files_expanded):
            continue
        listed = False
        for file in config.files:
            if os.path.join(home_dir, file) == path:
                listed = True
                break
        if not listed:
            not_listed_files.append(path)

    files = {}
    for file in not_listed_files:
        if file not in config.files:
            files[os.path.relpath(file, home_dir)] = file_info(file)

    return files

def get_listed_not_available_files(config: Config, exclude_files: [str]):
    not_available_files = []
    home_dir = os.path.expanduser('~')

    for path in config.files:
        path = os.path.join(home_dir, path)
        if check_is_excluded(path, exclude_files):
            not_available_files.append(path)
            continue

        if not file_exists(path, follow_symlinks=False):
            not_available_files.append(os.path.relpath(path, home_dir))

    return not_available_files

def check_is_excluded(path: str, exclude_files: [str]) -> bool:
    for exclude_file in exclude_files:
        if _path_match_pattern(path, exclude_file):
            if exclude_file not in used_exclude_files:
                used_exclude_files.append(exclude_file)
            logging.debug(f'Excluding {path} because of exclude rule {exclude_file}')
            return True
    return False

def _path_match_pattern(path, pattern):
    return re.match(f'^{re.sub(r"/$", "/.*", pattern.replace("*", ".*"))}$', path)

def get_listed_not_used_exclude_files(config, relative_path):
    listed_not_used_exclude_files = []
    for exclude_file in config.exclude_files:
        if os.path.join(relative_path, exclude_file) not in used_exclude_files:
            listed_not_used_exclude_files.append(exclude_file)
    return listed_not_used_exclude_files

def read_filtered_content(path: str, content_filters: [], relative_path: str):
    filtered_content = []
    matching_content_filters = []
    for content_filter in content_filters:
        filter_path = list(content_filter.keys())[0]
        if _path_match_pattern(path, os.path.join(relative_path, filter_path)):
            matching_content_filters.append(content_filter)
    try:
        with open(path, 'r') as file:
            content = file.read()
            for line in content.split('\n'):
                do_filter = False
                for content_filter in matching_content_filters:
                    if re.match(f'^{content_filter[list(content_filter.keys())[0]]}$', line):
                        if path not in used_content_filters:
                            used_content_filters.append(content_filter)
                        do_filter = True
                        break
                if not do_filter:
                    filtered_content.append(line)
        return '\n'.join(filtered_content)
    except UnicodeDecodeError:
        with open(path, 'rb') as file:
            return file.read()

def get_listed_not_used_content_filters(config):
    listed_not_used_content_filters = []
    for content_filter1 in config.content_filters:
        used = False
        path1 = list(content_filter1.keys())[0]
        for content_filter2 in used_content_filters:
            path2 = list(content_filter2.keys())[0]
            if path1 == path2 and content_filter1[path1] == content_filter2[path2]:
                used = True
        if not used:
            listed_not_used_content_filters.append(content_filter1)
    return listed_not_used_content_filters


# pip install bertopic

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

import nltk
import gensim
import gensim.corpora as corpora
import pandas as pd
import numpy as np
from bertopic import BERTopic
from sentence_transformers import SentenceTransformer
from umap import UMAP
from hdbscan import HDBSCAN
from nltk.stem import WordNetLemmatizer
from nltk.tokenize import word_tokenize
from sklearn.feature_extraction.text import CountVectorizer
from bertopic.vectorizers import ClassTfidfTransformer
from gensim.models.coherencemodel import CoherenceModel

nltk.download('wordnet')
nltk.download('omw-1.4')
nltk.download('punkt')

colnames=['name', 'description', 'price', 'link']
docs1 = pd.read_csv('appliance.csv',names=colnames, header=None)
docs2 = pd.read_csv('laptops.csv',names=colnames, header=None)
docs3 = pd.read_csv('phones.csv',names=colnames, header=None)
docs4 = pd.read_csv('sneakers.csv',names=colnames, header=None, sep=';')
docs5 = pd.read_csv('toys.csv',names=colnames, header=None, sep=';')
docs6 = pd.read_csv('makiyazh.csv',names=colnames, header=None, sep=';')
docs7 = pd.read_csv('reduced_data.csv',names=colnames, header=None)

docs = docs1.append(docs2, ignore_index=True)
docs = docs.append(docs3, ignore_index=True)
docs = docs.append(docs3, ignore_index=True)
docs = docs.append(docs4, ignore_index=True)
docs = docs.append(docs5, ignore_index=True)
docs = docs.append(docs6, ignore_index=True)
docs = docs.append(docs7, ignore_index=True)

# docs['comb'] = docs['name'] + " " + docs['description']
docs['comb'] = docs['name']

docs.head()

filtered_text = []
lemmatizer = WordNetLemmatizer()

for w in docs['comb']:
  filtered_text.append(lemmatizer.lemmatize(w))
print(filtered_text[:1])
docs['filtered'] = filtered_text

embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
umap_model = UMAP(n_neighbors=20, n_components=5, min_dist=0.0, metric='cosine')
hdbscan_model = HDBSCAN(min_cluster_size=40,metric='euclidean', cluster_selection_method='eom', prediction_data=True)
vectorizer_model = CountVectorizer()
ctfidf_model = ClassTfidfTransformer()

docs.head()

topic_model= BERTopic(language = "multilingual", top_n_words = 10, n_gram_range = (1,2), calculate_probabilities=True, verbose=True,
  embedding_model=embedding_model,
  umap_model=umap_model,
  hdbscan_model=hdbscan_model,
  vectorizer_model=vectorizer_model,
  ctfidf_model=ctfidf_model,
  nr_topics = "auto"
          )
topics, probs = topic_model.fit_transform(docs['filtered'])

# topic_model.update_topics(list(docs['filtered']), n_gram_range=(1, 2), top_n_words = 10)

hierarchical_topics = topic_model.hierarchical_topics(docs['filtered'])

hierarchical_topics

topic_model.visualize_hierarchy(hierarchical_topics=hierarchical_topics)

tree = topic_model.get_topic_tree(hierarchical_topics)
print(tree)

representative_docs = topic_model.get_representative_docs()
representative_docs

docs['Topic'] = topics

docs.head()

freq = topic_model.get_topic_info()
freq

topic_model.get_topic(0)

topic_model.visualize_hierarchy(top_n_topics=50)

topic_model.visualize_heatmap(n_clusters=20, width=1000, height=1000)

# from sklearn.metrics.pairwise import cosine_similarity
# sim_matrix = cosine_similarity(topic_model.c_tf_idf_)
# df = pd.DataFrame(sim_matrix, columns=topic_model.topic_labels_.values(), index=topic_model.topic_labels_.values())
# df
# # distance_matrix = cosine_similarity(np.array(topic_model.topic_embeddings)[1:, :])

topic_model.visualize_barchart(top_n_topics=len(freq))

topic_model.visualize_topics()

topic_model.save("my_model")

test_docs = "adidas кроссовки"
predicted_topics, predicted_probs = topic_model.transform(test_docs)
print(f'Predicted topic is {predicted_topics}, and the probability is {np.round(predicted_probs,2)}')

new_review ="adidas кроссовки"
similar_topics, similarity = topic_model.find_topics(new_review, top_n=5)
print(f'The most similar child topic is {similar_topics}, and the similarities is {np.round(similarity,2)}')