"""Independent volume-weighted replay of frozen fields; no evolution or Lab Engine."""
import argparse
import gzip
import hashlib
from io import BytesIO
import json
from pathlib import Path
import sys
import zipfile
import numpy as np


def sha(raw):
    return hashlib.sha256(raw).hexdigest()


def arrays(raw):
    with np.load(BytesIO(raw),allow_pickle=False) as saved:
        return {key:saved[key].copy() for key in saved.files}


def primitives(record,grid):
    state=record['snapshotConserved']
    rho=grid['rho0'][None,:]+state[:,:,0]
    velocity=state[:,:,1]/rho
    pressure=grid['p0'][None,:]+(5/3-1)*(state[:,:,2]-grid['phiMean'][None,:]*state[:,:,0]-.5*state[:,:,1]**2/rho)
    if not np.all(np.isfinite(state)) or np.any(rho<=0) or np.any(pressure<=0):
        raise ValueError('Nonphysical input')
    return np.stack([rho,velocity,pressure],axis=-1)


def metric(difference,normalization_signal,volumes,limit):
    norm=np.sum(abs(normalization_signal)*volumes[None,:,None],axis=1)
    numerator=np.sum(abs(difference)*volumes[None,:,None],axis=1)
    active=norm>0
    error=np.divide(numerator,norm,out=np.zeros_like(norm),where=active)
    zero_error=float(numerator[~active].max()) if np.any(~active) else 0.
    return {'pass':bool(zero_error<=1e-12 and error.max()<=limit),'maximum':float(error.max()),
        'limit':limit,'zeroSignalAbsoluteMaximum':zero_error,
        'perSnapshotRhoVelocityPressure':error.tolist(),'maximumSnapshotVariableIndices':list(map(int,np.unravel_index(error.argmax(),error.shape)))}


def main():
    parser=argparse.ArgumentParser()
    parser.add_argument('--materials',type=Path,default=Path(__file__).parent)
    parser.add_argument('--output',type=Path)
    parser.add_argument('--packet',type=Path,help='Optional location of the same hash-bound full packet')
    args=parser.parse_args();folder=args.materials
    manifest=json.loads((folder/'wave-refinement-manifest.json').read_text())
    def checked_file(row,override=None):
        raw=(override if override is not None else folder/row['name']).read_bytes()
        if len(raw)!=row['bytes'] or sha(raw)!=row['sha256']:
            raise ValueError('File identity: '+row['name'])
        return raw
    packet_path=args.packet if args.packet is not None else folder/manifest['packet']['name']
    packet_raw=packet_path.read_bytes()
    packet_identity=(len(packet_raw),sha(packet_raw))
    encoded=manifest['packet'];decoded=encoded['decodedVariant']
    if packet_identity==(encoded['bytes'],encoded['sha256']):
        packet_json=gzip.decompress(packet_raw)
        if (len(packet_json),sha(packet_json))!=(decoded['bytes'],decoded['sha256']):
            raise ValueError('Decoded packet identity differs')
    elif packet_identity==(decoded['bytes'],decoded['sha256']):
        packet_json=packet_raw
    else:
        raise ValueError('Packet is neither declared gzip nor declared HTTP-decoded JSON')
    packet=json.loads(packet_json)
    archive_raw=checked_file(manifest['archive'])
    data={}
    with zipfile.ZipFile(BytesIO(archive_raw)) as archive:
        if set(archive.namelist())!=set(manifest['members']):
            raise ValueError('Native archive member set differs')
        for name,row in manifest['members'].items():
            raw=archive.read(name)
            if len(raw)!=row['bytes'] or sha(raw)!=row['sha256']:
                raise ValueError('Native member identity: '+name)
            if name.endswith('.npz'):
                data[name]=arrays(raw)
    normal_pulse=np.array(packet['arrays']['pulseRhoVelocityPressureCode'],dtype='<f8')
    normal_cooling=np.array(packet['arrays']['coolingRhoVelocityPressureCode'],dtype='<f8')
    normal=normal_pulse-normal_cooling
    times=np.array(packet['arrays']['waveTimesCode'])
    for name,row in data.items():
        if 'snapshotTimesCode' in row and not np.allclose(times,row['snapshotTimesCode'],rtol=0,atol=1e-12):
            raise ValueError('Mismatched actual times: '+name)
    coarse_grid=data['coarse/grid.npz'];half_grid=data['half/grid.npz']
    fine_faces=np.array(packet['arrays']['waveFacesKpc'])
    coarse_faces=coarse_grid['faces']
    if not (len(fine_faces)==4097 and len(coarse_faces)==2049 and np.allclose(fine_faces[::2],coarse_faces,rtol=0,atol=1e-12)):
        raise ValueError('Native nested geometry differs')
    coarse=primitives(data['coarse/allocated-pulse.npz'],coarse_grid)-primitives(data['coarse/cooling-only.npz'],coarse_grid)
    half=primitives(data['half/allocated-pulse.npz'],half_grid)-primitives(data['half/cooling-only.npz'],half_grid)
    # Independent two-cell streaming integral instead of the producer's dense overlap matrix.
    def project_to_coarse(signal):
        projected=np.empty_like(coarse)
        for cell in range(2048):
            lower=coarse_faces[cell];upper=coarse_faces[cell+1]
            weights=[]
            for fine_cell in [2*cell,2*cell+1]:
                overlap=max(0.,min(upper,fine_faces[fine_cell+1])**3-max(lower,fine_faces[fine_cell])**3)
                weights.append(4*np.pi/3*overlap)
            projected[:,cell,:]=(signal[:,2*cell,:]*weights[0]+signal[:,2*cell+1,:]*weights[1])/coarse_grid['volumes'][cell]
        return projected
    projected=project_to_coarse(normal)
    projected_half=project_to_coarse(half)
    spatial=metric(coarse-projected,projected,coarse_grid['volumes'],.05)
    temporal=metric(half-normal,normal,half_grid['volumes'],.01)
    original_temporal=metric(projected_half-projected,projected,coarse_grid['volumes'],.01)
    expected=json.loads((folder/'gentle-wave-convergence.json').read_text())
    if abs(spatial['maximum']-expected['spatial']['maximum'])>1e-12 or abs(temporal['maximum']-expected['temporal']['maximum'])>1e-12:
        raise ValueError('Independent replay differs from actual saved comparison')
    if not spatial['pass'] or not temporal['pass'] or not original_temporal['pass']:
        raise ValueError('Original numerical gate failed')
    result={'status':'PASS_INDEPENDENT_FROZEN_FIELD_REFINEMENT_REPLAY','spatial':spatial,'temporalNative4096':temporal,'temporalOriginalCoarseProjected':original_temporal,
        'actualNumPy':np.__version__,'actualPython':sys.version,'newPDEOrODECalls':0,
        'AGNShockScope':expected['originalShockScopeStatus'],
        'boundary':'Recomputed original field norms, not a rerun of every PDE step, control, shock detector or feedback model.'}
    text=json.dumps(result,indent=2,allow_nan=False)+'\n'
    if args.output:
        with args.output.open('x') as stream:stream.write(text)
    print(text,end='')


if __name__=='__main__':
    main()