#!/usr/bin/env python3
"""Extract small real HST morphology data; requires NumPy, no image synthesis."""
from pathlib import Path
import hashlib
import json
import shutil
import warnings
import math
import argparse
import numpy as np

REPOSITORY = next((parent for parent in Path(__file__).resolve().parents if (parent/'src/features/labs').exists()), None)
ROOT = REPOSITORY or Path.cwd()
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--source-directory', type=Path, default=Path('/tmp/sandboxphysics-ac11-source'))
parser.add_argument('--output-directory', type=Path, default=ROOT/'src/features/labs/comet-dust/data' if REPOSITORY else ROOT)
parser.add_argument('--documentation-directory', type=Path, default=ROOT/'docs/astronomy/ac11' if REPOSITORY else ROOT)
arguments = parser.parse_args()
CACHE = arguments.source_directory
OUTPUT = arguments.output_directory
DOCS = arguments.documentation_directory
FITS_SHA = 'bd436ea87d7ab05ba8e10280e03ea04887db7055841260fb8b4bdf0b547e33f7'
JPEG_SHA = '96204fc6efbe4d13fe6ba9531fea8e50981bd8ee8c459f59c43ec1ea72bd989a'

def artifact(path, data):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, separators=(',', ':'), allow_nan=False) + '\n')
    raw = path.read_bytes()
    return {'path': str(path.relative_to(ROOT)) if path.is_relative_to(ROOT) else path.name, 'bytes': len(raw), 'sha256': hashlib.sha256(raw).hexdigest()}

