import hopsy
import matplotlib.pyplot as plt
import numpy as np

import x3cflux

# Load FluxML file with data from INST experiment
sim = simulator_inst = x3cflux.create_simulator_from_fml("spiralus.fml", "ms_INST")

# Configure simulation: solver + accuracy (optional)
sim.builder.set_solver("bdf")
sim.builder.solver.relative_tolerance = 1e-6
sim.builder.solver.absolute_tolerance = 1e-9

# Forward simulation of a specific parameter set at specific time points
params = [0.5, 1.2, 3.0, 1.1]  # p.n  # u.n  # G  # B
results = sim.compute_measurements(params=params, time_stamps=[0.0, 5.0, 10.0, 60])

# Multi-start fit (using external optimizer IPOPT)
optima, obj_vals = x3cflux.run_multi_optimization(
    sim, starting_points=x3cflux.run_uniform_sampling(sim, num_samples=100), num_procs=4
)
best = np.argsort(obj_vals)[0]
print(f"Optimum at {optima[:, best]}; Residual {obj_vals[best]}")

# Bayesian uncertainty analysis (using external sampler hopsy)
samples = x3cflux.run_non_uniform_sampling(
    sim,
    num_samples=5_000,
    proposal=hopsy.CSmMALAProposal,
    starting_point=np.array([0.5, 1.1, 5, 0.9]),
    progress_bar=True,
)
plt.hist(samples[0, 1, :])
plt.show()