controller_schema = {
  "type": "object",
  "properties": {
    "max_n_epochs": {"type": "integer"},
    "max_n_better_results": {"type": "integer"}
  }
}

optimizer_schema = {
  "type": "object",
  "properties": {
    "optimizer_type": {"type": "string", "enum": ["SGD"]},
    "momentum": {"type": "number"}
  }
}

model_schema = {
  "type": "object",
  "properties": {
    "name": {"type": "string", "enum": ["resnet34", "resnet", "convnet", "densenet",
                                        "basic_net", "toy_relu_net", "toy_res_net",
                                        "toy_svm", "toy_svm_poly", "toy_gaussian_naive",
                                        "toy_batch_relu_net", "toy_batch_res_net", "toy_batch_net",
                                        "toy_long_batch_relu_net", "toy_long_batch_res_net",
                                        "basic_net", "toy_relu_net", "toy_res_net", "toy_gaussian_naive"
                                        ]},

    "n_outputs": {"type": "integer"}
  }
}

LearnMapping_schema = {
  "type": "object",
  "properties": {
    "optimizer": optimizer_schema,
    "controller": controller_schema,
    "model": model_schema,
  },
  "required": ["optimizer", "controller", "model"]
}

Loss_schema = {
  "type": "object",
  "properties": {
    "loss_function": {"type": "object", "properties": {"threshold_high": {"type": "number"},
                                                       "threshold_low": {"type": "number"}}},
    "regularization_weight": {"type": "number"},
    "loss_aggregator": {"type": "object", "properties": {"name": {"type": "string"}}},
    "report_frequency": {"type": "integer"}
  }
}

SupervisedDataloader_schema = {
  "type": "object",
  "properties": {
    "batch_size": {"type": "integer"},
    "dataset": {"type": "string"},
    "grayscale": {"type": "boolean"},
    "shuffle": {"type": "boolean"},
    "transform": {"type": "object"}
  }
}


Domains_schema = {
  "type": "object",
  "properties": {
    "source": SupervisedDataloader_schema,
    "target": SupervisedDataloader_schema
  }
}


Task_schema = {
  "type": "object",
  "properties": {
    "task_type": {"type": "string", "enum": ["supervised-learning", "semi-supervised-learning", "domain-confusion"]},
    "model_to_train": {},
    "model_to_infer": {},
    "domain_defaults": SupervisedDataloader_schema,
    "domains": Domains_schema,
    "dataloaders": {"type": "object", "properties": {"labeled_dataloader": {"type": "string"},
                                                     "unlabeled_dataloader": {"type": "string"},
                                                     "inference_dataloader": {"type": "string"}}},
    "loss": Loss_schema,
    "inference_certainty_threshold": {"type": "number"}
  },
  "required": ["model_to_train", "dataloaders"]
}


TaskPathDefaults_schema = {
  "type": "object",
  "properties": {
    "domain_defaults": SupervisedDataloader_schema,
    "loss": Loss_schema
  }
}


TaskPath_schema = {
  "type": "object",
  "properties": {
    "defaults": TaskPathDefaults_schema,
    "tasks": {"type": "array", "items": Task_schema}
  },
  "required": ["defaults", "tasks"]
}


Experiment_schema = {
  "type": "object",
  "properties": {
    "task_path": TaskPath_schema,
    "learn_mapping": LearnMapping_schema
  },
  "required": ["task_path", "learn_mapping"]
}
