# Mushroom body inference for the M5Stack (pure Python: runs in MicroPython and CPython).
#
# The model file comes from experiments/export_for_m5stack.py. Feature calculation here
# must match code/enose/features.py, and the network must match enose/mushroom.py;
# tests/test_firmware_inference.py checks both on the laptop.
import json
import math

WINDOW_SECONDS = 60


def tail_mean(rows, key, seconds=WINDOW_SECONDS):
    """rows: list of dicts with 'second' and key; mean of the last `seconds` of valid values."""
    if not rows:
        return None
    end = rows[-1]["second"]
    values = [row[key] for row in rows if row["second"] > end - seconds and row[key] is not None]
    return sum(values) / len(values) if values else None


def relative_change(steady, baseline):
    if steady is None or baseline is None or baseline == 0:
        return None
    return (steady - baseline) / abs(baseline)


def features(purge_rows, sample_rows, heater_steps):
    """Same 8 odour features as enose.features.measurement_features, in the same order."""
    result = []
    for step in range(heater_steps):
        key = "gas{}".format(step)
        result.append(relative_change(tail_mean(sample_rows, key), tail_mean(purge_rows, key)))
    result.append(relative_change(tail_mean(sample_rows, "ethanol"), tail_mean(purge_rows, "ethanol")))
    result.append(relative_change(tail_mean(sample_rows, "h2"), tail_mean(purge_rows, "h2")))
    steady, baseline = tail_mean(sample_rows, "co2"), tail_mean(purge_rows, "co2")
    result.append(None if steady is None or baseline is None else (steady - baseline) / 1000.0)
    return result


class MushroomBody:
    def __init__(self, path):
        with open(path) as handle:
            model = json.load(handle)
        self.classes = model["classes"]
        self.scale = model["scale"]
        self.normalize = model["normalize"]
        self.active = model["active"]
        self.kc_inputs = model["kc_inputs"]
        # Random wiring is all 1.0, so the file leaves the weights out to save memory.
        self.kc_weights = model.get("kc_weights") or [[1.0] * len(inputs) for inputs in self.kc_inputs]
        self.weights = model["weights"]
        self.temperature = model["temperature"]
        # Recording schedule the training data was taken with; the device must use the same one.
        self.schedule = model["schedule"]
        self.intensity_centres = model.get("intensity_centres") or []
        self.intensity_width = model.get("intensity_width", 1.0)
        self.intensity_gain = model.get("intensity_gain", 1.0)
        self.adapt_mean = model.get("adapt_mean") or []
        self.adapt_std = model.get("adapt_std") or []

    @staticmethod
    def on_off(values):
        return [max(value, 0.0) for value in values] + [max(-value, 0.0) for value in values]

    def projection_neurons(self, feature_values):
        """Same steps as enose/mushroom.py: gain, normalisation, intensity, adaptation, ON/OFF."""
        values = [0.0 if value is None else value for value in feature_values]
        scaled = [value / scale for value, scale in zip(values, self.scale)]
        total = sum(abs(value) for value in scaled)
        pattern = [value / (total + 1e-9) for value in scaled] if self.normalize else scaled
        bumps = []
        if self.intensity_centres:
            log_total = math.log(total + 1e-6)
            size = self.intensity_gain / (2 * len(self.scale))
            bumps = [math.exp(-0.5 * ((log_total - centre) / self.intensity_width) ** 2) * size
                     for centre in self.intensity_centres]
        if self.adapt_mean:
            signal = pattern + bumps
            return self.on_off([(value - mean) / std for value, mean, std in zip(signal, self.adapt_mean, self.adapt_std)])
        return self.on_off(pattern) + bumps

    def predict(self, feature_values):
        """Return (class name, confidence 0-1, all scores)."""
        pn = self.projection_neurons(feature_values)
        drive = []
        for inputs, weights in zip(self.kc_inputs, self.kc_weights):
            drive.append(round(sum(pn[index] * weight for index, weight in zip(inputs, weights)), 9))
        # sorted() is stable: among equal drives the lower KC index wins, same as the laptop.
        winners = sorted(range(len(drive)), key=lambda kc: -drive[kc])[:self.active]
        scores = [sum(row[kc] for kc in winners) / self.active for row in self.weights]
        biggest = max(abs(score) for score in scores) + 1e-9
        exps = [pow(2.718281828, (score / biggest - 1) / self.temperature) for score in scores]
        best = scores.index(max(scores))
        return self.classes[best], exps[best] / sum(exps), scores