"""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, DT_BRACKETS,
)


@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 = 3.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 ────────────────────────────────────────

    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,
    ) -> dict[str, list[TunePoint]]:
        """Optimize deadtime for each current bracket.

        For each deadtime bracket, sets a load that puts the converter
        in that current range, then sweeps deadtime values to find the
        optimum.

        Args:
            load_values: Load setpoints to test (one per DT bracket).
                         If None, auto-selects based on bracket midpoints.
        """
        settle = settle_time or self.settle_time

        if load_values is None:
            # Auto-select load values targeting the middle of each bracket
            # Using CP mode: power ≈ voltage × current
            load_values = []
            for _, _, i_lo, i_hi in DT_BRACKETS:
                mid_i_A = (i_lo + i_hi) / 2 / 1000.0  # mA → A
                target_power = voltage * mid_i_A * 0.4  # rough vout/vin ratio
                load_values.append(max(10.0, target_power))

        print("=" * 80)
        print("DEADTIME OPTIMIZATION")
        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("=" * 80)

        all_results: dict[str, list[TunePoint]] = {}

        for i, (param_id, param_name, i_lo, i_hi) in enumerate(DT_BRACKETS):
            load_val = load_values[i] if i < len(load_values) else load_values[-1]
            unit = "A" if load_mode == "CC" else "W"

            print(f"\n── Bracket: {param_name} ({i_lo/1000:.0f}-{i_hi/1000:.0f}A) "
                  f"@ {load_mode}={load_val:.0f}{unit} ──")

            results = self.sweep_param(
                param_name=param_name,
                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[param_name] = results

            # Find and report best
            if results:
                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"  ★ Best: {param_name}={best.param_value:.0f} → "
                          f"EFF={best.meter_eff:.2f}%")

        # Summary
        print("\n" + "=" * 80)
        print("DEADTIME OPTIMIZATION SUMMARY")
        print(f"{'Bracket':<15} {'Best DT':>8} {'Efficiency':>12} {'Temp':>8}")
        print("-" * 45)
        for param_name, results in all_results.items():
            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"{param_name:<15} {best.param_value:>8.0f} "
                      f"{best.meter_eff:>11.2f}% {best.stm_etemp:>7.0f}°C")
            else:
                print(f"{param_name:<15} {'N/A':>8} {'N/A':>12} {'N/A':>8}")
        print("=" * 80)

        return all_results

    def apply_best_deadtimes(self, results: dict[str, list[TunePoint]]):
        """Apply the best deadtime from each bracket to the STM32."""
        print("\nApplying optimal deadtimes:")
        for param_name, points in results.items():
            valid = [p for p in points if 0 < p.meter_eff < 110]
            if valid:
                best = max(valid, key=lambda p: p.meter_eff)
                val = int(best.param_value)
                ack = self.link.write_param(param_name, val)
                status = "OK" if ack else "NO ACK"
                print(f"  {param_name} = {val} ({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