# Path: group_testing/group_selectors/group_selector.py
# class GroupSelector:
#   NEEDS_POSTERIOR = False
#   def __call__(self, rng, state):
#   def get_groups(self, rng, state):


import gin
import jax.numpy as np
import numpy as onp
from group_testing.group_selectors import group_selector
def origami_m2():
  """Origami M2."""
  new_groups = onp.zeros((22, 231))
  new_groups[0, [9, 17, 23, 32, 69, 71, 73, 79, 87, 107, 116, 126, 134, 138, 146, 171, 187, 190, 205, 209, 230]] = 1
  new_groups[1, [0, 1, 6, 41, 57, 71, 75, 97, 99, 103, 113, 148, 153, 164, 204, 208, 212, 214, 223, 228, 229]] = 1
  new_groups[2, [1, 25, 26, 32, 33, 35, 43, 44, 54, 56, 60, 61, 62, 67, 86, 123, 137, 183, 217, 218, 227]] = 1
  new_groups[3, [8, 11, 16, 23, 26, 38, 63, 77, 90, 99, 108, 142, 149, 175, 177, 181, 195, 198, 200, 210, 226]] = 1
  new_groups[4, [5, 12, 15, 16, 36, 39, 44, 47, 82, 89, 91, 107, 117, 153, 158, 172, 173, 178, 185, 186, 219]] = 1
  new_groups[5, [2, 7, 22, 30, 68, 76, 81, 104, 106, 113, 120, 123, 142, 144, 151, 152, 186, 187, 188, 192, 193]] = 1
  new_groups[6, [6, 18, 27, 28, 47, 84, 86, 96, 100, 108, 109, 133, 135, 159, 162, 167, 171, 176, 179, 188, 196]] = 1
  new_groups[7, [4, 20, 21, 27, 29, 30, 51, 62, 82, 87, 94, 114, 127, 139, 148, 150, 166, 174, 182, 203, 226]] = 1
  new_groups[8, [45, 53, 59, 60, 81, 90, 93, 110, 111, 112, 122, 145, 150, 158, 168, 176, 194, 206, 209, 212, 215]] = 1
  new_groups[9, [9, 13, 14, 22, 61, 84, 91, 103, 129, 143, 149, 161, 180, 189, 194, 201, 202, 203, 211, 216, 222]] = 1
  new_groups[10, [0, 3, 5, 18, 45, 55, 58, 66, 78, 85, 120, 128, 130, 138, 147, 163, 166, 181, 183, 184, 202]] = 1
  new_groups[11, [7, 35, 38, 58, 74, 98, 111, 114, 118, 131, 154, 157, 179, 190, 191, 199, 201, 213, 219, 225, 228]] = 1
  new_groups[12, [8, 14, 39, 49, 52, 56, 66, 70, 74, 76, 80, 95, 105, 110, 116, 119, 136, 139, 156, 164, 167]] = 1
  new_groups[13, [2, 15, 21, 31, 41, 52, 53, 54, 63, 64, 65, 72, 79, 102, 118, 129, 135, 140, 147, 160, 165]] = 1
  new_groups[14, [49, 77, 83, 85, 93, 100, 102, 106, 115, 117, 121, 132, 137, 170, 182, 191, 197, 205, 214, 216, 220]] = 1
  new_groups[15, [25, 28, 34, 40, 50, 55, 57, 73, 80, 89, 121, 122, 127, 141, 155, 165, 177, 189, 192, 207, 225]] = 1
  new_groups[16, [4, 12, 13, 19, 33, 48, 70, 72, 88, 104, 109, 124, 126, 128, 145, 154, 155, 175, 197, 208, 224]] = 1
  new_groups[17, [10, 24, 29, 48, 50, 64, 92, 112, 133, 134, 136, 152, 163, 169, 170, 173, 195, 199, 204, 222, 227]] = 1
  new_groups[18, [3, 20, 46, 65, 83, 97, 101, 125, 131, 151, 156, 169, 172, 196, 200, 206, 207, 211, 218, 224, 230]] = 1
  new_groups[19, [19, 34, 36, 37, 42, 67, 68, 69, 92, 98, 115, 119, 125, 130, 143, 159, 160, 174, 210, 215, 223]] = 1
  new_groups[20, [10, 17, 37, 46, 94, 105, 124, 132, 140, 141, 157, 162, 168, 180, 184, 185, 193, 198, 217, 221, 229]] = 1
  new_groups[21, [11, 24, 31, 40, 42, 43, 51, 59, 75, 78, 88, 95, 96, 101, 144, 146, 161, 178, 213, 220, 221]] = 1
  return np.array(new_groups, dtype=bool)
# pylint: enable=line-too-long


@gin.configurable

