# Copyright (c) 2024 - Zizhe Wang
# https://zizhe.wang

import json
import os
from optimize_main import run_optimization

class MOO4ModelicaWrapper:
    def __init__(self, orchestration_config_path, config_path):
        self.orchestration_config_path = orchestration_config_path
        self.config_path = config_path
        self.load_configs()
        self.update_config_file()

    def load_configs(self):
        with open(self.orchestration_config_path, 'r') as f:
            self.orchestration_config = json.load(f)

        with open(self.config_path, 'r') as f:
            self.config = json.load(f)

    def update_config_file(self):
        # Update config with orchestration config
        self.config['MODEL_NAME'] = self.orchestration_config['MODEL_NAME']
        self.config['MODEL_FILE'] = self.orchestration_config['MODEL_FILE']
        self.config['SIMULATION_STOP_TIME'] = self.orchestration_config['SIMULATION_TIME']
        self.config['PARAMETERS'] = self.orchestration_config['TUNABLE_PARAMETERS']['PARAMETERS']
        self.config['PARAM_BOUNDS'] = self.orchestration_config['TUNABLE_PARAMETERS']['PARAM_BOUNDS']
        self.config['OBJECTIVES'] = self.orchestration_config['OBJECTIVES']
        self.config['PLOT_CONFIG'] = self.orchestration_config['PLOT_CONFIG']
        self.config['N_JOBS'] = self.orchestration_config['N_JOBS']
        for key, value in self.orchestration_config['OPTIMIZATION_CONFIG'].items():
            self.config['OPTIMIZATION_CONFIG'][key] = value

        with open(self.config_path, 'w') as f:
            json.dump(self.config, f, indent=4)

    def update_config(self, simulation_inputs, simulation_time):
        # Update only the simulation-specific parameters
        self.config['SIMULATION_STOP_TIME'] = simulation_time
        for input_param, value in simulation_inputs.items():
            self.config[self.orchestration_config['INPUT_PARAMETERS'][input_param]] = value

        # Save the updated configuration to a file
        with open(self.config_path, 'w') as f:
            json.dump(self.config, f, indent=4)

    def run_optimization(self):
        run_optimization()

    def get_parameters(self):
        # Assuming the results are saved to a file "optimization_results.json"
        results_path = os.path.join(os.path.dirname(self.config_path), 'results', 'optimization_results.json')
        with open(results_path, 'r') as f:
            optimization_results = json.load(f)

        # The parameters are returned as they are in the results
        parameters = optimization_results["parameters"]

        return [
            {param: value for param, value in zip(self.config['PARAMETERS'], params)}
            for params in parameters
        ]