from pathlib import Path
from typing import cast

import numpy as np

from rock_physics_open.equinor_utilities import gen_utilities
from rock_physics_open.equinor_utilities.optimisation_utilities import (
    gen_mod_routine,
    gen_sub_routine,
    load_opt_params,
)
from rock_physics_open.equinor_utilities.various_utilities.types import Array1D

from .curvefit_t_matrix_min import curve_fit_2_inclusion_sets


def run_t_matrix_with_opt_params_petec(
    min_k: Array1D[np.float64],
    min_mu: Array1D[np.float64],
    min_rho: Array1D[np.float64],
    fl_k_orig: Array1D[np.float64],
    fl_rho_orig: Array1D[np.float64],
    fl_k_sub: Array1D[np.float64],
    fl_rho_sub: Array1D[np.float64],
    vp: Array1D[np.float64],
    vs: Array1D[np.float64],
    rhob: Array1D[np.float64],
    phi: Array1D[np.float64],
    angle: float,
    perm: float,
    visco: float,
    tau: float,
    freq: float,
    f_name: str | Path,
    fluid_sub: bool = True,
) -> tuple[
    Array1D[np.float64],
    Array1D[np.float64],
    Array1D[np.float64],
    Array1D[np.float64],
    Array1D[np.float64],
    Array1D[np.float64],
    Array1D[np.float64],
    Array1D[np.float64],
    Array1D[np.float64],
    Array1D[np.float64],
    Array1D[np.float64],
]:
    """
        Based on the input file with parameters for the optimally fitted model, the correct modelling version is run.
        Fluid substitution follows, in case it is selected. If not, the vp_sub and vs_sub will contain the same values as
        the input logs.
    def run_t_matrix_with_opt_params_exp(fl_k_orig, fl_rho_orig, fl_k_sub, fl_rho_sub,
                                     vp, vs, rhob,
                                     phi, vsh,
                                     angle, perm, visco, tau, freq,
                                     f_name, fluid_sub=True):
        Parameters
        ----------
        min_k : np.ndarray
            Effective mineral bulk modulus [Pa].
        min_mu : np.ndarray
            Effective mineral shear modulus [Pa].
        min_rho : np.ndarray
            Effective mineral density [kg/m^3].
        fl_k_orig : np.ndarray
            Effective in situ fluid bulk modulus [Pa].
        fl_rho_orig : np.ndarray
            Effective in situ fluid density [kg/m^3].
        fl_k_sub : np.ndarray
            Effective substituted fluid bulk modulus [Pa].
        fl_rho_sub : Effective substituted density [kg/m^3].
        vp : np.ndarray
            Compressional velocity [m/s].
        vs : np.ndarray
            Shear velocity [m/s].
        rhob : np.ndarray
            Bulk density [kg/m^3].
        phi : np.ndarray
            Porosity [fraction].
        angle : float
            Angle of symmetry plane [degrees]
        perm : float
            Permeability [mD].
        visco : float
            Viscosity [cP].
        tau : float
            Relaxation time constant [s].
        freq : float
            Signal frequency [Hz].
        f_name : str
            File name for parameter file for optimal parameters.
        fluid_sub : bool
            Boolean parameter to perform fluid substitution.

        Returns
        -------
        tuple
            Tuple of np.ndarrays: vp and vs for pressure substituted case, vp, vs and density for fluid substituted case, vp and vs for
            optimal fitted model, vp and vs residuals (observed logs minus modelled values).
    """
    _, opt_params, _ = load_opt_params(f_name)
    y_data = np.stack([vp, vs], axis=1)
    y_shape = y_data.shape
    phi, angle_, perm_, visco_, tau_, freq_, def_vpvs = cast(
        list[Array1D[np.float64]],
        gen_utilities.dim_check_vector((phi, angle, perm, visco, tau, freq, 1.0)),
    )

    rho_sub = rhob + (fl_rho_sub - fl_rho_orig) * phi
    # Set None values for inputs that will be defined in the different cases
    x_data_new = None

    opt_fcn = curve_fit_2_inclusion_sets
    # Generate x_data according to method min
    x_data = np.stack(
        (
            phi,
            min_k,
            min_mu,
            min_rho,
            fl_k_orig,
            fl_rho_orig,
            angle_,
            perm_,
            visco_,
            tau_,
            freq_,
            def_vpvs,
        ),
        axis=1,
    )
    if fluid_sub:
        x_data_new = np.stack(
            (
                phi,
                min_k,
                min_mu,
                min_rho,
                fl_k_sub,
                fl_rho_sub,
                angle_,
                perm_,
                visco_,
                tau_,
                freq_,
                def_vpvs,
            ),
            axis=1,
        )

        v_sub, v_mod, v_res = gen_sub_routine(
            opt_function=opt_fcn,
            xdata_orig=x_data,
            xdata_new=x_data_new,
            ydata=y_data,
            opt_params=opt_params,
        )
        vp_sub, vs_sub = [arr.flatten() for arr in np.split(v_sub, 2, axis=1)]
        vp_mod, vs_mod = [arr.flatten() for arr in np.split(v_mod, 2, axis=1)]
        vp_res, vs_res = [arr.flatten() for arr in np.split(v_res, 2, axis=1)]
    else:
        v_mod = gen_mod_routine(opt_fcn, x_data, y_shape, opt_params)
        vp_mod, vs_mod = [arr.flatten() for arr in np.split(v_mod, 2, axis=1)]
        vp_sub = vp
        vs_sub = vs
        vp_res = vp_mod - vp
        vs_res = vs_mod - vs

    rho_mod = min_rho * (1.0 - phi) + fl_rho_orig * phi
    rho_res = rho_mod - rhob
    ai_sub = vp_sub * rho_sub
    vpvs_sub = vp_sub / vs_sub

    return (
        vp_sub,
        vs_sub,
        rho_sub,
        ai_sub,
        vpvs_sub,
        vp_mod,
        vs_mod,
        rho_mod,
        vp_res,
        vs_res,
        rho_res,
    )
