# flake8: noqa T001, E501
# -*- coding: utf-8 -*-
"""
PyBaMM Integration Example

Demonstrates battery modeling with PyBaMM (Python Battery Mathematical Modeling).
Creates a battery simulation and displays the results in BT-Analysis.

Requirements:
    - pybamm (installs automatically on first import)
    - numpy (installs automatically on first import)

Features:
    - Thevenin equivalent circuit model
    - Custom charge/discharge experiment
    - Voltage and current visualization
    - Efficient data handling with Series
    - CSV export of results
"""

import BTAnalysisPyInterface as bta
import numpy as np
import pybamm


def run_pybamm_simulation():
    """Run PyBaMM battery simulation"""
    bta.print("Running PyBaMM simulation...")

    # Configure Thevenin model
    model = pybamm.equivalent_circuit.Thevenin()

    parameter_values = pybamm.ParameterValues("ECM_Example")
    parameter_values.update(
        {
            "Cell capacity [A.h]": 0.1,
            "Nominal cell capacity [A.h]": 0.1,
        }
    )

    # Define experiment
    experiment = [
        "Charge at 33mA for 24 hours or until 4V",
        (
            "Rest for 5 seconds",
            "Discharge at 20mA for 10 seconds or until 3.5V",
            "Rest for 10 seconds",
            "Charge at 20mA for 10 seconds or until 4V",
            "Rest for 10 seconds",
            "Discharge at 33mA for 1 hours or until 3.5V",
        )
        * 10,
    ]

    # Run simulation
    simulation = pybamm.Simulation(
        model, experiment=experiment, parameter_values=parameter_values
    )

    solution = simulation.solve()
    bta.print("PyBaMM simulation complete!")

    return solution


def main():
    # Run simulation
    simulation_solution = run_pybamm_simulation()

    # Create graph
    graph = bta.add_graph(
        title="PyBaMM Simulation: Voltage & Current",
        xAxis="Time",
        y1Axis="Voltage",
        y2Axis="Current",
    )

    # Extract arrays from PyBaMM
    time = simulation_solution.t
    voltage = simulation_solution["Voltage [V]"](simulation_solution.t)
    current = simulation_solution["Current [A]"](simulation_solution.t)

    # Add data to graph efficiently using Series
    time_series = bta.Series("Time")
    time_series.add_values(time.tolist())

    voltage_series = bta.Series("Voltage")
    voltage_series.add_values(voltage.tolist())

    current_series = bta.Series("Current")
    current_series.add_values(current.tolist())

    graph.add_series("Voltage", time_series, voltage_series, bta.GraphAxis.Y1)
    graph.add_series("Current", time_series, current_series, bta.GraphAxis.Y2)

    # Style the traces
    graph.configure_trace("Voltage", color="#2E86AB", line_width=2.0)

    graph.configure_trace("Current", color="#C73E1D", line_width=2.0)

    bta.print(f"Plotted {len(time)} points")

    # ========== EXPORT RESULTS TO TABLE AND CSV ==========
    results_table = bta.add_table(
        columns=["Time", "Voltage", "Current"], title="PyBaMM Simulation Results"
    )

    # Populate table
    for t, v, i in zip(time, voltage, current):
        results_table.add_row(
            {"Time": float(t), "Voltage": float(v), "Current": float(i)}
        )

    # Export to CSV
    results_table.export_to_csv("pybamm_simulation.csv")
    bta.print("Simulation data exported to pybamm_simulation.csv")
