from typing import Any, Dict, List, Optional, Callable
import os
import requests
import numpy as np

from .exceptions import APIUnavailableError, APIRequestError, LocalModelNotAvailable
from .utils.config import SDKConfig, api_is_available
from .models.surrogate import SurrogateModel
from .models.inverse_design import InverseDesign
from .models.benchmark import BenchmarkSuite


class RavanClient:
    def __init__(
        self,
        api_url: str = SDKConfig.api_url,
        api_key: Optional[str] = None,
        mode: str = "auto",
        model_dir: str = SDKConfig.model_dir,
        onnx_path: str = SDKConfig.onnx_path,
    ):
        self.api_url = api_url.rstrip("/")
        self.api_key = api_key
        self.headers = {"Content-Type": "application/json"}
        if api_key:
            self.headers["Authorization"] = f"Bearer {api_key}"
        self.mode = mode
        self.model_dir = model_dir
        self.onnx_path = onnx_path

        self._active_mode = self._select_mode()

        self.surrogate = SurrogateModel(self)
        self.inverse = InverseDesign(self)
        self.benchmark = BenchmarkSuite(self)

        self._onnx_session = None
        self._onnx_input_name = None
        self._pytorch_model = None

    def _select_mode(self) -> str:
        if self.mode == "http":
            if not api_is_available(self.api_url):
                raise APIUnavailableError(f"API unavailable at {self.api_url}")
            return "http"
        if self.mode == "local":
            return "local"
        return "http" if api_is_available(self.api_url) else "local"

    def predict(self, params: Dict[str, float] | List[float] | np.ndarray) -> Dict[str, Any]:
        parameters: List[List[float]]
        if isinstance(params, dict):
            vals = list(params.values())
            parameters = [vals] if vals and not isinstance(vals[0], list) else vals
        elif isinstance(params, np.ndarray):
            parameters = params.tolist() if params.ndim == 2 else [params.tolist()]
        else:
            parameters = [params]

        if self._active_mode == "http":
            try:
                resp = requests.post(f"{self.api_url}/predict", json={"parameters": parameters}, headers=self.headers, timeout=10)
                resp.raise_for_status()
                return resp.json()
            except requests.RequestException as e:
                raise APIRequestError(str(e))
        predictor = self._get_predict_function()
        arr = np.asarray(parameters, dtype=np.float32)
        preds = predictor(arr)
        return {"predictions": preds.tolist()}

    def explain(self, params: Dict[str, float] | List[float] | np.ndarray) -> Dict[str, Any]:
        if self._active_mode != "http":
            raise APIUnavailableError("Explain endpoint requires running API server")
        parameters: List[List[float]]
        if isinstance(params, dict):
            vals = list(params.values())
            parameters = [vals] if vals and not isinstance(vals[0], list) else vals
        elif isinstance(params, np.ndarray):
            parameters = params.tolist() if params.ndim == 2 else [params.tolist()]
        else:
            parameters = [params]
        try:
            resp = requests.post(f"{self.api_url}/explain", json={"parameters": parameters}, headers=self.headers, timeout=15)
            resp.raise_for_status()
            return resp.json()
        except requests.RequestException as e:
            raise APIRequestError(str(e))

    def _get_predict_function(self) -> Callable[[np.ndarray], np.ndarray]:
        if os.path.exists(self.onnx_path):
            return self._predict_with_onnx
        if api_is_available(self.api_url):
            def http_fn(x: np.ndarray) -> np.ndarray:
                resp = requests.post(f"{self.api_url}/predict", json={"parameters": x.astype(np.float32).tolist()}, headers=self.headers, timeout=10)
                resp.raise_for_status()
                data = resp.json()
                return np.asarray(data["predictions"], dtype=np.float32)
            return http_fn
        pt_dir = self.model_dir
        if os.path.isdir(pt_dir) and os.path.exists(os.path.join(pt_dir, "model.pt")):
            return self._predict_with_pytorch
        raise LocalModelNotAvailable("No available predictor: missing ONNX, HTTP API, and local PyTorch model")

    def _ensure_onnx(self):
        if self._onnx_session is not None:
            return
        import onnxruntime as ort
        providers = ['CPUExecutionProvider']
        sess = ort.InferenceSession(self.onnx_path, providers=providers)
        self._onnx_session = sess
        self._onnx_input_name = sess.get_inputs()[0].name

    def _predict_with_onnx(self, x: np.ndarray) -> np.ndarray:
        self._ensure_onnx()
        outputs = self._onnx_session.run(None, {self._onnx_input_name: x.astype(np.float32)})
        return outputs[0].astype(np.float32)

    def _ensure_pytorch(self):
        if self._pytorch_model is not None:
            return
        from mlp_regressor import MLPRegressor
        from model_base import ModelConfig
        config = ModelConfig(model_type='mlp', input_dim=3, output_dim=2, hyperparameters={})
        model = MLPRegressor(config)
        model.load(self.model_dir)
        self._pytorch_model = model

    def _predict_with_pytorch(self, x: np.ndarray) -> np.ndarray:
        self._ensure_pytorch()
        return self._pytorch_model.predict(x.astype(np.float32)).astype(np.float32)

    def simulate_quantum_circuit(self, params: Dict[str, Any]) -> Dict[str, Any]:
        if self._active_mode == "http":
            try:
                resp = requests.post(f"{self.api_url}/simulate/quantum_circuit", json=params, headers=self.headers, timeout=30)
                resp.raise_for_status()
                return resp.json()
            except requests.RequestException as e:
                raise APIRequestError(str(e))
        from quantum_circuit_simulator import QuantumCircuitSimulator
        sim = QuantumCircuitSimulator()
        result = sim.run(params)
        return {
            "results": getattr(result, "observables", result),
            "simulator_name": "quantum_circuit",
        }

    def simulate_schrodinger(self, params: Dict[str, Any]) -> Dict[str, Any]:
        if self._active_mode == "http":
            try:
                resp = requests.post(f"{self.api_url}/simulate/schrodinger", json=params, headers=self.headers, timeout=60)
                resp.raise_for_status()
                return resp.json()
            except requests.RequestException as e:
                raise APIRequestError(str(e))
        from schrodinger_solver import SchrodingerSolver
        sim = SchrodingerSolver()
        result = sim.run(params)
        return {
            "results": getattr(result, "observables", result),
            "simulator_name": "schrodinger",
        }

    def simulate_harmonic_oscillator(self, params: Dict[str, Any]) -> Dict[str, Any]:
        if self._active_mode == "http":
            try:
                resp = requests.post(f"{self.api_url}/simulate/harmonic_oscillator", json=params, headers=self.headers, timeout=30)
                resp.raise_for_status()
                return resp.json()
            except requests.RequestException as e:
                raise APIRequestError(str(e))
        from harmonic_oscillator import HarmonicOscillator
        sim = HarmonicOscillator()
        result = sim.run(params)
        return {
            "results": getattr(result, "observables", result),
            "simulator_name": "harmonic_oscillator",
        }


