import numpy as np
import pybamm

# logic copied from test_with_output_variables_and_sensitivities()

input_parameters = {  # Sensitivities dictionary
    "Current function [A]": 0.222,
    "Separator porosity": 0.3,
}

# construct model
model = pybamm.lithium_ion.DFN()
geometry = model.default_geometry
param = model.default_parameter_values
param.update({key: "[input]" for key in input_parameters})
param.process_model(model)
param.process_geometry(geometry)
var_pts = {"x_n": 50, "x_s": 50, "x_p": 50, "r_n": 5, "r_p": 5}
mesh = pybamm.Mesh(geometry, model.default_submesh_types, var_pts)
disc = pybamm.Discretisation(mesh, model.default_spatial_methods)
disc.process_model(model)

t_eval = [0, 100]
t_interp = np.linspace(t_eval[0], t_eval[-1], 5)

options = {
    "linear_solver": "SUNLinSol_KLU",
    "jacobian": "sparse",
    "num_threads": 4,
    "max_num_steps": 1000,
}

# Use a selection of variables of different types
output_variables = [
    "Voltage [V]",
    "Time [min]",
    "x [m]",  # 1d
    "Negative particle flux [mol.m-2.s-1]",  # 2d
    "Throughput capacity [A.h]",  # ExplicitTimeIntegral
]

sols = []
for out_vars in [True, False]:
    solver = pybamm.IDAKLUSolver(
        options=options,
        output_variables=output_variables if out_vars else None,
    )

    sol = solver.solve(
        model,
        t_eval,
        inputs=input_parameters,
        calculate_sensitivities=True,
        t_interp=t_interp,
    )
    sols.append(sol)

for varname in output_variables:
    try:
        np.testing.assert_allclose(
            sols[0][varname].sensitivities["Current function [A]"],
            sols[1][varname].sensitivities["Current function [A]"].flatten(),
            err_msg=f"Failed for {varname} sensitivities, variable type: {type(sols[1][varname])}",
        )
    except AssertionError as e:
        print(e)