"""XJTU battery dataset loading and preprocessing.

Data source: PINN4SOH repository (Wang et al., Nature Communications 2024),
github.com/wang-fujin/PINN4SOH. Each CSV = one battery; one row per
charge cycle: 16 statistical features from a partial charge window + capacity.
"""
import glob
import json
import os

import numpy as np
import pandas as pd

NOMINAL_CAPACITY = 2.0  # Ah, XJTU cells
T_SCALE = 500.0         # cycle-index normalization for the time coordinate
EOL_SOH = 0.80          # end-of-life threshold (test protocol stops at 80%)

FEATURE_COLS = 16  # first 16 columns are features, last is capacity


def load_battery(path):
    """Load one battery CSV -> dict(name, x [N,16], cycle [N], soh [N])."""
    df = pd.read_csv(path)
    df.insert(0, "cycle", np.arange(len(df), dtype=np.float32))
    df = df.replace([np.inf, -np.inf], np.nan).dropna().reset_index(drop=True)
    # 3-sigma outlier removal on features + capacity (as in PINN4SOH)
    keep = np.ones(len(df), dtype=bool)
    for col in df.columns[1:]:
        s = df[col]
        keep &= (s >= s.mean() - 3 * s.std()) & (s <= s.mean() + 3 * s.std())
    df = df[keep].reset_index(drop=True)
    return {
        "name": os.path.splitext(os.path.basename(path))[0],
        "crate": 2.0 if os.path.basename(path).startswith("2C") else 3.0,
        "x": df.iloc[:, 1:1 + FEATURE_COLS].values.astype(np.float32),
        "cycle": df["cycle"].values.astype(np.float32),
        "soh": (df["capacity"] / NOMINAL_CAPACITY).values.astype(np.float32),
    }


def load_dataset(data_dir, batches=("2C", "3C")):
    files = []
    for b in batches:
        files += sorted(glob.glob(os.path.join(data_dir, f"{b}_battery-*.csv")))
    return [load_battery(f) for f in files]


# Fixed battery-level split (no leakage across batteries)
TEST = ["2C_battery-7", "2C_battery-8", "3C_battery-14", "3C_battery-15"]
VAL = ["2C_battery-6", "3C_battery-13"]


def split_dataset(batteries):
    train = [b for b in batteries if b["name"] not in TEST + VAL]
    val = [b for b in batteries if b["name"] in VAL]
    test = [b for b in batteries if b["name"] in TEST]
    return train, val, test


def fit_norm(train):
    x = np.concatenate([b["x"] for b in train])
    return {"mean": x.mean(0), "std": x.std(0) + 1e-8}


def apply_norm(b, norm):
    return {
        **b,
        "x": (b["x"] - norm["mean"]) / norm["std"],
        "t": b["cycle"] / T_SCALE,
        "c": np.full(len(b["soh"]), (b["crate"] - 2.5) / 0.5, dtype=np.float32),
    }


def stack(batteries):
    """Concatenate normalized batteries into flat arrays."""
    x = np.concatenate([b["x"] for b in batteries])
    t = np.concatenate([b["t"] for b in batteries])
    y = np.concatenate([b["soh"] for b in batteries])
    c = np.concatenate([b["c"] for b in batteries])
    return x, t, y, c


def stack_pairs(batteries):
    """Consecutive-cycle pairs for the discrete physics loss."""
    cols = [[] for _ in range(7)]
    for b in batteries:
        arrs = (b["x"][:-1], b["t"][:-1], b["soh"][:-1],
                b["x"][1:], b["t"][1:], b["soh"][1:], b["c"][:-1])
        for c, a in zip(cols, arrs):
            c.append(a)
    import numpy as _np
    return tuple(_np.concatenate(c) for c in cols)


def rate_dataset(batteries):
    """(c, t, u, rate) degradation-rate targets for fitting the dynamics law.

    Per battery, a power-law fade curve SOH(cyc) = q0 - a*(cyc/100)**b is
    least-squares fitted (standard empirical battery-aging form; captures the
    knee). Targets are taken along the fitted curve with analytic rates,
    which are smooth and noise-free. Expects normalized batteries.
    """
    from scipy.optimize import curve_fit

    def fade(cyc, q0, a_, b_):
        return q0 - a_ * np.power(cyc / 100.0, b_)

    cs, ts, us, rs = [], [], [], []
    for b in batteries:
        cyc, y = b["cycle"], b["soh"]
        p, _ = curve_fit(fade, cyc, y, p0=[0.97, 0.02, 1.5],
                         bounds=([0.9, 1e-5, 0.5], [1.05, 1.0, 4.0]), maxfev=5000)
        q0, a_, b_ = p
        u_fit = fade(cyc, *p)
        # dSOH/dcycle * T_SCALE = rate in normalized-time units
        rate = -a_ * b_ / 100.0 * np.power(np.maximum(cyc, 1.0) / 100.0, b_ - 1.0) * T_SCALE
        cs.append(b["c"])
        ts.append((cyc / T_SCALE).astype(np.float32))
        us.append(u_fit.astype(np.float32))
        rs.append(rate.astype(np.float32))
    return (np.concatenate(cs), np.concatenate(ts),
            np.concatenate(us), np.concatenate(rs))


def save_norm(norm, path):
    with open(path, "w") as f:
        json.dump({k: v.tolist() for k, v in norm.items()}, f)


def load_norm(path):
    with open(path) as f:
        d = json.load(f)
    return {k: np.array(v, dtype=np.float32) for k, v in d.items()}
