"""Rebuild the two compositional feature archives from official raw recordings.

Use --hydraulic-archive and --temperature-archive to reuse downloaded archives.
Source scientific aggregation stays on CUDA. Raw archives are never republished.
"""
from pathlib import Path
import argparse,hashlib,io,itertools,json,urllib.request,zipfile
import numpy as np
import torch
from reproduce import ROOT,load

def digest(p):return hashlib.sha256(Path(p).read_bytes()).hexdigest()
def acquire(path,url,expected):
    path=Path(path)
    if not path.exists():path.parent.mkdir(parents=True,exist_ok=True);urllib.request.urlretrieve(url,path)
    assert digest(path)==expected,('archive hash mismatch',path)
    return path

def hydraulic(path,device):
    from unionmind.continuous_decision import source
    import tempfile,os
    config=json.loads((ROOT/'protocols/factorial_transfer_protocol_01.json').read_text())
    with tempfile.TemporaryDirectory(prefix='hydraulic-rebuild-') as tmp:
        tmp=Path(tmp);(tmp/'sources').mkdir();os.symlink(Path(path).resolve(),tmp/'sources/hydraulic.zip')
        source.prepare_hydraulic(tmp,config,device)
        rebuilt=torch.load(tmp/'prepared/hydraulic.pt',map_location=device,weights_only=True)
    expected=load('data/hydraulic/prepared',device)
    assert torch.allclose(rebuilt['x'],expected['x'],atol=1e-10,rtol=1e-12)
    assert torch.equal(rebuilt['h'],expected['h'])
    for k in expected['masks']:assert torch.equal(expected['masks'][k],rebuilt['masks'][k])
    return dict(records=len(rebuilt['x']),maximum_feature_error=float((rebuilt['x']-expected['x']).abs().max()),split_masks_identical=True)

def temperature(path,device):
    rows=[];labels=[];days=[];invalid=[]
    with zipfile.ZipFile(path) as outer:
        inner=zipfile.ZipFile(io.BytesIO(outer.read('gas-sensor-array-temperature-modulation.zip')))
        members=sorted(a.filename for a in inner.infolist() if a.filename.endswith('.csv'))
        assert len(members)==13
        for day,name in enumerate(members):
            data=np.loadtxt(io.BytesIO(inner.read(name)),delimiter=',',skiprows=1)
            raw=torch.as_tensor(data,device=device,dtype=torch.float64);t=raw[:,0]
            if not bool(torch.isfinite(raw).all()) or not bool((torch.diff(t)>=0).all()):invalid.append(name);continue
            for block in range(int(torch.floor((t[-1]-t[0])/900))+1):
                boundary=t[0]+torch.arange(5,device=device,dtype=torch.float64)*225+block*900
                ix=torch.searchsorted(t,boundary);a,b=int(ix[0]),min(int(ix[-1]),len(t))
                if b<=a or float(t[b-1]-t[a])<850:continue
                q=torch.stack([raw[int(ix[k]):int(ix[k+1]),6:20].mean(0) for k in range(4)])
                last=raw[int(ix[3]):b];co=torch.quantile(last[:,1],.5,interpolation='lower');rh=last[:,2].mean()
                levels=torch.linspace(0,20,10,device=device,dtype=torch.float64)
                rows.append(q.T.reshape(-1));labels.append(torch.stack(((levels-co).abs().argmin(),torch.bucketize(rh,rh.new_tensor([35.,55.])))));days.append(day)
    x=torch.stack(rows);h=torch.stack(labels);d=torch.tensor(days,device=device)
    grid=list(itertools.product(range(10),range(3)));held=sorted(grid,key=lambda v:hashlib.sha256(f'20821:{v[0]},{v[1]}'.encode()).hexdigest())[:6]
    unseen=torch.zeros(len(x),device=device,dtype=torch.bool)
    for binding in held:unseen|=(h==h.new_tensor(binding)).all(1)
    rebuilt=dict(x=x,h=h,days=d,source=(d<9)&~unseen,target=d>=9,unseen=unseen);expected=load('data/temperature/prepared',device)
    for k in expected:
        if k=='x':assert torch.allclose(rebuilt[k],expected[k],atol=1e-10,rtol=1e-12)
        else:assert torch.equal(rebuilt[k],expected[k]),k
    assert len(invalid)==1 and '20161003_085624.csv' in invalid[0]
    return dict(records=len(x),maximum_feature_error=float((x-expected['x']).abs().max()),invalid_days=invalid,original_primary_eligibility=False,secondary_split_identical=True)

def main():
    p=argparse.ArgumentParser(description=__doc__);p.add_argument('--device',default='cuda:0');p.add_argument('--cache',default='raw_downloads');p.add_argument('--hydraulic-archive');p.add_argument('--temperature-archive');p.add_argument('--output',default='raw_rebuild.json');a=p.parse_args()
    assert a.device.startswith('cuda') and torch.cuda.is_available();torch.cuda.set_device(a.device);torch.backends.cuda.matmul.allow_tf32=False
    cache=Path(a.cache);hyd=acquire(a.hydraulic_archive or cache/'hydraulic.zip','https://archive.ics.uci.edu/static/public/447/condition+monitoring+of+hydraulic+systems.zip','24128aad2ee45eea7e6b63ebbd9992cdf25d0483a2cebefbfc13bc69079af1f2')
    tmp=acquire(a.temperature_archive or cache/'temperature.zip','https://archive.ics.uci.edu/static/public/487/gas+sensor+array+temperature+modulation.zip','cb5ce4a6af1a51b933d1979952d7845f0e9baac54e15a81a1e3599e8b85905d4')
    report={'hydraulic':hydraulic(hyd,a.device),'temperature':temperature(tmp,a.device)};Path(a.output).write_text(json.dumps(report,indent=2)+'\n');print(json.dumps(report,indent=2))
if __name__=='__main__':main()