def transition_test():

    model = EnergyTrans(r"..\..\..\models\EnergyTrans", "fluCase")
    ensemble = ModelEnsemble()
    ensemble.set_model_structure(model)

    ensemble.perform_experiments(cases=10, callback=HDF5Callback)

def test_optimization():
    ema_logging.log_to_stderr(ema_logging.INFO)
    
    model = FluModel(r'..\data', "fluCase")
    ensemble = ModelEnsemble()
    
    ensemble.set_model_structure(model)
    ensemble.parallel=True
#    ensemble.processes = 12
        
    stats, pop  = ensemble.perform_outcome_optimization(obj_function = obj_function_multi,
                                                    reporting_interval=10, 
                                                    weights=(MAXIMIZE, MAXIMIZE),
                                                    pop_size=100,
                                                    nr_of_generations=5,
                                                    crossover_rate=0.5,
                                                    mutation_rate=0.05)
    res = stats.hall_of_fame.keys
    
    x = [entry.values[0] for entry in res]
    y = [entry.values[1] for entry in res]
    
    print len(x), len(y)
    
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.scatter(x,y)
    ax.set_ylabel("deceased population")
    ax.set_xlabel("infected fraction")
    
    plt.show()

def flu_test():

    model = FluModel(r"..\..\..\models\flu", "fluCase")
    ensemble = ModelEnsemble()
    ensemble.set_model_structure(model)

    ensemble.perform_experiments(cases=10, callback=HDF5Callback)

def test_inspect():
    import inspect_test
    model = FluModel(r'..\..\..\models\flu', "fluCase")
    ensemble = ModelEnsemble()
    ensemble.set_model_structure(model)
    
    ensemble.perform_experiments(cases = 10,
                                 callback=inspect_test.InspectCallback)

def perform_experiments():
    ema_logging.log_to_stderr(level=ema_logging.INFO)
    model = SalinizationModel(r"C:\workspace\EMA-workbench\models\salinization", "verzilting")
    model.step = 4
    
    ensemble = ModelEnsemble()
    ensemble.set_model_structure(model)
    
    ensemble.parallel = True
    nr_of_experiments = 10000
    results = ensemble.perform_experiments(nr_of_experiments)
    return results

def test_save_results():
    os.remove("test.h5")

    nrOfExperiments = 10
    fileName = "test.h5"
    experimentName = "one_exp_test"

    ensemble = ModelEnsemble()
    ensemble.set_model_structure(FluModel(r"..\..\..\models\flu", "fluCase"))

    ensemble.perform_experiments(
        nrOfExperiments, callback=HDF5Callback, fileName=fileName, experimentName=experimentName
    )

 def test_running_lookup_uncertainties(self):
     '''
     This is the more comprehensive test, given that the lookup
     uncertainty replaces itself with a bunch of other uncertainties, check
     whether we can successfully run a set of experiments and get results
     back. We assert that the uncertainties are correctly replaced by
     analyzing the experiments array. 
     
     '''
     model = LookupTestModel( r'../models/', 'lookupTestModel')
     
     #model.step = 4 #reduce data to be stored
     ensemble = ModelEnsemble()
     ensemble.set_model_structure(model)
     
     ensemble.perform_experiments(10)

 def test_vensim_model(self):
     #instantiate a model
     wd = r'../models'
     model = VensimExampleModel(wd, "simpleModel")
     
     #instantiate an ensemble
     ensemble = ModelEnsemble()
     
     #set the model on the ensemble
     ensemble.set_model_structure(model)
     
     nr_runs = 10
     experiments, outcomes = ensemble.perform_experiments(nr_runs)
     
     self.assertEqual(experiments.shape[0], nr_runs)
     self.assertIn('TIME', outcomes.keys())
     self.assertIn(model.outcomes[0].name, outcomes.keys())

def test_optimization():
    ema_logging.log_to_stderr(ema_logging.INFO)
    
    model = FluModel(r'..\data', "fluCase")
    ensemble = ModelEnsemble()
    
    ensemble.set_model_structure(model)
    ensemble.parallel=True
        
    stats, pop  = ensemble.perform_outcome_optimization(obj_function = obj_function_multi,
                                                    reporting_interval=100, 
                                                    weights=(MAXIMIZE, MAXIMIZE),
                                                    pop_size=100,
                                                    nr_of_generations=20,
                                                    crossover_rate=0.5,
                                                    mutation_rate=0.05,
                                                    caching=False)
    res = stats.hall_of_fame.keys
    
    print len(stats.tried_solutions.values())

