"""Download and read GALAH DR4 per-star spectra (observed, uncertainty, best-fit model, mask).

Two storage modes.  `read` keeps the full FITS file on disk (first-pass vetting sample).  The
revision steps (flagged-star vetting, injection-recovery) use `read_compact`, which downloads the
file into memory and keeps only the pixels inside the line segments of the neutron-capture masks
(about 50 kB per star instead of 0.5 MB), because the raw-data disk is nearly full."""
import io, os, time, urllib.request
import numpy as np
from astropy.io import fits
from common import RAW

URL = ('https://cloud.datacentral.org.au/teamdata/GALAH/public/GALAH_DR4/analysis_products_allstar/'
       '{date}/{sid}/{sid}_allstar_fit_spectrum.fits')


def path(sid):
    return RAW + 'spec/%d_allstar_fit_spectrum.fits' % sid


def fetch(sid, retries=3):
    sid = int(sid); p = path(sid)
    if os.path.exists(p) and os.path.getsize(p) > 10000:
        return p
    url = URL.format(date=str(sid)[:6], sid=sid)
    for k in range(retries):
        try:
            with urllib.request.urlopen(url, timeout=60) as r:
                buf = r.read()
            if len(buf) > 10000:
                with open(p, 'wb') as fh:
                    fh.write(buf)
                return p
        except Exception as ex:  # noqa
            err = ex
            time.sleep(1 + 2 * k)
    return None


def read(sid):
    p = fetch(sid)
    if p is None:
        return None
    with fits.open(p) as h:
        t = h[1].data
        return dict(wave=np.array(t['wave'], float), sob=np.array(t['sob'], float),
                    uob=np.array(t['uob'], float), smod=np.array(t['smod'], float),
                    mob=np.array(t['mob'], bool))


COMPACT_DIR = RAW + 'spec_compact/'
PAD = 0.6      # A kept on either side of each line segment


def compact_path(sid):
    return COMPACT_DIR + '%d.npz' % int(sid)


def _windows(masks):
    return sorted((m[3] - PAD, m[4] + PAD) for el in masks for m in masks[el])


def fetch_compact(sid, masks, retries=3):
    """Make sure a compact copy (line segments only) exists; returns its path, or the path of the
    full file if that is already on disk, or None."""
    sid = int(sid)
    if os.path.exists(path(sid)) and os.path.getsize(path(sid)) > 10000:
        return path(sid)
    cp = compact_path(sid)
    if os.path.exists(cp) and os.path.getsize(cp) > 1000:
        return cp
    os.makedirs(COMPACT_DIR, exist_ok=True)
    url = URL.format(date=str(sid)[:6], sid=sid)
    for k in range(retries):
        try:
            with urllib.request.urlopen(url, timeout=60) as r:
                buf = r.read()
            if len(buf) < 10000:
                raise IOError('short file')
            with fits.open(io.BytesIO(buf)) as h:
                t = h[1].data
                w = np.array(t['wave'], float)
                keep = np.zeros(len(w), bool)
                for a, b in _windows(masks):
                    keep |= (w >= a) & (w <= b)
                np.savez_compressed(cp + '.tmp.npz', wave=w[keep], sob=np.array(t['sob'], 'f4')[keep],
                                    uob=np.array(t['uob'], 'f4')[keep], smod=np.array(t['smod'], 'f4')[keep],
                                    mob=np.array(t['mob'], bool)[keep])
            os.replace(cp + '.tmp.npz', cp)
            return cp
        except Exception as ex:  # noqa
            time.sleep(1 + 2 * k)
    return None


def read_compact(sid, masks):
    p = fetch_compact(sid, masks)
    if p is None:
        return None
    if p.endswith('.fits'):
        return read(sid)
    z = np.load(p)
    return dict(wave=z['wave'].astype(float), sob=z['sob'].astype(float), uob=z['uob'].astype(float),
                smod=z['smod'].astype(float), mob=z['mob'].astype(bool))