"""Author-defined full-factorial source adapters and untouched binding splits."""
from __future__ import annotations
import hashlib,io,itertools,json,math,time,urllib.request,zipfile
from pathlib import Path
import numpy as np
import torch
from unionmind.predictive_transfer.data import member,digest_int
from unionmind.supported_observations.data import sha


def acquire(root,protocol):
    assert (root/'confirmation_lock.json').exists();out=root/'sources';out.mkdir(exist_ok=True);records=[]
    for name,spec in protocol['sources'].items():
        path=out/f'{name}.zip';assert not path.exists();start=time.perf_counter()
        reused=spec.get('opened_archive')
        if reused:path.symlink_to(Path(reused).resolve())
        else:urllib.request.urlretrieve(spec['url'],path)
        with zipfile.ZipFile(path) as z:inventory=[dict(name=a.filename,bytes=a.file_size,compressed_bytes=a.compress_size) for a in z.infolist()]
        records.append(dict(name=name,path=str(path),sha256=sha(path),bytes=path.stat().st_size,
            acquisition_seconds=time.perf_counter()-start,previously_opened=bool(reused),
            phase='opened_development' if reused else 'locked_confirmation',
            doi=spec['doi'],authors=spec['authors'],license=spec['license'],inventory=inventory))
        print(json.dumps(dict(acquired=name,bytes=path.stat().st_size)),flush=True)
    (out/'manifest.json').write_text(json.dumps(dict(archives=records),indent=2))


def parse(content,skip_header=False):
    if skip_header:content=content.split(b'\n',1)[1]
    content=content.strip()
    # CPU performs text parsing only. Feature statistics below are CUDA.
    lines=content.count(b'\n')+(not content.endswith(b'\n'))
    values=np.fromstring(content.decode(),sep=' ',dtype=np.float64)
    assert values.size%lines==0,(values.size,lines)
    return values.reshape(lines,-1)


def split(h,groups,dataset,sizes):
    grid=list(itertools.product(*[range(k) for k in sizes]))
    binding_order=sorted(grid,key=lambda row:hashlib.sha256(('20711:'+dataset+':'+','.join(map(str,row))).encode()).hexdigest())
    withheld=set(binding_order[:max(4,math.ceil(.2*len(grid)))])
    unseen=[tuple(row) in withheld for row in h.tolist()]
    seen=[not u and digest_int('20721:'+dataset+':'+str(g))%5==0 for u,g in zip(unseen,groups)]
    held_groups={g for g,a,b in zip(groups,unseen,seen) if a or b}
    forbidden={g+d for g in held_groups for d in (-1,0,1)}
    source=[not a and not b and g not in forbidden for a,b,g in zip(unseen,seen,groups)]
    device=h.device
    masks=dict(source=torch.tensor(source,device=device),unseen=torch.tensor(unseen,device=device),seen=torch.tensor(seen,device=device))
    source_count=sum(source);unseen_count=sum(unseen);seen_count=sum(seen)
    bindings=len({tuple(row) for row,u in zip(h.tolist(),unseen) if u})
    reasons=[]
    if source_count<64:reasons.append('source_less_than_64')
    if unseen_count<32:reasons.append('unseen_records_less_than_32')
    if bindings<3:reasons.append('unseen_bindings_less_than_3')
    if seen_count<16:reasons.append('seen_records_less_than_16')
    return masks,dict(withheld_ontology_bindings=[list(r) for r in sorted(withheld)],source_records=source_count,
        unseen_records=unseen_count,seen_records=seen_count,unseen_observed_bindings=bindings,
        embargoed_records=len(h)-source_count-unseen_count-seen_count,exclusions=reasons,
        primary_eligible=not reasons,secondary_eligible=not reasons and source_count>=128)


def encode_levels(profile,levels):
    encoded=[]
    for column,values in enumerate(levels):
        options=profile.new_tensor(values);match=profile[:,column,None]==options[None]
        assert torch.all(match.sum(-1)==1),(column,profile[:,column].unique().tolist())
        encoded.append(match.long().argmax(-1))
    return torch.stack(encoded,-1)


