import sys,os
import numpy as np
import time

from jcmoptimizer import Server, Client, Study, Obseravtion
server = Server()
client = Client(host=server.host)


#Amplitude of a driven harmonic oscillator. 
def amplitude(omega: float, F: float, omega0: float, gamma: float) -> float:
    return F/np.sqrt((2*omega*omega0*gamma)**2 + (omega0**2 - omega**2)**2)

# Definition of the search domain to tune the oscillator
design_space = [
    {'name': 'F', 'type': 'continuous', 'domain': (0.1, 40.0)}, 
    {'name': 'omega0', 'type': 'continuous', 'domain': (1.0, 4.0)},
    {'name': 'gamma', 'type': 'continuous', 'domain': (0.03, 0.3)},
]

#The amplitude is scaned for 10 different driving frequencies
omegas = np.linspace(1,3,10)

# Creation of the study object with study_id 'harmonic_oscillator_fit'
study = client.create_study(
    design_space=design_space,
    driver="ActiveLearning",
    study_name="Optimization of resonant system based on Gaussian fit",
    study_id="harmonic_oscillator_fit"
)

#configuration of study
study.configure(
    max_iter = 30,
    surrogates = [
        #A Gaussian process learns the dependence of the amplitudes for
        #all wavelengths on the design parameters.
        dict(type="GP", name="amplitudes", output_dim=len(omegas))
    ],    
    variables = [
        #The amplitudes are fitted to a Gaussian + linear expression with
        #amplitude A, resonance frequency tau, linear gradient B, and constant offset C.
        #The output are the fitted values and the mean-squared error (MSE).
        dict(type="Fit", name="fit", input="amplitudes", 
             expression="A*exp(-0.5*(omega-tau)^2/sigma^2) + B*omega + C", 
             output_names=["A", "B", "C", "tau", "sigma", "MSE"],
             output_dim=6, 
             model_variables=["omega"], variable_values=omegas[:,None].tolist(),
             initial_parameters=[1.0, 0.0, 0.0, 1.0, 0.5],
             prior_uncertainties=[10.0, 10.0, 10.0, 10.0, 10.0],
            ),  
        #The expression that defines the loss function:
        dict(type="Expression", name="loss",
             expression="(A-20)^2 + (tau - 2)^2 + sigma^2 + 0.001*MSE")
    ],
    objectives = [
        #The only objective of the study is to minimize the loss.
        dict(type="Minimizer", variable="loss"),
    ],
    acquisition_optimizer = dict(
        #Sample computation based on fits is more expensive. In this case
        #advanced sample computation is usually too laborious.
        compute_suggestion_in_advance = False
    ),    
)

# Evaluation of the black-box function for specified design parameters
def evaluate(study: Study, F: float, omega0: float, gamma: float) -> Observation:
    #The harmonic-oscillator amplitude is evaluated for all omega values and the
    #observed values are used as training input for the Gaussian process "amplitudes"
    observation = study.new_observation()
    observation.add([amplitude(omega, F, omega0, gamma) for omega in omegas])
    return observation

# Run the minimization
study.set_evaluator(evaluate)
study.run()

# Print result
best = study.get_state("driver.best_sample")
print(f"Best design parameters: F={best['F']:.3f}, omega0={best['omega0']:.3f}, "
      f"gamma={best['gamma']:.3f}")
print("Objective: (A-20)^2 + (tau - 2)^2 + sigma^2 + 0.001*MSE = "
      f"{study.get_state('driver.min_objective'):.3f}")