#!/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 logs/log_<date_time>_<name>.json
    python utils/replay_sim.py logs/log_<date_time>_<name>.json \\
        --heat-kp 0.1 --heat-ki 0.03 --heat-kt 1.0
    python utils/replay_sim.py logs/log_<date_time>_<name>.json \\
        --plant-M 25 --plant-L 0.15 --ambient 18

PID params and plant params are read from the log's embedded "Config" and
"PlantParams" sections when present (logs recorded since the config-in-log
feature was added).  Pass --config to supply a config file for older logs
that predate that feature.  Any --<section>-<gain> or --plant-* flag
overrides just that one value regardless of source.

Max heater power is inferred from the highest power_eff step in the log.
"""

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 import PidFactory
from components.plant.pot import Pot


GAIN_SECTIONS = (('Outer', 'outer'), ('Inner.Heat', 'inner-heat'),
                 ('Inner.Hold', 'inner-hold'), ('Inner.Cool', 'inner-cool'))


def _apply_gain_overrides(params, args):
    p = copy.deepcopy(params)
    for section, prefix in GAIN_SECTIONS:
        for gain in ('kp', 'ki', 'kd', 'kt'):
            val = getattr(args, '{}_{}'.format(prefix.replace('-', '_'), gain))
            if val is not None:
                if '.' in section:
                    outer, inner = section.split('.')
                    p[outer][inner][gain] = val
                else:
                    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_outer.get_y(); when the outer 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_type, pid_params, heatrate_soll_set, plant_params, param_events,
               ambient, dt):
    """Open-loop plant replay with parallel controller observation.

    The plant is driven by the real logged power_eff each tick so that
    sim_temp_ist tracks real_temp_ist when the plant model is accurate.
    The controller runs in parallel (observing sim_temp) to produce a
    comparable power signal — useful to spot gain issues — but its output
    does NOT feed back into the plant.

    param_events is a list of {t, params} dicts (from the log's PlantParams
    section) applied at their logged wall-clock t so the simulated plant
    tracks the same M/C changes (grain added, water boiling off) as the
    real run."""
    tc = PidFactory.create(pid_type, dt)
    tc.set_params(pid_params)
    tc.set_ambient_temperature(ambient)
    tc.set_model_plant_params(plant_params)
    tc.set_enabled(True)

    # The logged temp_ist is Smith-corrected (transport delay already removed
    # by the predictor), so the replay plant must not add another Td on top.
    # For Normal mode the raw sensor reading is logged so keep Td as-is.
    replay_plant_params = dict(plant_params, Td=0) if pid_type == 'Smith' else plant_params
    plant = Pot(dt)
    plant.set_plant_params(replay_plant_params)
    plant.set_ambient_temperature(ambient)
    plant.initial(samples[0]['temp_ist'])

    # Walk param_events in order, applying each when the sample t passes it.
    param_queue = list(param_events)
    param_idx = 0

    out = []
    for s in samples:
        # Apply any plant param changes that fall at or before this sample.
        while param_idx < len(param_queue) and param_queue[param_idx]['t'] <= s['t']:
            p = param_queue[param_idx]['params']
            plant.set_plant_params(dict(p, Td=0) if pid_type == 'Smith' else p)
            tc.set_model_plant_params(p)
            param_idx += 1

        plant.process()
        sim_temp = plant.get_temperature()

        # Controller observes sim_temp but its output doesn't drive the plant.
        tc.set_theta_ist(sim_temp)
        tc.set_theta_soll(s['temp_soll'])
        tc.set_heatrate_soll(heatrate_soll_set)
        tc.process()

        # Drive plant and Smith model with the real logged power.
        real_power_watts = s['power_eff']
        tc.set_model_power(real_power_watts)
        plant.set_power(real_power_watts)

        out.append({
            'sim_temp_ist': sim_temp,
            'power': max(0.0, 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: power_set (what the controller commanded) and power_eff (what
    # the heater actually delivered, quantised to its discrete steps).
    # The gap between the two is the heater's rounding/quantisation.
    pwr_set = [s['power_set'] / max_power for s in samples]
    pwr_eff = [s['power_eff'] / max_power for s in samples]
    ax_power.plot(t, pwr_set, '-b',  linewidth=1,       label='power_set (norm.)')
    ax_power.plot(t, pwr_eff, '-',   color='C1', linewidth=1, alpha=0.7, label='power_eff (norm.)')
    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'

    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument('log', help='Path to log_*.json file')
    parser.add_argument('--config', default=None,
                        help='Config file for PID params (default: embedded in log, or config.json)')
    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)')

    # Plant model params — all default to None so we can tell "not specified"
    # and fall back to the log's embedded PlantParams when available.
    parser.add_argument('--plant-M',  type=float, default=None, metavar='kg',    help='Plant mass [kg]')
    parser.add_argument('--plant-C',  type=float, default=None, metavar='J/kgK', help='Specific heat [J/(kg·K)]')
    parser.add_argument('--plant-L',  type=float, default=None, metavar='W/kgK', help='Heat loss coeff [W/(kg·K)]')
    parser.add_argument('--plant-Td', type=float, default=None, metavar='s',     help='Transport delay [s]')

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

    args = parser.parse_args()

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

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

    # Config: prefer embedded in log, then --config arg, then default config.json.
    config = log.get('Config')
    if config is None:
        config_path = Path(args.config) if args.config else default_config
        if config_path.exists():
            with open(config_path) as f:
                config = json.load(f)
        else:
            print('Error: no config embedded in log and config file not found: {}'.format(config_path),
                  file=sys.stderr)
            sys.exit(1)

    # Plant params: prefer log's PlantParams[0], then --plant-* flags, then hardcoded defaults.
    log_plant = log['PlantParams'][0]['params'] if log.get('PlantParams') else None
    plant_params = {
        'M':  args.plant_M  if args.plant_M  is not None else (log_plant['M']  if log_plant else 20.0),
        'C':  args.plant_C  if args.plant_C  is not None else (log_plant['C']  if log_plant else 4190.0),
        'L':  args.plant_L  if args.plant_L  is not None else (log_plant['L']  if log_plant else 0.2),
        'Td': args.plant_Td if args.plant_Td is not None else (log_plant['Td'] if log_plant else 30.0),
    }

    samples      = log['Samples']
    param_events = log.get('PlantParams', [])
    name         = log.get('Name', log_path.stem)
    pid_type     = config['TempCtrl'].get('pid_type', 'Normal')
    pid_params   = _apply_gain_overrides(config['TempCtrl'], args)
    # 'dt' is the physics timestep the controller/plant were created with —
    # distinct from the sample interval (wall-clock) in server logs.
    dt           = log.get('dt') or ((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)

    any_override = any(
        getattr(args, '{}_{}'.format(prefix.replace('-', '_'), g)) is not None
        for _, prefix in GAIN_SECTIONS
        for g in ('kp', 'ki', 'kd', 'kt')
    )
    config_source = 'log' if log.get('Config') else 'file'
    params_label  = 'candidate' if any_override else config_source

    print('Log:    {} ({} samples, dt={} s)'.format(name, len(samples), round(dt, 4)))
    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:.0f} J/kgK  L={L} W/kgK  Td={Td} s  ambient={amb} °C'.format(
        amb=ambient, **plant_params))
    print('Controller: {}'.format(pid_type))
    print('Params: {} ({})'.format('overridden' if any_override else 'from {}'.format(config_source),
                                   'log' if log_plant else 'defaults'))
    print()
    for section, _ in GAIN_SECTIONS:
        if '.' in section:
            outer, inner = section.split('.')
            p = pid_params[outer][inner]
        else:
            p = pid_params[section]
        print('  {}: kp={kp}  ki={ki}  kd={kd}  kt={kt}'.format(section, **p))

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


if __name__ == '__main__':
    main()