DynamisLab/src/CCD_analysis/real_ccd/core.py

"""Provenance-bound, mask-compressed streaming real-case Q=1 CCD."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterator, Mapping
import json
import numpy as np
from CCD_analysis.acquisition.contracts import ACTION_IDENTITIES
from CCD_analysis.direct_dq.analysis import coordinate_weights
from CCD_analysis.direct_dq.io import load_acquisition_artifact, load_result as load_direct_dq_result
from CCD_analysis.acquisition.artifacts import file_sha256
from CCD_analysis.direct_dq.schema import canonical_array_sha256

ACTION_CHANNEL_NAMES=("front","upper","lower")
ACTION_UNITS="native solver angular-velocity command units"
FLATTEN_ORDER="component-major ux then uy; C-order analysis-mask point order"
WEIGHT_RULE="direct_dq.coordinate_weights(x_D)*coordinate_weights(y_D), repeated ux then uy; not area-normalized"

@dataclass(frozen=True)
class MemoryBudget:
    ram_bytes:int
    scratch_bytes:int
    safety_margin:float=1.25
    def __post_init__(self):
        if type(self.ram_bytes) is not int or type(self.scratch_bytes) is not int or self.ram_bytes<=0 or self.scratch_bytes<0 or not np.isfinite(self.safety_margin) or self.safety_margin<1:
            raise ValueError("explicit positive RAM, nonnegative scratch, and safety_margin>=1 required")

@dataclass(frozen=True)
class StreamingConfig:
    chunk_size:int
    budget:MemoryBudget
    singular_block_rtol:float=1e-10
    singular_block_atol:float=0.0
    def __post_init__(self):
        if type(self.chunk_size) is not int or self.chunk_size<=0: raise ValueError("chunk_size must be positive")
        if min(self.singular_block_rtol,self.singular_block_atol)<0: raise ValueError("singular tolerances must be nonnegative")

@dataclass(frozen=True)
class RealCCDInput:
    case_id:str; direct_dq_root:Path; direct_manifest_sha256:str
    acquisition_identities:Mapping[str,Mapping[str,Any]]
    x_D:np.ndarray; y_D:np.ndarray; role_masks:Mapping[str,np.ndarray]; analysis_mask:np.ndarray
    selected_indices:np.ndarray; selected_relative_steps:np.ndarray
    q_ctl_absolute_steps:np.ndarray; dq_ctl:np.ndarray; actions:np.ndarray
    authoritative_mean_dq_ctl:np.ndarray

@dataclass(frozen=True)
class RealCCDResult:
    arrays:dict[str,np.ndarray]; config:dict[str,Any]; summary:dict[str,Any]; input_hashes:dict[str,Any]

def _read_json(path:Path)->dict[str,Any]:
    value=json.loads(path.read_text())
    if not isinstance(value,dict): raise ValueError(f"JSON object required: {path}")
    return value

def estimate_memory(*,m:int,n:int,nx:int,ny:int,full_count:int,chunk_size:int,budget:MemoryBudget)->dict[str,Any]:
    if min(m,n,nx,ny,full_count,chunk_size)<=0: raise ValueError("memory dimensions must be positive")
    c=min(chunk_size,n); grid=nx*ny
    terms={
      "validated_direct_result_float32_fields":6*n*2*grid*4,
      "validated_live_acquisition_float32_fields":3*full_count*2*grid*4,
      "loader_decompression_and_copy_allowance":(6*n+3*full_count)*2*grid*4,
      "field_mean_float64":m*8,"weights_and_roots_float64":2*m*8,
      "cross_and_modes_float64":(3*m+2*3*m)*8,
      "chunk_float64_working_set":3*m*c*8,
      "actions_coefficients_and_small_svd":(3*n+3*n+30)*8,
    }
    raw=sum(terms.values()); peak=int(np.ceil(raw*budget.safety_margin)); scratch_raw=0; scratch=int(np.ceil(scratch_raw*budget.safety_margin))
    decision=peak<=budget.ram_bytes and scratch<=budget.scratch_bytes
    out={"formula":"ceil(safety_margin * sum(terms)); loader residency explicitly included; no MxM or full float64 MxN term","terms_bytes":terms,"raw_peak_ram_bytes":raw,"estimated_peak_ram_bytes":peak,"raw_scratch_bytes":scratch_raw,"estimated_scratch_bytes":scratch,"ram_budget_bytes":budget.ram_bytes,"scratch_budget_bytes":budget.scratch_bytes,"safety_margin":budget.safety_margin,"decision":"PASS" if decision else "FAIL"}
    if not decision: raise MemoryError(f"real-CCD memory/scratch estimate exceeds explicit budget: {out}")
    return out

def inspect_direct_dq_dimensions(path:str|Path)->dict[str,int]:
    root=Path(path); summary=_read_json(root/'summary.json'); config=_read_json(root/'config.json')
    if config.get('schema_id')!='ccd-direct-dq-config/v2': raise ValueError('authoritative direct-dq config required')
    identities=_read_json(root/'input_hashes.json'); qctl=Path(identities['q_ctl']['path']); acq_manifest=_read_json(qctl/'manifest.json'); acq_config=_read_json(qctl/'config.json')
    frame=acq_config['runtime']['coordinate_frame']
    return {"n":int(summary['sample_count']),"nx":int(frame['x']['count']),"ny":int(frame['y']['count']),"full_count":int(acq_manifest['field_count'])}

def load_validated_input(path:str|Path,*,streaming_config:StreamingConfig)->tuple[RealCCDInput,dict[str,Any]]:
    root=Path(path).resolve(); dims=inspect_direct_dq_dimensions(root)
    estimate=estimate_memory(m=2*dims['nx']*dims['ny'],chunk_size=streaming_config.chunk_size,budget=streaming_config.budget,**dims)
    direct=load_direct_dq_result(root); a=direct['arrays']; cfg=direct['config']; recorded=direct['input_hashes']; case=cfg['case_id']
    mask=a['analysis_fluid_mask']; m=2*int(mask.sum())
    tight=estimate_memory(m=m,chunk_size=streaming_config.chunk_size,budget=streaming_config.budget,**dims)
    qctl=load_acquisition_artifact(recorded['q_ctl']['path'],expected_case=case,expected_role='q_ctl')
    idx=a['selected_timeline_indices']; rel=a['selected_acquisition_relative_lattice_steps']
    if not np.array_equal(qctl.fields['acquisition_relative_lattice_steps'][idx],rel): raise ValueError('q_ctl exact selected indices/timestamps mismatch')
    actions=qctl.fields['effective_applied_action'][idx,-3:].copy()
    if actions.dtype!=np.float32 or actions.shape!=(idx.size,3) or not np.isfinite(actions).all(): raise ValueError('q_ctl effective field-time actions invalid')
    dq=(a['q_ctl_instantaneous']-a['q_blk_instantaneous'])
    if not np.array_equal(dq,a['dq_ctl_instantaneous']): raise ValueError('authoritative direct_dq dq_ctl identity failed')
    inp=RealCCDInput(case,root,file_sha256(root/'manifest.json'),recorded,a['x_D'].copy(),a['y_D'].copy(),{r:a[f'{r}_solver_fluid_mask'].copy() for r in ('q_target','q_blk','q_ctl')},mask.copy(),idx.copy(),rel.copy(),qctl.fields['lattice_steps'][idx].copy(),dq,actions,a['dq_ctl_mean'].copy())
    estimate={**tight,"conservative_all_fluid_estimated_peak_ram_bytes":estimate['estimated_peak_ram_bytes'],"admission_basis":"all-fluid M estimate before authoritative loader; tight mask estimate also passed"}
    return inp,estimate

def _chunks(inp:RealCCDInput,chunk:int)->Iterator[tuple[slice,np.ndarray,np.ndarray]]:
    mask=inp.analysis_mask
    for start in range(0,inp.selected_indices.size,chunk):
        stop=min(start+chunk,inp.selected_indices.size); raw=inp.dq_ctl[start:stop]
        field=np.concatenate((raw[:,0][:,mask],raw[:,1][:,mask]),axis=1).T.astype(np.float64)
        actions=inp.actions[start:stop].T.astype(np.float64)
        if not np.isfinite(field).all() or not np.isfinite(actions).all(): raise ValueError('nonfinite streamed column')
        yield slice(start,stop),field,actions

def _classify(s:np.ndarray,rtol:float,atol:float)->tuple[np.ndarray,list[list[int]],list[int],float]:
    tol=atol+(rtol*float(s[0]) if s.size else 0.0); identifiable=s>tol; blocks=[]; boundaries=[]; stop=int(identifiable.sum()); start=0
    while start<stop:
        end=start+1
        while end<stop and abs(float(s[end]-s[start]))<=atol+rtol*max(float(s[start]),float(s[end])): end+=1
        if end-start>1: blocks.append([start,end])
        boundaries.append(end); start=end
    return identifiable,blocks,boundaries,tol

def decompose_streaming(inp:RealCCDInput,*,streaming_config:StreamingConfig,memory_estimate:dict[str,Any]|None=None)->RealCCDResult:
    if inp.actions.shape!=(inp.selected_indices.size,3): raise ValueError('exactly three action channels required')
    n=inp.selected_indices.size
    for name, values in (("selected indices", inp.selected_indices), ("selected relative timestamps", inp.selected_relative_steps), ("q_ctl absolute timestamps", inp.q_ctl_absolute_steps)):
        if values.dtype != np.int64 or values.shape != (n,) or np.any(np.diff(values) <= 0):
            raise ValueError(f'{name} must be exact strictly increasing int64 values')
    mask=inp.analysis_mask; points=int(mask.sum()); m=2*points
    estimate=memory_estimate or estimate_memory(m=m,n=n,nx=inp.x_D.size,ny=inp.y_D.size,full_count=inp.dq_ctl.shape[0],chunk_size=streaming_config.chunk_size,budget=streaming_config.budget)
    wx=coordinate_weights(inp.x_D); wy=coordinate_weights(inp.y_D); point_w=(wx[:,None]*wy[None,:])[mask]; weights=np.concatenate((point_w,point_w)); roots=np.sqrt(weights)
    fsum=np.zeros(m,np.float64); psum=np.zeros(3,np.float64); count=0
    for _,u,p in _chunks(inp,streaming_config.chunk_size): fsum+=u.sum(axis=1); psum+=p.sum(axis=1); count+=u.shape[1]
    if count!=n: raise ValueError('stream pass count mismatch')
    fmean=fsum/n; pmean=psum/n; cross=np.zeros((3,m),np.float64)
    for _,u,p in _chunks(inp,streaming_config.chunk_size): cross+=(p-pmean[:,None])@((u-fmean[:,None])*roots[:,None]).T
    cross/=n*np.sqrt(3.0)
    left,s,vh=np.linalg.svd(cross,full_matrices=False); weighted=vh.T
    for k in range(weighted.shape[1]):
        pivot=int(np.argmax(np.abs(weighted[:,k])))
        if weighted[pivot,k]<0: weighted[:,k]*=-1; left[:,k]*=-1
    modes=weighted/roots[:,None]; identifiable,blocks,boundaries,null_tol=_classify(s,streaming_config.singular_block_rtol,streaming_config.singular_block_atol)
    coeff=np.empty((3,n),np.float64); total_sq=0.0
    for sl,u,_ in _chunks(inp,streaming_config.chunk_size):
        x=(u-fmean[:,None])*roots[:,None]; coeff[:,sl]=weighted.T@x; total_sq+=float(np.sum(x*x))
    residual=np.asarray([max(total_sq-float(np.sum(coeff[:r]**2)),0.0) for r in boundaries],np.float64)
    residual=np.sqrt(residual/max(total_sq,np.finfo(float).tiny))
    arrays={"x_D":inp.x_D,"y_D":inp.y_D,"q_target_solver_fluid_mask":inp.role_masks['q_target'],"q_blk_solver_fluid_mask":inp.role_masks['q_blk'],"q_ctl_solver_fluid_mask":inp.role_masks['q_ctl'],"analysis_fluid_mask":mask,"selected_timeline_indices":inp.selected_indices,"selected_acquisition_relative_lattice_steps":inp.selected_relative_steps,"selected_q_ctl_absolute_lattice_steps":inp.q_ctl_absolute_steps,"coordinate_weights":weights,"field_mean":fmean,"action_mean":pmean,"effective_actions":inp.actions,"cross_correlation":cross,"left_functions":left,"singular_values":s,"physical_modes":modes,"coefficients":coeff,"identifiable_mode_mask":identifiable,"residual_block_boundaries":np.asarray(boundaries,np.int64),"weighted_relative_residuals":residual,"authoritative_mean_dq_ctl":inp.authoritative_mean_dq_ctl}
    input_hashes={"direct_dq":{"path":str(inp.direct_dq_root),"manifest_sha256":inp.direct_manifest_sha256},"acquisitions":{k:dict(v) for k,v in inp.acquisition_identities.items()},"canonical_arrays":{k:canonical_array_sha256(v) for k,v in arrays.items()}}
    config={"schema_id":"ccd-real-ccd-config/v1","case_id":inp.case_id,"Q":1,"tau":0,"observable_count":3,"channel_names":list(ACTION_CHANNEL_NAMES),"action_identities":list(ACTION_IDENTITIES),"action_units":ACTION_UNITS,"flatten_order":FLATTEN_ORDER,"weight_rule":WEIGHT_RULE,"center_snapshots":True,"center_observables":True,"standardization":False,"whitening":False,"chunk_size":streaming_config.chunk_size,"accumulation_dtype":"float64","input_field_dtype":"float32","singular_block_rtol":streaming_config.singular_block_rtol,"singular_block_atol":streaming_config.singular_block_atol,"memory":estimate,"full_reconstructions_persisted":False}
    summary={"schema_id":"ccd-real-ccd-summary/v1","sample_count":n,"spatial_dof_count":m,"numerical_rank":int(identifiable.sum()),"null_tolerance":null_tol,"degenerate_singular_blocks":blocks,"complete_block_boundaries":boundaries,"spectrum_label":"cross-correlation strength; not field energy, explained variance, or canonical coefficient","mean_context":"mean effective actions and authoritative mean dq_ctl are outside CCD","claim_boundary":"no CCD>POD, causal, mechanism, response-time, same-phase, independent-realization, uncertainty, or observable-prediction claim","passes":3,"provenance_status":"VERIFIED_LIVE_INPUTS_REQUIRED_ON_LOAD"}
    from .schema import validate_result
    arrays=validate_result(arrays=arrays,config=config,summary=summary,input_hashes=input_hashes)
    return RealCCDResult(arrays,config,summary,input_hashes)