"""Structurally matched custom track: classify unoriented line orientation."""
import numpy as np
META = {"name":"oriented_line_quotient", "domain":"vision_orientation", "description":"Small grayscale/RGB images containing apolar lines; classify their half-circle orientation bin."}

def get_dataset(seed, n_train, n_test):
    def make(n, s):
        rng=np.random.RandomState(s); H=W=16
        x=np.zeros((n,3,H,W),np.float32); y=np.empty(n,np.int64)
        yy,xx=np.mgrid[0:H,0:W].astype(np.float32); cx=cy=7.5
        for i in range(n):
            a=rng.uniform(0,np.pi); u=np.array([np.cos(a),np.sin(a)])
            # distance to an unoriented line through the center
            d=-(xx-cx)*u[1]+(yy-cy)*u[0]
            along=(xx-cx)*u[0]+(yy-cy)*u[1]
            img=np.exp(-(d*d)/1.1)*np.exp(-(along*along)/90.0)
            img += rng.normal(0,.035,(H,W)).astype(np.float32)
            img=np.clip(img,0,1)
            x[i]=img[None,:,:]
            x[i,1]=img; x[i,2]=img
            y[i]=int((a/np.pi*4)%4)
        return x,y
    xtr,ytr=make(n_train,seed); xte,yte=make(n_test,seed+5000)
    return {"xtr":xtr,"ytr":ytr,"xte":xte,"yte":yte,"task":"classification","metric":"err","out_dim":4}