use pyo3::prelude::*;
use numpy::{PyReadonlyArray1, PyReadonlyArray2, PyArray1, PyUntypedArrayMethods};
use csrk::GP as InnerGP;

/** Hybrid GP struct / class definition (Rust + Python)

Method names should ideally be the same as the inner GP methods

# Data
* `InnerGP` - The pure rust GP object from csrk

# Methods
* `new` - Constructor for Hybrid/Inner GP
* `predict_mean` - Predict GP interpolation mean on sample array
* `predict_mean_sgl` - Predict GP interpolation mean on single sample
* `predict_var` - Predict GP interpolation variance on sample array
* `predict_var_sgl` - Predict GP interpolation variance on single sample
* `predict_cov_sgl` - Predict GP covariance between two sample points
* `retrain` - Retrain GP with new hyperparameters
* `log_marginal_likelihood` - Evaluate the log marginal likelihood of the model
* `ndim` - Return the number of dimensions for the GP model
* `whitenoise` - Return the whitenoise hyperparameter for the GP model
* `scale` - Return the scale hyperparameters for the GP model
* `to_hdf5` - Save the GP to HDF5
* `from_hdf5` - Load a GP from HDF5
* `generate_ldl` - Generate the LDL decomposition for a loaded model

*/
#[pyclass]
pub struct HybridGP {
    inner: InnerGP,
}

