"""Precompute the demo's example reservoirs on the Space's GPU, through its /generate endpoint (same model and
sampler as the live demo). Needs a Hugging Face token (huggingface-cli login).
    python make_examples.py NAME     -> examples/NAME.npz"""
import io, json, sys, time
import urllib.request
import numpy as np
from huggingface_hub import get_token
from matplotlib.path import Path

SPACE = 'https://baghishov-resflow-demo.hf.space'
Z = 32
WELL_LOGS = {                                              # as in backend.py: base (z = 0) to top, 1 sand, 0 mud
    'Two sand bodies': [0] * 6 + [1] * 7 + [0] * 7 + [1] * 8 + [0] * 4,
    'Thick sand at the base': [1] * 14 + [0] * 18,
}
HEADERS = {'Authorization': f'Bearer {get_token()}'}


def fetch(url, data=None):
    req = urllib.request.Request(url, data=None if data is None else json.dumps(data).encode(),
                                 headers={**HEADERS, 'Content-Type': 'application/json'})
    return urllib.request.urlopen(req, timeout=900)


def call(api, data):
    """Run a Gradio API endpoint and return its output list."""
    eid = json.load(fetch(f'{SPACE}/gradio_api/call/{api}', {'data': data}))['event_id']
    event = None
    with fetch(f'{SPACE}/gradio_api/call/{api}/{eid}') as r:
        for line in (l.decode().rstrip('\n') for l in r):
            if line.startswith('event:'):
                event = line[6:].strip()
            elif line.startswith('data:') and event in ('complete', 'error'):
                if event == 'error':
                    raise RuntimeError(line)
                return json.loads(line[5:])
    raise RuntimeError('stream ended without a result')


def blob(nx, ny, seed=3):
    """An irregular reservoir outline: a smooth random closed curve."""
    rng = np.random.default_rng(seed)
    t = np.linspace(0, 2 * np.pi, 400)
    r = 0.36 + sum(rng.uniform(0.02, 0.06) * np.cos(k * t + rng.uniform(0, 6.3)) for k in range(2, 6))
    x, y = 0.5 + r * np.cos(t), 0.5 + 0.9 * r * np.sin(t)
    gx, gy = np.meshgrid((np.arange(nx) + .5) / nx, (np.arange(ny) + .5) / ny, indexing='ij')
    return Path(np.c_[x, y]).contains_points(np.c_[gx.ravel(), gy.ravel()]).reshape(nx, ny)


EX = {
    'lobe_outline': dict(title='Deep-water lobes, irregular outline, two wells', label='Lobe (deep-water)', nx=256, ny=256,
                         azimuth=45.0, outline=True,
                         wells=[(96, 120, 'Two sand bodies'), (168, 104, 'Thick sand at the base')]),
    'meander_strip': dict(title='Meander belt along flow', label='Meander (oxbow)', nx=384, ny=128, azimuth=0.0,
                          outline=False, wells=[]),
}
name = sys.argv[1]; e = EX[name]
wells = [(x, y, WELL_LOGS[log]) for x, y, log in e['wells']]
t = time.time()
out = call('generate', [e['label'], e['nx'], e['ny'], json.dumps(wells), 50, 0, json.dumps({'azimuth': e['azimuth']})])
z = np.load(io.BytesIO(fetch(out[0]['url']).read()))
f = z['facies']
params = {k: float(z[k]) for k in z.files if k not in ('facies', 'active', 'environment')}
for x, y, log in wells:                                    # wells must be honoured exactly
    assert (f[x, y] == np.array(log)).all(), f'well ({x}, {y}) not honoured'
active = blob(e['nx'], e['ny']) if e['outline'] else np.ones((e['nx'], e['ny']), bool)
np.savez_compressed(f'examples/{name}.npz', facies=f, active=active,
                    wells=np.array([[x, y] + list(log) for x, y, log in wells], np.int16).reshape(-1, 2 + Z),
                    meta=json.dumps({'title': e['title'], 'label': e['label'], 'nx': e['nx'], 'ny': e['ny'],
                                     'params': params, 'steps': 50, 'seed': 0}))
print(f'{name}: {f.shape}, sand {f[active].mean():.2f}, wells honoured {len(wells)}, {time.time() - t:.0f}s', flush=True)