#!/usr/bin/env python3
"""Estimate induction-motor parameters from inverter self-commissioning CSV.

The script intentionally uses only the Python standard library so it can run
on an engineering workstation without installing SciPy/Numpy.
"""

from __future__ import annotations

import csv
import json
import math
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable


EPS = 1e-12


@dataclass
class SweepPoint:
    freq_hz: float
    r_ohm: float
    x_ohm: float


@dataclass
class FitResult:
    rs_ohm: float
    rr_ohm: float
    lls_h: float
    llr_h: float
    lm_h: float
    rms_error_ohm: float


def parse_float(row: dict[str, str], key: str) -> float | None:
    value = row.get(key, "")
    if value is None or value.strip() == "":
        return None
    return float(value)


def load_rows(path: Path) -> list[dict[str, str]]:
    with path.open("r", encoding="utf-8-sig", newline="") as stream:
        return list(csv.DictReader(stream))


def linear_fit(xs: Iterable[float], ys: Iterable[float]) -> tuple[float, float]:
    x = list(xs)
    y = list(ys)
    n = len(x)
    if n < 2:
        raise ValueError("linear fit needs at least two points")

    sx = sum(x)
    sy = sum(y)
    sxx = sum(v * v for v in x)
    sxy = sum(a * b for a, b in zip(x, y))
    den = n * sxx - sx * sx
    if abs(den) < EPS:
        raise ValueError("linear fit has degenerate x values")

    slope = (n * sxy - sx * sy) / den
    offset = (sy - slope * sx) / n
    return slope, offset


def estimate_rs(rows: list[dict[str, str]]) -> tuple[float, float]:
    currents: list[float] = []
    voltages: list[float] = []

    for row in rows:
        if row.get("test") != "rs_dc":
            continue
        current = parse_float(row, "current_a")
        voltage = parse_float(row, "voltage_v")
        if current is None or voltage is None:
            continue
        currents.append(current)
        voltages.append(voltage)

    if len(currents) < 2:
        raise ValueError("need at least two rs_dc rows with current_a and voltage_v")

    rs, voltage_offset = linear_fit(currents, voltages)
    return max(rs, 0.0), voltage_offset


def estimate_sweep(rows: list[dict[str, str]]) -> list[SweepPoint]:
    points: list[SweepPoint] = []

    for row in rows:
        if row.get("test") != "ac_sweep":
            continue
        freq = parse_float(row, "freq_hz")
        v_rms = parse_float(row, "v_rms")
        i_rms = parse_float(row, "i_rms")
        power = parse_float(row, "p_w")
        if freq is None or v_rms is None or i_rms is None or power is None:
            continue
        if freq <= 0.0 or i_rms <= EPS:
            continue

        z_abs = v_rms / i_rms
        r = power / (i_rms * i_rms)
        x_sq = max(z_abs * z_abs - r * r, 0.0)
        points.append(SweepPoint(freq_hz=freq, r_ohm=r, x_ohm=math.sqrt(x_sq)))

    if len(points) < 2:
        raise ValueError("need at least two ac_sweep rows with freq_hz, v_rms, i_rms, p_w")
    return points


def estimate_lm_from_flux(rows: list[dict[str, str]], lls_h: float) -> float | None:
    values: list[float] = []
    for row in rows:
        if row.get("test") != "magnetizing":
            continue
        current = parse_float(row, "current_a")
        psi = parse_float(row, "psi_wb")
        if current is None or psi is None or abs(current) <= EPS:
            continue
        lm = psi / current - lls_h
        if lm > 0.0:
            values.append(lm)

    if not values:
        return None
    return sum(values) / len(values)


def z_model(freq_hz: float, rs: float, rr: float, lls: float, llr: float, lm: float) -> complex:
    w = 2.0 * math.pi * freq_hz
    z_lm = 1j * w * lm
    z_rotor = rr + 1j * w * llr
    z_parallel = (z_lm * z_rotor) / (z_lm + z_rotor)
    return rs + 1j * w * lls + z_parallel


