import sys,os
import numpy as np
import time
import pandas as pd
import torch

jcm_optimizer_path = r"<JCM_OPTIMIZER_PATH>"
sys.path.insert(0, os.path.join(jcm_optimizer_path, "interface", "python"))
from jcmoptimizer import Server, Client, Study, Observation
server = Server()
client = Client(server.host)

# Definition of the search domain
design_space = [
    {'name': 'b1', 'type': 'continuous', 'domain': (0,10)}, 
    {'name': 'b2', 'type': 'continuous', 'domain': (0.1,4)},
    {'name': 'b3', 'type': 'continuous', 'domain': (-4,-0.1)},
    {'name': 'b4', 'type': 'continuous', 'domain': (0.05,1)},
    {'name': 'b5', 'type': 'continuous', 'domain': (0.05,1)}
]
constraints = [
    {'name': 'test', 'expression': 'b2 + b3 <= 1.0'}
]

# Creation of the study object with study_id 'scipy_least_squares'
study = client.create_study(
    design_space=design_space,
    constraints=constraints,
    driver="ScipyLeastSquares",
    name="Solution of a non-expensive least-square problem based on a scipy implementation",
    study_id="scipy_least_squares"
)
#The vectorial model function of the MGH17 problem
def model(x: torch.Tensor) -> torch.Tensor:
    s = torch.arange(33)
    return x[0] + x[1]*torch.exp(-s*x[3]) + x[2]*torch.exp(-s*x[4])

#Target vector of the MGH17
target=torch.tensor([
    8.44E-01, 9.08E-01, 9.32E-01, 9.36E-01, 9.25E-01,
    9.08E-01, 8.81E-01, 8.50E-01, 8.18E-01, 7.84E-01,
    7.51E-01, 7.18E-01, 6.85E-01, 6.58E-01, 6.28E-01,
    6.03E-01, 5.80E-01, 5.58E-01, 5.38E-01, 5.22E-01,
    5.06E-01, 4.90E-01, 4.78E-01, 4.67E-01, 4.57E-01,
    4.48E-01, 4.38E-01, 4.31E-01, 4.24E-01, 4.20E-01,
    4.14E-01, 4.11E-01, 4.06E-01
])

study.configure(
    target_vector=target.tolist(),
    method="trf",
    max_iter=150,
    num_parallel=2,
    num_initial=2,
    jac=True
)

# Evaluation of the black-box function for specified design parameters
def evaluate(study: Study, b1: float, b2: float, b3: float, b4: float, b5: float) -> Observation:

    observation = study.new_observation()
    #tensor of design values to reconstruct
    x = torch.tensor([b1, b2, b3, b4, b5], requires_grad=True)
    
    observation.add(model(x).tolist())

    #determine Jacobian matrix
    jac = torch.autograd.functional.jacobian(
        func=model,
        inputs=x
    )

    for idx, param in enumerate(design_space):
        observation.add(jac[:, idx].tolist(), derivative=param["name"])
    return observation

# Run the minimization
study.set_evaluator(evaluate)
study.run()
best_sample = study.driver.best_sample
uncertainties = study.driver.uncertainties
print("Reconstructed parameters:")
for param in design_space:
    name = param['name']
    print(f"  {name} = {best_sample[name]:.3f} +/- {uncertainties[name]:.3f}")