from progpy.models import BatteryCircuit

batt = BatteryCircuit()

print("inputs:", batt.inputs)
print("outputs:", batt.outputs)

print("event(s): ", batt.events)
print("states: ", batt.states)

from pprint import pprint

print("Model configuration:")
pprint(batt.parameters)

config = {"save_freq": 100, "dt": 2, "t0": 700}


def future_loading(t, x=None):
    if t < 600:
        i = 2
    elif t < 900:
        i = 1
    elif t < 1800:
        i = 4
    elif t < 3000:
        i = 2
    else:
        i = 3
    return batt.InputContainer({"i": i})

simulated_results = batt.simulate_to_threshold(future_loading, **config)

fig = simulated_results.inputs.plot(
    xlabel="time (s)", ylabel="current draw (amps)", title="BatteryCircuit Input"
)

fig = simulated_results.outputs.plot(
    keys=["t"],
    xlabel="time (s)",
    ylabel="temperature (K)",
    figsize=(10, 4),
    title="BatteryCircuit Outputs",
)
fig2 = simulated_results.outputs.plot(
    keys=["v"], xlabel="time (s)", ylabel="voltage (V)", figsize=(10, 4)
)

from progpy.models import BatteryElectroChemEOD

batt = BatteryElectroChemEOD()

print("inputs:", batt.inputs)
print("outputs:", batt.outputs)
print("event(s): ", batt.events)
print("states:", batt.states)

config = {"save_freq": 100, "dt": 2, "events": ["EOD"]}

simulated_results = batt.simulate_to_threshold(future_loading, **config)

fig = simulated_results.inputs.plot(
    xlabel="time (s)", ylabel="current draw (amps)", title="BatteryElectroChemEOD Input"
)

fig = simulated_results.outputs.plot(
    keys=["v"],
    xlabel="time (s)",
    ylabel="voltage (V)",
    figsize=(10, 4),
    title="BatteryElectroChemEOD Outputs",
)
print("End of discharge voltage threshold:", batt.parameters["VEOD"])

fig2 = simulated_results.outputs.plot(
    keys=["t"], xlabel="time (s)", ylabel="temperature (°C)", figsize=(10, 4)
)

fig = simulated_results.event_states.plot(
    xlabel="time (s)",
    ylabel="event state",
    labels={"EOD"},
    title="BatteryElectroChemEOD Event State",
)

from progpy.models import BatteryElectroChemEOL

batt = BatteryElectroChemEOL()

print("inputs:", batt.inputs)
print("outputs:", batt.outputs)
print("event(s): ", batt.events)
print("states:", batt.states)

config = {"save_freq": 100, "dt": 2, "events": ["InsufficientCapacity"]}

simulated_results = batt.simulate_to_threshold(future_loading, **config)

fig = simulated_results.inputs.plot(
    xlabel="time (s)", ylabel="current draw (amps)", title="BatteryElectroChemEOL Input"
)

fig = simulated_results.event_states.plot(
    xlabel="time (s)",
    ylabel="event state",
    labels={"InsufficientCapacity"},
    title="BatteryElectroChemEOL Event State",
)

from progpy.models import BatteryElectroChem

batt = BatteryElectroChem()

print("inputs:", batt.inputs)
print("outputs:", batt.outputs)
print("event(s): ", batt.events)
print("states:", batt.states)

config = {
    "save_freq": 1000,
    "dt": 2,
    "events": "InsufficientCapacity",
}


def future_loading(t, x=None):
    load = 1

    if x is not None:
        event_state = batt.event_state(x)
        if event_state["EOD"] > 0.95:
            load = 1  # Discharge
        elif event_state["EOD"] < 0.05:
            load = -1  # Charge

    return batt.InputContainer({"i": load})

simulated_results = batt.simulate_to_threshold(future_loading, **config)

fig = simulated_results.inputs.plot(
    xlabel="time (s)", ylabel="current drawn (amps)", title="BatteryElectroChem Input"
)

fig = simulated_results.outputs.plot(
    keys=["v"],
    xlabel="time (s)",
    ylabel="voltage (V)",
    figsize=(10, 4),
    title="BatteryElectroChem Outputs",
)
fig2 = simulated_results.outputs.plot(
    keys=["t"], xlabel="time (s)", ylabel="temperature (°C)", figsize=(10, 4)
)

fig = simulated_results.event_states.plot(
    xlabel="time (s)",
    ylabel="event states",
    labels={"EOD", "InsufficientCapacity"},
    title="BatteryElectroChem Event States",
)

from progpy.models import SimplifiedBattery
from progpy.loading import Piecewise

batt = SimplifiedBattery()

print("inputs:", batt.inputs)
print("outputs:", batt.outputs)
print("event(s): ", batt.events)
print("states:", batt.states)

future_loading = Piecewise(
    dict, [600, 900, 1800, 3000, float("inf")], {"P": [25, 12, 50, 25, 33]}
)

simulated_results = batt.simulate_to(200, future_loading, {"v": 4.183})

fig = simulated_results.event_states.plot(
    xlabel="time (s)", ylabel="event state", title="SimplifiedBattery Event States"
)