import json
import os
import sys
import time
from pathlib import Path

import numpy as np
import torch
import tensorrt as trt
from huggingface_hub import snapshot_download

REPO=os.environ.get('CF_RUNTIME_REPO','patdev/Companion-Forge-L4-ONNX')
ROOT=Path(snapshot_download(REPO,repo_type='model',token=os.environ.get('HF_TOKEN'),local_dir='/tmp/cf-val',allow_patterns=['engines/l4-sm89/custom-ops/*.plan','plugins/tensorrt/companion_sparse_trt.py','plugins/reference/ops.py']))
sys.path.insert(0,str(ROOT/'plugins/tensorrt'));sys.path.insert(0,str(ROOT/'plugins/reference'))
import companion_sparse_trt  # noqa
import ops as ref

LOGGER=trt.Logger(trt.Logger.ERROR)
DT={trt.float32:torch.float32,trt.float16:torch.float16,trt.int32:torch.int32,trt.int64:torch.int64,trt.bool:torch.bool}
if hasattr(trt,'bfloat16'): DT[trt.bfloat16]=torch.bfloat16

class OutputAllocator(trt.IOutputAllocator):
    def __init__(self,dtype):
        trt.IOutputAllocator.__init__(self);self.dtype=dtype;self.tensor=None;self.shape=None;self.nbytes=0
    def reallocate_output(self,tensor_name,memory,size,alignment):
        self.nbytes=int(size); n=max(1,(int(size)+torch.tensor([],dtype=self.dtype).element_size()-1)//torch.tensor([],dtype=self.dtype).element_size())
        self.tensor=torch.empty(n,device='cuda',dtype=self.dtype);return int(self.tensor.data_ptr())
    def reallocate_output_async(self,tensor_name,memory,size,alignment,stream):
        return self.reallocate_output(tensor_name,memory,size,alignment)
    def notify_shape(self,tensor_name,dims):
        self.shape=tuple(int(x) for x in dims)

class Runner:
    def __init__(self,name):
        blob=(ROOT/'engines/l4-sm89/custom-ops'/f'{name}.plan').read_bytes();self.rt=trt.Runtime(LOGGER);self.eng=self.rt.deserialize_cuda_engine(blob);assert self.eng;self.ctx=self.eng.create_execution_context();self.name=name
    def run(self,feeds):
        keep=[]
        for n,x in feeds.items():
            dt=DT[self.eng.get_tensor_dtype(n)];x=x.to(dtype=dt).contiguous();feeds[n]=x;keep.append(x)
            if self.eng.get_tensor_mode(n)==trt.TensorIOMode.INPUT:
                self.ctx.set_input_shape(n,tuple(x.shape));self.ctx.set_tensor_address(n,int(x.data_ptr()))
        outs={};allocs={}
        for i in range(self.eng.num_io_tensors):
            n=self.eng.get_tensor_name(i)
            if self.eng.get_tensor_mode(n)!=trt.TensorIOMode.OUTPUT:continue
            dt=DT[self.eng.get_tensor_dtype(n)];shape=tuple(int(x) for x in self.ctx.get_tensor_shape(n))
            if any(x<0 for x in shape):
                a=OutputAllocator(dt);allocs[n]=a;self.ctx.set_output_allocator(n,a)
            else:
                y=torch.empty(shape if shape else (),device='cuda',dtype=dt);outs[n]=y;keep.append(y);self.ctx.set_tensor_address(n,int(y.data_ptr()))
        ok=self.ctx.execute_async_v3(stream_handle=torch.cuda.current_stream().cuda_stream);assert ok,f'{self.name} execute failed';torch.cuda.synchronize()
        for n,a in allocs.items():
            if a.tensor is None:raise RuntimeError(f'{self.name}/{n}: allocator not called')
            if a.shape is None:raise RuntimeError(f'{self.name}/{n}: shape notification missing')
            num=int(np.prod(a.shape)) if a.shape else 1
            outs[n]=a.tensor[:num].view(a.shape)
        return outs

def diff(a,b):
    a=a.float();b=b.float();d=(a-b).abs();return {'max_abs':float(d.max()) if d.numel() else 0.0,'mean_abs':float(d.mean()) if d.numel() else 0.0,'cos':float(torch.nn.functional.cosine_similarity(a.reshape(1,-1),b.reshape(1,-1)).item()) if d.numel() else 1.0}

def grid4():
    xyz=torch.cartesian_prod(torch.arange(4),torch.arange(4),torch.arange(4)).to(torch.int32);return torch.cat([torch.zeros((64,1),dtype=torch.int32),xyz],1).cuda()

report={}
torch.manual_seed(123)
# Subdivide
f=torch.randn(32,8,device='cuda',dtype=torch.float16);c=grid4()[:32].contiguous();r0,r1=ref.sparse_subdivide(f,c);o=Runner('SparseSubdivide').run({'feats':f,'coords':c});report['SparseSubdivide']={'feats':diff(o['out_feats'],r0),'coords_equal':bool(torch.equal(o['out_coords'],r1))}
# Downsample
f=torch.randn(64,8,device='cuda',dtype=torch.float16);c=grid4();r0,r1,ri=ref.sparse_downsample(f,c,(2,2,2));o=Runner('SparseDownsample').run({'feats':f,'coords':c});m=int(o['count'].item());report['SparseDownsample']={'count':m,'ref_count':int(r0.shape[0]),'feats':diff(o['out_feats'][:m],r0),'coords_equal':bool(torch.equal(o['out_coords'][:m],r1)),'inverse_equal':bool(torch.equal(o['inverse'],ri))}
# Upsample
f=torch.randn(8,8,device='cuda',dtype=torch.float16);tc=grid4();inv=torch.arange(64,device='cuda',dtype=torch.int32)%8;r0,r1=ref.sparse_upsample(f,tc,inv);o=Runner('SparseUpsample').run({'feats':f,'target_coords':tc,'inverse':inv});report['SparseUpsample']={'feats':diff(o['out_feats'],r0),'coords_equal':bool(torch.equal(o['out_coords'],r1))}
# Window attention
qkv=torch.randn(64,3,4,16,device='cuda',dtype=torch.float16);c=grid4();r=ref.sparse_window_attention(qkv,c,8,(0,0,0));o=Runner('SparseWindowAttention').run({'qkv':qkv,'coords':c});report['SparseWindowAttention']=diff(o['out'],r)
# Conv: create actual spconv reference and feed its native weight layout.
import spconv.pytorch as spconv
f=torch.randn(64,8,device='cuda',dtype=torch.float16);c=grid4();mod=spconv.SubMConv3d(8,16,3,bias=True,algo=spconv.ConvAlgo.Native).cuda().half().eval();st=spconv.SparseConvTensor(f,c,[64,64,64],1);ry=mod(st);w=mod.weight.detach().contiguous();b=mod.bias.detach().contiguous();print('SPCONV_WEIGHT_SHAPE',tuple(w.shape),flush=True)
try:
    o=Runner('SparseConv3D').run({'feats':f,'coords':c,'weight':w,'bias':b});m=int(o['count'].item());report['SparseConv3D']={'count':m,'ref_count':int(ry.features.shape[0]),'feats':diff(o['out_feats'][:m],ry.features),'coords_equal':bool(torch.equal(o['out_coords'][:m],ry.indices.to(torch.int32)))}
except Exception as e:
    report['SparseConv3D']={'error':repr(e),'weight_shape':tuple(w.shape)}
# Latency quick pass for successful fixed-shape ops
for name in ['SparseSubdivide','SparseUpsample','SparseWindowAttention']:
    pass
print('REPORT',json.dumps(report,indent=2),flush=True)
Path('/tmp/custom_validation.json').write_text(json.dumps(report,indent=2))