import matplotlib as mpl
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import pandas as pd
import statsmodels.api as sm
mpl.rcParams['figure.figsize'] = [25, 10]

# Number of individuals in the simulation. Used to generate a graph if none is passed.
N = 10000
# Properties of the connection graph. (2, 20, 0.2) makes top 0.1% = 15 * median.
# Minimal degree of nodes in the generated graph.
MIN_DEGREE = 2
# Mean degree of nodes in the generated graph.
MEAN_DEGREE = 20
# Parameter of degree distribution in generated graph.
# Higher gamma means a fatter tail of the degree distribution.
GAMMA = 0.2
# Number of simulation steps per one day of real life time.
STEPS_PER_DAY = 5
# Maximal number of steps in simulation. If this number is reached, simulation stops.
MAX_STEPS = 3000
# Number of individuals infected at t=0.
INITIAL_INFECTED_NUM = 10
# In a single interaction on a single day, what is the chance of an Infectious
# to infect a Susceptible. Fitted to produce doubling_time=3.1 For G with parameters
# (100000, 2, 20, 0.2) and 40% asymptomatic, along with all other defaults.
PROB_INFECT = 0.027
# Ratio of infectiousness between Exposed (in the final days before becoming
# Infected) and Infected. In other words,
# prob_infect_exposed = prob_infect * prob_infect_exposed_factor.
# Set this to 0 to make the Exposed non-infectious.
PROB_INFECT_EXPOSED_FACTOR = 0.5
# Ratio of infectiousness between Asymptomatic and symptomatic. In other words,
# prob_infect_asymptomatic = prob_infect * relative_infectiousness_asymptomatic.
# The same ratio holds for the final contagious part of the incubation period.
# Set this to 0 to make the Asymptomatic non-infectious.
# Source: Assumed in a few places. This paper might claim 0.5: https://science.sciencemag.org/content/368/6490/489
RELATIVE_INFECTIOUSNESS_ASYMPTOMATIC = 0.5
# How many days before developing symptoms (becoming Infected) an Exposed
# individual is contagious. Questionable evidence, but 1-2 produces a reasonable R0.
DURATION_EXPOSED_INFECTS = 2
# Incubation period duration distribution mean. In days.
# Source: 3rd Imperial College paper: https://www.imperial.ac.uk/media/imperial-college/medicine/sph/ide/gida-fellowships/Imperial-College-COVID19-Europe-estimates-and-NPI-impact-30-03-2020.pdf
INCUBATION_DURATION_MEAN = 5.1
# Incubation period duration distribution standard deviation. In days.
# Source: 3rd Imperial College paper: https://www.imperial.ac.uk/media/imperial-college/medicine/sph/ide/gida-fellowships/Imperial-College-COVID19-Europe-estimates-and-NPI-impact-30-03-2020.pdf
INCUBATION_DURATION_STD = 4.38
# Share of population who would not develop symptoms even when sick. They are
# not detected along with the symptomatic Infected in the testing step, but are
# detected by direct tests (mass testing or when neighbors are tested).
# Source for default: Vo, Italy study and others (See Eric Topol paper): https://www.scripps.edu/science-and-medicine/translational-institute/about/news/sarc-cov-2-infection/index.html
PROB_ASYMPTOMATIC = 0.40
# Probability to change from Infections to Recovered in a single day.
# Equals the inverse of the expected time to Recover.
# 3.5 = 0.3 * 0 + 0.56 * 5 + 0.1 * 5 + 0.04 * 5. Source: Pueyo spreadsheet: https://docs.google.com/spreadsheets/d/1uJHvBubps9Z2Iw_-a_xeEbr3-gci6c475t1_bBVkarc/edit#gid=0
PROB_RECOVER = 1 / 3.5
# Length of imposed quarantine, in days.
DAYS_IN_QUARANTINE = 14
# Probability that an Infected individual is tested and detected as a carrier in a single day.
# This encapsulates also the fact that many people are mildly symptomatic or will
# not be tested (or quarantined).
# Default: no tests at all.
PROB_INFECTED_DETECTED = 0
# Probability that a neighbor of an individual who tested positive is himself traced,
# and then tested / quarantined. The testing of neighbors happens once, once the
# original individual is tested positive, so this is *not* a probability per day.
# Default - no neighbors are traced.
PROB_NEIGHBOR_TRACED = 0
# Share of the general population which is tested on a single day. Exposed and
# Infected individuals who are tested are detected as carriers.
# Default - previous behavior, no exposed are detected.
PROB_EXPOSED_DETECTED = 0
# When an individual tests positive, whether their traced neighbors are quarantined.
QUARANTINE_NEIGHBORS = False
# When an individual tests positive, whether their traced neighbors are tested.
TEST_NEIGHBORS = False
# Delay between the time a test is performed and the time results become known
# to the subject, and contact tracing is complete. Days.
TEST_DELAY_TIME = 0

_EPSILON = 1e-10
# USA deaths data (worldometers.info): 11 -> 2220 from March 4th to March 28th.
DOUBLING_DAYS = float((28 - 4) / np.log2(2220 / 11))  # About 3.13
# Names of columns in SimulationResults DataFrame.
MAIN_GROUPS = ['susceptible', 'exposed', 'recovered', 'infected', 'quarantined']
# Column list out of the columns in the SimulationResults DataFrame.
ALL_COLUMNS = MAIN_GROUPS + ['test_rate']
# Plotting color conventions.
GROUP2COLOR = dict(susceptible='blue', exposed='orange', recovered='green',
                   quarantined='purple', infected='red', test_rate='brown')
# Attributes to print in SimulationResults summary.
SUMMARY_ATTRS = ['duration', 'fraction_infected', 'doubling_days', 'fraction_quarantine_time', 'peak_infected_time', 'peak_fraction_infected', 'fraction_tests', 'peak_test_rate']

class SimulationResults(object):
  """Simulation run results: metadata, aggregates and time series."""

  def __init__(self, results_df, G=None, **kwargs):
    """Initialize with simulation results, graph and hyper-parameters."""
    self.df = results_df
    self.hyperparams = kwargs
    for name, value in kwargs.items():
      setattr(self, name, value)
    if G is None:
      self.N = results_df[['susceptible', 'exposed', 'infected', 'recovered']].iloc[0].sum()
    else:
      self.N = len(G)
      # Extract graph parameters from its name.
      if G.name.startswith('power_law'):
        self.G_attrs = dict(zip(['gamma', 'min_degree', 'mean_degree'], map(float, G.name.split('_')[-3:])))
        self.G_attrs['N'] = self.N

    if not hasattr(self, 'steps_per_day'):
      self.steps_per_day = ((results_df['step'].iloc[1] - results_df['step'].iloc[0]) /
                            (results_df['day'].iloc[1] - results_df['day'].iloc[0]))

    self.analyze_results_df()

  def calculate_doubling_time(self):
    """Returns doubling time of initial spread (in days)."""
    results_df = self.df
    # Exponential regime before leveling off.
    idx_end = (results_df['exposed'] > results_df['exposed'].max() * 0.5).to_numpy().nonzero()[0][0]
    if self.peak_exposed_time < 3 or idx_end == 0:
      # 3 days peak time, or peak value is less than twice initial value - probably containment.
      return np.inf

    # Don't start from the very beginning, to avoid the initial dip in the
    # number of exposed, due to imperfect starting conditions.
    # Start when number of exposed > 2 * minimum (after passing the minimum).
    exposed_min = results_df['exposed'][:idx_end].min()
    idx_min = results_df['exposed'][:idx_end].idxmin()
    start_candidates = ((results_df.index >= idx_min) &
                        (results_df.index < idx_end) &
                        (results_df['exposed'] > exposed_min * 2)).to_numpy().nonzero()[0]
    if not start_candidates.size:
      # Empty - no candidates. Probably containment.
      return np.inf
    idx_start = start_candidates[0]

    # Linear regression to find doubling time: log2(exposed) ~ day + const
    try:
      X = sm.add_constant(results_df[idx_start:idx_end][['day']], prepend=False)
      log2_exposed = np.log2(results_df[idx_start:idx_end]['exposed'])
      regression_results = sm.OLS(log2_exposed, X).fit()
      # Days for doubling is the inverse of the doubling effect of one day.
      doubling_days = 1 / regression_results.params['day']
    except ValueError:
      doubling_days = None
    return doubling_days

  def calculate_halving_time(self):
    """Returns halving time of spread after peak (in days)."""
    results_df = self.df
    # Find peak.
    idx_peak = results_df['exposed'].idxmax()
    # Find end point for calculation, not right at the peak, but not at the noisy end.
    end_candidates = ((results_df.index >= idx_peak) &
                      (results_df['exposed'] < self.peak_exposed / 5) &
                      (results_df['exposed'] > 5)).to_numpy().nonzero()[0]
    if not end_candidates.size:
      # Halving is too short/noisy to calculate the halving time.
      return None

    idx_end = end_candidates[0]
    idx_start = idx_peak
    if idx_end - idx_start < 20:
      # Halving is too short to calculate the halving time.
      return None

    # Linear regression to find halving time: log2(exposed) ~ day + const
    try:
      X = sm.add_constant(results_df[idx_start:idx_end][['day']], prepend=False)
      log2_exposed = np.log2(results_df[idx_start:idx_end]['exposed'])
      regression_results = sm.OLS(log2_exposed, X).fit()
      # Days for halving is the inverse of the halving effect of one day.
      halving_days = -1 / regression_results.params['day']
    except ValueError:
      halving_days = None
    return halving_days

  def analyze_results_df(self):
    """Calculate various summary stats."""
    results_df = self.df
    self.duration = results_df['day'].iloc[-1]
    # Find peak infections.
    self.peak_infected_time = results_df['day'].iloc[results_df['infected'].idxmax()]
    self.peak_infected = results_df['infected'].max()
    self.peak_fraction_infected = results_df['infected'].max() / self.N
    self.peak_exposed_time = results_df['day'].iloc[results_df['exposed'].idxmax()]
    self.peak_exposed = results_df['exposed'].max()
    self.doubling_days = self.calculate_doubling_time()
    self.halving_days = self.calculate_halving_time()
    # Other result summary stats.
    self.fraction_infected = results_df['recovered'].iloc[-1] / self.N
    # Units: [steps] * [fraction of population]
    fraction_quarantine_steps = results_df['quarantined'].sum() / self.N
    # Units: [days] * [fraction of population]
    self.fraction_quarantine_time = fraction_quarantine_steps / self.steps_per_day
    total_tests = results_df['test_rate'].sum() / self.steps_per_day
    # Number of tests performed, as fraction of the population.
    self.fraction_tests = total_tests / self.N
    self.peak_test_rate = results_df['test_rate'].max() / self.N

  def plot_trends(self, fraction_of_population=True, hyperparams=True, G_attrs=False, columns=None, vertical=False):
    """Plots the time series of self.df, with the specified columns."""
    if columns is None:
      columns = MAIN_GROUPS
    title = ''
    if hyperparams:
      title += str({k: round(v, 3) for k, v in self.hyperparams.items()})
    if G_attrs:
      title += str(self.G_attrs)

    if fraction_of_population:
      scale = self.N
      ylabel = 'Fraction of population'
    else:
      scale = 1
      ylabel = 'Individuals'

    if vertical:
      fig, ax_arr = plt.subplots(2, 1, figsize=(10, 20))
      ax_arr[0].set_title('Epidemic Simulation')
      ax_arr[1].set_title('log scale')
    else:
      fig, ax_arr = plt.subplots(1, 2)
    fig.suptitle(title, fontsize=18)
    results_to_plot = self.df.drop('step', axis=1).set_index('day') / scale
    results_to_plot = results_to_plot[columns]
    for pane_ind, logy in enumerate([False, True]):
      ax = results_to_plot.plot(ax=ax_arr[pane_ind], logy=logy)
      ax.set_ylabel(ylabel)
      if logy:
        ax.get_legend().remove()
    return ax_arr

  def summary(self, hyperparams=False, G_attrs=False, plot=False, **plot_kwargs):
    """Print several summary stats, and potentially plot the trends."""
    summary_dict = {attr_name: getattr(self, attr_name) for attr_name in SUMMARY_ATTRS}
    summary_series_list = [pd.Series(summary_dict)]
    if hyperparams:
      summary_series_list.append(pd.Series(self.hyperparams))
    if G_attrs:
      summary_series_list.append(pd.Series(self.G_attrs))
    print(pd.concat(summary_series_list))
    if plot:
      self.plot_trends(hyperparams=hyperparams, G_attrs=G_attrs, **plot_kwargs)

