"""Causal sampled control track for inverse-gain structured distillation."""
import numpy as np

META = {
    "name": "inverse_gain_control",
    "domain": "dynamics",
    "description": "Causal windows of regulated state, velocity, and reference; target is privileged inverse-gain expert action."
}

def get_dataset(seed, n_train=400, n_test=400):
    def make(n, offset):
        rng = np.random.RandomState(seed + offset)
        dt, k1 = 0.05, 1.5
        X = np.empty((n, 24), np.float32); Y = np.empty((n, 1), np.float32)
        for i in range(n):
            a = rng.choice([0.25, 0.5, 1.0, 2.0])
            d = rng.normal(0, 0.35)
            x1, x2 = rng.normal(0, .3), rng.normal(0, .2)
            hist = []
            for k in range(8):
                r = 1.0*np.sin(.045*k + rng.uniform(0, 6.28)) + .25*np.sin(.13*k)
                prev = x1
                # sampled plant: x1 evolves under unknown gain and disturbance
                x1 = x1 + dt*(a*x2 + d)
                delta = (x1-prev)/dt
                e = r-x1
                expert = x2 + (k1*e-delta)/a
                hist.append([x1, x2, r])
                # low-level actuator tracks the expert reference
                x2 = x2 + dt*3.0*(expert-x2)
            X[i] = np.asarray(hist, np.float32).reshape(-1)
            Y[i, 0] = expert
        return X, Y
    xtr, ytr = make(n_train, 0); xte, yte = make(n_test, 5000)
    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
            "task": "regression", "metric": "mse"}