def objective(points: list[SweepPoint], rs: float, rr: float, lls: float, llr: float, lm: float) -> float:
    err = 0.0
    for point in points:
        z = z_model(point.freq_hz, rs, rr, lls, llr, lm)
        scale = max(abs(complex(point.r_ohm, point.x_ohm)), 1e-3)
        dr = (z.real - point.r_ohm) / scale
        dx = (z.imag - point.x_ohm) / scale
        err += dr * dr + dx * dx
    return err / len(points)


def initial_from_sweep(points: list[SweepPoint], rs: float, lm_hint: float | None) -> tuple[float, float, float, float]:
    r_slope, r0 = linear_fit([p.freq_hz for p in points], [p.r_ohm for p in points])
    _ = r_slope
    rr = max(r0 - rs, 1e-4)

    leakage_values = []
    for point in points:
        w = 2.0 * math.pi * point.freq_hz
        leakage_values.append(max(point.x_ohm / w, 1e-7))
    ll_total = sum(leakage_values) / len(leakage_values)
    lls = max(0.5 * ll_total, 1e-7)
    llr = max(0.5 * ll_total, 1e-7)
    lm = max(lm_hint if lm_hint is not None else 8.0 * ll_total, 5.0 * ll_total, 1e-6)
    return rr, lls, llr, lm


def fit_t_model(points: list[SweepPoint], rs: float, lm_hint: float | None) -> FitResult:
    rr, lls, llr, lm = initial_from_sweep(points, rs, lm_hint)
    params = [rr, lls, llr, lm]
    factors = [2.0, 2.0, 2.0, 2.0]
    fitted_count = 3 if lm_hint is not None else 4
    best = objective(points, rs, *params)

    for _ in range(80):
        improved = False
        for idx in range(fitted_count):
            for direction in (1.0, -1.0):
                trial = params[:]
                factor = factors[idx] if direction > 0.0 else 1.0 / factors[idx]
                trial[idx] = max(trial[idx] * factor, 1e-9)
                score = objective(points, rs, *trial)
                if score < best:
                    params = trial
                    best = score
                    improved = True

        if not improved:
            factors = [math.sqrt(f) for f in factors]
            if max(factors) < 1.005:
                break

    raw_error = 0.0
    for point in points:
        z = z_model(point.freq_hz, rs, *params)
        raw_error += (z.real - point.r_ohm) ** 2 + (z.imag - point.x_ohm) ** 2
    rms_error = math.sqrt(raw_error / len(points))

    return FitResult(
        rs_ohm=rs,
        rr_ohm=params[0],
        lls_h=params[1],
        llr_h=params[2],
        lm_h=params[3],
        rms_error_ohm=rms_error,
    )


def derived(result: FitResult) -> dict[str, float]:
    ls = result.lm_h + result.lls_h
    lr = result.lm_h + result.llr_h
    sigma = 1.0 - (result.lm_h * result.lm_h) / max(ls * lr, EPS)
    tr = lr / result.rr_ohm if result.rr_ohm > EPS else math.inf
    return {
        "Ls_H": ls,
        "Lr_H": lr,
        "Ll_total_H": result.lls_h + result.llr_h,
        "sigma": sigma,
        "Tr_s": tr,
    }


def main(argv: list[str]) -> int:
    if len(argv) != 2:
        print(f"usage: {Path(argv[0]).name} measurements.csv", file=sys.stderr)
        return 2

    rows = load_rows(Path(argv[1]))
    rs, voltage_offset = estimate_rs(rows)
    points = estimate_sweep(rows)

    rough_lls = initial_from_sweep(points, rs, None)[1]
    lm_hint = estimate_lm_from_flux(rows, rough_lls)
    fit = fit_t_model(points, rs, lm_hint)

    payload = {
        "parameters": {
            "Rs_ohm": fit.rs_ohm,
            "Rr_ohm": fit.rr_ohm,
            "Lls_H": fit.lls_h,
            "Llr_H": fit.llr_h,
            "Lm_H": fit.lm_h,
            **derived(fit),
        },
        "diagnostics": {
            "rs_voltage_offset_V": voltage_offset,
            "sweep_points": len(points),
            "t_model_rms_error_ohm": fit.rms_error_ohm,
            "lm_hint_used": lm_hint is not None,
            "note": "Use as commissioning seed; validate against current-loop response before enabling torque control.",
        },
    }
    print(json.dumps(payload, indent=2, ensure_ascii=False))
    return 0


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