import json, os, time
import torch
import emt_lab as emt

def run(ctx):
    assert torch.cuda.is_available() and str(ctx.device).startswith('cuda'), 'Real CUDA execution required'
    # Keep initial claims occupied briefly so both nodes exercise their GPU workers.
    time.sleep(45)
    n=256
    scale=int(ctx.params['scale'])
    a=torch.arange(n*n,device=ctx.device,dtype=torch.float64).reshape(n,n)
    identity=torch.eye(n,device=ctx.device,dtype=torch.float64)*scale
    result=a@identity
    torch.cuda.synchronize()
    actual=float(result.sum().item())
    expected=scale*(n*n)*(n*n-1)/2
    max_error=float((result-scale*a).abs().max().item())
    assert max_error == 0 and actual == expected
    with open(ctx.output('matrix-result.json'),'w') as stream:
        json.dump({'n':n,'scale':scale,'input_definition':'A[i,j]=i*n+j; B=s*I',
                   'result':result.detach().cpu().tolist(),'sum':actual,'expected_sum':expected},stream)
    ctx.explain('Matrix multiplication on CUDA','C = A × (sI) = sA',
      inputs={'n':n,'scale':scale,'A_definition':'A[i,j] = i*n+j','device':str(ctx.device)},
      result={'sum_C':actual,'maximum_element_error':max_error},
      explanation='Each element was computed on the assigned GPU; compared against s*A on the same device.')
    ctx.explain('Independent analytical sum','sum(C) = s × N × (N−1) / 2, N=n²',
      inputs={'n':n,'N':n*n,'scale':scale},result=expected,
      explanation='Closed-form Python scalar calculation independently checks the sum of the CUDA matrix result.')
    return emt.result(scene=emt.particles([0],[[0,0,0]],[emt.body(1,'#8aceff')]),
      summary=[emt.value('CUDA matrix sum',actual,''),emt.value('Analytical sum',expected,''),emt.value('CUDA device',torch.cuda.current_device(),'')],
      checks=[emt.check('Maximum element error',max_error,0),emt.check('Analytical sum error',abs(actual-expected),0)],
      integrator='CUDA float64 matrix multiplication with closed-form sum comparison',steps=1)