#!/usr/bin/env python3
"""Closed-loop replay simulation for the BrewPi temp controller.

Loads a real run log and drives a simulated Pot plant with a fresh
TempController (Normal mode) using the same setpoints (temp_soll,
heatrate_soll_set) as the original run.  The controller's power output
feeds back into the plant model each tick, generating a simulated
temperature trajectory (sim_temp_ist) that can be compared against the
real one from the log (real_temp_ist).

Use this to verify PID params: tweak gains and see whether the simulated
plant reaches setpoint faster, overshoots less, or oscillates less than
the original logged run, all without touching real hardware.

Usage:
    python utils/replay_sim.py <date_time> <sud_name>
    python utils/replay_sim.py <date_time> <sud_name> \\
        --heat-kp 0.1 --heat-ki 0.03 --heat-kt 1.0
    python utils/replay_sim.py <date_time> <sud_name> \\
        --plant-M 25 --plant-L 0.15 --ambient 18

Plant params default to water (C=4190 J/kg·K), M=20 kg, L=0.2 W/kg·K,
Td=30 s — pass --plant-* to match the actual batch.  Max heater power is
inferred from the highest power_eff step in the log.

By default reads PID params from config.json in the repo root.  Any
--<section>-<gain> flag overrides just that one value.
"""

import argparse
import copy
import json
import sys
from pathlib import Path

import matplotlib
matplotlib.use("TkAgg")
import matplotlib.pyplot as plt
from matplotlib.ticker import AutoMinorLocator

sys.path.insert(0, str(Path(__file__).parent.parent))

from components.pid.temp_controller import TempController
from components.plant.pot import Pot


def _apply_gain_overrides(params, args):
    p = copy.deepcopy(params)
    for section, prefix in (('Hold', 'hold'), ('Heat', 'heat'), ('Cool', 'cool')):
        for gain in ('kp', 'ki', 'kd', 'kt'):
            val = getattr(args, '{}_{}'.format(prefix, gain))
            if val is not None:
                p[section][gain] = val
    return p


def _infer_max_power(samples):
    """Largest discrete heater step seen in the log — heater's effective max."""
    return max((s['power_eff'] for s in samples), default=1.0) or 1.0


def _infer_heatrate_soll_set(samples):
    """Estimate heatrate_soll_set from the log.

    rate_soll = heatrate_soll_set * pid_hold.get_y(); when the hold PID is
    saturated at 1.0 during active heating the two are equal, so the max
    rate_soll seen while the heater is on is a good upper bound."""
    candidates = [s['rate_soll'] for s in samples if s['power_set'] > 0]
    return max(candidates, default=1.0) or 1.0


def run_replay(samples, pid_params, heatrate_soll_set, plant_params, ambient, max_power, dt):
    """Closed-loop simulation: controller drives plant, plant feeds temp back.

    The plant is seeded at the real starting temperature from the log so the
    simulation begins at a realistic operating point.  Setpoints follow the
    real run's temp_soll schedule sample by sample."""
    tc = TempController(dt)
    tc.set_params(pid_params)
    tc.set_enabled(True)

    plant = Pot(dt)
    plant.set_plant_params(plant_params)
    plant.set_ambient_temperature(ambient)
    plant.initial(samples[0]['temp_ist'])

    out = []
    for s in samples:
        plant.process()
        sim_temp = plant.get_temperature()

        tc.set_theta_ist(sim_temp)
        tc.set_theta_soll(s['temp_soll'])
        tc.set_heatrate_soll(heatrate_soll_set)
        tc.process()

        power_watts = max(0.0, tc.get_power() * max_power)
        plant.set_power(power_watts)

        out.append({
            'sim_temp_ist': sim_temp,
            'power': tc.get_power(),
            'state': tc.state.name,
            'heatrate_ist': tc.get_heatrate_ist(),
            'heatrate_soll': tc.get_heatrate_soll(),
        })
    return out


def _style_axis(ax):
    ax.set_facecolor('white')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.xaxis.set_minor_locator(AutoMinorLocator())
    ax.yaxis.set_minor_locator(AutoMinorLocator())
    ax.tick_params(which='major', direction='out', labelsize='small')
    ax.tick_params(which='minor', direction='out', length=2)
    ax.grid(which='major', linestyle='-', linewidth=0.5, alpha=0.3)
    ax.grid(which='minor', linestyle=':', linewidth=0.5, alpha=0.15)


