from gearup import expand

def test_expand_library_basic():
  library = {
    'experiment': {
      '_template': {
        'model': '$model',
        'lr': '$lr'
      },
      '_values': {
        'model': ['resnet', 'vgg'],
        'lr': [0.01, 0.001]
      }
    }
  }

  name, expanded = expand(library)
  assert name is None
  assert len(expanded) == 4
  assert 'experiment-0-0' in expanded
  assert 'experiment-1-1' in expanded
  assert expanded['experiment-0-0']['model'] == 'resnet'
  assert expanded['experiment-0-0']['lr'] == 0.01
  assert expanded['experiment-1-1']['model'] == 'vgg'
  assert expanded['experiment-1-1']['lr'] == 0.001

def test_expand_library_with_name():
  library = {
    '_name': 'my_library',
    'config': {
      'value': 42
    }
  }

  name, expanded = expand(library)
  assert name == 'my_library'
  assert len(expanded) == 1
  assert 'config' in expanded
  assert expanded['config']['value'] == 42

def test_expand_library_non_template():
  library = {
    'config1': {
      'foo': 'bar',
      'num': 123
    },
    'config2': {
      'baz': 'qux'
    }
  }

  name, expanded = expand(library)
  assert name is None
  assert len(expanded) == 2
  assert expanded['config1']['foo'] == 'bar'
  assert expanded['config1']['num'] == 123
  assert expanded['config2']['baz'] == 'qux'

def test_expand_library_mixed():
  library = {
    'static_config': {
      'type': 'static'
    },
    'template_config': {
      '_template': {
        'type': 'dynamic',
        'value': '$val'
      },
      '_values': {
        'val': [1, 2]
      }
    }
  }

  name, expanded = expand(library)
  assert len(expanded) == 3
  assert 'static_config' in expanded
  assert expanded['static_config']['type'] == 'static'
  assert 'template_config-0' in expanded
  assert 'template_config-1' in expanded
  assert expanded['template_config-0']['value'] == 1
  assert expanded['template_config-1']['value'] == 2

def test_expand_library_collision():
  library = {
    'experiment-0-0': {
      'static': True
    },
    'experiment': {
      '_template': {
        'model': '$model',
        'lr': '$lr'
      },
      '_values': {
        'model': ['rI want esnet'],
        'lr': [0.01]
      }
    }
  }

  try:
    expand(library)
    assert False, 'Expected ValueError for name collision'
  except ValueError as e:
    assert 'name collision' in str(e)

def test_expand_library_dict_values():
  library = {
    'experiment': {
      '_template': {
        'optimizer': '$opt'
      },
      '_values': {
        'opt': {
          'adam': {'type': 'Adam', 'lr': 0.001},
          'sgd': {'type': 'SGD', 'lr': 0.01}
        }
      }
    }
  }

  name, expanded = expand(library)
  assert len(expanded) == 2
  assert 'experiment-adam' in expanded
  assert 'experiment-sgd' in expanded
  assert expanded['experiment-adam']['optimizer']['type'] == 'Adam'
  assert expanded['experiment-sgd']['optimizer']['type'] == 'SGD'