{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generalized Multi-Step Charge Rate Optimization\n",
    "\n",
    "This example demonstrates how to optimize the charge rate in battery experiments using PyBaMM's new InputParameter functionality while maximizing the negative electrode potential constraint. This approach allows dynamic adjustment of charging currents during optimization.\n",
    "\n",
    "Key features:\n",
    "- Uses PyBaMM InputParameters for dynamic charge current control\n",
    "- Optimizes charge rate to maximize charging speed\n",
    "- Ensures negative electrode potential stays within safe operating limits\n",
    "- Demonstrates advanced experiment design with parameterized protocols\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ionworkspipeline as iwp\n",
    "from ionworkspipeline.data_fits.models.metrics import Minimum, Time\n",
    "from ionworkspipeline.data_fits.models.metrics.actions import GreaterThan, Minimize\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pybamm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup Model and Base Parameters\n",
    "\n",
    "We'll use a DFN model with the Chen2020 parameter set as our baseline configuration.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup base model and parameters\n",
    "model = pybamm.lithium_ion.DFN()\n",
    "params_baseline = pybamm.ParameterValues(\"Chen2020\")\n",
    "\n",
    "print(\"Base model and parameter set configured\")\n",
    "print(f\"Model: {model.name}\")\n",
    "print(\"Parameter set: Chen2020\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Experiment with InputParameters\n",
    "\n",
    "We'll create an experiment protocol that uses PyBaMM's InputParameter functionality to make the charge current optimizable. This allows us to dynamically adjust the charging rate during optimization.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Number of charge steps with corresponding potential transitions\n",
    "num_charge_steps = 6\n",
    "potentials = [3.6, 3.7, 3.9, 4.0, 4.1, 4.2]\n",
    "\n",
    "# Define optimization parameters for charge current\n",
    "parameters = {}\n",
    "for i in range(num_charge_steps):\n",
    "    parameters[f\"I_charge_{i}\"] = iwp.Parameter(\n",
    "        f\"I_charge_{i}\", initial_value=-4.0, bounds=(-10.0, -0.1)\n",
    "    )\n",
    "\n",
    "# Define experiment steps with parameterized charge current\n",
    "charge_steps = []\n",
    "i = 0\n",
    "for _, value in parameters.items():\n",
    "    charge_steps.append(pybamm.step.current(value, termination=f\"> {potentials[i]} V\"))\n",
    "    i += 1\n",
    "\n",
    "experiment_steps = [\n",
    "    pybamm.step.current(5.0, termination=\"< 3.0 V\"),  # Discharge at 5A (fixed)\n",
    "    pybamm.step.rest(duration=\"10 minutes\"),  # Rest period\n",
    "    *charge_steps,\n",
    "    pybamm.step.voltage(4.20, termination=\"< C/20\"),  # CV hold\n",
    "]\n",
    "\n",
    "# Create experiment with the parameterized steps\n",
    "experiment = pybamm.Experiment(experiment_steps, period=\"10 seconds\")\n",
    "\n",
    "print(\"Experiment protocol with InputParameter defined:\")\n",
    "print(experiment_steps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Optimization Metrics\n",
    "\n",
    "We'll create metrics that focus on:\n",
    "1. Minimizing charge time (faster charging)\n",
    "2. Maximizing negative electrode potential (safety constraint)\n",
    "3. Monitoring overall charging performance\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create optimization actions\n",
    "actions = {\n",
    "    # Minimize charge time (maximize charge rate)\n",
    "    \"Final time [s]\": Minimize(Time(variable=\"Time [s]\", value=-1), weight=0.01),\n",
    "}\n",
    "\n",
    "# Constraints define bounds that must be satisfied (using GreaterThan/LessThan actions)\n",
    "constraints = {\n",
    "    # Penalize negative electrode potential below 0 V (safety constraint)\n",
    "    \"Anode potential constraint violation [V]\": GreaterThan(\n",
    "        Minimum(\n",
    "            \"Negative electrode surface potential difference at separator interface [V]\"\n",
    "        ),\n",
    "        value=0.0,\n",
    "    ),\n",
    "}\n",
    "\n",
    "print(\n",
    "    f\"Defined {len(actions)} optimization metrics and {len(constraints)} constraints:\"\n",
    ")\n",
    "print(\"  1. Charge Time (minimize) - Faster charging\")\n",
    "print(\"  2. Min Negative Electrode Potential constraint - Safety constraint\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up Custom Cost Function\n",
    "\n",
    "We'll create a weighted cost function that balances fast charging with safety constraints.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Custom cost function created with variable weights\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run Charge Rate Optimization\n",
    "\n",
    "Now we'll run the optimization to find the optimal charge current that maximizes charging speed while maintaining safe electrode potentials.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up design objective with InputParameter support\n",
    "objective_options = {\n",
    "    \"model\": model,\n",
    "    \"simulation_kwargs\": {\n",
    "        \"experiment\": experiment,\n",
    "    },\n",
    "}\n",
    "\n",
    "objective = iwp.objectives.DesignObjective(\n",
    "    options=objective_options,\n",
    "    actions=actions,\n",
    "    constraints=constraints,\n",
    ")\n",
    "\n",
    "# Create and run optimization\n",
    "optim = iwp.optimizers.Pints(\n",
    "    max_iterations=10,  # Increase for better results (~200)\n",
    "    max_unchanged_iterations=80,\n",
    "    max_unchanged_iterations_threshold=1e-4,\n",
    "    log_to_screen=True,\n",
    "    sigma0=0.2,\n",
    "    # population_size=30, # Enable during parallel optimisation\n",
    ")\n",
    "design_datafit = iwp.DataFit(\n",
    "    objective,\n",
    "    parameters=parameters,\n",
    "    options={\"seed\": 42},\n",
    "    optimizer=optim,\n",
    "    # parallel=True,\n",
    ")\n",
    "\n",
    "print(\"Running charge rate optimization...\")\n",
    "print(\"This may take a few minutes as we optimize for safe fast charging...\")\n",
    "\n",
    "results = design_datafit.run(params_baseline)\n",
    "\n",
    "print(\"\\nCharge rate optimization completed!\")\n",
    "print(f\"Final cost: {results.costs:.6f}\")\n",
    "print(\"\\nOptimal parameters:\")\n",
    "for name in parameters:\n",
    "    value = results.parameter_values[name]\n",
    "    print(f\"  Optimal charge current: {value:.3f}A\")\n",
    "\n",
    "optimal_charge_current = results.parameter_values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validate and Compare Results\n",
    "\n",
    "Let's compare the optimized charge rate with a baseline charge rate to see the improvements.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper function to extract charging metrics\n",
    "def extract_charge_metrics(solution, charge_step_index=2):\n",
    "    \"\"\"Extract key charging performance metrics from simulation solution.\"\"\"\n",
    "    charge_sol = solution.sub_solutions[charge_step_index:]\n",
    "\n",
    "    # Extract charging metrics\n",
    "    charge_time = (charge_sol[-1].t[-1] - charge_sol[0].t[0]) / 3600  # Convert to hours\n",
    "\n",
    "    # Negative electrode potential analysis\n",
    "    anode_potential = []\n",
    "    for sol in charge_sol:\n",
    "        anode_potential.extend(\n",
    "            sol[\n",
    "                \"Negative electrode surface potential difference at separator interface [V]\"\n",
    "            ]()\n",
    "        )\n",
    "    min_anode_potential = min(anode_potential)\n",
    "    mean_anode_potential = np.asarray(anode_potential).mean()\n",
    "\n",
    "    # Voltage analysis\n",
    "    voltage = charge_sol[-1][\"Voltage [V]\"]()\n",
    "    final_voltage = voltage[-1]\n",
    "\n",
    "    return {\n",
    "        \"charge_time\": charge_time,\n",
    "        \"min_anode_potential\": min_anode_potential,\n",
    "        \"mean_anode_potential\": mean_anode_potential,\n",
    "        \"final_voltage\": final_voltage,\n",
    "        \"solution\": solution,\n",
    "        \"charge_sol\": charge_sol,\n",
    "    }\n",
    "\n",
    "\n",
    "# Run simulations for comparison\n",
    "input_names = list(parameters.keys())\n",
    "baseline_current = -4.0  # Baseline charge current\n",
    "baseline_inputs = {name: baseline_current for name in input_names}\n",
    "optimal_inputs = {}\n",
    "for name in input_names:\n",
    "    optimal_inputs[name] = results.parameter_values[name]\n",
    "\n",
    "print(\"Running comparison simulations...\")\n",
    "print(f\"Baseline charge current: {baseline_current}A\")\n",
    "\n",
    "# Baseline simulation\n",
    "sim_baseline = iwp.Simulation(\n",
    "    model=model, parameter_values=params_baseline, experiment=experiment\n",
    ")\n",
    "sol_baseline = sim_baseline.solve(\n",
    "    initial_soc=1.0,\n",
    "    inputs=baseline_inputs,\n",
    ")\n",
    "\n",
    "# Optimal simulation\n",
    "sim_optimal = iwp.Simulation(\n",
    "    model=model, parameter_values=params_baseline, experiment=experiment\n",
    ")\n",
    "sol_optimal = sim_optimal.solve(\n",
    "    initial_soc=1.0,\n",
    "    inputs=optimal_inputs,\n",
    ")\n",
    "\n",
    "# Extract metrics\n",
    "baseline_metrics = extract_charge_metrics(sol_baseline)\n",
    "optimal_metrics = extract_charge_metrics(sol_optimal)\n",
    "\n",
    "print(\"\\nCharging Performance Comparison:\")\n",
    "print(\"=\" * 75)\n",
    "print(f\"{'Metric':<35} {'Baseline':<15} {'Optimal':<15} {'Improvement'}\")\n",
    "print(\"-\" * 75)\n",
    "\n",
    "# Calculate improvements\n",
    "time_improvement = (\n",
    "    (baseline_metrics[\"charge_time\"] - optimal_metrics[\"charge_time\"])\n",
    "    / baseline_metrics[\"charge_time\"]\n",
    "    * 100\n",
    ")\n",
    "min_anode_change = (\n",
    "    optimal_metrics[\"min_anode_potential\"] - baseline_metrics[\"min_anode_potential\"]\n",
    ")\n",
    "\n",
    "print(\n",
    "    f\"{'Charge Time [h]':<35} {baseline_metrics['charge_time']:<15.4f} {optimal_metrics['charge_time']:<15.4f} {time_improvement:+.1f}%\"\n",
    ")\n",
    "print(\n",
    "    f\"{'Min Anode Potential [V]':<35} {baseline_metrics['min_anode_potential']:<15.4f} {optimal_metrics['min_anode_potential']:<15.4f} {min_anode_change:+.4f}\"\n",
    ")\n",
    "print(\n",
    "    f\"{'Mean Anode Potential [V]':<35} {baseline_metrics['mean_anode_potential']:<15.4f} {optimal_metrics['mean_anode_potential']:<15.4f} {(optimal_metrics['mean_anode_potential'] - baseline_metrics['mean_anode_potential']):+.4f}\"\n",
    ")\n",
    "print(\"=\" * 75)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameter Analysis and Visualization\n",
    "\n",
    "Let's analyze the optimized parameters and create comprehensive visualizations of the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract parameter comparison data\n",
    "print(\"Parameter Optimization Results:\")\n",
    "print(\"=\" * 60)\n",
    "for name, param in parameters.items():\n",
    "    baseline_val = param.initial_value\n",
    "    optimal_val = results.parameter_values[name]\n",
    "    change = optimal_val - baseline_val\n",
    "    percent_change = (change / baseline_val) * 100\n",
    "\n",
    "    short_name = name.replace(\"I_charge_\", \"Charge Current \")\n",
    "    print(f\"\\n{short_name}:\")\n",
    "    print(f\"  Baseline:       {baseline_val:.4f}A\")\n",
    "    print(f\"  Optimal:        {optimal_val:.4f}A\")\n",
    "    print(f\"  Change:         {change:+.4f}A ({percent_change:+.2f}%)\")\n",
    "\n",
    "print(\"\\n\" + \"=\" * 60)\n",
    "\n",
    "# Create comprehensive visualization\n",
    "fig, axes = plt.subplots(2, 2, figsize=(12.5, 10))\n",
    "\n",
    "# Plot 1: Voltage comparison\n",
    "t_baseline = sol_baseline.t / 3600\n",
    "t_optimal = sol_optimal.t / 3600\n",
    "\n",
    "axes[0, 0].plot(\n",
    "    t_baseline,\n",
    "    sol_baseline[\"Voltage [V]\"](sol_baseline.t),\n",
    "    label=\"Baseline\",\n",
    "    linewidth=2,\n",
    "    alpha=0.8,\n",
    ")\n",
    "axes[0, 0].plot(\n",
    "    t_optimal,\n",
    "    sol_optimal[\"Voltage [V]\"](sol_optimal.t),\n",
    "    label=\"Optimal\",\n",
    "    linewidth=2,\n",
    "    alpha=0.8,\n",
    "    linestyle=\"--\",\n",
    ")\n",
    "axes[0, 0].set_xlabel(\"Time [h]\")\n",
    "axes[0, 0].set_ylabel(\"Voltage [V]\")\n",
    "axes[0, 0].set_title(\"Cell Voltage: Baseline vs Optimal\")\n",
    "axes[0, 0].legend()\n",
    "axes[0, 0].grid(True, alpha=0.3)\n",
    "\n",
    "# Plot 2: Current comparison\n",
    "axes[0, 1].plot(\n",
    "    t_baseline,\n",
    "    sol_baseline[\"Current [A]\"](sol_baseline.t),\n",
    "    label=\"Baseline\",\n",
    "    linewidth=2,\n",
    "    alpha=0.8,\n",
    ")\n",
    "axes[0, 1].plot(\n",
    "    t_optimal,\n",
    "    sol_optimal[\"Current [A]\"](sol_optimal.t),\n",
    "    label=\"Optimal\",\n",
    "    linewidth=2,\n",
    "    alpha=0.8,\n",
    "    linestyle=\"--\",\n",
    ")\n",
    "axes[0, 1].set_xlabel(\"Time [h]\")\n",
    "axes[0, 1].set_ylabel(\"Current [A]\")\n",
    "axes[0, 1].set_title(\"Current Profile: Baseline vs Optimal\")\n",
    "axes[0, 1].legend()\n",
    "axes[0, 1].grid(True, alpha=0.3)\n",
    "\n",
    "# Plot 3: Anode potential during charge\n",
    "# Combine charge steps for both solutions\n",
    "charge_steps_baseline = sol_baseline.sub_solutions[2:-1]\n",
    "charge_steps_optimal = sol_optimal.sub_solutions[2:-1]\n",
    "\n",
    "# For baseline\n",
    "charge_sol_baseline = charge_steps_baseline[0]\n",
    "for step in charge_steps_baseline[1:]:\n",
    "    charge_sol_baseline = charge_sol_baseline + step\n",
    "\n",
    "# For optimal\n",
    "charge_sol_optimal = charge_steps_optimal[0]\n",
    "for step in charge_steps_optimal[1:]:\n",
    "    charge_sol_optimal = charge_sol_optimal + step\n",
    "\n",
    "t_charge_baseline = charge_sol_baseline.t / 3600 - charge_sol_baseline.t[0] / 3600\n",
    "t_charge_optimal = charge_sol_optimal.t / 3600 - charge_sol_optimal.t[0] / 3600\n",
    "\n",
    "anode_baseline = charge_sol_baseline[\n",
    "    \"Negative electrode surface potential difference at separator interface [V]\"\n",
    "](charge_sol_baseline.t)\n",
    "anode_optimal = charge_sol_optimal[\n",
    "    \"Negative electrode surface potential difference at separator interface [V]\"\n",
    "](charge_sol_optimal.t)\n",
    "\n",
    "axes[1, 0].plot(\n",
    "    t_charge_baseline, anode_baseline, label=\"Baseline\", linewidth=2, alpha=0.8\n",
    ")\n",
    "axes[1, 0].plot(\n",
    "    t_charge_optimal,\n",
    "    anode_optimal,\n",
    "    label=\"Optimal\",\n",
    "    linewidth=2,\n",
    "    alpha=0.8,\n",
    "    linestyle=\"--\",\n",
    ")\n",
    "axes[1, 0].axhline(y=0.0, color=\"r\", linestyle=\":\", alpha=0.7, label=\"Safety limit\")\n",
    "axes[1, 0].set_xlabel(\"Charge Time [h]\")\n",
    "axes[1, 0].set_ylabel(\"Electrode - Electrolyte Potential [V]\")\n",
    "axes[1, 0].set_title(\"Negative Electrode - Electrolyte Potential During Charge\")\n",
    "axes[1, 0].legend()\n",
    "axes[1, 0].grid(True, alpha=0.3)\n",
    "\n",
    "# Plot 4: Performance metrics comparison\n",
    "metrics_names = [\n",
    "    \"Charge Time [h]\",\n",
    "    \"Min Anode\\nPotential [V]\",\n",
    "    \"Mean Anode\\nPotential [V]\",\n",
    "]\n",
    "baseline_values = [\n",
    "    baseline_metrics[\"charge_time\"],\n",
    "    baseline_metrics[\"min_anode_potential\"],\n",
    "    baseline_metrics[\"mean_anode_potential\"],\n",
    "]\n",
    "optimal_values = [\n",
    "    optimal_metrics[\"charge_time\"],\n",
    "    optimal_metrics[\"min_anode_potential\"],\n",
    "    optimal_metrics[\"mean_anode_potential\"],\n",
    "]\n",
    "\n",
    "x = np.arange(len(metrics_names))\n",
    "width = 0.35\n",
    "\n",
    "bars1 = axes[1, 1].bar(\n",
    "    x - width / 2,\n",
    "    baseline_values,\n",
    "    width,\n",
    "    label=\"Baseline\",\n",
    "    alpha=0.8,\n",
    "    color=\"lightblue\",\n",
    ")\n",
    "bars2 = axes[1, 1].bar(\n",
    "    x + width / 2, optimal_values, width, label=\"Optimal\", alpha=0.8, color=\"orange\"\n",
    ")\n",
    "\n",
    "# Add value labels on bars\n",
    "for i, (baseline, optimal) in enumerate(\n",
    "    zip(baseline_values, optimal_values, strict=False)\n",
    "):\n",
    "    y_offset = abs(baseline) * 0.02 if baseline != 0 else 0.005\n",
    "    axes[1, 1].text(\n",
    "        i - width / 2,\n",
    "        baseline + y_offset,\n",
    "        f\"{baseline:.3f}\",\n",
    "        ha=\"center\",\n",
    "        va=\"bottom\",\n",
    "        fontsize=9,\n",
    "    )\n",
    "    axes[1, 1].text(\n",
    "        i + width / 2,\n",
    "        optimal + y_offset,\n",
    "        f\"{optimal:.3f}\",\n",
    "        ha=\"center\",\n",
    "        va=\"bottom\",\n",
    "        fontsize=9,\n",
    "    )\n",
    "\n",
    "axes[1, 1].set_xlabel(\"Performance Metrics\")\n",
    "axes[1, 1].set_ylabel(\"Value\")\n",
    "axes[1, 1].set_ylim(top=3)\n",
    "axes[1, 1].set_title(\"Performance Metrics: Baseline vs Optimal\")\n",
    "axes[1, 1].set_xticks(x)\n",
    "axes[1, 1].set_xticklabels(metrics_names)\n",
    "axes[1, 1].legend()\n",
    "axes[1, 1].grid(True, alpha=0.3)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(\"\\nCharge rate optimization analysis completed!\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}