package rlpark.plugin.rltoys.algorithms.functions.policydistributions.structures;

import java.util.Random;

import rlpark.plugin.rltoys.algorithms.functions.policydistributions.BoundedPdf;
import rlpark.plugin.rltoys.algorithms.functions.policydistributions.PolicyDistribution;
import rlpark.plugin.rltoys.envio.actions.Action;
import rlpark.plugin.rltoys.envio.actions.ActionArray;
import rlpark.plugin.rltoys.envio.actions.Actions;
import rlpark.plugin.rltoys.math.vector.RealVector;
import rlpark.plugin.rltoys.math.vector.implementations.PVector;
import zephyr.plugin.core.api.monitoring.abstracts.LabeledCollection;
import zephyr.plugin.core.api.monitoring.annotations.Monitor;

@Monitor
public abstract class AbstractNormalDistribution implements PolicyDistribution, LabeledCollection, BoundedPdf {
  private static final long serialVersionUID = -6707070542157254303L;
  @Monitor(level = 4)
  protected PVector u_mean;
  @Monitor(level = 4)
  protected PVector u_stddev;
  protected double mean = 0;
  protected double stddev = 1;
  protected final Random random;
  public double a_t;
  protected double meanStep;
  protected double stddevStep;

  protected RealVector x;

  public AbstractNormalDistribution(Random random) {
    this.random = random;
  }

  @Override
  public PVector[] createParameters(int nbFeatures) {
    u_mean = new PVector(nbFeatures);
    u_stddev = new PVector(nbFeatures);
    return new PVector[] { u_mean, u_stddev };
  }

  public double stddev() {
    return stddev;
  }

  public double mean() {
    return mean;
  }

  @Override
  final public void update(RealVector x) {
    this.x = x.copy();
    updateDistribution();
  }

  abstract protected void updateDistribution();

  @Override
  public double pi(Action a) {
    assert Actions.isOneDimension(a);
    return pi_s(ActionArray.toDouble(a));
  }

  public abstract double pi_s(double a);

  @Override
  public String label(int index) {
    return index == 0 ? "mean" : "stddev";
  }

  @Override
  public int nbParameterVectors() {
    return 2;
  }
}

@@ -541,24 +541,156 @@
 ns.PVector;%0A
+import zephyr.plugin.core.api.internal.monitoring.wrappers.Abs;%0Aimport zephyr.plugin.core.api.internal.monitoring.wrappers.Squared;%0A
 import zephy
@@ -819,16 +819,49 @@
 Monitor%0A
+@SuppressWarnings(%22restriction%22)%0A
 public a
@@ -1256,34 +1256,126 @@
 ;%0A  
-protected double meanStep;
+@Monitor(wrappers = %7B Squared.ID, Abs.ID %7D)%0A  protected double meanStep;%0A  @Monitor(wrappers = %7B Squared.ID, Abs.ID %7D)
 %0A  p

Add monitor wrappers