#!/usr/bin/env python3
"""Replay a BrewPi log through the temp controller with candidate PID params.

Feeds real temp_ist readings from a log back into a fresh TempController
(Normal mode) tick by tick and compares the resulting power output to the
original logged values.  Because the real sensor measurements are injected at
each tick, no plant model is needed: the replay shows exactly how a different
set of PID gains would have responded to the same temperature trajectory.

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> --config alt_config.json

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

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


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 — used to normalise power_set."""
    return max((s['power_eff'] for s in samples), default=1.0) or 1.0


def _infer_heatrate_soll_set(samples):
    """Heat-rate setpoint from the log: max rate_soll seen while power > 0.

    rate_soll = heatrate_soll_set * pid_hold.get_y(); when the hold PID is
    saturated at 1.0 (full-power heating) the two are equal, so the maximum
    over all heating samples is a good estimate of heatrate_soll_set."""
    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, dt):
    tc = TempController(dt)
    tc.set_params(pid_params)
    tc.set_enabled(True)

    out = []
    for s in samples:
        tc.set_theta_ist(s['temp_ist'])
        tc.set_theta_soll(s['temp_soll'])
        tc.set_heatrate_soll(heatrate_soll_set)
        tc.process()
        out.append({
            '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), fontsize='small')

    # Temperature: identical in both runs (real sensor readings are injected)
    ax_temp.plot(t, [s['temp_ist']  for s in samples], '-b', linewidth=1, label='theta_ist')
    ax_temp.plot(t, [s['temp_soll'] for s in samples], '-r', linewidth=1, label='theta_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 (log)')
    ax_rate.plot(t, [r['heatrate_ist'] for r in replay],     '--',  color='C1', linewidth=1, label='rate_ist (replay)')
    ax_rate.plot(t, [s['rate_soll'] for s in samples],       '-r',  linewidth=1, alpha=0.4, label='rate_soll (log)')
    ax_rate.set_ylabel('Rate [K/min]')
    ax_rate.legend(loc='upper left', fontsize='small')

    # Power (both normalised to 0–1)
    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 (log, norm.)')
    ax_power.plot(t, replay_pwr, '-', color='C1', linewidth=1, label='power ({})'.format(params_label))
    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] fed to the replay controller '
                             '(default: inferred from log as max rate_soll while heating)')

    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)

    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('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, dt)
    plot_replay(samples, replay, name, max_power, params_label)
    plt.show()


if __name__ == '__main__':
    main()