from typing import Any, Dict
import numpy as np
from scipy.optimize import basinhopping

from ..exceptions import InverseDesignError


class InverseDesign:
    def __init__(self, client: "RavanClient"):
        self.client = client

    def run(self, target_entropy: float, target_fidelity: float, x0: list[float] | None = None, niter: int = 25) -> Dict[str, Any]:
        predictor = self.client._get_predict_function()

        def objective(p):
            apply_h = max(0.0, min(1.0, float(p[0])))
            apply_cnot = max(0.0, min(1.0, float(p[1])))
            shots_norm = max(0.05, min(1.0, float(p[2])))
            x = np.array([[apply_h, apply_cnot, shots_norm]], dtype=np.float32)
            pred = predictor(x)[0]
            return (pred[0] - target_entropy) ** 2 + (pred[1] - target_fidelity) ** 2

        x0 = x0 or [1.0, 1.0, 0.5]
        minimizer_kwargs = {
            'method': 'L-BFGS-B',
            'bounds': [(0, 1), (0, 1), (0.05, 1.0)],
            'options': {'maxiter': 50, 'ftol': 1e-6, 'disp': False}
        }

        try:
            res = basinhopping(
                func=objective,
                x0=np.array(x0, dtype=np.float64),
                niter=niter,
                T=1.0,
                stepsize=0.1,
                minimizer_kwargs=minimizer_kwargs,
            )
        except Exception as e:
            raise InverseDesignError(str(e))

        solution = {
            'x': res.x.tolist(),
            'fun': float(res.fun),
            'success': True,
            'message': getattr(res, 'message', ''),
        }
        return solution