def maxmin_optimize():
    ema_logging.log_to_stderr(ema_logging.INFO)
 
    model = TestModel("", 'simpleModel') #instantiate the model
    ensemble = ModelEnsemble() #instantiate an ensemble
    ensemble.set_model_structure(model) #set the model on the ensemble
    ensemble.parallel = True
    ensemble.processes = 12
    
    def obj_function1(outcomes):
        return outcomes['y']
    
    policy_levers = { "L1": (0,1),
                      "L2": (0,1)}
    
    
    
    results = ensemble.perform_maximin_optimization(obj_function1 = obj_function1, 
                                                policy_levers = policy_levers,
                                                minimax1='minimize',
                                                nrOfGenerations1=50,
                                                nrOfPopMembers1=200,
                                                minimax2 = "maximize",                                   
                                                nrOfGenerations2 = 50,
                                                nrOfPopMembers2 = 100,
                                                )

    graph_errorbars_raw(results['stats'])
    plt.show()

def robust_optimize():
    ema_logging.log_to_stderr(ema_logging.INFO)
 
    model = TestModel("", 'simpleModel') #instantiate the model
    ensemble = ModelEnsemble() #instantiate an ensemble
    ensemble.set_model_structure(model) #set the model on the ensemble
    
    
    policy_levers = { "L1": (0,1),
                      "L2": (0,1)}
    
    def obj_func(results):
        return np.average(results['y'])        
    
    results = ensemble.perform_robust_optimization(cases=1000, 
                                                   obj_function=obj_func, 
                                                   policy_levers=policy_levers, 
                                                   minimax='minimize', 
                                                   nrOfGenerations=50, 
                                                   nrOfPopMembers=20 )
    graph_errorbars_raw(results['stats'])
    plt.show()

def test_multiple_models():
    class Model1(ModelStructureInterface):
        uncertainties = [ParameterUncertainty((0,1),"a"),
                         ParameterUncertainty((0,1),"b")]
        
        outcomes = [Outcome("test")]
        
        def model_init(self, policy, kwargs):
            pass
        
        def run_model(self, case):
            self.output['test'] = 1

    class Model2(ModelStructureInterface):
        uncertainties = [ParameterUncertainty((0,1),"b"),
                         ParameterUncertainty((0,1),"c")]
        
        outcomes = [Outcome("test")]
        
        def model_init(self, policy, kwargs):
            pass
        
        def run_model(self, case):
            self.output['test'] = 1
    
#    os.remove('test.h5')

    nrOfExperiments = 10
    fileName = 'test.h5'
    experimentName = "one_exp_test"
    
    ensemble = ModelEnsemble()
    ensemble.add_model_structure(Model1('', "test1"))
    ensemble.add_model_structure(Model2('', "test2"))
    
    ensemble.perform_experiments(nrOfExperiments,
                                 callback=HDF5Callback,
                                 fileName=fileName,
                                 experimentName=experimentName)

def test_optimization():
    ema_logging.log_to_stderr(ema_logging.INFO)
    
    model = FluModel(r'../models', "fluCase")
    ensemble = ModelEnsemble()
    
    ensemble.set_model_structure(model)
    ensemble.parallel=True
    
    pop_size = 8
    nr_of_generations = 10
    eps = np.array([1e-3, 1e6])

    stats, pop  = ensemble.perform_outcome_optimization(obj_function = obj_function_multi,
                                                    algorithm=epsNSGA2,
                                                    reporting_interval=100, 
                                                    weights=(MAXIMIZE, MAXIMIZE),
                                                    pop_size=pop_size,          
                                                    nr_of_generations=nr_of_generations,
                                                    crossover_rate=0.8,
                                                    mutation_rate=0.05,
                                                    eps=eps)
    fn = '../data/test optimization save.bz2'

def test_tree():
    
    log_to_stderr(level= INFO)
        
    model = FluModel(r'..\..\models\flu', "fluCase")
    ensemble = ModelEnsemble()
    ensemble.parallel = True
    ensemble.set_model_structure(model)
    
    policies = [{'name': 'no policy',
                 'file': r'\FLUvensimV1basecase.vpm'},
                {'name': 'static policy',
                 'file': r'\FLUvensimV1static.vpm'},
                {'name': 'adaptive policy',
                 'file': r'\FLUvensimV1dynamic.vpm'}
                ]
    ensemble.add_policies(policies)
    
    results = ensemble.perform_experiments(10)
   
    a_tree = tree(results, classify)

def test_feature_selection():
    log_to_stderr(level= INFO)
        
    model = FluModel(r'..\..\models\flu', "fluCase")
    ensemble = ModelEnsemble()
    ensemble.parallel = True
    ensemble.set_model_structure(model)
    
    policies = [{'name': 'no policy',
                 'file': r'\FLUvensimV1basecase.vpm'},
                {'name': 'static policy',
                 'file': r'\FLUvensimV1static.vpm'},
                {'name': 'adaptive policy',
                 'file': r'\FLUvensimV1dynamic.vpm'}
                ]
    ensemble.add_policies(policies)
    
    results = ensemble.perform_experiments(5000)
   
    results = feature_selection(results, classify)
    for entry in results:
        print entry[0] +"\t" + str(entry[1])