#[pymethods]
impl HybridGP {
    /**
    # Arguments
    * `training_points` - a 2D NumPy array of shape (n_points, n_dimensions)
    * `training_values` - a 1D NumPy array of shape (n_points,)
    * `training_error` - a 1D NumPy array of shape (n_points,)
    * `scale` - a 1D NumPy array of shape (n_dimensions,) (sparsity hyperparameter)
    * `whitenoise` - a floating point whitenoise hyperparameter
    * `msq_order` - i32: the mean square differentiability order hyperparameter
    
    # Returns
    * A Hybrid GP struct / object
    */
    #[new]
    pub fn new (
        training_points: PyReadonlyArray2<f64>,
        training_values: PyReadonlyArray1<f64>,
        training_error: PyReadonlyArray1<f64>,
        scale: PyReadonlyArray1<f64>,
        whitenoise: f64,
        msq_order: i32,
    ) -> PyResult<Self> {
        // Convert inputs to array view
        let training_points = training_points.as_array();
        let training_values = training_values.as_array();
        let training_error = training_error.as_array();
        let scale = scale.as_array();
        // Check ndim and npts
        let shape = training_points.shape();
        let npts = shape[0];
        let ndim = shape[1];
        if shape.len() != 2 {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                "training_points must be a 2D array",
            ));
        }

        if training_values.len() != npts {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                "training_points and training_values must have matching dimensions!",
            ));
        }
        if training_error.len() != npts {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                "training_points and training_error must have matching dimensions!",
            ));
        }
        if scale.len() != ndim {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                "scale must have length equal to the number of dimensions"
            ));
        }
        // This is a zero-copy: ndarray -> ndarray::ArrayView
        let gp = InnerGP::train(
            training_points,
            training_values,
            training_error,
            scale,
            whitenoise,
            msq_order,
        );

        // Return new HybridGP object
        Ok(Self{ inner: gp })
    }
    /** Predict the mean interpolation on some evaluation samples

    # Arguments
    * `eval_points` - A 2D NumPy array of shape (n_points, n_dimensions)

    # Returns
    A 1D NumPy array of shape (n_points,) containing the predicted mean values
    */
    #[pyo3(text_signature = "($self, eval_points)")]
    pub fn predict_mean<'py> (
        &self,
        py: Python<'py>, // GIL token
        eval_points: PyReadonlyArray2<f64> // Zero-copy input
    ) -> PyResult<Bound<'py, PyArray1<f64>>> { // Owned output
        // Establish Array-view interpretation of NumPy array
        let eval_points_view = eval_points.as_array();
        // Check dimensions
        // Check ndim and npts
        let shape = eval_points.shape();
        if shape.len() != 2 {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                "eval_points must be a 2D array",
            ));
        }
        let ndim = shape[1];
        if eval_points_view.ncols() != ndim {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                format!(
                    "eval_points represents {} dimensions but GP was trained with {}",
                    eval_points_view.ncols(),
                    ndim
                )
            ));
        }
        // Call InnerGP predict_mean method
        let means = self.inner.predict_mean(eval_points_view);
        // Generate NumPy representation
        let numpy_means = PyArray1::from_vec(py, means.to_vec());
        // Return NumPy array
        Ok(numpy_means)
    }
    /** Wrapper for predict mean
    */
    #[pyo3(signature = (eval_points))]
    pub fn __call__<'py> (
        &self,
        py: Python<'py>, // GIL token
        eval_points: PyReadonlyArray2<f64> // Zero-copy input
    ) -> PyResult<Bound<'py, PyArray1<f64>>> { // Owned output
        self.predict_mean(py, eval_points)
    }
    /** Predict the mean interpolation on a single evaluation sample

    # Arguments
    * `eval_point` - A 1D NumPy array of shape (n_dimensions)

    # Returns
    A double precision float for the evaluated mean value
    */
    #[pyo3(text_signature = "($self, eval_point)")]
    pub fn predict_mean_sgl(
        &self,
        eval_point: PyReadonlyArray1<f64>,
    ) -> PyResult<f64> {
        // Establish Array-view interpretation of NumPy array
        let eval_point_view = eval_point.as_array();
        // Validate dimensions
        let shape = eval_point.shape();
        if shape.len() != 1 {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                "eval_point must be a 1D array (did you mean predict_mean?)",
            ));
        }
        let ndim = shape[0];
        if eval_point_view.len() != ndim {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                format!(
                    "eval_point represents {} dimensions but GP was trained with {}",
                    eval_point_view.len(),
                    ndim,
                )
            ));
        }
        // Call the InnerGP predict_mean_sgl method
        Ok(self.inner.predict_mean_sgl(eval_point_view))
    }
    /** Predict the variance at given evaluation points
    
    # Arguments
    * `eval_points` - A 2D NumPy array of shape (n_points, n_dimensions)

    # Returns
    A 1D NumPy array of shape (n_points,) containing the predicted variances
    */
    #[pyo3(text_signature = "($self, eval_points)")]
    pub fn predict_var<'py> (
        &self,
        py: Python<'py>, // GIL token
        eval_points: PyReadonlyArray2<f64> // Zero-copy input
    ) -> PyResult<Bound<'py, PyArray1<f64>>> { // Owned output
        // Establish Array-view interpretation of NumPy array
        let eval_points_view = eval_points.as_array();
        // Check dimensions
        // Check ndim and npts
        let shape = eval_points.shape();
        if shape.len() != 2 {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                "eval_points must be a 2D array",
            ));
        }
        let ndim = shape[1];
        if eval_points_view.ncols() != ndim {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                format!(
                    "eval_points represents {} dimensions but GP was trained with {}",
                    eval_points_view.ncols(),
                    ndim
                )
            ));
        }
        // Call InnerGP predict_var method
        let vars = self.inner.predict_var(eval_points_view);
        // Generate NumPy representation
        let numpy_vars = PyArray1::from_vec(py, vars.to_vec());
        // Return NumPy array
        Ok(numpy_vars)
    }
    /** Predict the variance at a single evaluation point
    
    # Arguments
    * `eval_point` - A 1D NumPy array of shape (n_dimensions,)

    # Returns
    A float containing the predicted variance
    */
    #[pyo3(text_signature = "($self, eval_point)")]
    pub fn predict_var_sgl(
        &self,
        eval_point: PyReadonlyArray1<f64>,
    ) -> PyResult<f64> {
        // Establish Array-view interpretation of NumPy array
        let eval_point_view = eval_point.as_array();
        // Validate dimensions
        let shape = eval_point.shape();
        if shape.len() != 1 {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                "eval_point must be a 1D array (did you mean predict_var?)",
            ));
        }
        let ndim = shape[0];
        if eval_point_view.len() != ndim {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                format!(
                    "eval_point represents {} dimensions but GP was trained with {}",
                    eval_point_view.len(),
                    ndim,
                )
            ));
        }
        // Call the InnerGP predict_var_sgl method
        Ok(self.inner.predict_var_sgl(eval_point_view))
    }
    /** Predict Covariance between two evaluation points

    # Arguments
    * `eval_point_1` - First evaluation point (1D array)
    * `eval_point_2` - Second evaluation point (2D array)

    # Returns
    The covariance between the two points
    */
    #[pyo3(text_signature = "($self, eval_point_1, eval_point_2)")]
    pub fn predict_cov_sgl(
        &self,
        eval_point_1: PyReadonlyArray1<f64>,
        eval_point_2: PyReadonlyArray1<f64>,
    ) -> PyResult<f64> {
        // Establish Array-view interpretation of NumPy array
        let eval_point_1_view = eval_point_1.as_array();
        let eval_point_2_view = eval_point_2.as_array();
        // Validate dimensions
        let shape = eval_point_1.shape();
        if shape.len() != 1 {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                "eval_point_1 must be a 1D array",
            ));
        }
        let ndim = shape[0];
        if eval_point_1_view.len() != ndim {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                format!(
                    "eval_point_1 represents {} dimensions but GP was trained with {}",
                    eval_point_1_view.len(),
                    ndim,
                )
            ));
        }
        let shape = eval_point_2.shape();
        if shape.len() != 1 {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                "eval_point_2 must be a 1D array",
            ));
        }
        let ndim = shape[0];
        if eval_point_2_view.len() != ndim {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                format!(
                    "eval_point_2 represents {} dimensions but GP was trained with {}",
                    eval_point_2_view.len(),
                    ndim,
                )
            ));
        }
        // Call the InnerGP predict_cov_sgl method
        Ok(self.inner.predict_cov_sgl(eval_point_1_view, eval_point_2_view))
    }
    /** Retrain the GP with new hyperparameters

    # Arguments
    * `scale` - New length scale parameters (1D array of shape (n_dimensions,))
    * `whitenoise` - New white noise parameter
    * `msq_order` - New Wendland Kernel order

    # Note
    This retrains the GP using the same training data but with new hyperparameters
    */
    #[pyo3(text_signature = "($self, scale, whitenoise, msq_order)")]
    pub fn retrain(
        &mut self,
        scale: PyReadonlyArray1<f64>,
        whitenoise: f64,
        msq_order: i32,
    ) -> PyResult<()> {
        // Convert scale to array view
        let scale_view = scale.as_array();
        // Validate dimensions
        if scale_view.len() != self.inner.ndim {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                format!(
                    "scale has {} dimensions but GP was trained with {}",
                    scale_view.len(),
                    self.inner.ndim,
                )
            ));
        }
        self.inner.retrain(scale_view, whitenoise, msq_order);
        Ok(())
    }
    /** Get the log marginal likelihood of the current model

    # Returns
    The log marginal likelihood (f64)

    # Note
    Higher values indicate better fit. Useful for hyperparameter optimization
    */
    #[pyo3(text_signature = "($self)")]
    pub fn log_marginal_likelihood(
        &mut self,
    ) -> f64 {
        self.inner.log_marginal_likelihood()
    }
    /// Get the number of dimensions
    #[getter]
    pub fn ndim(&self) -> usize {
        self.inner.ndim
    }
    /// Get the whitenoise parameter
    #[getter]
    pub fn whitenoise(&self) -> f64 {
        self.inner.kernel.whitenoise()
    }
    /// Get the scale parameter
    #[getter]
    pub fn scale<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1<f64>> {
        PyArray1::from_vec(py, self.inner.kernel.scale().to_vec())
    }

    /// String representation
    fn __repr__(&mut self) -> String {
        #[allow(non_snake_case)]
        let lnL = self.inner.log_marginal_likelihood();
        format!(
            "HybridGP(ndim={}, whitenoise={:.2e}, lnL={:.4e})",
            self.inner.ndim,
            self.inner.kernel.whitenoise(),
            lnL,
        )
    }

    /** Save the GP model to an HDF5 file
    
    # Arguments
    * `filename` - Path to the HDF5 file (will be created if it doesn't exist)
    * `group_name` - Name of the HDF5 group to store the model in

    # Example
    ```
    gp = HybridGP(x_train, y_train, y_err, scale, whitenoise, msq_order)
    gp.to_hdf5("model.hdf5", "my_gp")
    ```

    # Note
    Will raise an error if the group already exists
    */
    #[pyo3(text_signature = "($self, filename, group_name)")]
    pub fn to_hdf5(
        &self,
        filename: &str,
        group_name: &str,
    ) -> PyResult<()> {
        self.inner.save_to_hdf5(filename, group_name)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(
                format!("Failed to save GP to HDF5: {}", e)
            ))
    }
    /** Load a GP model from an HDF5 file

    # Arguments
    * `filename` - Path to the HDF5 file
    * `group_name` - Name of the HDF5 group containing the model

    Returns
    A new HybridGP instance loaded from the file

    # Example
    ```
    gp = HybridGP.from_hdf5("model.hdf5", "my_gp")
    gp.generate_ldl()
    ```

    # Note
    After loading, you MUST call `generate_ldl()` before using methods
    which require access to the LDL decomposition.
    The LDL decomposition is not serialized and must be regenerated.
    For the time being though, none of the included methods in the Python
    interface require the LDL composition to be in scope.
    */
    #[staticmethod]
    #[pyo3(text_signature = "(filename, group_name)")]
    pub fn from_hdf5(
        filename: &str,
        group_name: &str,
    ) -> PyResult<Self> {
        let inner = InnerGP::load_from_hdf5(filename, group_name)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(
                format!("Failed to load GP from HDF5: {}", e)
            ))?;
        Ok(Self {inner})
    }

    /** Regenerate the LDL decomposition after loading from an HDF5

    # Note
    This must be called after loading a GP from HDF5 before using methods
        which require access to the LDL decomposition.
    The LDLD decomposition is not saved to the disk and must be recomputed.

    # Example
    ```
    gp = HybridGP.from_hdf5("model.hdf5", "my_gp")
    gp.generate_ldl()
    ```

    */
    #[pyo3(text_signature = "($self)")]
    pub fn generate_ldl(&mut self) {
        self.inner.generate_ldl();
    }
}
