from components.plant import Pot
from components.pid import PidFactory
from components.sud import Sud, SudState

# Safety cap so a schedule whose target a step can never actually reach
# (e.g. a "hold" colder than ambient with no active cooling) can't hang the
# simulation forever - the estimate is simply cut off there.
MAX_TICKS = 200000

# Mirrors tasks/heater.py's own pulse_period_s - HeaterTask duty-cycles its
# device's discrete power steps over a rolling window this many *real*
# seconds wide, regardless of warp factor (see estimate()'s actuate()
# closure below).
PULSE_PERIOD_S = 10


def _next_smaller_power(powers, power):
	"""Mirrors HeaterTask.get_next_smaller_power() (tasks/heater.py) -
	duplicated rather than imported (components/ has no business depending
	on tasks/) - keep the two in sync if HeaterTask's PWM logic changes."""
	p_list = [p for p in powers if p <= power]
	return p_list[-1] if p_list else powers[0]


def _next_greater_power(powers, power):
	"""Mirrors HeaterTask.get_next_greater_power() (tasks/heater.py) - see
	_next_smaller_power()."""
	p_list = [p for p in powers if p > power]
	return p_list[0] if p_list else powers[0]


class SudForecastEstimator:
	"""Predicts how long a Sud schedule will actually take by simulating it
	with the same machinery (and params) the real server's brewpi.py wires
	up - a fresh Pot and temperature controller of the configured pid_type,
	driven through the schedule exactly as tasks/sud.py's SudTask would.

	This is deliberately independent of wall-clock/asyncio time: it just
	iterates dt-sized ticks as fast as the CPU allows (a multi-hour brew
	simulates in well under a second), so it can be run synchronously
	whenever a client needs an estimate - the naive "abs(delta)/rate" model
	the GUI used to compute itself has no way to see the real PID cascade's
	spin-up/settling lag, which is exactly why its estimate drifted so far
	from reality (see README.md's "Forecast vs. actual duration").

	estimate() also duty-cycles the PID's continuous output across the
	heater's own discrete power steps (see its actuate() closure), exactly
	like the real run's HeaterTask/device chain does - feeding tc.get_power()
	straight to the plant instead lets pid_heat's own oscillatory tendency
	reach the plant undamped, producing a forecast far more jagged than any
	real run actually is (the discrete steps end up duty-cycling it back
	down to something close to the commanded average)."""

	def __init__(self, dt, theta_amb, pid_type, tempctrl_params, heater_powers, sim_warp_factor):
		self.dt = dt
		self.theta_amb = theta_amb
		self.pid_type = pid_type
		self.tempctrl_params = tempctrl_params
		# The heater's own discrete output levels (AHeater.get_powers()),
		# not just its max - duty-cycling needs the actual steps to pick the
		# pair straddling the commanded power.
		self.heater_powers = heater_powers
		self.sim_warp_factor = sim_warp_factor

	def set_ambient_temperature(self, theta_amb):
		self.theta_amb = theta_amb

	def estimate(self, doc, start_theta=None):
		"""Returns (t, theta, final_state, step_starts): t/theta are
		parallel lists of elapsed simulated seconds and temperature,
		covering doc['steps'] from the start all the way to the end
		(final_state is SudState.DONE), or, in the pathological case of
		a step whose target can never actually be reached, wherever
		MAX_TICKS cut the simulation off. step_starts maps each step's
		index in doc['steps'] (0-based, local to this doc - the caller
		rebases onto its own absolute schedule indices/timeline, same as
		it does for t/theta themselves) to the simulated second it began
		at - consulted by tasks/sud.py's SudTask to show each step's
		predicted total/remaining duration (the GUI's Progress tab).

		A step requiring user confirmation doesn't stop the simulation
		either - a human's response time genuinely can't be forecast,
		so it's modeled as zero delay (auto-confirmed the instant that
		step's hold completes) rather than leaving the estimate stuck
		there forever. That optimistic guess gets corrected for real once
		the caller (tasks/sud.py's SudTask) sees the schedule actually
		reach the next step boundary - see SudTask._reanchor_forecast(),
		called from on_step_changed() on every real transition, not just
		confirmations.

		start_theta defaults to the configured ambient temperature - i.e.
		a cold start, same as the GUI's static estimate."""
		if start_theta is None:
			start_theta = self.theta_amb

		sud = Sud()
		if not sud.load(doc) or not sud.schedule:
			return [0.0], [start_theta], SudState.DONE, {}

		# Plant params (M/C/L/Td) are deliberately *not* seeded here from
		# any default - sud.start() below synchronously fires the first
		# step's on_step_changed callback (assigning Sud.step triggers
		# its callbacks immediately, before sud.start() even returns),
		# which sets them from this doc's own derive_plant_params() - see
		# there - before any pot.process()/tc.process() tick ever runs.
		pot = Pot(self.dt)
		pot.set_ambient_temperature(self.theta_amb)
		pot.initial(start_theta)
		tc = PidFactory.create(self.pid_type, self.dt)
		tc.set_params(self.tempctrl_params)
		tc.set_ambient_temperature(self.theta_amb)
		tc.set_enabled(True)
		tc.set_theta_ist(pot.get_temperature())
		# Seed the target at start_theta - this tc is a fresh, throwaway
		# instance (unlike the real run's persistent one), so without this
		# its theta_soll_set defaults to 0 until a step pushes its own.
		# Steps without their own 'temperature' (common now that ramping
		# isn't gated by a 'ramp' key - see components/sud.py) rely on
		# inheriting whatever target was already running, which for a
		# schedule starting mid-brew (the dynamic remaining forecast) is
		# start_theta, not 0 - without this, such a schedule's first step
		# would have the simulated controller chase 0 degrees indefinitely,
		# hitting MAX_TICKS and producing a needlessly huge result.
		tc.set_theta_soll(start_theta)

		# Ticks here are dt simulated seconds each, same as HeaterTask's own
		# loop (one DT_TASK real seconds == dt simulated seconds, by warp
		# factor's definition) - so the PWM period, in ticks, is the same
		# pulse_period_s * sim_warp_factor regardless of dt.
		pulse_period_count = PULSE_PERIOD_S * self.sim_warp_factor
		pulse_counter = 0

		def actuate(y):
			"""Duty-cycles y (tc.get_power(), in -1..1) across self.
			heater_powers exactly like HeaterTask.on_process()'s loop does
			with power_actor - see SudForecastEstimator's own docstring."""
			nonlocal pulse_counter
			power = max(0, self.heater_powers[-1] * y)
			power_low = _next_smaller_power(self.heater_powers, power)
			power_high = _next_greater_power(self.heater_powers, power)
			power_step = power_high - power_low
			duty = (power - power_low) / power_step if power_step else 0.0
			on_count = pulse_period_count * duty
			pulse_counter += 1
			if pulse_counter >= pulse_period_count:
				pulse_counter = 0
			return power_low if (power == 0 or pulse_counter >= on_count) else power_high

		step_starts = {}

		def on_step_changed(step):
			if step is None:
				return
			# setdefault: this callback also re-fires for the same step's
			# ramp->hold phase switch (components/sud.py's temp_reached()
			# re-assigns self.step to retrigger it) - only the *first* call
			# for a given index is its actual start.
			step_starts.setdefault(sud.index, t[-1])
			step_pot = step.get('pot', {})
			params = sud.derive_plant_params(step_pot.get('grain_mass', 0), step_pot.get('water_mass', 0))
			pot.set_plant_params(params)
			tc.set_model_plant_params(params)
			if sud.state == SudState.RAMPING and step['temperature'] is not None:
				tc.set_theta_soll(step['temperature'])
				tc.set_heatrate_soll(step['ramp']['rate'])
		sud.set_on_changed('step', on_step_changed)

		t = [0.0]
		theta = [pot.get_temperature()]

		sud.start()
		ticks = 0
		while sud.state != SudState.DONE and ticks < MAX_TICKS:
			if sud.state == SudState.WAIT_USER:
				sud.confirm()
				continue

			pot.process()
			tc.set_theta_ist(pot.get_temperature())
			tc.process()
			power = actuate(tc.get_power())
			pot.set_power(power)
			tc.set_model_power(power)

			if sud.state == SudState.RAMPING:
				if tc.is_holding():
					sud.temp_reached()
			sud.tick(self.dt)

			t.append(t[-1] + self.dt)
			theta.append(pot.get_temperature())
			ticks += 1

		return t, theta, sud.state, step_starts