def outcome_optimize():
    ema_logging.log_to_stderr(ema_logging.INFO)
 
    model = TestModel("", 'simpleModel') #instantiate the model
    ensemble = ModelEnsemble() #instantiate an ensemble
    ensemble.set_model_structure(model) #set the model on the ensemble
    policy = {"name": "test",
              "L1": 1,
              "L2": 1}
    ensemble.add_policy(policy)
    
    def obj_func(results):
        return results['y']        
    
    results = ensemble.perform_outcome_optimization(obj_function=obj_func, 
                                          minimax = 'minimize', 
                                          nrOfGenerations = 1000, 
                                          nrOfPopMembers = 10)

    graph_errorbars_raw(results['stats'])
    plt.show()

                             "normal contact rate region 2")]
                         
    def model_init(self, policy, kwargs):
        '''initializes the model'''
        
        try:
            self.model_file = policy['file']
        except KeyError:
            ema_logging.warning("key 'file' not found in policy")
        super(FluModel, self).model_init(policy, kwargs)
        
if __name__ == "__main__":
    ema_logging.log_to_stderr(ema_logging.INFO)
        
    model = FluModel(r'..\..\models\flu', "fluCase")
    ensemble = ModelEnsemble()
    ensemble.set_model_structure(model)
    
    #add policies
    policies = [{'name': 'no policy',
                 'file': r'\FLUvensimV1basecase.vpm'},
                {'name': 'static policy',
                 'file': r'\FLUvensimV1static.vpm'},
                {'name': 'adaptive policy',
                 'file': r'\FLUvensimV1dynamic.vpm'}
                ]
    ensemble.add_policies(policies)

    ensemble.parallel = True #turn on parallel processing

    results = ensemble.perform_experiments(1000)

                             "susceptible to immune population delay time region 1"),
        ParameterUncertainty((0.5,2), 
                             "susceptible to immune population delay time region 2"),
        ParameterUncertainty((0.01, 5), 
                             "root contact rate region 1"),
        ParameterUncertainty((0.01, 5), 
                             "root contact ratio region 2"),
        ParameterUncertainty((0, 0.15), 
                             "infection ratio region 1"),
        ParameterUncertainty((0, 0.15), 
                             "infection rate region 2"),
        ParameterUncertainty((10, 100), 
                             "normal contact rate region 1"),
        ParameterUncertainty((10, 200), 
                             "normal contact rate region 2")]
                         
        
if __name__ == "__main__":
    ema_logging.log_to_stderr(ema_logging.INFO)
        
    model = FluModel(r'./models/flu', "fluCase")
    ensemble = ModelEnsemble()
    ensemble.set_model_structure(model)
    
    ensemble.parallel = True #turn on parallel processing

    nr_experiments = 1000
    results = ensemble.perform_experiments(nr_experiments)
    
    fh =  r'./data/{} flu cases no policy.tar.gz'.format(nr_experiments)
    save_results(results, fh)

                     ParameterUncertainty((0.45,0.55),
                                          "tH"),
                     ParameterUncertainty((0.1,0.3),
                                          "kk")]
    
    #specification of the outcomes
    outcomes = [Outcome("B4:B1076", time=True),  #we can refer to a range in the normal way
                Outcome("P_t", time=True)] # we can also use named range
    
    #name of the sheet
    sheet = "Sheet1"
    
    #relative path to the Excel file
    workbook = r'\excel example.xlsx'
    

if __name__ == "__main__":    
    ema_logging.log_to_stderr(level=ema_logging.INFO)
    
    model = ExcelModel(r"./models/excelModel", "predatorPrey")
    
    ensemble = ModelEnsemble()
    ensemble.set_model_structure(model) 

    ensemble.parallel = True #turn on parallel computing
    ensemble.processes = 2 #using only 2 cores 
    
    #run 100 experiments
    nr_experiments = 100
    results = ensemble.perform_experiments(nr_experiments) 

    a_mean = np.mean(a)
    b_mean = np.mean(b)

    if a_mean < 0.5 or b_mean < 0.5:
        return (np.inf,) * 2
    else:
        return a_mean, b_mean

if __name__ == "__main__":
    ema_logging.log_to_stderr(ema_logging.INFO)
    model = DummyModel(r"", "dummy")
    
    np.random.seed(123456789)
       
    ensemble = ModelEnsemble()
    ensemble.set_model_structure(model)

    
    policy_levers = {'Trigger a': {'type':'list', 'values':[0, 0.25, 0.5, 0.75, 1]},
                     'Trigger b': {'type':'list', 'values':[0, 0.25, 0.5, 0.75, 1]},
                     'Trigger c': {'type':'list', 'values':[0, 0.25, 0.5, 0.75, 1]}}
    
    cases = ensemble._generate_samples(10, UNION)[0]
    ensemble.add_policy({"name":None})
    experiments = [entry for entry in ensemble._generate_experiments(cases)]
    for entry in experiments:
        entry.pop("model")
        entry.pop("policy")
    cases = experiments