def get_gamma_distribution_params(mean, std):
  """Turn mean and std of Gamma distribution into parameters k and theta."""
  # mean = k * theta
  # var = std**2 = k * theta**2
  theta = std**2 / mean
  k = mean / theta
  return k, theta

def edges2graph(edges, N=N):
  """Creates MultiGraph from list of edges."""
  H = nx.MultiGraph()
  H.add_nodes_from(np.arange(N))
  H.add_edges_from(edges)
  return H

def generate_power_law_degrees(N=N, min_degree=MIN_DEGREE, mean_degree=MEAN_DEGREE, gamma=GAMMA):
  """Generate array of degrees according to a power law distribution.

  If U~U[0,1] then generate power law values as
  X = U^(-gamma) - 1
  and then stretch them linearly to provide mean and min as required.
  Then convert them to integers with probability of being rounded up equal to
  the fractional part of the number.
  """
  assert gamma >= 0
  power_law_values = 1 / np.random.random(N)**gamma - 1 + _EPSILON
  value_mean = np.mean(power_law_values)
  frac_degrees = min_degree + power_law_values / value_mean * (
      mean_degree - min_degree)
  int_degrees = frac_degrees.astype(int) + (
      np.random.random(N) < (frac_degrees % 1))
  return int_degrees

def generate_scale_free_graph(N=N, min_degree=MIN_DEGREE, mean_degree=MEAN_DEGREE, gamma=GAMMA):
  """Generates a graph with power-law degree distribution.

  - Draws random degrees according to a power-law distribution.
  - Creates one-sided edges for each node according to its degree.
  - Randomly matches these one-sided edges to another one-sided edge to create
    a single edge.
  - Genereates a graph from this edge list.

  Higher gamma means a fatter tail of the degree distribution.
  Name of returned graph is 'power_law_{gamma}_{min_degree}_{mean_degree}'
  """
  degrees = generate_power_law_degrees(N, min_degree, mean_degree, gamma)
  # pd.Series(degrees).value_counts().sort_index().to_frame().plot(loglog=True)
  nodes_multiple = np.concatenate([np.full(degree, i) for i, degree in enumerate(degrees)])
  np.random.shuffle(nodes_multiple)
  if nodes_multiple.shape[0] % 2 == 1:
    nodes_multiple = nodes_multiple[:-1]
  edges = nodes_multiple.reshape((nodes_multiple.shape[0] // 2, 2))
  H = edges2graph(edges, N)
  # In Graph name, include the parameters.
  H.name = f'power_law_{round(gamma, 3)}_{min_degree}_{mean_degree}'
  return H

def genereate_local_edges(N=N, n_divisions=4, mean_degree=MEAN_DEGREE):
  """Generates edges with high clustering coefficient and constant degree.

  Smaller n_divisions -> higher clustering.
  """
  all_nodes = np.arange(N)
  division_sizes = np.diff(np.linspace(0, mean_degree/2, n_divisions+1).astype(int))
  edges_groups = []
  for division_size in division_sizes:
    np.random.shuffle(all_nodes)
    for i in range(1, division_size + 1):
      # In the random ordering, create edges between each node to the one
      # that is `i` places later.
      edges_groups.append(np.stack([all_nodes[i:], all_nodes[:-i]], axis=1))

  edges = np.concatenate(edges_groups)
  return edges

def generate_local_graph(N=N, n_divisions=4, mean_degree=MEAN_DEGREE):
  """Generates a graph with a much higher clustering coefficient.

  Smaller n_divisions -> higher clustering.
  """
  edges = genereate_local_edges(N, n_divisions, mean_degree)

  H = edges2graph(edges, N)
  # In Graph name, include the parameters.
  H.name = f'local_graph_{n_divisions}_{mean_degree}'
  return H

def genereate_local_edges_from_degrees(degrees, n_divisions=4):
  """Generates edges with high clustering coefficient, using list of degrees.

  Doesn't exactly preserve this list of degrees. It becomes narrower.

  Smaller n_divisions -> higher clustering: clustering coefficient ~ 1/n_divisions
  """
  N = len(degrees)
  all_nodes = np.arange(N)
  divisions_per_node = [np.diff(np.linspace(0, degree, n_divisions+1).astype(int))
                        for degree in degrees]
  edges_groups = []
  for division_ind in range(n_divisions):
    np.random.shuffle(all_nodes)
    for j in range(N):
      node = all_nodes[j]
      n_neighbors = divisions_per_node[node][division_ind]
      # In the random ordering, create edges between each node to the ones
      # that are immediately after it.
      neighbors = all_nodes[j + 1:j + 1 + n_neighbors]
      node_edges = np.stack([np.full_like(neighbors, node), neighbors], axis=1)
      edges_groups.append(node_edges)

  edges = np.concatenate(edges_groups)
  # Drop half of edges, to make the average degree right.
  edges = edges[np.random.random(len(edges)) < 0.5]
  return edges

def generate_local_scale_free_graph(N=N, n_divisions=4, min_degree=MIN_DEGREE,
                                    mean_degree=MEAN_DEGREE, gamma=GAMMA):
  """Generates a scale-free graph with a much higher clustering coefficient.

  WARNING: Resulting graph distribution is not the same as generate_scale_free_graph
    with the sama parameters. This one has a thinner tail.
  Smaller n_divisions -> higher clustering: clustering coefficient ~ 1/n_divisions
  Graph name is 'local_power_law_{n_divisions}_{gamma}_{min_degree}_{mean_degree}'
  """
  degrees = generate_power_law_degrees(N, min_degree, mean_degree, gamma)
  edges = genereate_local_edges_from_degrees(degrees, n_divisions=n_divisions)

  H = edges2graph(edges, N)
  # In Graph name, include the parameters.
  H.name = f'local_power_law_{n_divisions}_{round(gamma, 3)}_{min_degree}_{mean_degree}'
  return H

def random_subset(indicator_arr, sample_probs):
  """In array of 0,1 leave subset of 1s with probabilities sample_probs.

  Args:
    indicator_arr: An array of 0,1 values.
    sample_probs: A number or an array with same shape as indicator_arr, for
      individual heterogeneous sampling probabilities.
  """
  subset_arr = (np.random.random(indicator_arr.shape) < sample_probs) & indicator_arr
  return subset_arr

def init_states(N=N,
                initial_infected_num=INITIAL_INFECTED_NUM,
                incubation_duration_mean=INCUBATION_DURATION_MEAN,
                incubation_duration_std=INCUBATION_DURATION_STD,
                prob_recover=PROB_RECOVER,
                prob_asymptomatic=PROB_ASYMPTOMATIC):
  """Init the state variables of the simulation.

  Each group has an array of 0/1 indicators, and potentially days left array.
  Susceptible (S_arr), Exposed (E_arr, E_left), Infected (I_arr), Asymptomatic (A_arr),
  Recovered (R_arr), Quarantined (Q_arr, Q_left), Tested Positive (TP_arr),
  Time until test results return (T_result_left), what the test results wil be
  when they return (T_result_positive_arr).
  i-th indicator value of 1 means the i-th node is a member of that group.
  *_left are arrays of number of days left until transition to next state. A
  value of -1 or less means Not Available. For example an individual not in the
  E group (i.e. has E_arr[i]==0) does not have a meaningful E_left[i], so it
  will be -1 or less.

  Args:
    prob_recover: per day. Inverse of expected duration of being Infected.
    prob_asymptomatic: Share of population who would not develop symptoms even
      when sick. They are not detected along with the symptomatic Infected in
      the testing step, but are detected by direct tests (mass testing or when
      neighbors are tested).
  """
  all_nodes_arr = np.arange(N)
  S_arr = np.ones_like(all_nodes_arr)
  # Use expected ratio of exposed to infected.
  initial_exposed_num = int(initial_infected_num * incubation_duration_mean * prob_recover)
  initial_infected_num = int(initial_infected_num)
  # Draw E & I, mutually exclusive.
  exposed_and_infected_inds = np.random.choice(all_nodes_arr, initial_exposed_num + initial_infected_num, replace=False)
  exposed_inds = exposed_and_infected_inds[:initial_exposed_num]
  infected_inds = exposed_and_infected_inds[initial_exposed_num:]
  E_arr = np.zeros_like(S_arr)
  E_arr[exposed_inds] = 1
  I_arr = np.zeros_like(S_arr)
  I_arr[infected_inds] = 1
  A_arr = random_subset(np.ones_like(S_arr), prob_asymptomatic)
  S_arr -= (I_arr + E_arr)
  R_arr = np.zeros_like(S_arr)
  Q_arr = np.zeros_like(S_arr)
  # Time left to exit quarantine.
  Q_left = np.full_like(S_arr, -1, dtype='float')
  # Time left to become infected.
  E_left = np.full_like(S_arr, -1, dtype='float')
  # Seed with incubation periods.
  incubation_k, incubation_theta = get_gamma_distribution_params(incubation_duration_mean, incubation_duration_std)
  incubation_durations = np.random.gamma(incubation_k, incubation_theta, N)
  np.copyto(E_left, incubation_durations, where=(E_arr > 0), casting='safe')
  TP_arr = np.zeros_like(S_arr)
  # Time left until test results come back.
  T_result_left = np.full_like(S_arr, -1)
  # What the test results will be, when they come back.
  T_result_positive_arr = np.zeros_like(S_arr)
  return S_arr, E_arr, E_left, I_arr, A_arr, R_arr, Q_arr, Q_left, TP_arr, T_result_left, T_result_positive_arr

def number_of_edges_to_group(group_arr, adj_mat):
  """Calculates number of edges from each node to nodes in group_arr.

  Returns array of integers representing number of edges for each node.
  Multiple edges are counted as their multiplication.

  Args:
    group_arr: Indicator array - 1 indicates group membership.
    adj_mat: Adjacency matrix of the graph.

  Returns:
    Array indicating for each member of the population the number of neighbors
    they have in the group indicated by group_arr.
  """
  if (group_arr == 0).all():
    # Performance optimization. If the group is empty, don't count edges.
    return np.zeros_like(group_arr)
  # Adjacency matrix filtered down to neighbors of group.
  adj_mat_group = adj_mat.multiply(group_arr)
  # number of connections to infected (could be multiple connections from the same infected)
  connections_to_group = np.array(adj_mat_group.sum(axis=1))[:, 0]
  return connections_to_group

def spread_to_neighbors(from_group_arr, prob_spread, adj_mat, to_group=None):
  """Neighbors of the group are marked with a certain probability for each edge.

  Each member of the group denoted by from_group_arr spreads to each one of its
  neighbors with probability prob_spread. This can be used for infection,
  tracing or anything else.

  Args:
    from_group_arr: Indicator array of seed group, from which the spread emanates.
    prob_spread: Probability that each member of the group "marks" or spreads
      to each of its neighbors.
    adj_mat: Adjacency matrix of the graph.
    to_group: Indicator array. Restrict spread to members of this group. Default
      is no restriction - the entire population.

  Returns:
    Indicator array of 1s for those who were marked.
  """
  spread_probs = individual_infection_probabilities(from_group_arr, prob_spread, adj_mat)
  if to_group is None:
    to_group = np.ones_like(from_group_arr)
  return random_subset(to_group, spread_probs)

def individual_infection_probabilities(I_arr, prob_infect, adj_mat):
  """Individual probabilities of being infected by the infectious in I_arr.

  Counts the number of infectious neighbors, and calculates individual
  probability of being infected.
  Probability of being infected by each connection is prob_infect.

  Args:
    prob_infect: Probability of infecting a single neighbor.
    adj_mat: Adjacency matrix of the graph.
  """
  if prob_infect == 0:
    # Performance optimization - Don't count edges when no chance of infection.
    return np.zeros_like(I_arr)
  # Number of connections to infected (could be multiple connections from the same infected)
  connections_to_I = number_of_edges_to_group(I_arr, adj_mat)
  # Individual probability of not being infected: (1 - p_infect)**n_connections
  no_infection_probs = np.exp(np.log(1 - prob_infect) * connections_to_I)
  return 1 - no_infection_probs

def infection_step(S_arr, E_arr, E_left, I_arr, A_arr, Q_arr, adj_mat,
                   incubation_k, incubation_theta, prob_infect,
                   prob_infect_exposed_factor=PROB_INFECT_EXPOSED_FACTOR,
                   relative_infectiousness_asymptomatic=RELATIVE_INFECTIOUSNESS_ASYMPTOMATIC,
                   duration_exposed_infects=DURATION_EXPOSED_INFECTS):
  """S -> E. Infectious individuals infect their neighbors with some probability.

  Individuals who are infectious are those who are in I but not in Quarantine
  and those who are Exposed (during final infectiousness duration) but not in Quarantine.
  Infectious individuals interact with each of their neighbors once (or multiple
  times if there are multiple edges between them).
  Each interaction with a neighbor infects that neighbor with probability
  prob_infect (for Infected) and prob_infect * prob_infect_exposed_factor (for Exposed),
  but only if that neighbor is Susceptible and not Quarantined.
  For Asymptomatics, their probabilities of infecting (in both Infected and Exposed)
  are multiplied by relative_infectiousness_asymptomatic.

  All infected individuals become Exposed, and are given an incubation time.

  Args:
    prob_infect: Probability of infection *per step* (not per day).
    incubation_k, incubation_theta: parameters of incubation period length
      in *days*, not steps.
    prob_infect_exposed_factor: Ratio of infectiousness between Exposed (in the
      final days before becoming Infected) and Infected. In other words,
      prob_infect_exposed = prob_infect * prob_infect_exposed_factor.
      Set this to 0 to make the Exposed non-infectious.
    relative_infectiousness_asymptomatic: Ratio of infectiousness between
      Asymptomatic and symptomatic. In other words,
      prob_infect_asymptomatic = prob_infect * relative_infectiousness_asymptomatic.
      The same ratio holds for the final contagious part of the incubation period.
      Set this to 0 to make the Asymptomatic non-infectious.
    duration_exposed_infects: How many days before developing symptoms (becoming
      Infected) an Exposed individual is contagious.
  """
  # TODO: Can unite these 4 matrix multiplications into 1, to optimize.
  # Those who can be infected - Quarantined cannot be.
  susceptible_not_quarantined = S_arr & (1 - Q_arr)

  # Newly infected by symptomatic Infected who are not Quarantined.
  new_infected = spread_to_neighbors(
      I_arr & (1 - A_arr) & (1 - Q_arr), prob_infect, adj_mat, susceptible_not_quarantined)

  # Newly infected by Asymptomatic Infected who are not Quarantined.
  prob_infect_asymptomatic = prob_infect * relative_infectiousness_asymptomatic
  new_infected |= spread_to_neighbors(
      I_arr & A_arr & (1 - Q_arr), prob_infect_asymptomatic, adj_mat, susceptible_not_quarantined)

  # Newly infected by the Exposed who are infectious, and not Quarantined.
  # Only infectious a few days before developing symptoms.
  E_infectious_arr = E_arr & (E_left <= duration_exposed_infects)
  # Exposed are less infectious.
  prob_infect_exposed = prob_infect * prob_infect_exposed_factor
  # Infections by Exposed who will become symptomatic.
  new_infected |= spread_to_neighbors(
      E_infectious_arr & (1 - A_arr) & (1 - Q_arr), prob_infect_exposed,
      adj_mat, susceptible_not_quarantined)
  # Infections by Exposed who will become Asymptomatic.
  prob_infect_exposed_asymptomatic = prob_infect_exposed * relative_infectiousness_asymptomatic
  new_infected |= spread_to_neighbors(
      E_infectious_arr & A_arr & (1 - Q_arr), prob_infect_exposed_asymptomatic,
      adj_mat, susceptible_not_quarantined)

  # new_infected are now Exposed.
  E_arr = E_arr | new_infected
  incubation_durations = np.random.gamma(incubation_k, incubation_theta, S_arr.shape[0])
  np.copyto(E_left, incubation_durations, where=(new_infected > 0), casting='safe')
  S_arr = S_arr & (1 - new_infected)
  return S_arr, E_arr

def incubation_step(E_arr, E_left, I_arr, steps_per_day=STEPS_PER_DAY):
  """E -> I. Incubation period ends for some individuals. They become Infected.

  Args:
    E_left: Array of days left in incubation. When reaches 0, become infectious.
  """
  E_left = E_left - 1 / steps_per_day  # Decrease time left in incubation.
  # Those whose incubation period ends.
  become_infected = E_arr & (E_left < _EPSILON)
  E_arr = E_arr - become_infected
  I_arr = I_arr + become_infected
  return E_arr, E_left, I_arr

def recovery_step(I_arr, R_arr, prob_recover):
  """I -> R. Infected individuals Recover with given step-wise probability."""
  new_recovered = (np.random.random(I_arr.shape) < prob_recover) & I_arr
  R_arr = R_arr + new_recovered
  I_arr = I_arr - new_recovered
  return I_arr, R_arr

def tests_analyzed(T_result_left, T_result_positive_arr, steps_per_day):
  """Time left for results to come back decreased, some results come back.

  Args:
    T_result_left: Array. For each agent, number days left until test result comes back.
    T_result_positive_arr: Whether test results (When they come back) will be positive.

  Returns:
    new_TP_arr: New positives.
    T_result_left: Updated, with less time left.
  """
  # Decrease time left for test results to come back.
  T_result_left = T_result_left - 1. / steps_per_day
  # Test results come back for some agents.
  test_comes_back_arr = ((T_result_left < _EPSILON) & (T_result_left > -1)).astype(int)
  new_TP_arr = test_comes_back_arr & T_result_positive_arr
  # No longer waiting for test results.
  T_result_left = np.where(test_comes_back_arr, -1, T_result_left)
  return new_TP_arr, T_result_left

def test_group(tested_arr, carrier_arr, T_result_left, T_result_positive_arr, test_delay_time):
  """Tests a group, returning test results, results time, number tested.

  Args:
    tested_arr: Indicator array for those to be tested.
    carrier_arr: Of tested, those who will test positive. Only its logical and
      with tested_arr is used.
    T_result_left: Array. For each agent, number days left until test result comes back.
    T_result_positive_arr: Whether test results (When they come back) will be positive.
    test_delay_time: Delay between the time a test is performed and the time
      results return and contact-traced neighbors are tested/quarantined. Days.
  """
  # Record positive results for the positives (when they come back)
  T_result_positive_arr = T_result_positive_arr | (tested_arr & carrier_arr)
  # Start the clock on new tests being analyzed.
  T_result_left = np.where(tested_arr, test_delay_time, T_result_left)
  # Assumes 0 negatives (false & true). Divide by P(test=positive) for more realistic estimate.
  n_tested = tested_arr.sum()

  return T_result_left, T_result_positive_arr, n_tested

def symptomatics_tested(I_arr, A_arr, TP_arr, new_TP_arr, T_result_left, T_result_positive_arr, prob_infected_detected, test_delay_time):
  """Some fraction of symptomatics is tested. Asymptomatic patients aren't tested.

  Args:
    I_arr: Indicator array of Infected.
    A_arr: Asymptomatics. Those who would not develop symptoms even if infected.
    TP_arr: Those who tested positive previously. They aren't retested.
    new_TP_arr: Those whose positive test results came back in this step.
    T_result_left: Array. For each agent, number days left until test result comes back.
    T_result_positive_arr: Whether test results (When they come back) will be positive.
    prob_infected_detected: probability *per step*, not per day.
    test_delay_time: Delay between the time a test is performed and the time
      results return and contact-traced neighbors are tested/quarantined. Days.

  Returns:
    T_result_left, T_result_positive_arr, n_infected_tested.
  """
  # Infected group is tested with some probability. Except Asymptomatic.
  new_symptomatic_tested_arr = random_subset(I_arr & (1 - A_arr), prob_infected_detected)
  # Known positives and those pending test results aren't retested.
  pending_test_results_arr = (T_result_left > 0).astype(int)
  new_symptomatic_tested_arr &= (1 - (TP_arr | new_TP_arr | pending_test_results_arr))

  T_result_left, T_result_positive_arr, n_infected_tested = test_group(
      new_symptomatic_tested_arr, I_arr, T_result_left, T_result_positive_arr, test_delay_time)

  return T_result_left, T_result_positive_arr, n_infected_tested, new_symptomatic_tested_arr

def general_population_tested(E_arr, I_arr, TP_arr, new_TP_arr, T_result_left, T_result_positive_arr, prob_exposed_detected, test_delay_time):
  """Some fraction of the general population is tested.

  Args:
    E_arr: Indicator array of Exposed.
    I_arr: Indicator array of Infected.
    TP_arr: Those who tested positive previously. They aren't retested.
    new_TP_arr: Those whose positive test results came back in this step.
    T_result_left: Array. For each agent, number days left until test result comes back.
    T_result_positive_arr: Whether test results (When they come back) will be positive.
    prob_exposed_detected: Share of general population which is tested, *per step*.
    test_delay_time: Delay between the time a test is performed and the time
      results return and contact-traced neighbors are tested/quarantined. Days.

  Returns:
    T_result_left, T_result_positive_arr, n_infected_tested.
  """
  # Known positives and those pending test results aren't retested.
  pending_test_results_arr = (T_result_left > 0).astype(int)
  # Random subset of entire population tested,
  # except known positives and those pending test results.
  new_genpop_tested_arr = random_subset(1 - (TP_arr | new_TP_arr | pending_test_results_arr),
                                        prob_exposed_detected)

  # Carriers (Exposed or Infected) will have positive test results (when back).
  carrier_arr = E_arr | I_arr
  T_result_left, T_result_positive_arr, n_general_tested = test_group(
      new_genpop_tested_arr, carrier_arr, T_result_left, T_result_positive_arr, test_delay_time)

  return T_result_left, T_result_positive_arr, n_general_tested

def contact_tracing(E_arr, I_arr, TP_arr, new_TP_arr, T_result_left,
                    T_result_positive_arr, new_Q_arr, prob_neighbor_traced,
                    quarantine_neighbors, test_neighbors, test_delay_time, adj_mat):
  """Contact tracing. Neighbors of positives are traced and tested / quarantined.

  For each positive, each neighbor is traced with probability prob_neighbor_traced.

  Args:
    quarantine_neighbors: When an individual tests positive, whether their
      traced neighbors are quarantined.
    test_neighbors: When an individual tests positive, whether their traced
      neighbors are tested.

  Returns:
    new_Q_arr, T_result_left, T_result_positive_arr, n_neighbors_traced, n_neighbors_tested.
  """
  # Check if no neighbors are tested, to avoid matrix multiplication if not required.
  if prob_neighbor_traced == 0:
    n_neighbors_tested = 0
    n_neighbors_traced = 0
    return new_Q_arr, T_result_left, T_result_positive_arr, n_neighbors_traced, n_neighbors_tested

  neighbors_traced = spread_to_neighbors(new_TP_arr, prob_neighbor_traced, adj_mat)
  n_neighbors_traced = neighbors_traced.sum()
  n_neighbors_tested = 0

  if quarantine_neighbors:
    # Quarantine traced neighbors.
    # Known positives aren't newly quarantined (either covered elsewhere or recovered).
    new_Q_arr = new_Q_arr | (neighbors_traced & (1 - (new_TP_arr | TP_arr)))

  if test_neighbors:
    # Test the traced neighbors.
    # Carriers (Exposed or Infected) will have positive test results (when back).
    carrier_arr = E_arr | I_arr
    # Known positives and those pending test results aren't retested.
    pending_test_results_arr = (T_result_left > 0).astype(int)
    neighbors_tested = neighbors_traced & (1 - (TP_arr | new_TP_arr | pending_test_results_arr))
    T_result_left, T_result_positive_arr, n_neighbors_tested = test_group(
        neighbors_tested, carrier_arr, T_result_left, T_result_positive_arr, test_delay_time)

  return new_Q_arr, T_result_left, T_result_positive_arr, n_neighbors_traced, n_neighbors_tested

def testing_step(E_arr, I_arr, A_arr, TP_arr, T_result_left, T_result_positive_arr,
                 prob_infected_detected,
                 prob_neighbor_traced,
                 prob_exposed_detected,
                 quarantine_neighbors,
                 test_neighbors,
                 test_delay_time,
                 adj_mat=None,
                 steps_per_day=STEPS_PER_DAY):
  """Testing of subsets of Infected, the general population. Contact tracing.

  A test comes out positive if a node is either Exposed or Infected.
  Test symptomatics: a random subset of Infected people are tested (and found
    positive). Symptomatics are quarantined even before test results return.
  Mass-testing: a random subset of the population is tested (Exposed & Infected
    people test positive.)
  Contact tracing: a random subset of neighbors of those who tested positive,
  are traced, and then either tested or quarantined without being tested.
  Counts number of tests performed, assuming 0 false negative tests.
  Known positives and those pending test results are not retested.

  Args:
    A_arr: Asymptomatics. Those who would not develop symptoms even if infected.
    TP_arr: Those who tested positive previously. They aren't retested.
    T_result_left: Array. For each agent, number days left until test result comes back.
    T_result_positive_arr: Whether test results (When they come back) will be positive.
    prob_infected_detected: probability *per step*, not per day.
    prob_neighbor_traced: absolute probability, not per day / step.
    prob_exposed_detected: Share of general population which is tested, *per step*.
    quarantine_neighbors: When an individual tests positive, whether their
      traced neighbors are quarantined.
    test_neighbors: When an individual tests positive, whether their traced
      neighbors are tested.
    test_delay_time: Delay between the time a test is performed and the time
      results return and contact-traced neighbors are tested/quarantined. Days.
    adj_mat: not required if prob_neighbor_traced=0.

  Returns:
    TP_arr, new_Q_arr, T_result_left, T_result_positive_arr, n_infected_tested,
    n_neighbors_traced, n_neighbors_tested, n_general_tested
  """
  # TODO: Testing for other populations (high degree nodes).
  new_TP_arr, T_result_left = tests_analyzed(
      T_result_left, T_result_positive_arr, steps_per_day)

  (T_result_left, T_result_positive_arr, n_infected_tested,
   new_symptomatic_tested_arr) = symptomatics_tested(
       I_arr, A_arr, TP_arr, new_TP_arr, T_result_left, T_result_positive_arr,
       prob_infected_detected, test_delay_time)

  T_result_left, T_result_positive_arr, n_general_tested = general_population_tested(
      E_arr, I_arr, TP_arr, new_TP_arr, T_result_left, T_result_positive_arr,
      prob_exposed_detected, test_delay_time)

  # Indicator for those who need to enter Quarantine. New positives and symptomatic.
  new_Q_arr = new_TP_arr | new_symptomatic_tested_arr
  # TODO: Add non-covid symptomatics who are quarantined, and released on negative test results.

  (new_Q_arr, T_result_left, T_result_positive_arr, n_neighbors_traced,
   n_neighbors_tested) = contact_tracing(E_arr, I_arr, TP_arr, new_TP_arr,
                                         T_result_left, T_result_positive_arr,
                                         new_Q_arr, prob_neighbor_traced,
                                         quarantine_neighbors, test_neighbors,
                                         test_delay_time, adj_mat)
  TP_arr = TP_arr | new_TP_arr
  return TP_arr, new_Q_arr, T_result_left, T_result_positive_arr, n_infected_tested, n_neighbors_traced, n_neighbors_tested, n_general_tested

def quarantine_step(Q_arr, Q_left, R_arr, TP_arr, new_Q_arr,
                    days_in_quarantine=DAYS_IN_QUARANTINE,
                    steps_per_day=STEPS_PER_DAY):
  """Individuals exit quarantine. Those who tested positive enter quarantine.

  Quarantined people whose time is up exit Quarantine.
  Those who tested positive and haven't recovered stay in Quarantine.
  Those who tested positive and have recovered exit Quarantine, even if their time isn't up.
  Those who should (tested positive or traced neighbors) enter Quarantine.
  Quarantined individuals can neither infect nor be infected.

  Args:
    Q_left: Array of days left in quarantine. When reaches 0, quarantine ends.
    TP_arr: array indicating those who tested positive.
    new_Q_arr: array indicating those who should enter Quarantine in this step.
  """
  # Exiting quarantine
  # ------------------
  Q_left = Q_left - 1. / steps_per_day  # Decrease time left in Quarantine.
  # Known positives who haven't recovered.
  known_positives = TP_arr & (1 - R_arr)
  known_recovered = TP_arr & R_arr
  # Release those whose Quarantine time is up.
  exit_quarantine = Q_arr & (Q_left < _EPSILON)
  # Don't release known positives who haven't recovered.
  exit_quarantine &= (1 - known_positives)
  # Release known positives who have recovered (even if their 14 days aren't up).
  exit_quarantine |= known_recovered & Q_arr
  # Apply exiting quarantine.
  Q_arr &= (1 - exit_quarantine)
  # Entering quarantine
  # ------------------
  new_in_quarantine = new_Q_arr & (1 - Q_arr)
  Q_arr |= new_in_quarantine
  # Start the clock on time in Quarantine.
  np.copyto(Q_left, days_in_quarantine, where=(new_in_quarantine > 0), casting='safe')
  return Q_arr, Q_left

def create_counter(I_arr, E_arr, R_arr, Q_arr, S_arr, n_infected_tested, n_neighbors_traced, n_neighbors_tested, n_general_tested):
  """Counts number of individuals in each group.

  Returns a dict from group name to sum of its indicator array.
  """
  return dict(infected=I_arr.sum(),
              exposed=E_arr.sum(),
              recovered=R_arr.sum(),
              quarantined=Q_arr.sum(),
              susceptible=S_arr.sum(),
              n_infected_tested=n_infected_tested,
              n_neighbors_traced=n_neighbors_traced,
              n_neighbors_tested=n_neighbors_tested,
              n_general_tested=n_general_tested)

def simulation(N=N,
               G=None,
               initial_infected_num=INITIAL_INFECTED_NUM,
               prob_infect=PROB_INFECT,
               incubation_duration_mean=INCUBATION_DURATION_MEAN,
               incubation_duration_std=INCUBATION_DURATION_STD,
               prob_asymptomatic=PROB_ASYMPTOMATIC,
               prob_recover=PROB_RECOVER,
               quarantine_neighbors=QUARANTINE_NEIGHBORS,
               test_neighbors=TEST_NEIGHBORS,
               prob_infected_detected=PROB_INFECTED_DETECTED,
               prob_neighbor_traced=PROB_NEIGHBOR_TRACED,
               prob_exposed_detected=PROB_EXPOSED_DETECTED,
               test_delay_time=TEST_DELAY_TIME,
               days_in_quarantine=DAYS_IN_QUARANTINE,
               prob_infect_exposed_factor=PROB_INFECT_EXPOSED_FACTOR,
               relative_infectiousness_asymptomatic=RELATIVE_INFECTIOUSNESS_ASYMPTOMATIC,
               duration_exposed_infects=DURATION_EXPOSED_INFECTS,
               steps_per_day=STEPS_PER_DAY,
               min_degree=MIN_DEGREE,
               mean_degree=MEAN_DEGREE,
               gamma=GAMMA,
               max_steps=MAX_STEPS,
               verbose=False):
  """Simulates SEQIR model on a network.

  A Susceptible-Exposed-Infected-Recovered (SEIR) model for disease progression.
  This is an agent-based model, where each agent (person) is represented as a
  node in a graph, and infection happens between contacts, represented by graph
  edges. It also models testing and a Quarantine group of those who test positive,
  and contact tracing of their neighbors in the graph.
  At any given step, each node belongs to exactly one of the following groups:

  Susceptible (S): Nodes which weren’t infected yet. These are the only nodes
    which can be infected. When infected, they become Exposed.

  Exposed (E): Nodes which were infected and are now in their pre-symptomatic
    incubation period. They are infectious starting 2 days before the end of
    their (Gamma distributed) incubation period, at which point they develop
    symptoms (and become Infected).

  Infected (I): Nodes which are now symptomatic. More infectious than during the
    final Exposed period. When their (exponentially distributed) Infected period
    ends, they become Recovered.

  Recovered (R): Nodes who have recovered (or died). They are no longer
    infectious and cannot be infected.

  In addition to those groups, a node can also belong to additional groups:

  Tested Positive (TP): Nodes who tested positive for the disease. At each
    simulation step, a certain share of Infected nodes are tested and a certain
    share of all nodes are tested. Nodes which are Exposed or Infected test
    positive. They become Quarantined, and potentially some share of their
    neighbors is traced. This does not model false positives.

  Quarantined (Q): Nodes which are in quarantine. They can neither infect nor
    get infected. Quarantine ends after a set period. Their disease progression
    continues irrespective of Quarantine.

  The simulation keeps a state in the form of indicator arrays, which describe
  for each node which groups they are a member of: S_arr, E_arr, I_arr, R_arr,
  Q_arr, TP_arr. The time left before transition to the next state is kept in
  additional arrays Q_left (time left to exit Quarantine) and E_left (incubation
  time left before becoming Infected).

  All probabilities in args are *daily* probabilities, unless stated otherwise.

  Args:
    N: Number of nodes in the simulation.
    G: Graph on which to run the simulation. If None, generated using N and the
      args of graph attribute below.
    initial_infected_num: Number of individuals infected at t=0.
    prob_infect: In a single interaction on a single day, what is the chance of
     an Infectious to infect a Susceptible.
    incubation_duration_mean: Incubation period duration distribution mean. In
      days.
    incubation_duration_std: Incubation period duration distribution standard
      deviation. In days.
    prob_asymptomatic: Share of population who would not develop symptoms even
      when sick. They are not detected along with the symptomatic Infected in
      the testing step, but are detected by direct tests (mass testing or when
      neighbors are tested).
    prob_recover: Probability to change from Infections to Recovered in a single
      day.
    quarantine_neighbors: When an individual tests positive, whether their
      traced neighbors are quarantined.
    test_neighbors: When an individual tests positive, whether their traced
      neighbors are tested.
    prob_infected_detected: Probability that an Infected individual is tested
      and detected as a carrier in a single day. This encapsulates also the fact
      that many people are asymptomatic or will not be tested (or are not
      quarantined).
    prob_neighbor_traced: Probability that a neighbor of an individual who
      tested positive is himself traced and tested/quarantined. The testing of
      neighbors happens once, once the original individual is tested positive,
      so this is *not* a probability per day.
    prob_exposed_detected: Share of the general population which is tested on a
      single day. Exposed and Infected individuals who are tested are detected
      as carriers.
    test_delay_time: Delay between the time a test is performed and the time
      results return and contact-traced neighbors are tested/quarantined. Days.
    days_in_quarantine: Length of imposed quarantine, in days.
    prob_infect_exposed_factor: Ratio of infectiousness between Exposed (in the
      final days before becoming Infected) and Infected. In other words,
      prob_infect_exposed = prob_infect * prob_infect_exposed_factor.
      Set this to 0 to make the Exposed non-infectious.
    relative_infectiousness_asymptomatic: Ratio of infectiousness between
      Asymptomatic and symptomatic. In other words,
      prob_infect_asymptomatic = prob_infect * relative_infectiousness_asymptomatic.
      The same ratio holds for the final contagious part of the incubation period.
      Set this to 0 to make the Asymptomatic non-infectious.
    duration_exposed_infects: How many days before developing symptoms (becoming
      Infected) an Exposed individual is contagious. Questionable evidence, but
      1-2 produces a reasonable R0.
    steps_per_day: Number of simulation steps per one day of real life time.
    min_degree: Minimal degree of nodes in a generated graph.
    mean_degree: Mean degree of nodes in a generated graph.
    gamma: Parameter of degree distribution in generated graph.
      Higher gamma means a fatter tail of the degree distribution.
    max_steps: Maximal number of steps in simulation. If this number is reached,
      simulation stops.
    verbose: Whether to print some progress indicators.

  Returns:
    SimulationResults object with the results of the simulation. This holds the
    time series of total number of members in each group (Susceptible, Exposed
    etc.) as well as other variables of interest, such as number of tests
    performed.
  """
  # Process input.
  if G is None:
    if verbose:
      print('Generating graph...')
    G = generate_scale_free_graph(N, min_degree=min_degree, mean_degree=mean_degree, gamma=gamma)
    if verbose:
      print('Done!')
  else:
    N = G.number_of_nodes()
  adj_mat = nx.adjacency_matrix(G)

  initial_infected_num = int(initial_infected_num)  # Make sure it's an integer.
  incubation_k, incubation_theta = get_gamma_distribution_params(incubation_duration_mean, incubation_duration_std)

  (S_arr, E_arr, E_left, I_arr, A_arr, R_arr, Q_arr, Q_left, TP_arr,
   T_result_left, T_result_positive_arr) = init_states(
       N, initial_infected_num=initial_infected_num, incubation_duration_mean=incubation_duration_mean,
       incubation_duration_std=incubation_duration_std, prob_recover=prob_recover,
       prob_asymptomatic=prob_asymptomatic)
  n_infected_tested = n_neighbors_traced = n_neighbors_tested = n_general_tested = 0
  # Step-wise probabilities.
  prob_infect_per_step = prob_infect / steps_per_day
  prob_infected_detected_per_step = prob_infected_detected / steps_per_day
  prob_exposed_detected_per_step = prob_exposed_detected / steps_per_day
  prob_recover_per_step = prob_recover / steps_per_day

  max_step_wise_prob = max(prob_infect_per_step, prob_infected_detected_per_step,
                           prob_exposed_detected_per_step, prob_recover_per_step)
  if max_step_wise_prob > 0.5:
    print('WARNING: steps_per_day too small? Maximal step-wise probability =', max_step_wise_prob)

  counters = []

  # Main loop.
  for step_num in range(max_steps):
    # Save metrics.
    counters.append(create_counter(I_arr, E_arr, R_arr, Q_arr, S_arr,
                                   n_infected_tested, n_neighbors_traced, n_neighbors_tested, n_general_tested))

    # Infection step. S -> E
    S_arr, E_arr = infection_step(S_arr, E_arr, E_left, I_arr, A_arr, Q_arr, adj_mat,
                                  incubation_k=incubation_k,
                                  incubation_theta=incubation_theta,
                                  prob_infect=prob_infect_per_step,
                                  prob_infect_exposed_factor=prob_infect_exposed_factor,
                                  relative_infectiousness_asymptomatic=relative_infectiousness_asymptomatic,
                                  duration_exposed_infects=duration_exposed_infects)

    # Incubation step. E -> I
    E_arr, E_left, I_arr = incubation_step(E_arr, E_left, I_arr,
                                           steps_per_day=steps_per_day)

    # Testing step. Update tested_positive and newly quarantined.
    (TP_arr, new_Q_arr, T_result_left, T_result_positive_arr, n_infected_tested,
     n_neighbors_traced, n_neighbors_tested, n_general_tested) = testing_step(
         E_arr, I_arr, A_arr, TP_arr, T_result_left, T_result_positive_arr,
         prob_infected_detected=prob_infected_detected_per_step,
         prob_neighbor_traced=prob_neighbor_traced,
         prob_exposed_detected=prob_exposed_detected_per_step,
         quarantine_neighbors=quarantine_neighbors,
         test_neighbors=test_neighbors,
         test_delay_time=test_delay_time,
         adj_mat=adj_mat,
         steps_per_day=steps_per_day)

    # Quarantine step. Update Q
    Q_arr, Q_left = quarantine_step(
        Q_arr, Q_left, R_arr, TP_arr, new_Q_arr,
        days_in_quarantine=days_in_quarantine,
        steps_per_day=steps_per_day)

    # Recovery step. I -> R
    I_arr, R_arr = recovery_step(I_arr, R_arr, prob_recover_per_step)

    # Stop condition. No more Infected or Exposed.
    if I_arr.sum() == 0 and E_arr.sum() == 0:
      break
  else:
    print('WARNING: max number of steps reached.')

  # Add final counters.
  counters.append(create_counter(I_arr, E_arr, R_arr, Q_arr, S_arr,
                                 n_infected_tested, n_neighbors_traced, n_neighbors_tested, n_general_tested))

  # Convert to DataFrame.
  results_df = pd.DataFrame(counters)
  tested = results_df.filter(regex='_tested').sum(axis=1)
  # Test rate: number of tests per day. Smoothe using day-length sliding window.
  results_df['test_rate'] = np.convolve(tested, np.ones((steps_per_day,)), mode='same')
  results_df['step'] = np.arange(len(results_df))
  results_df['day'] = results_df['step'] / steps_per_day
  results = SimulationResults(results_df, G, prob_infect=prob_infect,
                              prob_infected_detected=prob_infected_detected,
                              prob_neighbor_traced=prob_neighbor_traced,
                              prob_exposed_detected=prob_exposed_detected,
                              quarantine_neighbors=quarantine_neighbors,
                              test_neighbors=test_neighbors,
                              steps_per_day=steps_per_day)
  if verbose:
    print('Simulation finished!')
  return results


import json
import os
import traceback
from collections import UserDict
from functools import partial

from configobj import ConfigObj
from yaml import load, dump
from yaml import Loader, Dumper

def format_json(content):
    if isinstance(content, (str, bytes)):
        try:
            content = json.loads(content)
        except:
            return content

    result = json.dumps(content, sort_keys=True, indent=4, separators=(',', ': ')). \
        encode('latin-1').decode('unicode_escape')

    return result

def get_root_dir(flag_file_name, cur_dir):
    if os.path.isfile(cur_dir):
        cur_dir = os.path.dirname(cur_dir)
    if flag_file_name in os.listdir(cur_dir):
        return cur_dir

    up_dir = os.path.dirname(cur_dir)
    if up_dir == cur_dir:
        return None
    return get_root_dir(flag_file_name, up_dir)

class WDict(UserDict):
    """
    It's hard to name！！！
    注意：正常使用key中不能再出现 "."，否则会被拆分掉，如a['a.b.c']等于a['a']['b']['c']
    示例：> d = {'a': {'b': 'c': 1}}
         > print(WDict(**d)['a.b.c']) # -> 1
    """

    def __getitem__(self, item: str):
        key_list = item.split('.')
        content = self.data
        for key in key_list:
            key = int(key) if str.isdigit(key) else key
            try:
                content = content[key]
            except (KeyError, IndexError):
                return None
        return content

class DDict:
    """
    It's hard to name！！！
    示例: > d = {'a': {'b': 'c': 1}}
         > print(d.a.b.c) # 1
    """

    def __init__(self, d):
        self.__d = d

    def __getattr__(self, item):
        try:
            return DDict(self.__d[item])
        except TypeError:
            raise ValueError('找不到%s配置，请检查' % item)

    def __str__(self):
        return str(self.__d)

def merge_dict(d1: dict, d2: dict):
    """
    merge d2 to d1
    d2有d1没有，则新增
    d1、d2同时有，但不全是字典，则替换
    d1、d2同时有，但都是字典，则递归调用合并函数
    """
    for k in d2:
        if k in d1:
            if isinstance(d1[k], dict) and isinstance(d2[k], dict):
                merge_dict(d1[k], d2[k])
                continue
        d1[k] = d2[k]

class ConfigFileParser:
    def __init__(self, suffix, file_path):
        if not isinstance(file_path, str):
            raise ValueError('file_path error, is not a valid str')

        if not os.path.isfile(file_path):
            raise ValueError('file_path error, is not a valid file path')

        self.file_path = file_path
        self.content = None
        self.parser = self.get_parser(suffix)

        try:
            with open(self.file_path, encoding='utf-8') as f:
                self.content = self.parser(f) or {}
        except Exception:
            traceback.print_exc()
            raise ValueError('%s is not a valid %s file' % (file_path, suffix))

    @staticmethod
    def get_parser(suffix):
        if suffix == 'yaml':
            return partial(load, Loader=Loader)
        elif suffix == 'json':
            return json.load
        elif suffix == 'ini':
            return ConfigObj
        else:
            raise ValueError('不支持的配置文件类型')

    def as_dict(self):
        return self.content


from treq import get

def done(response):
    print(response.code)
    reactor.stop()

def eb(e):
    print(e)
    reactor.stop()

get("https://httpbin.org/get").addCallback(done).addErrback(eb)

from twisted.internet import reactor

reactor.run()


import matplotlib.pyplot as plt

#displays error graph
def errorgraph(data):
    error = data.errors
    x = list(range(0, len(data.errors)))
    plt.scatter(x, error, alpha=0.8, edgecolors='none', s=5)
    plt.title('Graph of error for ' + data.kind + " " + str(data.params[len(data.params) - 2]))
    plt.axis([0, 1000, 0, max(error) + max(error)/10])
    plt.show()


# -*- coding: utf-8 -*-
import os
from Qt.QtWidgets import *
from Qt.QtCore import *
from Qt.QtGui import *
import maya.cmds as mc
from miraLibs.pipeLibs import pipeFile, pipeMira
from miraLibs.pyLibs import join_path

def get_proxy_dir():
    obj = pipeFile.PathDetails.parse_path()
    project = obj.project
    asset_name = obj.asset_name
    root_dir = pipeMira.get_root_dir(project)
    proxy_dir = join_path.join_path2(root_dir, project, "assets", "proxy", asset_name, "static")
    return proxy_dir

def export_ass():
    proxy_dir = get_proxy_dir()
    meshes = mc.ls(type="mesh", long=1)
    progress_dialog = QProgressDialog("Ass Exporting...", 'Cancel', 0, len(meshes))
    progress_dialog.setWindowModality(Qt.WindowModal)
    progress_dialog.setMinimumWidth(350)
    progress_dialog.show()
    for index, mesh in enumerate(meshes):
        progress_dialog.setValue(index)
        if progress_dialog.wasCanceled():
            break
        mc.select(mesh, r=1)
        new_name = mesh.replace("|", "@")
        mc.file(os.path.join(proxy_dir, "%s.ass" % new_name),
                options="-mask 8;-lightLinks 0;-boundingBox;-shadowLinks 0",
                f=1, type="ASS Export", shader=0, exportSelected=1)

if __name__ == "__main__":
    export_ass()


# snakeRules
IS_ORACLE = False
MOVE_PENALTY = 1
MICE_TO_WIN = 30
MOUSE_REWARD_MULTIPLIER = 10
SNAKE_SPEED = 2

# game
DEFAULT_DIMENSONS = (10, 10)
DEFAULT_NUM_MICE = 2
DEFAULT_NUM_GAMES = 1

# what do each of these items look like in the graphics?
SNAKE_HEAD = "0"
MOUSE = "*"
SNAKE_BODY = "O"
GRID_SPACE = " "

DISPLAY = True
SLEEP_TIME = 0.08
SHOW_EACH_AGENT_MOVE = False #if True, you can see each agent's move every turn


from typing import List

class Solution:
    def largestSubmatrix(self, matrix: List[List[int]]) -> int:
        if not matrix or not matrix[0]:
            return 0

        rows, cols = len(matrix), len(matrix[0])
        count = [0] * cols
        max_area = 0

        for i in range(rows):
            for j in range(cols):
                if matrix[i][j] == 1:
                    count[j] = count[j] + 1 if i > 0 else 1
                else:
                    count[j] = 0

            sorted_count = sorted(count)
            for k in range(cols):
                max_area = max(max_area, sorted_count[k] * (cols - k))

        return max_area


import pytest
import contextlib
import os
import release_version_manager as ReleaseVersion
import random as real_random
import subprocess
import urllib.parse
import atexit
import yaml
from typing import NamedTuple
import requests
from pathlib import Path

from selenium import webdriver
from appium import webdriver as appiumdriver
from appium.options.android import UiAutomator2Options
from appium.options.ios import XCUITestOptions
from selenium.common.exceptions import WebDriverException
from selenium.webdriver.chrome.options import Options as ChromeOptions
from selenium.webdriver.firefox.options import Options as FirefoxOptions
from selenium.webdriver.safari.options import Options as SafariOptions
from datetime import datetime, timedelta, timezone
from urllib.parse import urlencode

import sentry_sdk

def get_system_user():
    return subprocess.check_output(['id', '-un']).strip().decode()

def is_running_in_gcp():
    try:
        # Google Compute Engine instances have a metadata server available at this URL
        metadata_url = "http://metadata.google.internal"
        headers = {'Metadata-Flavor': 'Google'}
        response = requests.get(metadata_url, headers=headers, timeout=2)
        return response.status_code == 200
    except requests.exceptions.RequestException:
        return False

def is_running_in_gha():
    return os.getenv('GITHUB_ACTIONS') == 'true'

def is_production():
    return is_running_in_gcp() or is_running_in_gha()

# Get current timestamp, with +/- adjustment if desired, as string in ISO format (for Sentry URLs)
# round seconds down (floor): adjust_seconds = 0
# round seconds up (ceil): adjust_seconds = 1
def get_current_time_iso_utc(adjust_seconds=0):
    dt = datetime.now(timezone.utc).replace(microsecond=0) + timedelta(seconds = adjust_seconds)
    return dt.strftime('%Y-%m-%dT%H:%M:%S')

# Example command usage:
#
#   RUN_ID=test_desktop_web_13 BATCH_SIZE=10 py.test -n 6 desktop_web
#

start_time = get_current_time_iso_utc()

class Browser(NamedTuple):
    remote: bool
    browserName: str
    browserVersion: str | None = None
    platformName: str | None = None
    param_display: str | None = None

class Config(NamedTuple):
    # TDA can operate in two modes, 'mock' and 'direct'
    # mock - captures a small sample of data using mini-relay, stores it, then clones those events
    # fudging timestamps and other fields to make it look like they were generated by the app
    # direct - just sends data directly to Sentry
    mode: str
    browsers: tuple[Browser, ...]
    dsn: str
    react_endpoints: tuple[str, ...]
    vue_endpoints: tuple[str, ...]

def _config() -> Config:
    cfg_filename = os.environ.get('TDA_CONFIG', 'config.yaml')
    with open(cfg_filename) as f:
        contents = yaml.safe_load(f)
    return Config(
        mode=contents['mode'],
        browsers=tuple(Browser(**d) for d in contents['browsers']),
        dsn=contents['dsn'],
        react_endpoints=tuple(contents['react_endpoints']),
        vue_endpoints=tuple(contents['vue_endpoints']),
    )

CONFIG = _config()

# optional, for tda-mock no need to specify as it will be set to request nodeid
RUN_ID = os.getenv("RUN_ID")

# We don't really need true randomness for anything. Graphs could look similar after TDA restarts due
# to same batch size sequence, but no one would notice. Saucelabs introduces some randomness anyway.
REPEATABLE_RANDOM = False

# BATCH_SIZE can be either <NUMBER> or random_<NUMBER>
# The later means each time a test is run the inner steps will be repeated a random
# number of times between 0 and <NUMBER> - 1
# e.g. BATCH_SIZE=5, or BATCH_SIZE=random_100
# IS_CANARY always overrides BATCH_SIZE
BATCH_SIZE = os.getenv("IS_CANARY") and "1" or (os.getenv("BATCH_SIZE") or "1")

SLEEP_LENGTH = os.getenv("SLEEP_LENGTH") or "random_2_1"

# Currently only used in desktop_web/ tests. Mobile apps have it hardcoded.
BACKENDS = (os.getenv("BACKENDS") or "flask,express,springboot,ruby,laravel,rails,aspnetcore").split(',')

import urllib3
urllib3.disable_warnings()

sentry_sdk.init(
    dsn=CONFIG.dsn,
    traces_sample_rate=0,
    # we don't really use environment for anything, as we'd have to propagate it to all
    # the other services, so instead out of laziness we effectively just encode it as the
    # first part of SE tag ("prod-", or system $USER when running locally).
    environment="production",
)

_browser2class = {
    'chrome': ChromeOptions,
    'firefox': FirefoxOptions,
    'safari': SafariOptions
}

def pytest_addoption(parser):
    parser.addoption("--dc", action="store", default='us', help="Set Sauce Labs Data Center (US or EU)")

@pytest.fixture
def data_center(request):
    return request.config.getoption('--dc')

@pytest.fixture
def random(request):
    if REPEATABLE_RANDOM:
        # Ensure tests produce repeatable outcomes when re-run, see:
        cwd = os.path.dirname(os.path.realpath(__file__)) + '/'
        test_relpath = str(request.path).split(cwd)[1]
        # e.g: desktop_web/test_mytest.py:test_myfunction[desktop_web_driver0]
        seed = f'{test_relpath}:{request.node.name}'
        # Use a new Random object doesn't share state with global random generator
        # this is just in case some package uses global random in a way that would
        # upset repeatability (e.g. numer of random() calls varies b/w executions)
        return real_random.Random(seed)
    else:
        return real_random

# random_1.5
# random_3_1.0
# 5.5
@pytest.fixture
def sleep_length(random):
    def random_sleep_length():
        if SLEEP_LENGTH.startswith("random_"):
            spl = SLEEP_LENGTH.split('_')
            r = random.randrange(float(spl[1]))
            if len(spl) == 3:
                r += float(spl[2])
            return r
        else:
            r = random.random() # unused, to make sure we call random same number of times
            return float(SLEEP_LENGTH)
    return random_sleep_length

@pytest.fixture
def batch_size(random):
    if BATCH_SIZE.startswith("random_"):
        return random.randrange(int(BATCH_SIZE.split('_')[1]))
    else:
        r = random.random() # unused, to make sure we call random same number of times
        return int(BATCH_SIZE)

@pytest.fixture
def backend(random):
    def random_backend(exclude=[], include=[]):
        if not include:
            include = BACKENDS
        else:
            include = [include] if isinstance(include, str) else include

        exclude = [exclude] if isinstance(exclude, str) else exclude
        # Must sort to get same order across processes to ensure that seeded random.sample()
        # always returns the same values.
        # python3
        #   >>> list(set(['b', 'a', 'f', 'd', 'e', 'c']) - set(['c', 'a']))
        #   ['b', 'd', 'e', 'f']
        #   >>> list(set(['b', 'a', 'f', 'd', 'e', 'c']) - set(['c', 'a']))
        #   ['b', 'd', 'e', 'f']
        #   >>> exit()
        # python3
        #   >>> list(set(['b', 'a', 'f', 'd', 'e', 'c']) - set(['c', 'a']))
        #   ['d', 'f', 'b', 'e']
        backends = sorted(list(set(include) - set(exclude)))
        if 'flask' in backends:
            backends += ['flask'] # make flask 2x more likely than others
        return random.sample(backends, 1)[0]
    return random_backend

@pytest.fixture
def endpoints():
    return CONFIG

class _ExtraParams:
    def __init__(self, *args, extra_params, **kwargs):
        if extra_params.startswith(('&', '?')):
            raise ValueError('extra_params must be in format: "param1=value1&param2=value2..."')
        self.extra_params = extra_params
        super().__init__(*args, **kwargs)

    def get(self, url):
        url += ('&' if '?' in url else '?') + self.extra_params
        return super().get(url)

# Automatically append a set of extra parameters to all URLs
#
#   remote = RemoteWithExtraUrlParams(..., extra_params='se=prod-tda-...')
#   remote.get('https://google.com/')
#       -> webdriver.Remote.get(https://google.com/?se=tda)
class RemoteWithExtraUrlParams(_ExtraParams, webdriver.Remote): pass
class ChromeWithExtraUrlParams(_ExtraParams, webdriver.Chrome): pass

@pytest.fixture
def selenium_endpoint(data_center):
    username = os.environ['SAUCE_USERNAME']
    access_key = os.environ['SAUCE_ACCESS_KEY']

    if data_center and data_center.lower() == 'eu':
        return "https://{}:{}@ondemand.eu-central-1.saucelabs.com/wd/hub".format(username, access_key)
    else:
        return "https://{}:{}@ondemand.us-west-1.saucelabs.com/wd/hub".format(username, access_key)

def sanitize_se_tag_component(c):
    if not c:
        return ''
    elif not isinstance(c, str):
        c = str(c)
    # so that we can parse the se tag into components later by splitting on '-'
    return c.replace('-', '_').replace('\n', '_')

# `se` tag captures all the information necessary to identify the process, user, parameters and
# particular test file that generated the event. It will be both:
#   1. Passed on to desktop_web as ?se= URL parameter so will show up in
#   demo.sentry.io project issues generated as a result of running the test
#   2. Used for internal reporting of this test automation's own errors (to DSN above,
#   i.e. "empower-tda" project)
# The 'base' part just identifies the who's running the entire suite, in what mode and with what
# global parameters. When run in production (i.e. in GCP or GHA) user will be "prod".
# Example: kosty-tda-direct-testing_my_branch_3
def se_base_tag():
    user = 'prod' if is_production() else get_system_user()
    s = sanitize_se_tag_component
    return f'{s(user)}-tda-{s(CONFIG.mode)}-{s(RUN_ID)}'

# 'prefix' additionally captures the particular test file
# Example se_prefix (no RUN_ID):    prod-tda-mock--desktop_web/test_checkout.py
# Example final se tag:             prod-tda-mock--desktop_web/test_checkout.py-sauce-chrome_windows
@pytest.fixture
def se_prefix(request):
    test_filepath = request.path
    if request.node.name.split('[')[0] != test_filepath.stem:
        raise ValueError(f'TDA expects only 1 function in each test file with the same name as the \
                        file. Please fix {test_filepath}.')
    rel_test_filepath = test_filepath.relative_to(Path(__file__).parent)
    s = sanitize_se_tag_component
    return f'{se_base_tag()}-{s(rel_test_filepath)}'

@contextlib.contextmanager
def _sauce_browser(request, selenium_endpoint, se):
    try:
        test_name = request.node.name
        build_tag = os.environ.get('BUILD_TAG', "Application-Monitoring-TDA")

        options = _browser2class[request.param.browserName]()

        # Testing out if this helps surface web vitals
        if request.param.browserName == 'chrome':
            options.add_argument("--enable-logging")
            options.add_argument("--enable-benchmarking")

        options.set_capability('platformName', request.param.platformName)
        options.set_capability('browserVersion', request.param.browserVersion)
        options.set_capability('sauce:options', {
            'seleniumVersion': '4.8.0',
            'build': build_tag,
            'name': test_name
        })

        browser = RemoteWithExtraUrlParams(
            command_executor=selenium_endpoint,
            options=options,
            keep_alive=True,
            # Note: these tags might not be supported by some frontends, e.g. Vue
            extra_params=urlencode({'se': se})
        )

        browser.implicitly_wait(10)

        sentry_sdk.set_tag("seleniumSessionId", browser.session_id)

        # This is specifically for SauceLabs plugin.
        # In case test fails after selenium session creation having this here will help track it down.
        if browser is not None:
            print("SauceOnDemandSessionID={} job-name={}".format(browser.session_id, test_name))
        else:
            raise WebDriverException("Never created!")

        try:
            yield browser
        finally:
            # Teardown starts here
            # report results
            # use the test result to send the pass/fail status to Sauce Labs
            sauce_result = "failed" if request.node.rep_call.failed else "passed"

            # Handler failure scenario, send to Sentry empower-tda
            if sauce_result == "failed":
                sentry_sdk.capture_message("Sauce Result: %s" % (sauce_result))

            browser.execute_script("sauce:job-result={}".format(sauce_result))
            browser.quit()

            # desktop_web tests finished, send to Sentry empower-tda, look for tags pytestName, pytestPlatform, seleniumSessionId
            sentry_sdk.capture_message("Selenium Session Done")

    except Exception as err:
        sentry_sdk.capture_exception(err)

@contextlib.contextmanager
def _local_browser(request, se):
    assert request.param.browserName == 'chrome'  # TODO: add others
    options = webdriver.ChromeOptions()
    options.add_argument("no-sandbox")
    options.add_argument("disable-gpu")
    options.add_argument("disable-dev-shm-usage")
    options.add_argument("headless")
    with ChromeWithExtraUrlParams(
            options=options,
            extra_params=urlencode({'se': se}),
    ) as driver:
        yield driver

@pytest.fixture(params=CONFIG.browsers, ids=[b.param_display for b in CONFIG.browsers])
def desktop_web_driver(request, se_prefix):
    if request.param.remote:
        se = f'{se_prefix}-sauce-{request.param.param_display}'
        sentry_sdk.set_tag("se", se)
        with _sauce_browser(request, request.getfixturevalue('selenium_endpoint'), se) as b:
            yield b
    else:
        se = f'{se_prefix}-local-{request.param.param_display}'
        sentry_sdk.set_tag("se", se)
        with _local_browser(request, se) as b:
            yield b

@pytest.fixture
def android_react_native_emu_driver(request, selenium_endpoint, se_prefix):

    se = f'{se_prefix}-sauce-android10'
    sentry_sdk.set_tag('se', se)
    try:
        release_version = ReleaseVersion.latest_react_native_github_release()

        options = UiAutomator2Options().load_capabilities({
            'deviceName': 'Android GoogleAPI Emulator',
            'platformVersion': '10.0',
            'platformName': 'Android',
            'app': f'https://github.com/sentry-demos/sentry_react_native/releases/download/{release_version}/app-release.apk',
            'sauce:options': {
                'appiumVersion': '1.20.2',
                'build': 'RDC-Android-Python-Best-Practice',
                'name': request.node.name
            },
            'appWaitForLaunch': False
        })

        driver = appiumdriver.Remote(selenium_endpoint, options=options)
        driver.implicitly_wait(20)

        sentry_sdk.set_tag("seleniumSessionId", driver.session_id)

        yield driver
        sauce_result = "failed" if request.node.rep_call.failed else "passed"
        driver.execute_script("sauce:job-result={}".format(sauce_result))
        driver.quit()

        # android_react_native tests finished, send to Sentry empower-tda, look for tags pytestName, pytestPlatform, seleniumSessionId
        sentry_sdk.capture_message("Selenium Session Done")

    except Exception as err:
        sentry_sdk.capture_exception(err)

@pytest.fixture
def android_emu_driver(request, selenium_endpoint, se_prefix):

    se = f'{se_prefix}-sauce-android10'
    sentry_sdk.set_tag('se', se)
    try:
        release_version = ReleaseVersion.latest_android_github_release()

        options = UiAutomator2Options().load_capabilities({
            'deviceName': 'Android GoogleAPI Emulator',
            'platformVersion': '10.0',
            'platformName': 'Android',
            'app': f'https://github.com/sentry-demos/android/releases/download/{release_version}/app-release.apk',
            'sauce:options': {
                'appiumVersion': '1.20.2',
                'build': 'RDC-Android-Python-Best-Practice',
                'name': request.node.name
            },
            'appWaitForLaunch': False
        })

        driver = appiumdriver.Remote(selenium_endpoint, options=options)
        driver.implicitly_wait(20)

        sentry_sdk.set_tag("seleniumSessionId", driver.session_id)

        yield driver
        sauce_result = "failed" if request.node.rep_call.failed else "passed"
        driver.execute_script("sauce:job-result={}".format(sauce_result))
        driver.quit()

        # android tests finished, send to Sentry empower-tda, look for tags pytestName, pytestPlatform, seleniumSessionId
        sentry_sdk.capture_message("Selenium Session Done")

    except Exception as err:
        sentry_sdk.capture_exception(err)
        raise

@pytest.fixture
def ios_react_native_sim_driver(request, selenium_endpoint, se_prefix):

    se = f'{se_prefix}-sauce-ios14.5'
    sentry_sdk.set_tag('se', se)
    try:
        release_version = ReleaseVersion.latest_react_native_github_release()

        options = XCUITestOptions().load_capabilities({
            'appium:deviceName': 'iPhone 11 Simulator',
            'platformName': 'iOS',
            'appium:platformVersion': '14.5',

            'sauce:options': {
                'appiumVersion': '1.21.0',
                'build': 'RDC-iOS-Python-Best-Practice',
                'name': request.node.name,
            },
            'appium:app': f'https://github.com/sentry-demos/sentry_react_native/releases/download/{release_version}/sentry_react_native.app.zip',
        })

        driver = appiumdriver.Remote(selenium_endpoint, options=options)
        driver.implicitly_wait(20)

        sentry_sdk.set_tag("seleniumSessionId", driver.session_id)

        yield driver
        sauce_result = "failed" if request.node.rep_call.failed else "passed"
        driver.execute_script("sauce:job-result={}".format(sauce_result))
        driver.quit()

        # ios_react_native tests finished, send to Sentry empower-tda, look for tags pytestName, pytestPlatform, seleniumSessionId
        sentry_sdk.capture_message("Selenium Session Done")

    except Exception as err:
        sentry_sdk.capture_exception(err)

@pytest.fixture
def ios_sim_driver(request, selenium_endpoint, se_prefix):

    se = f'{se_prefix}-sauce-ios15.5'
    sentry_sdk.set_tag('se', se)
    try:
        release_version = ReleaseVersion.latest_ios_github_release()

        options = XCUITestOptions().load_capabilities({
            'appium:deviceName': 'iPhone 13 Simulator',
            'platformName': 'iOS',
            'appium:platformVersion': '15.5',

            'sauce:options': {
                'appiumVersion': '1.22.3',
                'build': 'RDC-iOS-Mobile-Native',
                'name': request.node.name,
            },
            'appium:app': f'https://github.com/sentry-demos/ios/releases/download/{release_version}/EmpowerPlant_release.zip',
        })

        driver = appiumdriver.Remote(selenium_endpoint, options=options)
        driver.implicitly_wait(20)

        sentry_sdk.set_tag("seleniumSessionId", driver.session_id)

        yield driver
        sauce_result = "failed" if request.node.rep_call.failed else "passed"
        driver.execute_script("sauce:job-result={}".format(sauce_result))
        driver.quit()

        # ios tests finished, send to Sentry empower-tda, look for tags pytestName, pytestPlatform, seleniumSessionId
        sentry_sdk.capture_message("Selenium Session Done")

    except Exception as err:
        sentry_sdk.capture_exception(err)

@pytest.hookimpl(hookwrapper=True, tryfirst=True)
def pytest_runtest_makereport(item, call):
    # this sets the result as a test attribute for Sauce Labs reporting.
    # execute all other hooks to obtain the report object
    outcome = yield
    rep = outcome.get_result()

    # set an report attribute for each phase of a call, which can
    # be "setup", "call", "teardown"
    setattr(item, "rep_" + rep.when, rep)

if CONFIG.mode == 'direct':
    def final_report():
        end_time = get_current_time_iso_utc(adjust_seconds=1)
        start = urllib.parse.quote(start_time)
        end = urllib.parse.quote(end_time)
        project = CONFIG.dsn.split('/')[-1]
        print()
        print(f'GENERATED errors: https://demo.sentry.io/issues/?query=se%3A{se_base_tag()}%2A+%21project%3Aempower-tda&start={start}&end={end}')
        print(f'OWN errors:       https://demo.sentry.io/issues/?project={project}&query=se%3A{se_base_tag()}%2A&start={start}&end={end}')

    atexit.register(final_report)


import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

metrics = pd.read_csv('metrics.csv')
metrics.set_index('Trait', inplace=True)
# print(metrics)

traits = ['O', 'C', 'E', 'A', 'N']
Y = ['Image','Audio']
Z = ['False', 'True']

metrics_dict = {}
for y in Y:
	metrics_dict[y] = {}
	for z in Z:
		metrics_dict[y][z] = {}
		for trait in traits:
				metrics_dict[y][z][trait]= metrics.loc[trait, str(y)+' ' + z]

print(metrics_dict)

sns.set()
for y in metrics_dict.keys():
	count= 0
	fig, axes = plt.subplots(nrows=1, ncols=1)
	# fig.suptitle("Audio Metrics for Y = "+str(y))
	# for z in metrics_dict[y].keys():
		# print(y,z,metrics_dict[y][z])
	df = pd.DataFrame(index=['Debiased Model', 'Traditional (Biased) Model'], columns=traits)
	df.loc['Debiased Model', :] = metrics_dict[y]['False']
	df.loc['Traditional (Biased) Model',:] = metrics_dict[y]['True']
	print(df)
		# df_T = df.transpose()
		# df_T = df_T.iloc[::-1]
	df.plot.barh()
	plt.show()
	# a = df.plot.barh(ax=axes[count], color = ['r','g'], legend=False)
	# axes[count].axvline(x=1)
	# 	# if z=='E':
	# 	# 	axes[count].set_title('Ethnicity')
	# 	# else:
	# 	# 	axes[count].set_title('Gender')
	# count+=1
	# for container in a.containers:
	# 		# print(container)
	# 	a.bar_label(container, fmt='%.2f')

	# lines = []
	# labels = []

	# for ax in fig.axes:
	#     Line, Label = ax.get_legend_handles_labels()
	#     lines.extend(Line)
	#     labels.extend(Label)

	# fig.legend(lines[0:3], labels[0:3], loc='upper right')
	# plt.show()


""" 2. Da se napravi programa vo koja korisnikot ke moze da vnese proizvolen broj na elementi vo set,
 da se pomine niz elementite so setot i parnite da se dodadat vo drug set,
  a ne patnite vo drug"""

setBasic = set()
setP = set()
setN = set()
i=0
vkupnoElementi = int(input("kOLKU ELEMENTI SAKATE DA VNESETE?"))

"""ako sakame da koristime while loop, togas treba da vneseme opcija dali korisnikot da prodolzi da vnesuvva
so prodolzi == "da"""

for i in range(vkupnoElementi):
    i = int(input("Vnesete vkupno elementi"))
    elementi = setBasic.add(i)
    if i %2 == 0:
        setP.add(i)
    else:
        setN.add(i)

print(setP)
print(setN)


import socket
from _thread import *
import pandas as pd
import numpy as np
import argparse
import requests
from requests.auth import HTTPBasicAuth
from datetime import datetime
from ml.LinearRegression import DistributedLinearRegression

parser = argparse.ArgumentParser()

parser.add_argument("--email", type=str)
parser.add_argument("--password", type=str)
parser.add_argument("--exp_name", type=str)
parser.add_argument("--exp_id", type=int)
parser.add_argument("--desc", type=str)

args = parser.parse_args()
dateformat = r"%d-%m-%Y %H:%M:%S"

body = {
    "device_type": 1,
    "id": args.exp_id,
    "name": args.exp_name,
    "description": args.desc,
    "active": 1,
    "start": datetime.now().strftime(dateformat)
}

url = "http://172.18.132.255:5000/experiments/commit"
response = requests.post(
    url,
    json = body,
    auth = HTTPBasicAuth(
        args.email, args.password
    )
)

if response.content != b'1':
    print("Unsuccessful !")
    print("Run the script again !")
    exit()

def recvall(sock):
    BUFF_SIZE = 4096 # 4 KiB
    data = b''
    while True:
        part = sock.recv(BUFF_SIZE)
        data += part
        if len(part) < BUFF_SIZE:
            # either 0 or end of data
            break
    return data

ServerSideSocket = socket.socket()
host = '127.0.0.1'
port = 2004
connectport = [12345, 12346]  #, 12346, 12347
nworkers = len(connectport)
ThreadCount = 0

try:
    ServerSideSocket.bind((host, port))
except socket.error as e:
    print(str(e))

print('Socket is listening..')
ServerSideSocket.listen(5)

# data = pd.read_csv("master/cleaned.csv").to_json(orient = 'table', index = False)
df = pd.read_csv("cleaned_normalised.csv")
df_shuffled = df.sample(frac=1)
df_splits = np.array_split(df_shuffled, nworkers)
df_splits = [
    df_split.to_json(
        orient = 'table',
         index = False
    ) for df_split in df_splits]

thetas = []
completed = [nworkers]

def multi_threaded_client(connection, data, thetas, completed):
    connection.sendall(data.encode())
    print("Distributed")
    data = recvall(connection).decode()
    print("received")
    thetas.append(data)
    connection.close()
    completed[0] -= 1

i = 1
while(i <= nworkers):
    conn, addr = ServerSideSocket.accept()
    print(addr)
    if addr[1] not in connectport:
        conn.close()
        i -= 1
    else:
        # do what ever you want here
        print('Connected to: ' + addr[0] + ':' + str(addr[1]))
        start_new_thread(multi_threaded_client, (conn, df_splits[i-1], thetas, completed))
        ThreadCount += 1
        print('Thread Number: ' + str(ThreadCount))
    i += 1

while(completed[0] != 0):
    pass

lr = DistributedLinearRegression()

def literal_eval(lst):
    lst = lst[1:len(lst)-1]    # Removes '[', ']'
    lst = list(map(float, lst.split()))
    return lst

thetas = np.array(list(map(literal_eval, thetas)))
print(thetas)

lr.from_weights(thetas)
theta = lr.get_weights()
print(theta)

print("Closing server")
ServerSideSocket.close()

body["active"] = 0
body["end"] = datetime.now().strftime(dateformat)

response = requests.post(
    url,
    json = body,
    auth = HTTPBasicAuth(
        args.email, args.password
    )
)

print("Execution successful !")


from ...common.utils import Singleton, Mkdir, Touch
from ...common import config
import json

class Table(Singleton):
    def __init__(self, filename:str) -> None:
        self.__dir = config.make_data_path('work_table')
        Mkdir(self.__dir)
        self.__path = self.__dir + '/' + filename
        Touch(self.__path, '{}')
        self.__data = dict()

        self.__load_from_db()

    def __load_from_db(self):
        with open(self.__path, 'r') as file:
            self.__data = json.loads(file.read())

    def __write_to_db(self):
        with open(self.__path, 'w') as file:
            file.write(json.dumps(self.__data))

    def add(self, group_id:str, obj:str) -> bool:
        "Return True if successfully add a new obj, False if exists obj for the group"
        if group_id in self.__data.keys():
            return False

        self.__data[group_id] = obj
        self.__write_to_db()
        return True

    def delete(self, group_id:str) -> bool:
        "Return True if successfully delete obj for the group, False if not exists obj"
        if group_id not in self.__data.keys():
            return False

        del self.__data[group_id]
        self.__write_to_db()
        return True

    def get(self, group_id:str) -> str:
        if group_id not in self.__data.keys():
            return ""

        return self.__data[group_id]

    def update(self, group_id:str, obj:str) -> bool:
        "Return True if successfully update obj for the group, False if not exists obj"
        if group_id not in self.__data.keys():
            return False

        self.__data[group_id] = obj
        self.__write_to_db()
        return True


from collections import defaultdict
from typing import List, Dict

def merge_dicts(lst_of_dicts: List[Dict], keys_to_merge: List[str]) -> dict:
    res = defaultdict(list)
    for dictionary in lst_of_dicts:
        for key in keys_to_merge:
            res[key] += dictionary[key]
    return res


# make a function that accepts a string inut, and then hashes it returning a unique input

class HashMap:

    def __init__(self):
        self.hash = []

        #Please teach me list comp cuz I tried and didnt know how
        for x in range(10000):
            self.hash.append(0)

    def mapAdd(self, key, value):
        numVal = 0

        for car in key:
            numVal += ord(car) + len(key)

            while self.hash[numVal]:
                numVal += 1

            self.hash[numVal] = value

        return self.hash

    def mapFind(self, key):
        numVal = 0

        #car is short for character lmao
        for car in key:
            # ord returns the ASCII value of the char
            numVal += ord(car) + len(key)

        while self.hash[numVal] != key:
            if self.hash[numVal] == 0:
                return "None"
            numVal += 1

        return self.hash[numVal]

    def mapDraw(self):
        print(self.hash)

hashmap1 = HashMap()
hashmap1.mapAdd("A", "Yoyoyoyoyoyoyoyoyo")
hashmap1.mapDraw()

# yoyoyoyo don't be mad at me for this code below, I just wanted to preserve it
# This was the original before making it into a class

def hashing(sLlist):
    hashList = []
    numVal = 0
    numValList = []
    check = False
    slist = []

    for x1 in sLlist:
        slist.append(x1[0])

    for x in range(10000):
        hashList.append(0)

    for i, car in enumerate(slist):
        for nerd in car:
            print(i)
            numVal += ord(nerd) + len(car)

        #havent implemented checking this yet
        while hashList[numVal]:
            numValList.append(numVal)
            hashList[numVal] = slist[i]
            numVal = 0
        #

        if not check:
            numValList.append(numVal)
            hashList[numVal] = slist[i]
            numVal = 0

#Need to check if the current hashList index is 0

    return hashList


import asyncio
from message import Message

@asyncio.coroutine
def run(message, matches, chat_id, step):
    return [Message(chat_id).set_text(
        "Hi :-) \n*Bot Search Music!*\n✔Fast and accurate✔\n\n_/Music <Artist name>_\n*Creator* = @LullSky", parse_mode="markdown")]

plugin = {
    "name": "About",
    "desc": "Get Information About The Developer!",
    "usage": ["/about"],
    "run": run,
    "sudo": False,
    "patterns": ["^[/]start$"]
}


def merge_sort(array):
    if len(array) > 1:
        mid = len(array) // 2
        left_array = array[:mid]
        right_array = array[mid:]

        merge_sort(left_array)
        merge_sort(right_array)

        i = j = k = 0
        while i < len(left_array) and j < len(right_array):
            if left_array[i] < right_array[j]:
                array[k] = left_array[i]
                i += 1
            else:
                array[k] = right_array[j]
                j += 1
            k += 1

        while i < len(left_array):
            array[k] = left_array[i]
            i += 1
            k += 1

        while j < len(right_array):
            array[k] = right_array[j]
            j += 1
            k += 1

#array = [1,8,6,3,9,99,121,2,3]
#merge_sort(array)
#print(array)


import cv2
from visualize_cv2 import model, display_instances, class_names, apply_mask_without_color
import sys
import numpy as np

def maakHistogram(image):
    histogramLijst = []
    color = ('b', 'g', 'r')
    # Maak voor elke kleurcomponent een histogram = matrix met vorm 256 x 1
    # Hou histogrammen bij in histogramLijst ==> histogramLijst.append
    for i, col in enumerate(color):
        histogram = cv2.calcHist([image], [i], None, [256], [10, 256])
        histogramLijst.append(histogram)
    return histogramLijst

def interpreteerHistogram(histogramLijst):
    # Voor elk histogram (3 in totaal) in de histogramLijst:
    tellerBlauw = 0
    tellerGroen = 0
    tellerRood = 0
    tellerTot = 0
    for i in range(150, 256):
        tellerBlauw = tellerBlauw + histogramLijst[0][i][0] * i / 100
        tellerGroen = tellerGroen + histogramLijst[1][i][0] * i / 100
        tellerRood = tellerRood + histogramLijst[2][i][0] * i / 100
    tellerTot = tellerBlauw + tellerGroen + tellerRood
    verhoudingGroen = tellerGroen / tellerTot
    print("B: " + str(tellerBlauw) + " G: " + str(tellerGroen) + " R: " + str(tellerRood) + " Tot: " + str(tellerTot))
    # print("Verhouding G/Tot: " + str(verhoudingGroen))
    if tellerRood > tellerGroen and tellerRood > tellerBlauw and verhoudingGroen > 0.35:
        print("Op basis van histogram: gele banaan")
    elif tellerGroen > tellerBlauw and tellerGroen > tellerRood:
        print("Op basis van histogram: Groene banaan")
    elif tellerRood > tellerGroen and tellerRood > tellerBlauw and verhoudingGroen < 0.35:
        print("Op basis van histogram: Bruine banaan")
    else:
        print("Op basis van histogram: onzeker")

args = sys.argv
if (len(args) < 2):
    print("run command: python video_demo.py 0 or video file name")
    sys.exit(0)
name = args[1]
if (len(args[1]) == 1):
    name = int(args[1])

stream = cv2.VideoCapture(name)

while True:
    ret, frame = stream.read()
    if not ret:
        print("unable to fetch frame")
        break
    results = model.detect([frame], verbose=0)

    # Visualize results
    r = results[0]
    masked_image = display_instances(frame, r['rois'], r['masks'], r['class_ids'],
                                     class_names, r['scores'])
    N = r['rois'].shape[0]
    masks = r['masks']
    resultaten = []
    for i in range(N):
        oppervlakteROI = (abs(r['rois'][i][0] - r['rois'][i][2])) * (abs(r['rois'][i][1] - r['rois'][i][3]))
        verhoudingROI = oppervlakteROI / (frame.shape[0] * frame.shape[1])

        if (r['class_ids'][i] == 47 and r['scores'][i] >= 0.90 and verhoudingROI > 0.10):
            print("Banaan gevonden (Verhouding ROI/Afbeelding = " + str(verhoudingROI) + ", Score: " + str(r['scores'][i]) + ")")

            im = frame.astype(np.uint32).copy()
            mask = masks[:, :, i]
            im = apply_mask_without_color(im, mask)
            for j in range(im.shape[2]):
                im[:, :, j] = im[:, :, j] * mask
            # im = cv2.cvtColor(im.astype(np.uint8), cv2.COLOR_BGR2RGB)
            resultaten.append(im.astype(np.uint8))
            cv2.imshow("Resultaat van Mask R-CNN", im.astype(np.uint8))

        elif (r['class_ids'][i] == 47 and r['scores'][i] < 0.90 and verhoudingROI > 0.10):
            print("Banaan gevonden maar met onvoldoende zekerheid")

    for res in resultaten:
        res = cv2.resize(res, (200, 150), interpolation = cv2.INTER_AREA)
        histogramLijst = maakHistogram(res)
        interpreteerHistogram(histogramLijst)

    cv2.imshow("masked_image", masked_image)
    if (cv2.waitKey(1) & 0xFF == ord('q')):
        break
stream.release()
cv2.destroyWindow("masked_image")


from heapq import heappop, heappush, heapify
class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:

        heap = []
        heapify(heap)
        for num in nums:
            heappush(heap, -num)

        kthLargestElement = float('-inf')
        while k > 0:
            maxElement = heappop(heap)
            kthLargestElement = - maxElement
            k -= 1

        return kthLargestElement