"""Memory-bounded construction of paired and phase-conditioned correction products."""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import scipy.signal

from .artifacts import atomic_save_npz, atomic_write_json


def _phase(signal: np.ndarray) -> np.ndarray:
    centered=np.asarray(signal,float)-np.mean(signal)
    return np.mod(np.angle(scipy.signal.hilbert(centered)),2*np.pi)


def build_products(run_dir: Path, *, x0: int, x1: int, y0: int, y1: int, stride: int, bins: int=24) -> dict:
    destination=run_dir/"aligned"
    if destination.exists(): raise FileExistsError(destination)
    destination.mkdir(parents=True)
    controlled=run_dir/"controlled/fields.npz"; baseline=run_dir/"time_paired_baseline/fields.npz"
    with np.load(controlled,allow_pickle=False) as c,np.load(baseline,allow_pickle=False) as b:
        steps=c["lattice_steps"].copy(); indices=c["control_indices"].copy()
        if not np.array_equal(steps,b["lattice_steps"]): raise ValueError("paired timelines differ")
        parts=[]
        for key in ("ux","uy"):
            ca=c[key][:,x0:x1:stride,y0:y1:stride].astype(np.float32)
            ba=b[key][:,x0:x1:stride,y0:y1:stride].astype(np.float32)
            parts.append((ca-ba).reshape(len(steps),-1))
        correction=np.hstack(parts).astype(np.float32)
    with np.load(run_dir/"controlled/telemetry.npz",allow_pickle=False) as t:
        sensors=t["sensors"].copy(); forces=t["forces"].copy(); actions=t["actions"].copy()
    phase_control=_phase(sensors[:,3] if sensors.shape[1]>3 else sensors[:,0])
    phase_field=phase_control[np.minimum(indices,len(phase_control)-1)]
    edges=np.linspace(0,2*np.pi,bins+1); labels=np.minimum(np.digitize(phase_field,edges)-1,bins-1)
    phase_means=np.full((bins,correction.shape[1]),np.nan,np.float32); counts=[]
    for index in range(bins):
        selected=labels==index; counts.append(int(selected.sum()))
        if selected.any(): phase_means[index]=correction[selected].mean(axis=0)
    product=atomic_save_npz(destination/"correction_reduced.npz",correction=correction,
        lattice_steps=steps,control_indices=indices,phase=phase_field,phase_bins=labels,
        phase_conditioned_mean=phase_means,phase_counts=np.asarray(counts,np.int64),
        sensors=sensors[indices],forces=forces[indices],actions=actions[indices])
    speed=np.sqrt(np.sum(correction**2,axis=1)); ux_width=((x1-x0+stride-1)//stride)*((y1-y0+stride-1)//stride)
    mean_ux=correction[:,:ux_width].mean(axis=0)
    report={"schema_version":"oid-v2-correction-v1","crop":{"x":[x0,x1],"y":[y0,y1],"stride":stride},
        "samples":len(correction),"features":correction.shape[1],"timeline_exact":True,
        "finite":bool(np.isfinite(correction).all()),"correction_rms":float(np.sqrt(np.mean(correction**2))),
        "snapshot_speed_rms_mean":float(speed.mean()),"mean_ux_positive_fraction":float(np.mean(mean_ux>0)),
        "phase_bins":bins,"phase_bin_min_count":min(counts),"phase_bin_max_count":max(counts),
        "product":str(product.name)}
    atomic_write_json(destination/"correction_report.json",report); return report


def main(argv=None):
    p=argparse.ArgumentParser(); p.add_argument("run_dir",type=Path); p.add_argument("--x0",type=int,required=True); p.add_argument("--x1",type=int,required=True)
    p.add_argument("--y0",type=int,required=True); p.add_argument("--y1",type=int,required=True); p.add_argument("--stride",type=int,default=4); p.add_argument("--bins",type=int,default=24)
    a=p.parse_args(argv); print(json.dumps(build_products(a.run_dir,x0=a.x0,x1=a.x1,y0=a.y0,y1=a.y1,stride=a.stride,bins=a.bins),indent=2)); return 0
if __name__=="__main__": raise SystemExit(main())
