"""Automated tuning routines combining testbench instruments + STM32 link.

Uses the power analyzer (HIOKI) as ground truth for efficiency while
adjusting converter parameters via the STM32 debug protocol.
"""

from __future__ import annotations

import csv
import time
from dataclasses import dataclass, field
from pathlib import Path

from testbench.bench import MPPTTestbench, IDLE_VOLTAGE
from testbench.stm32_link import (
    STM32Link, Telemetry, PARAM_BY_NAME,
)


@dataclass
class TunePoint:
    """One measurement during a tuning sweep."""
    param_name: str
    param_value: float
    voltage_set: float
    load_setpoint: float
    load_mode: str

    # HIOKI measurements (ground truth)
    meter_pin: float = 0.0
    meter_pout: float = 0.0
    meter_eff: float = 0.0

    # STM32 telemetry
    stm_vin: float = 0.0
    stm_vout: float = 0.0
    stm_iin: float = 0.0
    stm_iout: float = 0.0
    stm_eff: float = 0.0
    stm_vfly: float = 0.0
    stm_etemp: float = 0.0

    timestamp: float = field(default_factory=time.time)


class Tuner:
    """Combines MPPTTestbench + STM32Link for automated tuning.

    Usage::

        tuner = Tuner(bench, link)
        results = tuner.sweep_param(
            "dt_10_20A", start=14, stop=40, step=1,
            voltage=60.0, current_limit=20.0,
            load_mode="CP", load_value=300.0,
        )
        tuner.print_results(results)
    """

    def __init__(
        self,
        bench: MPPTTestbench,
        link: STM32Link,
        settle_time: float = 5.0,
        stm_avg_samples: int = 10,
    ):
        self.bench = bench
        self.link = link
        self.settle_time = settle_time
        self.stm_avg_samples = stm_avg_samples

    def _measure(
        self,
        param_name: str,
        param_value: float,
        voltage: float,
        load_value: float,
        load_mode: str,
    ) -> TunePoint:
        """Take one combined measurement from HIOKI + STM32."""
        point = TunePoint(
            param_name=param_name,
            param_value=param_value,
            voltage_set=voltage,
            load_setpoint=load_value,
            load_mode=load_mode,
        )

        # HIOKI measurement
        meter_vals = self.bench._wait_meter_ready(max_retries=10, retry_delay=1.0)
        point.meter_pin = meter_vals.get("P5", 0.0)
        point.meter_pout = meter_vals.get("P6", 0.0)
        point.meter_eff = meter_vals.get("EFF1", 0.0)

        # STM32 telemetry (averaged)
        t = self.link.read_telemetry_avg(n=self.stm_avg_samples)
        if t:
            point.stm_vin = t.vin_V
            point.stm_vout = t.vout_V
            point.stm_iin = t.iin_A
            point.stm_iout = t.iout_A
            point.stm_eff = t.efficiency
            point.stm_vfly = t.vfly / 1000.0
            point.stm_etemp = t.etemp

        return point

    # ── Parameter sweep ──────────────────────────────────────────────

    def sweep_param(
        self,
        param_name: str,
        start: float,
        stop: float,
        step: float,
        voltage: float,
        current_limit: float,
        load_mode: str = "CC",
        load_value: float = 5.0,
        settle_time: float | None = None,
    ) -> list[TunePoint]:
        """Sweep a single STM32 parameter while measuring efficiency.

        Sets up the testbench at the given operating point, then steps
        the parameter from start to stop, measuring at each step.

        Returns list of TunePoints with both HIOKI and STM32 data.
        """
        if param_name not in PARAM_BY_NAME:
            raise ValueError(f"Unknown parameter: {param_name!r}")

        settle = settle_time or self.settle_time

        if step == 0:
            raise ValueError("step cannot be zero")
        if start > stop and step > 0:
            step = -step
        elif start < stop and step < 0:
            step = -step

        # Count steps
        n_steps = int(abs(stop - start) / abs(step)) + 1
        unit = "A" if load_mode == "CC" else "W"
        print(f"Parameter sweep: {param_name} = {start} → {stop} (step {step})")
        print(f"  Operating point: V={voltage:.1f}V, {load_mode}={load_value:.1f}{unit}")
        print(f"  {n_steps} points, settle={settle:.1f}s")
        print()

        # Set up testbench
        self.bench.supply.set_current(current_limit)
        self.bench.supply.set_voltage(voltage)
        self.bench.supply.output_on()
        self.bench.load.set_mode(load_mode)
        self.bench._apply_load_value(load_mode, load_value)
        self.bench.load.load_on()

        # Initial settle
        print("  Settling...")
        time.sleep(settle * 2)

        results: list[TunePoint] = []
        val = start
        n = 0

        try:
            while True:
                if step > 0 and val > stop + step / 2:
                    break
                if step < 0 and val < stop + step / 2:
                    break

                # Write parameter
                ack = self.link.write_param(param_name, val)
                if not ack:
                    print(f"  WARNING: No ACK for {param_name}={val}")

                time.sleep(settle)

                # Measure
                point = self._measure(param_name, val, voltage, load_value, load_mode)
                results.append(point)
                n += 1

                print(
                    f"  [{n:>3d}/{n_steps}] {param_name}={val:>6.1f}  "
                    f"HIOKI: Pin={point.meter_pin:7.1f}W Pout={point.meter_pout:7.1f}W "
                    f"EFF={point.meter_eff:5.2f}%  "
                    f"STM32: EFF={point.stm_eff:5.1f}% T={point.stm_etemp:.0f}°C"
                )

                val += step

        finally:
            self.bench.load.load_off()
            self.bench.supply.set_voltage(IDLE_VOLTAGE)
            print(f"\n  Load OFF. Supply at {IDLE_VOLTAGE:.0f}V.")

        return results

    # ── Deadtime optimization ────────────────────────────────────────

    # Firmware now uses a single global dead-time (`dt_normal`, 0x60) rather
    # than per-current-bracket values, so we sweep the one parameter — optionally
    # at several load points to expose any load dependence — and pick the best.

    DT_PARAM = "dt_normal"

    def tune_deadtime(
        self,
        dt_start: int = 14,
        dt_stop: int = 50,
        dt_step: int = 1,
        voltage: float = 60.0,
        current_limit: float = 20.0,
        load_mode: str = "CP",
        load_values: list[float] | None = None,
        settle_time: float | None = None,
    ) -> list[TunePoint]:
        """Optimize the single global dead-time (`dt_normal`).

        Sweeps `dt_normal` from `dt_start` to `dt_stop` at each requested load
        point and returns a flat list of measurements. The best value can be
        applied with :meth:`apply_best_deadtime`.

        Args:
            load_values: Load setpoints to test. If None, a single mid-range
                         load is used.
        """
        settle = settle_time or self.settle_time

        if not load_values:
            load_values = [200.0 if load_mode == "CP" else 5.0]

        unit = "A" if load_mode == "CC" else "W"
        print("=" * 80)
        print("DEAD-TIME OPTIMIZATION (dt_normal)")
        print(f"  DT range: {dt_start} → {dt_stop} (step {dt_step})")
        print(f"  V={voltage:.0f}V, I_limit={current_limit:.0f}A, mode={load_mode}")
        print(f"  Loads: {', '.join(f'{lv:.0f}{unit}' for lv in load_values)}")
        print("=" * 80)

        all_results: list[TunePoint] = []
        best_per_load: list[tuple[float, TunePoint]] = []

        for load_val in load_values:
            print(f"\n── {self.DT_PARAM} @ {load_mode}={load_val:.0f}{unit} ──")
            results = self.sweep_param(
                param_name=self.DT_PARAM,
                start=dt_start,
                stop=dt_stop,
                step=dt_step,
                voltage=voltage,
                current_limit=current_limit,
                load_mode=load_mode,
                load_value=load_val,
                settle_time=settle,
            )
            all_results.extend(results)

            valid = [p for p in results if 0 < p.meter_eff < 110]
            if valid:
                best = max(valid, key=lambda p: p.meter_eff)
                best_per_load.append((load_val, best))
                print(f"  ★ Best: {self.DT_PARAM}={best.param_value:.0f} → "
                      f"EFF={best.meter_eff:.2f}%")

        # Summary
        print("\n" + "=" * 80)
        print("DEAD-TIME OPTIMIZATION SUMMARY")
        print(f"{'Load':<12} {'Best DT':>8} {'Efficiency':>12} {'Temp':>8}")
        print("-" * 42)
        for load_val, best in best_per_load:
            print(f"{load_val:<11.0f}{unit} {best.param_value:>8.0f} "
                  f"{best.meter_eff:>11.2f}% {best.stm_etemp:>7.0f}°C")
        if not best_per_load:
            print("  (no valid points)")
        print("=" * 80)

        return all_results

    def apply_best_deadtime(self, results: list[TunePoint]):
        """Apply the single best dead-time (highest efficiency) to the STM32."""
        valid = [p for p in results if 0 < p.meter_eff < 110]
        if not valid:
            print("\nNo valid points — dead-time not applied.")
            return
        best = max(valid, key=lambda p: p.meter_eff)
        val = int(best.param_value)
        ack = self.link.write_param(self.DT_PARAM, val)
        status = "OK" if ack else "NO ACK"
        print(f"\nApplying best dead-time: {self.DT_PARAM} = {val} "
              f"(EFF={best.meter_eff:.2f}%) ({status})")

    # ── Multi-point sweep ────────────────────────────────────────────

    def sweep_param_multi(
        self,
        param_name: str,
        start: float,
        stop: float,
        step: float,
        voltages: list[float],
        current_limit: float,
        load_mode: str = "CP",
        load_values: list[float] | None = None,
        settle_time: float | None = None,
    ) -> list[TunePoint]:
        """Sweep a parameter across multiple voltage/load combinations.

        Produces a comprehensive dataset showing how the parameter
        affects efficiency across the full operating range.
        """
        if load_values is None:
            load_values = [200.0]  # default: 200W

        all_results: list[TunePoint] = []

        for v in voltages:
            for lv in load_values:
                unit = "A" if load_mode == "CC" else "W"
                print(f"\n── V={v:.0f}V, {load_mode}={lv:.0f}{unit} ──")
                results = self.sweep_param(
                    param_name=param_name,
                    start=start, stop=stop, step=step,
                    voltage=v, current_limit=current_limit,
                    load_mode=load_mode, load_value=lv,
                    settle_time=settle_time,
                )
                all_results.extend(results)

        return all_results

    # ── Output ───────────────────────────────────────────────────────

    @staticmethod
    def print_results(results: list[TunePoint]):
        """Print a summary table of tuning results."""
        if not results:
            print("No results.")
            return

        valid = [p for p in results if 0 < p.meter_eff < 110]
        if valid:
            best = max(valid, key=lambda p: p.meter_eff)
            print(f"\nBest: {best.param_name}={best.param_value:.1f} → "
                  f"EFF={best.meter_eff:.2f}% "
                  f"(Pin={best.meter_pin:.1f}W Pout={best.meter_pout:.1f}W)")

    @staticmethod
    def write_csv(results: list[TunePoint], path: str):
        """Write tuning results to CSV."""
        if not results:
            return

        with open(path, "w", newline="") as f:
            w = csv.writer(f)
            w.writerow([
                "param_name", "param_value",
                "voltage_set", "load_setpoint", "load_mode",
                "meter_pin", "meter_pout", "meter_eff",
                "stm_vin", "stm_vout", "stm_iin", "stm_iout",
                "stm_eff", "stm_vfly", "stm_etemp",
            ])
            for p in results:
                w.writerow([
                    p.param_name, f"{p.param_value:.4f}",
                    f"{p.voltage_set:.4f}", f"{p.load_setpoint:.4f}", p.load_mode,
                    f"{p.meter_pin:.4f}", f"{p.meter_pout:.4f}", f"{p.meter_eff:.4f}",
                    f"{p.stm_vin:.4f}", f"{p.stm_vout:.4f}",
                    f"{p.stm_iin:.4f}", f"{p.stm_iout:.4f}",
                    f"{p.stm_eff:.4f}", f"{p.stm_vfly:.4f}", f"{p.stm_etemp:.4f}",
                ])
        print(f"Results saved to {path}")

    @staticmethod
    def plot_sweep(results: list[TunePoint], show: bool = True):
        """Plot parameter sweep results."""
        import numpy as np
        import matplotlib.pyplot as plt

        if not results:
            return

        param_name = results[0].param_name
        vals = np.array([p.param_value for p in results])
        eff_hioki = np.array([p.meter_eff for p in results])
        eff_stm = np.array([p.stm_eff for p in results])
        temp = np.array([p.stm_etemp for p in results])

        # Filter valid
        valid = (eff_hioki > 0) & (eff_hioki < 110)

        # Group by operating point
        ops = sorted(set((p.voltage_set, p.load_setpoint) for p in results))

        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8), sharex=True)

        cmap = plt.cm.viridis
        for i, (v, l) in enumerate(ops):
            color = cmap(i / max(len(ops) - 1, 1))
            mask = np.array([
                (p.voltage_set == v and p.load_setpoint == l and
                 0 < p.meter_eff < 110)
                for p in results
            ])
            if not np.any(mask):
                continue
            x = vals[mask]
            order = np.argsort(x)
            unit = results[0].load_mode
            ax1.plot(x[order], eff_hioki[mask][order], "o-", color=color,
                     markersize=4, label=f"{v:.0f}V/{l:.0f}{unit}")
            ax2.plot(x[order], temp[mask][order], "o-", color=color,
                     markersize=4, label=f"{v:.0f}V/{l:.0f}{unit}")

        # Mark best
        valid_pts = [p for p in results if 0 < p.meter_eff < 110]
        if valid_pts:
            best = max(valid_pts, key=lambda p: p.meter_eff)
            ax1.axvline(best.param_value, color="red", linestyle="--", alpha=0.5)
            ax1.plot(best.param_value, best.meter_eff, "*", color="red",
                     markersize=15, zorder=10,
                     label=f"Best: {best.param_value:.0f} → {best.meter_eff:.2f}%")

        ax1.set_ylabel("Efficiency (%)", fontsize=12)
        ax1.set_title(f"Parameter Sweep: {param_name}", fontsize=14)
        ax1.legend(fontsize=8)
        ax1.grid(True, alpha=0.3)

        ax2.set_xlabel(param_name, fontsize=12)
        ax2.set_ylabel("Temperature (°C)", fontsize=12)
        ax2.legend(fontsize=8)
        ax2.grid(True, alpha=0.3)

        fig.tight_layout()
        if show:
            plt.show()
        return fig
