"""
battery_runtime.py: will the battery finish the mission?

EECE 5554 Robotics Sensing and Navigation, demo for Lecture 06.

The pack is S cells in series and P in parallel. Each step of one second:

    1. the load asks for power P (watts), from the mission profile
    2. the pack voltage sags under that load:  V = OCV(SOC) - I * R
       with I = P / V, so V solves  V^2 - OCV * V + P * R = 0
    3. the charge drops by I * dt, which lowers SOC and the open-circuit voltage

The run ends when the loaded voltage hits the BMS cutoff, which can happen
long before the pack is empty if a high-power burst makes the voltage sag.
Cold and aging both raise the internal resistance and shrink the capacity.

All cell numbers are typical values, not one product. Change them to match a
datasheet.

Run it from a terminal:
    python3 battery_runtime.py
"""
import json
import math
import sys

# open-circuit voltage per cell versus state of charge (typical shapes)
SOC_PTS = [0, 5, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
OCV = {
    "nmc": [3.00, 3.30, 3.45, 3.55, 3.62, 3.68, 3.74, 3.81, 3.90, 3.99, 4.08, 4.20],
    "lfp": [2.50, 2.90, 3.10, 3.20, 3.24, 3.26, 3.28, 3.29, 3.30, 3.31, 3.33, 3.40],
}
CUTOFF = {"nmc": 3.00, "lfp": 2.50}           # BMS under-voltage limit per cell, under load
NOMINAL = {"nmc": 3.6, "lfp": 3.2}
WIRING_OHM = 0.015                            # connectors, fuse, and wires

# cold makes the cell slower: resistance up, usable capacity down (rule-of-thumb values)
TEMPERATURE = {"25": (1.0, 1.00), "0": (2.0, 0.85), "-10": (3.0, 0.75)}

# mission profiles: repeating lists of (seconds, watts)
PROFILES = {
    "patrol": {"label": "Patrol robot", "cycle": [(60, 120), (30, 55), (10, 320)]},
    "delivery": {"label": "Delivery robot", "cycle": [(5, 260), (115, 90), (60, 40), (30, 60)]},
    "turtlebot": {"label": "Small indoor robot", "cycle": [(20, 15), (20, 24)]},
    "constant": {"label": "Constant power", "cycle": [(60, 100)]},
}


def ocv_cell(chem, soc):
    """Linear interpolation in the OCV table; soc in percent."""
    pts, v = SOC_PTS, OCV[chem]
    soc = max(0.0, min(100.0, soc))
    for i in range(len(pts) - 1):
        if soc <= pts[i + 1]:
            a = (soc - pts[i]) / (pts[i + 1] - pts[i])
            return v[i] + a * (v[i + 1] - v[i])
    return v[-1]


def power_at(t, cycle):
    period = sum(d for d, _ in cycle)
    x = t % period
    for d, w in cycle:
        if x < d:
            return w
        x -= d
    return cycle[-1][1]


def simulate(chem="nmc", s=7, p=3, cell_ah=3.0, cell_mohm=35.0, temp="25", soh=100.0,
             profile="patrol", const_w=100.0, max_hours=12.0):
    r_mult, c_mult = TEMPERATURE[str(temp)]
    age = soh / 100.0
    cap_ah = p * cell_ah * c_mult * age
    r_pack = (s * cell_mohm / 1000.0 / p) * r_mult * (1 + 2.5 * (1 - age)) + WIRING_OHM
    cycle = PROFILES[profile]["cycle"] if profile != "constant" else [(60, const_w)]
    v_cut = s * CUTOFF[chem]
    q = cap_ah * 3600.0                   # coulombs available
    used = 0.0
    t, dt = 0, 1.0
    rows = []
    reason = "time limit"
    energy = 0.0
    peak_i, min_v = 0.0, 1e9
    while t < max_hours * 3600:
        soc = 100.0 * (1 - used / q)
        ocv = s * ocv_cell(chem, soc)
        pw = power_at(t, cycle)
        disc = ocv * ocv - 4 * pw * r_pack
        if disc < 0:
            reason = "collapse"           # the pack cannot deliver this much power at all
            rows.append((t, ocv / 2, pw / (ocv / 2), soc, pw, ocv))
            break
        v = (ocv + math.sqrt(disc)) / 2
        i = pw / v
        rows.append((t, v, i, soc, pw, ocv))
        peak_i, min_v = max(peak_i, i), min(min_v, v)
        if v <= v_cut:
            reason = "sag" if soc >= 5 else "empty"
            break
        used += i * dt
        energy += pw * dt / 3600.0
        if used >= q:
            reason = "empty"
            break
        t += 1
    return {
        "rows": rows, "reason": reason, "runtime_s": t, "energy_wh": energy,
        "peak_a": peak_i, "min_v": min_v, "v_cut": v_cut, "r_pack": r_pack, "cap_ah": cap_ah,
        "nameplate_wh": s * p * NOMINAL[chem] * cell_ah, "soc_end": rows[-1][3] if rows else 100.0,
        "c_rate_peak": peak_i / cap_ah if cap_ah else 0.0,
        "left_wh": (rows[-1][3] / 100.0) * cap_ah * s * NOMINAL[chem] if rows else 0.0,
        "avg_w": sum(w * d for d, w in cycle) / sum(d for d, _ in cycle),
    }


def decimate(rows, n=700):
    """Keep the shape for plotting: per bucket, the lowest-voltage row and the last row."""
    if len(rows) <= n:
        return rows
    out, size = [], len(rows) / n
    for k in range(n):
        b = rows[int(k * size):int((k + 1) * size)] or [rows[-1]]
        lo = min(b, key=lambda r: r[1])
        out.extend(sorted({lo, b[-1]}, key=lambda r: r[0]))
    return out


def demo_api(params_json):
    """Web page entry point. Input: pack, conditions, profile, and mission length (minutes)."""
    p = json.loads(params_json or "{}")
    r = simulate(p.get("chem", "nmc"), int(p.get("s", 7)), int(p.get("p", 3)), float(p.get("cell_ah", 3.0)),
                 float(p.get("cell_mohm", 35.0)), str(p.get("temp", "25")), float(p.get("soh", 100)),
                 p.get("profile", "patrol"), float(p.get("const_w", 100)))
    mission_s = float(p.get("mission_min", 90)) * 60
    rows = decimate(r.pop("rows"))
    r["series"] = {"t_min": [x[0] / 60 for x in rows], "v": [x[1] for x in rows], "i": [x[2] for x in rows],
                   "soc": [x[3] for x in rows], "w": [x[4] for x in rows], "ocv": [x[5] for x in rows]}
    r["mission_s"] = mission_s
    r["completes"] = r["runtime_s"] >= mission_s
    r["profile_label"] = PROFILES[p.get("profile", "patrol")]["label"]
    return json.dumps(r)


def main(argv):
    for temp in ("25", "0", "-10"):
        r = simulate(temp=temp)
        print("{:>4} C: runtime {:5.1f} min, stopped by {:8s}  min V {:5.2f}  peak {:5.1f} A ({:.1f}C)  SOC left {:4.1f}%".format(
            temp, r["runtime_s"] / 60, r["reason"], r["min_v"], r["peak_a"], r["c_rate_peak"], r["soc_end"]))


if __name__ == "__main__":
    main(sys.argv)