#!/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.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, heater_efficiency, 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 * heater_efficiency)
        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'

    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]')
    parser.add_argument('--heater-efficiency', type=float, default=1.0, metavar='0-1',
                        help='Fraction of rated heater power delivered to the plant (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)
    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']
    name       = log.get('Name', log_path.stem)
    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)

    any_override = any(
        getattr(args, '{}_{}'.format(p, g)) is not None
        for p in ('hold', 'heat', 'cool')
        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={:.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:.0f} J/kgK  L={L} W/kgK  Td={Td} s  ambient={amb} °C  heater_efficiency={eff}'.format(
        amb=ambient, eff=args.heater_efficiency, **plant_params))
    print('Params: {} ({})'.format('overridden' if any_override else 'from {}'.format(config_source),
                                   'log' if log_plant else 'defaults'))
    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, args.heater_efficiency, dt)
    plot_replay(samples, replay, name, max_power, params_label)
    plt.show()


if __name__ == '__main__':
    main()