import  itertools
import  shutil
import  os
from    pathlib import Path
import  sys
import  logging
import  yaml

class niqcs_para_sweep:
    """
    Generates a parametric sweep over several parameters.
    """

    parameters = ['Current']
    parameter_values = []
    output_values = ['nzpv', 'hot-spot']
    itervalues = []
    input_vars = []
    statics = None

    # Filenames
    pybbq_path = 'pyBBQ.py'
    yaml_param = ''             # Yaml file with the parameters to sweep
    yaml_init = ''              # Yaml file with all input parameters for pyBBQ
    yaml_sweep = ''             # Working yaml file copied from yaml_init
    output_csv = ''
    output_folder = 'Output' + os.sep + 'sweeper'
    result_file = ''
    mode = 'Sweep'  # 'Sweep' is standard, 'Sweep-Full' saves the full output
    mode2 = 'full' # partial or full, partial sweeps over range(len(x)) of all arrays, full makes all combinations
    name = None
    optimize = False
    delimiter = ','

    def __init__(self, sweeper_input_yaml = None, original_yaml=None):
        self.yaml_input = sweeper_input_yaml
        self.yaml_original = original_yaml

        if self.yaml_input is None and self.yaml_original is None:
            self.make_folders()
        else:
            # Loading the vars for the parametric sweep
            with open(self.yaml_input) as file:
                contents = yaml.load(file, Loader=yaml.FullLoader)

            if 'name' in contents.keys():
                setattr(self, 'name', contents['name'])

            if 'sweep' in contents.keys():
                setattr(self, 'sweep', contents['sweep'])

            # Making the output folder
            self.make_folders()

            keys = [i for i in contents['sweep']]
            unpacked_values = []
            for i in keys:
                unpacked_values.append([(i, j) for j in contents['sweep'][i]])

            # Creating the sweeper list
            if self.mode2 == 'full':
                self.itervalues = list(itertools.product(*unpacked_values))
            elif self.mode2 == 'partial':
                self.itervalues = list(zip(*unpacked_values))

            self.start_sweep()

    def uniquify(self):
        # Creates a unique foldername for the simulation
        self.unique_number = 1
        path = self.file_folder / 'sweeper' / 'sweeper_output' / (self.name + "_" + str(self.unique_number))
        while os.path.isdir(path):
            self.unique_number += 1
            path = self.file_folder / 'sweeper' / 'sweeper_output' / (self.name + "_" + str(self.unique_number))
        return path

    def make_folders(self):
        """
            Making the sweeper output folders
        """

        # Creating the sweeper folder
        self.file_folder = Path(os.path.dirname(os.path.basename(__file__)))
        try:
            os.mkdir(self.file_folder / 'sweeper')
        except OSError:
            print("Creation of the directory %s failed" % (self.file_folder / 'sweeper'))
        else:
            print("Successfully created the directory %s " % (self.file_folder / 'sweeper'))

        # Creating the sweeper input folder
        try:
            os.mkdir(self.file_folder / 'sweeper' / 'sweeper_input_yaml')
        except OSError:
            print("Creation of the directory %s failed" % (self.file_folder / 'sweeper'))
        else:
            print("Successfully created the directory %s " % (self.file_folder / 'sweeper'))

        # Creating the sweeper output folder
        try:
            os.mkdir(self.file_folder / 'sweeper' / 'sweeper_output')
        except OSError:
            print("Creation of the directory %s failed" % (self.file_folder / 'sweeper'))
        else:
            print("Successfully created the directory %s " % (self.file_folder / 'sweeper'))

        if self.yaml_input is not None and self.yaml_original is not None and self.name is not None:
            self.output_folder = self.uniquify()

            try:
                os.mkdir(self.output_folder)
            except OSError:
                print("Creation of the directory %s failed" % self.output_folder)
            else:
                print("Successfully created the directory %s " % self.output_folder)

            # Copying the input yaml file.
            self.yaml_original_out = self.output_folder / 'original.yaml'
            try:
                shutil.copy(self.yaml_original, self.yaml_original_out)
                print("File copied successfully.")
            # If source and destination are same
            except shutil.SameFileError:
                print("Source and destination represents the same file.")
            # If there is any permission issue
            except PermissionError:
                print("Permission denied.")
            # For other errors
            except:
                print("Error occurred while copying file.")

            # Copying the input yaml file.
            self.yaml_sweep_out = self.output_folder / 'sweeper_input.yaml'
            try:
                shutil.copy(self.yaml_input, self.yaml_sweep_out)
                print("File copied successfully.")
            # If source and destination are same
            except shutil.SameFileError:
                print("Source and destination represents the same file.")
            # If there is any permission issue
            except PermissionError:
                print("Permission denied.")
            # For other errors
            except:
                print("Error occurred while copying file.")

            # creating a python file to run the sweep:
            with open(self.output_folder / 'run_sweep.py', 'w+') as the_file:
                the_file.write('# !/usr/bin/env python3\n')
                the_file.write('import subprocess\n\n')

    def start_sweep(self):
        """ Starts the parametric sweep """
        print(self.itervalues)
        for i in range(len(self.itervalues)):
            yaml_sweep_n = 'sweep_' + str(i) + '.yaml'
            self.set_state(self.itervalues[i], yaml_sweep_n)
            self.input_vars.append(yaml_sweep_n)


    def set_state(self, change_params, output_yaml):
        """ Changes the parameters in the working yaml file
            Input:  zip of the parameters, for example: [('Current', 10), ('T1', 50)]
        """

        logging.info('Preparing: ' + output_yaml)

        with open(self.yaml_original_out) as f:
            doc = yaml.safe_load(f)

        for i in change_params:
            doc[i[0]] = i[1]

        if 'tag' in doc.keys():
            doc['tag'] = doc['tag'] + '_' + self.name + '_' + str(self.unique_number) + '_' + output_yaml.split('.')[0]

        if self.statics is not None:
            for i in self.statics.keys():
                doc[i] = self.statics[i]

        with open(self.output_folder / output_yaml, 'w') as f:
            yaml.safe_dump(doc, f, default_flow_style=False)

        # creating a python file to run the sweep:
        with open(self.output_folder / 'run_sweep.py', 'a+') as the_file:
            the_file.write('subprocess.call([\'python\', \'NI_Coil_Solver_Torch.py\', \'' + str(self.output_folder / output_yaml) + '\'])\n')

if __name__ == "__main__":
    print(sys.argv)
    if len(sys.argv) == 3:
        print('Loading the provided yaml files: ' + sys.argv[1] + ' and ' + sys.argv[2])
        x = niqcs_para_sweep(sys.argv[1], sys.argv[2])
    else:
        x = niqcs_para_sweep()