#!/usr/bin/env python3
"""Offline analysis of a log (and optional forecast) produced by the BrewPi server.

Usage:
    python analyze_log.py <date_time> <sud_name> [--log-dir <dir>]
    python analyze_log.py <path/to/log_*.json>

The two-argument form loads logs/<date_time>_<sud_name>.json; the
single-argument form accepts a direct path to any log_*.json (e.g. a
server-session log without a Sud name). The forecast-vs-actual panel is
shown only for a Sud log (one carrying a 'Doc', 'Config' and 'HeaterPowers'
- see tasks/sud_log.py), by re-simulating that run's own schedule with
SudForecastEstimator, the same as the live server does at Play.

Renders the 3-panel realtime strip chart and, when available, the
forecast-vs-actual comparison.
"""

import argparse
import bisect
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.sud import SudState
from components.sud_forecast import SudForecastEstimator


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_realtime(samples, name):
    """3-panel strip chart mirroring RealtimePlot in the PyQt GUI."""
    t          = [s['t']         for s in samples]
    temp_ist   = [s['temp_ist']  for s in samples]
    temp_soll  = [s['temp_soll'] for s in samples]
    rate_ist   = [s['rate_ist']  for s in samples]
    rate_soll  = [s['rate_soll'] for s in samples]
    power_set  = [s['power_set'] for s in samples]
    power_eff  = [s['power_eff'] for s in samples]

    fig, (ax_temp, ax_rate, ax_power) = plt.subplots(3, 1, sharex=True, facecolor='white')
    fig.suptitle(name, fontsize='small')

    ax_temp.plot(t, temp_ist,  '-b', linewidth=1, label='theta_ist')
    ax_temp.plot(t, temp_soll, '-r', linewidth=1, label='theta_soll')
    ax_temp.set_ylabel('Temp [°C]')
    ax_temp.legend(loc='upper left', fontsize='small')

    ax_rate.plot(t, rate_ist,  '-b', linewidth=1, label='heatrate_ist')
    ax_rate.plot(t, rate_soll, '-r', linewidth=1, label='heatrate_soll')
    ax_rate.set_ylabel('Rate [K/min]')
    ax_rate.legend(loc='upper left', fontsize='small')

    ax_power.plot(t, power_set, '-r', linewidth=1, label='power_set')
    ax_power.plot(t, power_eff, '-b', linewidth=1, label='power_eff')
    ax_power.set_ylabel('Power [W]')
    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 plot_forecast(forecast, samples, name):
    """Forecast-vs-actual comparison mirroring SudForecastPlot in the PyQt GUI."""
    t_min_fc = [t / 60.0 for t in forecast['T']]
    theta_fc  = forecast['Theta']
    finished  = forecast['Finished']

    t_min_actual = [s['t'] / 60.0  for s in samples]
    theta_actual  = [s['temp_ist'] for s in samples]

    fig, ax = plt.subplots(1, 1, facecolor='white')

    ax.plot(t_min_actual, theta_actual, '-b', linewidth=1.5, zorder=2, label='actual')
    ax.plot(t_min_fc,    theta_fc,      '-b', linewidth=1.5, alpha=0.5, zorder=2, label='forecast')

    total = t_min_fc[-1] if t_min_fc else 0.0
    suffix = '' if finished else ' so far'
    ax.set_title('{} - {:.0f} min{} total'.format(name, total, suffix), fontsize='small')
    ax.set_xlabel('t [min]')
    ax.set_ylabel('Temp [°C]')
    ax.legend(loc='upper left', fontsize='small')
    _style_axis(ax)

    fig.tight_layout()
    return fig


