"""Orthographic geometry metrics; silhouettes do not depend on object colour."""
import numpy as np
from PIL import Image, ImageDraw
from scipy.spatial import cKDTree
from scipy.spatial.transform import Rotation

VIEWS = ((45,15),(135,15),(225,15),(315,15))

def rotation(yaw, pitch=0):
    y,x=np.deg2rad([yaw,pitch])
    ry=np.array([[np.cos(y),0,np.sin(y)],[0,1,0],[-np.sin(y),0,np.cos(y)]])
    rx=np.array([[1,0,0],[0,np.cos(x),-np.sin(x)],[0,np.sin(x),np.cos(x)]])
    return (rx@ry).astype(np.float32)

def rotate_splats(splats, matrix, scale=1.):
    from image3d_student.splats import Splats
    q=splats.quat[:,[1,2,3,0]]
    out=(Rotation.from_matrix(matrix)*Rotation.from_quat(q)).as_quat()[:,[3,0,1,2]]
    return Splats(splats.xyz@matrix.T*scale,splats.rgb,splats.scale*scale,out,splats.opacity)

def chamfer(pred, reference):
    if not len(pred):
        return 2.0 # fixed maximum-error sentinel, also report empty=True
    # Symmetric mean Euclidean nearest-neighbour distance, in normalized object units.
    return float((cKDTree(reference).query(pred)[0].mean()+cKDTree(pred).query(reference)[0].mean())/2)

def mesh_mask(vertices, faces, view, size=96):
    points=vertices@view.T
    points=points[:,:2]*np.array([1,-1])*(size/1.25)+size/2
    image=Image.new("L",(size,size))
    draw=ImageDraw.Draw(image)
    for triangle in points[faces]:
        draw.polygon([tuple(v) for v in triangle],fill=255)
    return np.asarray(image)>0

def splat_mask(splats, view, size=96):
    from image3d_student.splats import rotation_matrices
    if splats is None:
        return np.zeros((size,size),bool)
    xyz=splats.xyz@view.T
    rot=rotation_matrices(splats.quat)
    cov=(rot*(splats.scale**2)[:,None,:])@rot.transpose(0,2,1)
    cov=view[None]@cov@view.T[None]
    alpha=np.zeros((size,size),np.float32)
    flip=np.diag([1.,-1.])
    for i in range(len(xyz)):
        center=xyz[i,:2]*np.array([1,-1])*(size/1.25)+size/2
        c=flip@cov[i,:2,:2]@flip*(size/1.25)**2+np.eye(2)*.3
        radius=min(size,int(np.ceil(3*np.sqrt(np.linalg.eigvalsh(c).max()))))
        lo=np.maximum(0,np.floor(center-radius).astype(int))
        hi=np.minimum(size,np.ceil(center+radius+1).astype(int))
        if np.any(hi<=lo): continue
        yy,xx=np.mgrid[lo[1]:hi[1],lo[0]:hi[0]]
        delta=np.stack([xx-center[0],yy-center[1]],axis=-1)
        a=(splats.opacity[i,0]*np.exp(-.5*np.einsum("...i,ij,...j->...",delta,np.linalg.inv(c),delta))).clip(0,.99)
        patch=alpha[lo[1]:hi[1],lo[0]:hi[0]]
        patch[:]=1-(1-patch)*(1-a)
    return alpha>=.5

def iou(a,b):
    return float(np.logical_and(a,b).sum()/max(1,np.logical_or(a,b).sum()))