def main():
    source = CACHE / 'ic7301ooq_drz.fits'
    assert hashlib.sha256(source.read_bytes()).hexdigest() == FITS_SHA
    assert hashlib.sha256((CACHE / 'ic7301ooq_drz.jpg').read_bytes()).hexdigest() == JPEG_SHA
    # Read the source FITS headers directly so re-extraction needs only the original file.
    raw_fits = source.read_bytes()
    offset = 0
    hdus = []
    while offset < len(raw_fits):
        header = {}
        ended = False
        while not ended:
            block = raw_fits[offset:offset+2880].decode('ascii')
            offset += 2880
            for start in range(0,2880,80):
                card = block[start:start+80]
                if card[:8].strip() == 'END':
                    ended = True
                    break
                if card[8:10] != '= ':
                    continue
                value = card[10:].strip()
                if value.startswith("'"):
                    value = value.split("'")[1].strip()
                else:
                    value = value.split('/')[0].strip()
                    try:
                        value = float(value.replace('D','E')) if any(c in value for c in '.ED') else int(value)
                    except ValueError:
                        pass
                header[card[:8].strip()] = value
        size = abs(int(header['BITPIX']))//8*math.prod(int(header['NAXIS'+str(i)]) for i in range(1,int(header['NAXIS'])+1)) if header['NAXIS'] else 0
        hdus.append({'header':header,'offset':offset})
        offset += math.ceil(size/2880)*2880
    header = hdus[1]['header']
    image = np.memmap(source, dtype='>f4', offset=hdus[1]['offset'], shape=(1087, 1030))
    weights = np.memmap(source, dtype='>f4', offset=hdus[2]['offset'], shape=(1087, 1030))
    blocks = np.where(weights[:1080, :1024] > 0, image[:1080, :1024], np.nan).reshape(135,8,128,8)
    with warnings.catch_warnings():
        warnings.simplefilter('ignore')
        medians = np.nanmedian(blocks, axis=(1,3))
    coverage = np.sum(np.isfinite(blocks), axis=(1,3)) / 64
    y, x = np.unravel_index(np.nanargmax(medians), medians.shape)
    global_y, global_x = np.unravel_index(np.nanargmax(image), image.shape)
    observations = {'source': {'dataset':'ic7301ooq', 'program':13198, 'uri':'mast:HST/product/ic7301ooq_drz.fits',
       'sourceBytes':source.stat().st_size,'sha256':FITS_SHA,'scienceHdu':1,'scienceDataOffsetBytes':hdus[1]['offset'],
       'calibrationDate':hdus[0]['header']['DATE'],'calibrationVersion':hdus[0]['header']['CAL_VER'],
       'acknowledgement':'HST(WFC3), program13198, PI J.-Y. Li; public data from MAST, STScI operated by AURA under NASA contract NAS5-26555'},
       'observation':{key:hdus[0]['header'][key] for key in ['DATE-OBS','TIME-OBS','EXPSTART','EXPEND','EXPTIME']},
       'wcs':{key:header[key] for key in ['NAXIS1','NAXIS2','CTYPE1','CTYPE2','CRPIX1','CRPIX2','CRVAL1','CRVAL2','CD1_1','CD1_2','CD2_1','CD2_2','ORIENTAT','BUNIT']},
       'processing':{'operation':'Median of finite, positive-coverage native8x8 blocks; morphology diagnostic, not aperture photometry',
         'binSize':8,'retainedNativeShape':[1080,1024],'coordinateConvention':'FITS native pixels are1-based with+y upward; array rows increase toward+y; block centres use zero-based native pixel coordinate8i+3.5',
         'nativeOmittedEdgeRows':7,'nativeOmittedEdgeColumns':6,'width':128,'height':135},
       'values':[[float(v) if np.isfinite(v) else None for v in row] for row in medians],
       'coverage':coverage.tolist(), 'referenceBlockNucleus':[int(x),int(y)],
       'naiveBrightestNativePixel':[int(global_x),int(global_y)],
       'kind':'Measured calibrated HST image, explicitly rebinned with a median; no generated tail pixels'}
    ephemeris = json.loads((CACHE / 'ephemeris.json').read_text())
    artifacts = [artifact(OUTPUT / 'ison-image.json',observations),artifact(OUTPUT / 'ison-ephemeris.json',ephemeris),artifact(OUTPUT / 'ison-perturbers.json',json.loads((CACHE/'perturber-sample.json').read_text()))]
    shutil.copy2(CACHE/'ic7301ooq_drz.jpg',OUTPUT/'ic7301ooq_drz.jpg')
    raw_queries = []
    raw_contents = {}
    for path in sorted(CACHE.glob('horizons-*.json')):
        if 'test' in path.name:
            continue
        target = DOCS / 'source-queries' / path.name
        target.parent.mkdir(parents=True,exist_ok=True)
        shutil.copy2(path,target)
        raw = target.read_bytes()
        raw_contents[path.name] = raw.decode('utf-8')
        raw_queries.append({'path':str(target.relative_to(ROOT)) if target.is_relative_to(ROOT) else target.name,'bytes':len(raw),'sha256':hashlib.sha256(raw).hexdigest()})
    (DOCS/'horizons-source-queries.json').write_text(json.dumps(raw_contents,indent=2)+'\n')
    manifest={'retrievedDateUtc':'2026-10-28','artifacts':artifacts,
       'imageSource':{'url':'https://mast.stsci.edu/api/v0.1/Download/file?uri=mast:HST/product/ic7301ooq_drz.fits','bytes':source.stat().st_size,'sha256':FITS_SHA},
       'previewSource':{'url':'https://mast.stsci.edu/api/v0.1/Download/file?uri=mast:HST/product/ic7301ooq_drz.jpg','bytes':190967,'sha256':JPEG_SHA},
       'dataUsePolicy':'https://archive.stsci.edu/publishing/data-use','dataUse':'Public mission science data; HST/MAST/program attribution preserved. No paper or archive web text is redistributed.',
       'horizonsRawQueries':raw_queries,'frameAndTime':'ICRF; exact HST UTC midpoint converted with Horizons TDB-UTC. Heliocentric release states, barycentric Earth at reception and Sun at emission.',
       'frozenVersions':'HST2026-06-01,CAL_VER3.7.3; JPL#57, solution2025-Nov-21_13:42:48, DE441, SB441-N16, EOP eop.261027.p270123'}
    (DOCS/'provenance.json').write_text(json.dumps(manifest,indent=2)+'\n')
    print(json.dumps({'imageShape':[135,128],'nucleusBlock':[int(x),int(y)],'artifacts':artifacts},indent=2))

if __name__ == '__main__':
    main()