def build_forecast(log):
    """Re-simulates the schedule a Sud log's own 'Doc' recorded, the same
    way the live server does at Play (see components/sud_forecast.py) -
    returns a {'T', 'Theta', 'Finished'} dict matching plot_forecast()'s
    expected shape, or None if log isn't a Sud log (no 'Doc'/'Config'/
    'HeaterPowers' - e.g. a server-session log with no run playing).

    Replays every real _reanchor_forecast() trigger recorded in the log's
    'ForecastAnchors' (see tasks/sud.py's SudTask.set_on_reanchor()), in
    order, splicing a freshly anchored simulation of the remaining
    schedule onto the truncated forecast at each one - the exact same
    correction a live client's forecast plot goes through as a run
    progresses. Without this, the single cold-start estimate() below is
    all a client ever saw at Load - it can't have anticipated things like
    a WAIT_USER confirmation's real timing (estimate() can only ever model
    that as zero-delay - see its own docstring), so a plot comparing it
    straight against Samples would show a small, spurious gap at every
    such step that isn't actually a forecast error."""
    if 'Doc' not in log or 'Config' not in log or 'HeaterPowers' not in log:
        return None
    doc = log['Doc']
    config = log['Config']
    tempctrl = config['TempCtrl']
    # Old logs predate 'SimWarpFactor' being recorded - 1.0 (real time) is
    # the right fallback for those, since it's also what every real-hardware
    # run already used (see server/brewpi.py's plant_sim check).
    estimator = SudForecastEstimator(
        log['dt'], config['ambient_temperature'], tempctrl['pid_type'], tempctrl,
        log['HeaterPowers'], log.get('SimWarpFactor', 1.0), config.get('Pot', {}))

    t, theta, state, step_starts = estimator.estimate(doc)
    finished = (state == SudState.DONE)

    for anchor in log.get('ForecastAnchors', []):
        real_elapsed = anchor['t']
        index = anchor['index']
        if not (0 <= index < len(doc['steps'])):
            continue
        # Cut at step_starts[index], the old forecast's own (possibly very
        # wrong) belief of where step `index` begins - not at real_elapsed
        # directly: a step whose real timing blew way past its zero-delay
        # guess (a long WAIT_USER confirm, above all) leaves the old
        # forecast's entire speculative remainder sitting at timestamps
        # still numerically less than real_elapsed, so bisecting against
        # real_elapsed itself would find nothing to discard and just tack
        # the fresh, correct simulation on after it - see tasks/sud.py's
        # SudTask._reanchor_forecast(), which this mirrors exactly.
        cut = bisect.bisect_right(t, step_starts.get(index, real_elapsed))
        t, theta = t[:cut], theta[:cut]
        step_starts = {i: tt for i, tt in step_starts.items() if i < index}
        sub_doc = {**doc, 'steps': doc['steps'][index:]}
        sub_t, sub_theta, sub_state, sub_step_starts = estimator.estimate(sub_doc, anchor['theta_ist'])
        if t and t[-1] < real_elapsed:
            t.append(real_elapsed)
            theta.append(anchor['theta_ist'])
        t.extend(real_elapsed + seconds for seconds in sub_t)
        theta.extend(sub_theta)
        step_starts.update((index + local_index, real_elapsed + local_t)
                            for local_index, local_t in sub_step_starts.items())
        finished = (sub_state == SudState.DONE)

    return {'T': t, 'Theta': theta, 'Finished': finished}


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('date_time_or_path',
                        help='Timestamp part of the filename (e.g. 20260628T184903) '
                             'or a direct path to a log_*.json file')
    parser.add_argument('sud_name', nargs='?', default=None,
                        help='Sud name part of the filename (e.g. Sud-0010); '
                             'omit when passing a direct file path')
    default_log_dir = Path(__file__).parent.parent / 'logs'
    parser.add_argument('--log-dir', default=str(default_log_dir),
                        help='Directory containing the log files (default: %(default)s)')
    args = parser.parse_args()

    if args.sud_name is None:
        log_path = Path(args.date_time_or_path)
        stem = log_path.stem[len('log_'):] if log_path.stem.startswith('log_') else log_path.stem
    else:
        log_dir = Path(args.log_dir)
        stem = '{}_{}'.format(args.date_time_or_path, args.sud_name)
        log_path = log_dir / 'log_{}.json'.format(stem)

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

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

    name = log.get('Name', stem)
    samples = log['Samples']

    plot_realtime(samples, name)
    forecast = build_forecast(log)
    if forecast is not None:
        plot_forecast(forecast, samples, name)

    plt.show()


if __name__ == '__main__':
    main()