# Path: group_testing/group_selectors/group_selector.py
# class GroupSelector:
#   NEEDS_POSTERIOR = False
#   def __call__(self, rng, state):
#   def get_groups(self, rng, state):


# Path: group_testing/group_selectors/origami.py
import gin
import jax.numpy as np
import numpy as onp
from group_testing.group_selectors import group_selector

def origami_m2():
  """Origami M2."""
  new_groups = onp.zeros((22, 231))
  new_groups[0, [9, 17, 23, 32, 69, 71, 73, 79, 87, 107, 116, 126, 134, 138, 146, 171, 187, 190, 205, 209, 230]] = 1
  new_groups[1, [0, 1, 6, 41, 57, 71, 75, 97, 99, 103, 113, 148, 153, 164, 204, 208, 212, 214, 223, 228, 229]] = 1
  new_groups[2, [1, 25, 26, 32, 33, 35, 43, 44, 54, 56, 60, 61, 62, 67, 86, 123, 137, 183, 217, 218, 227]] = 1
  new_groups[3, [8, 11, 16, 23, 26, 38, 63, 77, 90, 99, 108, 142, 149, 175, 177, 181, 195, 198, 200, 210, 226]] = 1
  new_groups[4, [5, 12, 15, 16, 36, 39, 44, 47, 82, 89, 91, 107, 117, 153, 158, 172, 173, 178, 185, 186, 219]] = 1
  new_groups[5, [2, 7, 22, 30, 68, 76, 81, 104, 106, 113, 120, 123, 142, 144, 151, 152, 186, 187, 188, 192, 193]] = 1
  new_groups[6, [6, 18, 27, 28, 47, 84, 86, 96, 100, 108, 109, 133, 135, 159, 162, 167, 171, 176, 179, 188, 196]] = 1
  new_groups[7, [4, 20, 21, 27, 29, 30, 51, 62, 82, 87, 94, 114, 127, 139, 148, 150, 166, 174, 182, 203, 226]] = 1
  new_groups[8, [45, 53, 59, 60, 81, 90, 93, 110, 111, 112, 122, 145, 150, 158, 168, 176, 194, 206, 209, 212, 215]] = 1
  new_groups[9, [9, 13, 14, 22, 61, 84, 91, 103, 129, 143, 149, 161, 180, 189, 194, 201, 202, 203, 211, 216, 222]] = 1
  new_groups[10, [0, 3, 5, 18, 45, 55, 58, 66, 78, 85, 120, 128, 130, 138, 147, 163, 166, 181, 183, 184, 202]] = 1
  new_groups[11, [7, 35, 38, 58, 74, 98, 111, 114, 118, 131, 154, 157, 179, 190, 191, 199, 201, 213, 219, 225, 228]] = 1
  new_groups[12, [8, 14, 39, 49, 52, 56, 66, 70, 74, 76, 80, 95, 105, 110, 116, 119, 136, 139, 156, 164, 167]] = 1
  new_groups[13, [2, 15, 21, 31, 41, 52, 53, 54, 63, 64, 65, 72, 79, 102, 118, 129, 135, 140, 147, 160, 165]] = 1
  new_groups[14, [49, 77, 83, 85, 93, 100, 102, 106, 115, 117, 121, 132, 137, 170, 182, 191, 197, 205, 214, 216, 220]] = 1
  new_groups[15, [25, 28, 34, 40, 50, 55, 57, 73, 80, 89, 121, 122, 127, 141, 155, 165, 177, 189, 192, 207, 225]] = 1
  new_groups[16, [4, 12, 13, 19, 33, 48, 70, 72, 88, 104, 109, 124, 126, 128, 145, 154, 155, 175, 197, 208, 224]] = 1
  new_groups[17, [10, 24, 29, 48, 50, 64, 92, 112, 133, 134, 136, 152, 163, 169, 170, 173, 195, 199, 204, 222, 227]] = 1
  new_groups[18, [3, 20, 46, 65, 83, 97, 101, 125, 131, 151, 156, 169, 172, 196, 200, 206, 207, 211, 218, 224, 230]] = 1
  new_groups[19, [19, 34, 36, 37, 42, 67, 68, 69, 92, 98, 115, 119, 125, 130, 143, 159, 160, 174, 210, 215, 223]] = 1
  new_groups[20, [10, 17, 37, 46, 94, 105, 124, 132, 140, 141, 157, 162, 168, 180, 184, 185, 193, 198, 217, 221, 229]] = 1
  new_groups[21, [11, 24, 31, 40, 42, 43, 51, 59, 75, 78, 88, 95, 96, 101, 144, 146, 161, 178, 213, 220, 221]] = 1
  return np.array(new_groups, dtype=bool)
# pylint: enable=line-too-long


@gin.configurable

class Origami(group_selector.GroupSelector):