def plot_replay(samples, replay, name, max_power, params_label):
    t = [s['t'] for s in samples]

    fig, (ax_temp, ax_rate, ax_power) = plt.subplots(3, 1, sharex=True, facecolor='white')
    fig.suptitle('{} — replay simulation ({})'.format(name, params_label), fontsize='small')

    # Temperature: real log vs simulated plant trajectory
    ax_temp.plot(t, [s['temp_ist']          for s in samples], '-b',       linewidth=1,       label='real_temp_ist')
    ax_temp.plot(t, [r['sim_temp_ist']       for r in replay],  '-',        color='C1', linewidth=1, label='sim_temp_ist')
    ax_temp.plot(t, [s['temp_soll']          for s in samples], '-r',       linewidth=1, alpha=0.5, label='temp_soll')
    ax_temp.set_ylabel('Temp [°C]')
    ax_temp.legend(loc='upper left', fontsize='small')

    # Heat rate
    ax_rate.plot(t, [s['rate_ist']           for s in samples], '-b',       linewidth=1, alpha=0.6, label='rate_ist (real)')
    ax_rate.plot(t, [r['heatrate_ist']       for r in replay],  '-',        color='C1', linewidth=1, label='rate_ist (sim)')
    ax_rate.plot(t, [s['rate_soll']          for s in samples], '-r',       linewidth=1, alpha=0.4, label='rate_soll (real)')
    ax_rate.set_ylabel('Rate [K/min]')
    ax_rate.legend(loc='upper left', fontsize='small')

    # Power: log (normalised) vs sim controller output
    log_norm   = [s['power_set'] / max_power for s in samples]
    replay_pwr = [r['power']                 for r in replay]
    ax_power.plot(t, log_norm,   '-b',        linewidth=1,       label='power (real, norm.)')
    ax_power.plot(t, replay_pwr, '-', color='C1', linewidth=1,   label='power (sim)')
    ax_power.set_ylabel('Power [0–1]')
    ax_power.set_xlabel('t [s]')
    ax_power.legend(loc='upper left', fontsize='small')

    for ax in (ax_temp, ax_rate, ax_power):
        _style_axis(ax)

    ax_temp.label_outer()
    ax_rate.label_outer()
    fig.tight_layout()
    return fig


def main():
    repo_root       = Path(__file__).parent.parent
    default_config  = repo_root / 'config.json'
    default_log_dir = repo_root / 'logs'

    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument('date_time', help='Timestamp part of log filename, e.g. 20260628T184903')
    parser.add_argument('sud_name',  help='Sud name part of log filename, e.g. Sud-0010')
    parser.add_argument('--log-dir', default=str(default_log_dir),
                        help='Log directory (default: %(default)s)')
    parser.add_argument('--config', default=str(default_config),
                        help='Config file for base PID params (default: %(default)s)')
    parser.add_argument('--rate-soll', type=float, default=None,
                        help='Heat-rate setpoint [K/min] (default: inferred from log)')
    parser.add_argument('--ambient', type=float, default=None,
                        help='Ambient temperature [°C] (default: from config ambient_temperature)')

    # Plant model params
    parser.add_argument('--plant-M',  type=float, default=20.0,  metavar='kg',
                        help='Plant mass [kg] (default: %(default)s)')
    parser.add_argument('--plant-C',  type=float, default=4190.0, metavar='J/kgK',
                        help='Specific heat capacity [J/(kg·K)] (default: %(default)s — water)')
    parser.add_argument('--plant-L',  type=float, default=0.2,   metavar='W/kgK',
                        help='Heat loss coefficient [W/(kg·K)] (default: %(default)s)')
    parser.add_argument('--plant-Td', type=float, default=30.0,  metavar='s',
                        help='Transport delay [s] (default: %(default)s)')

    for section, prefix in (('Hold', 'hold'), ('Heat', 'heat'), ('Cool', 'cool')):
        for gain in ('kp', 'ki', 'kd', 'kt'):
            parser.add_argument(
                '--{}-{}'.format(prefix, gain),
                type=float, default=None,
                dest='{}_{}'.format(prefix, gain),
                metavar='VAL',
                help='Override TempCtrl.{}.{}'.format(section, gain),
            )

    args = parser.parse_args()

    log_path = Path(args.log_dir) / 'log_{}_{}.json'.format(args.date_time, args.sud_name)
    if not log_path.exists():
        print('Error: log file not found: {}'.format(log_path), file=sys.stderr)
        sys.exit(1)

    config_path = Path(args.config)
    if not config_path.exists():
        print('Error: config file not found: {}'.format(config_path), file=sys.stderr)
        sys.exit(1)

    with open(log_path) as f:
        log = json.load(f)
    with open(config_path) as f:
        config = json.load(f)

    samples     = log['Samples']
    name        = log.get('Name', '{}_{}'.format(args.date_time, args.sud_name))
    pid_params  = _apply_gain_overrides(config['TempCtrl'], args)
    dt          = (samples[1]['t'] - samples[0]['t']) if len(samples) > 1 else 1.0
    max_power   = _infer_max_power(samples)
    rate_soll   = args.rate_soll if args.rate_soll is not None else _infer_heatrate_soll_set(samples)
    ambient     = args.ambient   if args.ambient   is not None else config.get('ambient_temperature', 20.0)

    plant_params = {'M': args.plant_M, 'C': args.plant_C, 'L': args.plant_L, 'Td': args.plant_Td}

    any_override = any(
        getattr(args, '{}_{}'.format(p, g)) is not None
        for p in ('hold', 'heat', 'cool')
        for g in ('kp', 'ki', 'kd', 'kt')
    )
    params_label = 'candidate' if any_override else 'config'

    print('Log:    {} ({} samples, dt={:.1f} s)'.format(name, len(samples), dt))
    print('Power:  max={:.0f} W'.format(max_power))
    print('Rate:   heatrate_soll_set={:.2f} K/min'.format(rate_soll))
    print('Plant:  M={M} kg  C={C} J/kgK  L={L} W/kgK  Td={Td} s  ambient={amb} °C'.format(
        amb=ambient, **plant_params))
    print('Params: {}'.format('overridden' if any_override else 'from config'))
    print()
    for section in ('Hold', 'Heat', 'Cool'):
        p = pid_params[section]
        print('  {}: kp={kp}  ki={ki}  kd={kd}  kt={kt}'.format(section, **p))

    replay = run_replay(samples, pid_params, rate_soll, plant_params, ambient, max_power, dt)
    plot_replay(samples, replay, name, max_power, params_label)
    plt.show()


if __name__ == '__main__':
    main()