def prepare_hydraulic(root,protocol,device='cuda:1'):
    out=root/'prepared';out.mkdir(exist_ok=True);series=[]
    spec=protocol['sources']['hydraulic']
    with zipfile.ZipFile(root/'sources/hydraulic.zip') as z:
        profiles=torch.as_tensor(parse(member(z,'profile.txt')),device=device,dtype=torch.float64)
        h=encode_levels(profiles,spec['levels']);features=[];feature_names=[];units=[];rates=[]
        for sensor,(unit,rate) in spec['physical_sensors'].items():
            raw=torch.as_tensor(parse(member(z,sensor+'.txt')),device=device,dtype=torch.float64)
            assert raw.shape==(len(h),60*rate)
            values=raw.reshape(len(h),4,15*rate).mean(-1);features.append(values)
            feature_names.extend([f'{sensor}:seconds_{q*15}_{(q+1)*15}' for q in range(4)])
            units.extend([unit]*4);rates.extend([rate]*4)
            del raw
        x=torch.cat(features,-1);groups=[i//5 for i in range(len(h))]
        masks,counts=split(h,groups,'hydraulic',[len(v) for v in spec['levels']])
        path=out/'hydraulic.pt';torch.save(dict(x=x,h=h,masks=masks,groups=torch.tensor(groups,device=device),stable_flag=profiles[:,4].long()),path)
        series.append(dict(name='hydraulic',system='hydraulic_rig',file=path.name,sha256=sha(path),rows=len(h),dimensions=x.shape[-1],
            factor_names=spec['factor_names'],levels=spec['levels'],features=feature_names,units=units,rates_hz=rates,
            records=[dict(id=f'hydraulic:cycle{i}',cycle_zero_based=i,source_file='profile.txt',source_row=i,group=groups[i]) for i in range(len(h))],**counts))
    print(json.dumps(dict(prepared='hydraulic',**counts)),flush=True)
    return series[0]


def prepare(root,protocol,device='cuda:1'):
    start=time.perf_counter();out=root/'prepared';out.mkdir(exist_ok=True);series=[];exclusions=[]
    series.append(prepare_hydraulic(root,protocol,device))
    spec=protocol['sources']['gas']
    with zipfile.ZipFile(root/'sources/gas.zip') as z:
        for name,file_spec in spec['files'].items():
            raw=torch.as_tensor(parse(member(z,file_spec['file']),True),device=device,dtype=torch.float64)
            assert raw.shape[1]==19 and torch.all(raw[1:,0]>raw[:-1,0]) and torch.isfinite(raw).all()
            starts=torch.cat((torch.zeros(1,device=device,dtype=torch.long),torch.where((raw[1:,1:3]!=raw[:-1,1:3]).any(-1))[0]+1))
            ends=torch.cat((starts[1:],starts.new_tensor([len(raw)])))
            rows=[];labels=[];records=[];groups=[]
            for episode,(a,b) in enumerate(zip(starts.tolist(),ends.tolist())):
                first=float(raw[a,0]);last=float(raw[b-1,0]);duration=last-first
                if duration<75:
                    exclusions.append(dict(series=name,episode=episode,first_row=a,last_row_exclusive=b,duration_seconds=duration,reason='duration_less_than_75_seconds'));continue
                early_end=int(torch.searchsorted(raw[:,0],raw[a,0]+20,right=False))
                late_start=int(torch.searchsorted(raw[:,0],raw[b-1,0]-20,right=False))
                early=raw[a:early_end,3:].mean(0);late=raw[late_start:b,3:].mean(0)
                rows.append(torch.stack((late,late-early),-1).flatten())
                primary=raw[a,1];ethylene=raw[a,2]
                hp=(primary>primary.new_tensor([0.,file_spec['primary_max_ppm']/4,file_spec['primary_max_ppm']/2])).sum()
                he=(ethylene>ethylene.new_tensor([0.,5.,10.])).sum()
                labels.append(torch.stack((hp,he)));groups.append(episode)
                records.append(dict(id=f'{name}:episode{episode}',episode=episode,source_file=file_spec['file'],
                    first_row_zero_based=a,last_row_exclusive=b,start_seconds=first,end_seconds=last,group=episode,
                    early_last_row_exclusive=early_end,late_first_row_zero_based=late_start))
            x=torch.stack(rows);h=torch.stack(labels).long();masks,counts=split(h,groups,name,[4,4])
            path=out/f'{name}.pt';torch.save(dict(x=x,h=h,masks=masks,groups=torch.tensor(groups,device=device)),path)
            series.append(dict(name=name,system='gas_array_rig',file=path.name,sha256=sha(path),rows=len(h),raw_rows=len(raw),
                dimensions=x.shape[-1],factor_names=spec['factor_names'],levels=spec['levels'],
                features=[f'sensor{j}:{stat}' for j in range(16) for stat in ('last20_mean','last20_minus_first20')],
                units=['source_sensor_reading']*32,rates_hz=[100]*32,records=records,**counts))
            print(json.dumps(dict(prepared=name,**counts)),flush=True);del raw
    torch.cuda.synchronize(device)
    result=dict(series=series,episode_exclusions=exclusions,preparation_seconds=time.perf_counter()-start,
                all_three_primary_eligible=all(s['primary_eligible'] for s in series),physical_systems=2,
                source_units_retained_without_resistance_conversion=True,source_receipt_times_not_available=True)
    (out/'manifest.json').write_text(json.dumps(result,indent=2))