"""Canonical boundary conversion primitives; cohort/geometry audit must precede use.

Source depths are stored pixel coordinates. RWVM is retained for provenance and is
not blindly applied: Zeiss metadata can carry lateral rather than axial scale.
"""
from pathlib import Path
import numpy as np
import pydicom

def code_key(segment):
 codes=segment.get('SegmentedPropertyTypeCodeSequence',[])
 if len(codes)!=1:raise ValueError('Boundary needs one unambiguous anatomical code')
 return str(codes[0].CodingSchemeDesignator)+'_'+str(codes[0].CodeValue)

FRAME_MAPPING_VERSION='aireadi-heightmap-anatomy-v2'

def heightmap_frame_numbers(ds,device=None):
 """Return effective anatomy IDs for stored frames, with a Topcon export erratum.

 AI-READI Topcon payloads are in anatomical depth order. Their sequential frame
 IDs use a different SegmentSequence order for OPL/ELM/EZ/RPE. The fixed mapping
 below is supported by train/validation raw B-scan overlays, layer definitions,
 and depth ordering; it does not sort or otherwise alter per-location values.
 Heidelberg and Zeiss follow their explicit frame IDs. See discovery evidence.
 """
 declared=[int(f.SegmentIdentificationSequence[0].ReferencedSegmentNumber) for f in ds.get('PerFrameFunctionalGroupsSequence',[])]
 if device not in ['topcon_maestro2','topcon_triton']:return declared
 expected_n=9 if device=='topcon_maestro2' else 10
 expected_codes=['SCT_280677004','DCM_128289','DCM_128290','DCM_128291','DCM_128320','DCM_128297','DCM_128299','DCM_128293','SCT_76710003']+(['DCM_128301'] if expected_n==10 else [])
 segments={int(s.SegmentNumber):s for s in ds.SegmentSequence}
 if sorted(segments)!=list(range(1,expected_n+1)) or [code_key(segments[n]) for n in sorted(segments)]!=expected_codes:raise ValueError('Unsupported Topcon codebook')
 effective=[1,2,3,4,8,9,5,6,7]+([10] if expected_n==10 else [])
 if declared==effective:return declared
 if declared!=list(range(1,expected_n+1)):raise ValueError('Unsupported Topcon frame identity pattern')
 return effective

def read_heightmap(path,height,device=None):
 ds=pydicom.dcmread(path)
 if str(ds.get('SegmentationType',''))!='HEIGHTMAP':raise ValueError('Unsupported segmentation representation')
 raw=ds.pixel_array.astype(np.float32)
 if raw.ndim==2:raw=raw[None]
 frames=ds.get('PerFrameFunctionalGroupsSequence',[])
 if len(frames)!=len(raw):raise ValueError('Missing frame-to-segment identities')
 frame_numbers=heightmap_frame_numbers(ds,device)
 segments={int(s.SegmentNumber):s for s in ds.SegmentSequence}
 if len(frame_numbers)!=len(set(frame_numbers)) or set(frame_numbers)!=set(segments):raise ValueError('Ambiguous or incomplete heightmap frame mapping')
 keys=[code_key(segments[n]) for n in sorted(segments)]
 if len(keys)!=len(set(keys)):raise ValueError('Duplicate anatomical boundary code')
 values=raw[[frame_numbers.index(n) for n in sorted(segments)]]
 valid=np.isfinite(values)
 padding=ds.get('FloatPixelPaddingValue',None)
 if padding is not None:
  low=float(padding);high=float(ds.get('FloatPixelPaddingRangeLimit',padding))
  valid &= ~((values>=min(low,high))&(values<=max(low,high)))
 # Zero is a real location. Values outside the imaged axial support cannot be
 # scored in original-volume coordinates and remain source-invalid locations.
 valid &= (values>=0)&(values<=height-1)
 values=np.where(valid,values,0).astype(np.float32)
 return values,np.asarray(keys),valid,ds

def source_support_statistics(ds,height,device=None):
 """Audit every source location without changing reference measurement support."""
 raw=np.asarray(ds.pixel_array,dtype=np.float32)
 if raw.ndim==2:raw=raw[None]
 numbers=heightmap_frame_numbers(ds,device)
 segments={int(s.SegmentNumber):s for s in ds.SegmentSequence}
 padding=ds.get('FloatPixelPaddingValue',None)
 result={}
 for number in sorted(segments):
  values=raw[numbers.index(number)];finite=np.isfinite(values);pad=np.zeros(values.shape,bool)
  if padding is not None:
   low=float(padding);high=float(ds.get('FloatPixelPaddingRangeLimit',padding))
   pad=(values>=min(low,high))&(values<=max(low,high))
  outside=finite&~pad&((values<0)|(values>height-1));valid=finite&~pad&~outside
  result[code_key(segments[number])]={'locations':int(values.size),'valid_points':int(valid.sum()),
    'nonfinite_points':int((~finite).sum()),'padding_points':int((finite&pad).sum()),
    'outside_imaged_axial_support':int(outside.sum()),'valid_zero_depth_points':int(((values==0)&valid).sum()),
    'b_scans_with_support':int(valid.any(-1).sum())}
 return result

def axial_spacing_um(headers,device):
 spacings=[]
 for path,ds in headers:
  shared=ds.get('SharedFunctionalGroupsSequence',[])
  base=shared[0] if shared else ds
  for fr in ds.get('PerFrameFunctionalGroupsSequence',[ds]):
   pm=fr.get('PixelMeasuresSequence',base.get('PixelMeasuresSequence',[]))
   if not pm:return None
   spacing=float(pm[0].PixelSpacing[0])
   if not np.isfinite(spacing) or spacing<=0:return None
   spacings.append(spacing)
 if not spacings or not np.allclose(spacings,spacings[0],rtol=1e-5,atol=1e-9):return None
 return spacings[0]*1000

def convert_case(row,root,headers,geometry,with_audit=False):
 b,h,a=geometry
 if len(row['segmentation_paths'])!=1:raise ValueError('Multiple segmentation assembly requires explicit frozen mapping')
 values,keys,valid,ds=read_heightmap(Path(root)/row['segmentation_paths'][0],h,row['device'])
 if values.shape[1:]!=(b,a):raise ValueError('Heightmap and structural grid differ')
 # Frame order and axes are verified by frozen metadata audit plus source overlays.
 spacing=axial_spacing_um(headers,row['device'])
 data={'positions':values,'boundary_keys':keys,'valid':valid,'height':np.asarray(h,dtype=np.int32),**({'axial_spacing_um':np.asarray(spacing)} if spacing else {})}
 if with_audit:return data,source_support_statistics(ds,h,row['device'])
 return data