import numpy as np
import torch

META = {
    'name': 'piola_surface_operator',
    'domain': 'pde',
    'description': 'Manufactured vector-valued surface operator on a fixed reference chart with affine deformations and contravariant Piola transport.'
}

N = 6
q = np.linspace(0.08, 0.92, N).astype(np.float32)
xx, yy = np.meshgrid(q, q, indexing='ij')
X = np.stack([xx.ravel(), yy.ravel()], axis=1)
M = len(X)
DSX = np.pi*np.cos(np.pi*X[:,0])*np.sin(np.pi*X[:,1])
DSY = np.pi*np.sin(np.pi*X[:,0])*np.cos(np.pi*X[:,1])


def geometry(y):
    a, b, c = [float(v) for v in y]
    F = np.zeros((M, 3, 2), dtype=np.float32)
    F[:,0,0] = 1 + a*DSX; F[:,0,1] = a*DSY
    F[:,1,0] = b*DSX; F[:,1,1] = 1 + b*DSY
    F[:,2,0] = c*DSX; F[:,2,1] = c*DSY
    G = np.einsum('nki,nkj->nij', F, F)
    J = np.sqrt(np.maximum(np.linalg.det(G), 1e-8)).astype(np.float32)
    return F, J


def piola(F, J, u):
    return np.einsum('nki,ni->nk', F, u) / J[:, None]


def inverse_piola(F, J, v):
    G = np.einsum('nki,nkj->nij', F, F)
    rhs = np.einsum('nki,nk->ni', F, J[:, None] * v)
    return np.einsum('nij,nj->ni', np.linalg.inv(G), rhs)


def _sample(rng, n, physical):
    ys = rng.uniform(-0.35, 0.35, (n, 3)).astype(np.float32)
    out = []
    for y in ys:
        F, J = geometry(y)
        u = rng.normal(size=(M, 2)).astype(np.float32)
        u += 0.45*np.stack([np.sin(2*np.pi*X[:,0]), np.cos(2*np.pi*X[:,1])], axis=1)
        target = 0.7*u + 0.3*u.mean(axis=0, keepdims=True)
        inpvec = piola(F, J, u) if physical else np.concatenate([u, np.zeros((M, 1), dtype=np.float32)], axis=1)
        # Same architecture and tensor shape; baseline sees physical components,
        # idea sees the inverse-Piola fixed-reference components.
        geomfeat = np.broadcast_to(y, (M, 3))
        meanfeat = np.broadcast_to(inpvec.mean(axis=0), (M, 3))
        feats = np.concatenate([X, geomfeat, inpvec, meanfeat], axis=1)
        out.append((feats.reshape(-1), target.reshape(-1)))
    return np.stack([z[0] for z in out]).astype(np.float32), np.stack([z[1] for z in out]).astype(np.float32)



def math_check():
    y=np.array([0.18,-0.13,0.16], dtype=np.float32)
    F,J=geometry(y)
    u=np.stack([X[:,0]**2+X[:,1], X[:,0]-X[:,1]**2], axis=1).astype(np.float32)
    v=np.stack([np.sin(2*np.pi*X[:,0]), np.cos(2*np.pi*X[:,1])], axis=1).astype(np.float32)
    lhs=float(np.mean(np.sum(piola(F,J,u)*piola(F,J,v),axis=1)*J))
    pu, pv = np.einsum('nki,ni->nk',F,u), np.einsum('nki,ni->nk',F,v)
    rhs=float(np.mean(np.sum(pu*pv,axis=1)/J))
    return {'area_cancellation_relative_error': abs(lhs-rhs)/(abs(rhs)+1e-12),
            'J_min': float(J.min()), 'J_max': float(J.max())}

def get_dataset(seed, n_train, n_test, physical=True):
    xtr, ytr = _sample(np.random.RandomState(seed), n_train, physical)
    xte, yte = _sample(np.random.RandomState(seed + 5000), n_test, physical)
    return {'xtr': torch.from_numpy(xtr), 'ytr': torch.from_numpy(ytr),
            'xte': torch.from_numpy(xte), 'yte': torch.from_numpy(yte),
            'task': 'regression', 'metric': 'mse', 'input_shape': (xtr.shape[1],),
            'out_dim': ytr.shape[1]}