from azureml.core import Workspace
from azureml.core import VERSION
from azureml.core import Dataset
from azureml.core import Model
from azureml.core.compute import ComputeTarget, AmlCompute
from azureml.core.compute_target import ComputeTargetException
from azureml.core import Experiment, ScriptRunConfig, Environment
from azureml.train.hyperdrive import GridParameterSampling, HyperDriveConfig, PrimaryMetricGoal, choice

import os
import time
import json

from database import db_management as mldb


ws = Workspace.from_config() 
print('Ready to use Azure ML {} to work with {}'.format(VERSION, ws.name))

default_ds = ws.get_default_datastore()
cwd = os.getcwd()
dataset_name = 'diabetes dataset'
dataset1_file = 'data/diabetes.csv'
dataset2_file = 'data/diabetes2.csv'

if dataset_name not in ws.datasets:
    ds1 = os.path.join(cwd,dataset1_file)
    ds2 = os.path.join(cwd,dataset2_file)
    default_ds.upload_files(files=[ds1, ds2], # Upload the diabetes csv files in /data
                        target_path='diabetes-data/', # Put it in a folder path in the datastore
                        overwrite=True, # Replace existing files of the same name
                        show_progress=True)

    #Create a tabular dataset from the path on the datastore (this may take a short while)
    tab_data_set = Dataset.Tabular.from_delimited_files(path=(default_ds, 'diabetes-data/*.csv'))

    # Register the tabular dataset
    try:
        tab_data_set = tab_data_set.register(workspace=ws, 
                                name=dataset_name,
                                description='diabetes data',
                                tags = {'format':'CSV'},
                                create_new_version=True)
        print('Dataset registered.')
    except Exception as ex:
        print(ex)
else:
    print('Dataset already registered.')

cluster_name = "qa-cluster"

try:
    # Check for existing compute target
    training_cluster = ComputeTarget(workspace=ws, name=cluster_name)
    print('Found existing cluster, use it.')
except ComputeTargetException:
    # If it doesn't already exist, create it
    try:
        compute_config = AmlCompute.provisioning_configuration(vm_size='STANDARD_DS11_V2', max_nodes=2)
        training_cluster = ComputeTarget.create(ws, cluster_name, compute_config)
        training_cluster.wait_for_completion(show_output=True)
    except Exception as ex:
        print(ex)



def create_run_experiment(p1, p2, p3, p4):
    experiment_folder = os.path.join(cwd,'experiment')

    # Create a Python environment for the experiment
    hyper_env = Environment.from_conda_specification("experiment_env", experiment_folder + "/hyperdrive_env.yml")

    # Get the training dataset
    diabetes_ds = ws.datasets.get("diabetes dataset")

    # Create a script config
    script_config = ScriptRunConfig(source_directory=experiment_folder,
                                    # New change
                                    # script='diabetes_training.py',
                                    script='diabetes_training_decision_tree.py',
                                    # Add non-hyperparameter arguments -in this case, the training dataset
                                    arguments = ['--input-data', diabetes_ds.as_named_input('training_data')],
                                    environment=hyper_env,
                                    compute_target = training_cluster)

    # Sample a range of parameter values
    params = GridParameterSampling(
        {
            # Hyperdrive will try 4 combinations, adding these as script arguments
            ### Sixth change
            #'--penalty_strength': choice(p1, p2, p3, p4)
            '--max_depth': choice(p1, p2),
            '--min_samples_leaf': choice(p3, p4)
        }
    )

    # Configure hyperdrive settings
    hyperdrive = HyperDriveConfig(run_config=script_config, 
                            hyperparameter_sampling=params, 
                            policy=None, # No early stopping policy
                            # New change: AUC works only with binary class
                            primary_metric_name='Accuracy', # Find the highest Accuracy metric ### relevant for discussion/mention
                            primary_metric_goal=PrimaryMetricGoal.MAXIMIZE, 
                            max_total_runs=6, # Restict the experiment to 6 iterations
                            max_concurrent_runs=2) # Run up to 2 iterations in parallel

    # Run the experiment
    experiment = Experiment(workspace=ws, name='diabetes-logistic-regression')
    run = experiment.submit(config=hyperdrive)

    job_id = mldb.add_job(run.id, 0, 0, '', '', 1)
    
    while run.get_status() not in ['Completed', 'Failed']: # For example purposes only, not exhaustive
       print('Run {} not in terminal state'.format(run.id))
       time.sleep(30)
       mldb.update_job(job_id, run.id, 0, 0, '', '', 2)
    
    # Get the best run, and its metrics and arguments
    best_run = run.get_best_run_by_primary_metric()
    best_run_metrics = best_run.get_metrics()
    script_arguments = best_run.get_details() ['runDefinition']['arguments']

    run_id = best_run.id
    auc = best_run_metrics['AUC']
    accuracy = best_run_metrics['Accuracy']
    arguments = json.dumps(script_arguments)

    # Register model
    best_run.register_model(model_path='outputs/diabetes_model.pkl', model_name='diabetes_model',
                            ### Seventh change
                            tags={'Training context':'Logistic Regression'},
                            properties={'AUC': best_run_metrics['AUC'], 'Accuracy': best_run_metrics['Accuracy']})

    ### New change
    ###mldb.update_job(job_id, run_id, auc, accuracy, arguments, 'Logistic Regression', 3)
    mldb.update_job(job_id, run_id, auc, accuracy, arguments, 'Decision Tree Classifier', 3)

    return job_id