lf_mri_platform/services/seq-interp/api.py

"""
seq_interp FastAPI service.

Endpoints:
  GET  /health              - liveness probe
  POST /interpret/          - upload .seq file, run full pipeline, return task_id
  GET  /status/             - status of all submitted tasks
  GET  /result/{task_id}    - full result: xml_text, post_json, metadata, waveforms
"""
from __future__ import annotations

import asyncio
import io
import json
import os
import shutil
import uuid
import zipfile
from typing import Any

from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import Response

from src.config import config
from src.hardware.constraints import HardwareConstraints
from src.interfaces.pulseq_adapter import PulseqLoader
from src.core.synchronizer import Synchronizer
from src.interfaces.xml_generator import XMLGenerator
from src.interfaces.rf_exporter import RFExporter
from src.interfaces.gradient_exporter import GradientExporter
from src.interfaces.picoscope_exporter import PicoScopeExporter
from src.interfaces.post_request_generator import PostRequestGenerator

app = FastAPI(title="seq-interp", version="1.0.0")

UPLOAD_DIR = config.get("upload_dir", "data/input")
OUTPUT_DIR = config.get("output_dir", "data/output")
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)

# in-memory cache: task_id -> {"status": str, "result": dict | None}
_tasks: dict[str, dict] = {}

_MAX_WAVEFORM_POINTS = 8000   # downsample for JSON transport

_STATUS_FILE = "_status.json"
_RESULT_FILE = "_result.json"


# -- disk persistence --------------------------------------------------------

def _task_dir(task_id: str) -> str:
    return os.path.join(OUTPUT_DIR, task_id)


def _persist_task(task_id: str, entry: dict) -> None:
    """Write task status (and result if present) to disk."""
    d = _task_dir(task_id)
    os.makedirs(d, exist_ok=True)
    with open(os.path.join(d, _STATUS_FILE), "w", encoding="utf-8") as f:
        json.dump({"status": entry["status"]}, f)
    if entry.get("result") is not None:
        with open(os.path.join(d, _RESULT_FILE), "w", encoding="utf-8") as f:
            json.dump(entry["result"], f, ensure_ascii=False)


def _load_task_from_disk(task_id: str) -> dict | None:
    """Restore task entry from disk (used after container restart)."""
    d = _task_dir(task_id)
    status_path = os.path.join(d, _STATUS_FILE)
    result_path = os.path.join(d, _RESULT_FILE)
    if not os.path.exists(status_path):
        return None
    with open(status_path, encoding="utf-8") as f:
        status = json.load(f).get("status", "unknown")
    result = None
    if os.path.exists(result_path):
        with open(result_path, encoding="utf-8") as f:
            result = json.load(f)
    return {"status": status, "result": result}


# -- helpers ----------------------------------------------------------------

def _downsample(arr, max_pts: int) -> list:
    """Convert numpy array to list, downsampling if too long."""
    try:
        import numpy as np
        a = np.asarray(arr, dtype=float).flatten()
        if len(a) > max_pts:
            idx = np.linspace(0, len(a) - 1, max_pts, dtype=int)
            a = a[idx]
        return [float(x) for x in a]
    except Exception:
        return []


def _extract_waveforms(seq_data: dict, sync_data: dict) -> dict:
    """Extract waveform arrays from seq_data for GUI display."""
    try:
        import numpy as _np
    except ImportError:
        _np = None

    waveforms: dict[str, list] = {}

    # Gradients
    for key in ("gx", "gy", "gz", "t_gx", "t_gy", "t_gz"):
        if key in seq_data:
            waveforms[key] = _downsample(seq_data[key], _MAX_WAVEFORM_POINTS)

    # RF — seq_data stores a complex array under "rf" / "t_rf".
    # Split into amplitude and phase for JSON transport (complex is not serialisable).
    if "rf" in seq_data and "t_rf" in seq_data and _np is not None:
        rf = _np.asarray(seq_data["rf"])
        waveforms["rf_amp"]   = _downsample(_np.abs(rf),    _MAX_WAVEFORM_POINTS)
        waveforms["rf_phase"] = _downsample(_np.angle(rf),  _MAX_WAVEFORM_POINTS)
        waveforms["t_rf"]     = _downsample(seq_data["t_rf"], _MAX_WAVEFORM_POINTS)

    # Sync gate arrays
    for key in ("gate_adc", "gate_rf", "gate_tr_switch", "blocks_duration"):
        if key in sync_data:
            waveforms[key] = _downsample(sync_data[key], _MAX_WAVEFORM_POINTS)

    return waveforms


async def _run_pipeline(file_path: str, task_id: str,
                        hw_overrides: dict | None = None) -> None:
    """Run the full interpretation pipeline and store result in _tasks."""
    out_dir = os.path.join(OUTPUT_DIR, task_id)
    os.makedirs(out_dir, exist_ok=True)
    try:
        hw = HardwareConstraints(json_path=config.get("hw_config_path"))
        # GUI-supplied hardware overrides (delay values, raster, and the
        # TR/RF/START enable flags) are applied on top of the service config,
        # so the GUI controls work without a local interpreter copy.
        for key, value in (hw_overrides or {}).items():
            setattr(hw, key, value)
        hw_cfg = config.hw_config

        loader = PulseqLoader(hw)
        seq_data = await asyncio.to_thread(loader.load, file_path)
        params = seq_data.get("params", {})

        sync = Synchronizer(hw)
        sync_data = await asyncio.to_thread(sync.process, seq_data["sequence"])

        xml_gen = XMLGenerator()
        xml_path = os.path.join(out_dir, "sync_v2.xml")
        adc_values, adc_starts = await asyncio.to_thread(
            xml_gen.generate, sync_data, xml_path, hw
        )
        with open(xml_path, encoding="utf-8") as fh:
            xml_text = fh.read()

        export_tasks = [
            asyncio.to_thread(RFExporter().export, seq_data, params, out_dir)
        ]
        if all(k in seq_data for k in ("gx", "gy", "gz")):
            export_tasks.append(
                asyncio.to_thread(GradientExporter().export, seq_data, params, out_dir)
            )
        iadc = hw_cfg.get("iadc", {})
        export_tasks.append(asyncio.to_thread(
            PicoScopeExporter().generate,
            adc_values, adc_starts, out_dir, hw,
            sampling_freq=iadc.get("srate", 8e6),
            num_channels=iadc.get("n_channels", 3),
        ))
        await asyncio.gather(*export_tasks)

        post_gen = PostRequestGenerator()
        post_payload = post_gen.build(
            seq_data=seq_data,
            adc_values=adc_values,
            sequence_path=file_path,
            output_dir=out_dir,
            hw_cfg=hw_cfg,
            rf_raster_time=params.get("rf_raster_time", 1e-6),
        )
        post_gen.write(post_payload, out_dir)

        blocks = seq_data.get("blocks", [])
        total_s = sum(sync_data.get("blocks_duration", []))
        adc_blocks = [b for b in blocks if b.get("has_adc")]

        # Serialize only the fields needed by the GUI block-row builder
        blocks_slim = [
            {"has_adc": bool(b.get("has_adc")),
             "type":    list(b.get("type", []))}
            for b in blocks
        ]

        result: dict[str, Any] = {
            "task_id": task_id,
            "status": "completed",
            "output_dir": out_dir,
            "xml_text": xml_text,
            "post_json": post_payload,
            "metadata": {
                "block_count": len(blocks),
                "sync_block_count": sync_data.get("number_of_blocks", 0),
                "adc_count": len(adc_blocks),
                "adc_windows": len(adc_values),
                "total_duration_ms": round(total_s * 1e3, 4),
                "rf_raster_us": params.get("rf_raster_time", 1e-6) * 1e6,
                "grad_raster_us": params.get("grad_raster_time", 1e-5) * 1e6,
            },
            "blocks":   blocks_slim,
            "waveforms": _extract_waveforms(seq_data, sync_data),
        }
        entry = {"status": "completed", "result": result}
        _tasks[task_id] = entry
        _persist_task(task_id, entry)

    except Exception as exc:
        entry = {"status": f"failed: {exc}", "result": None}
        _tasks[task_id] = entry
        _persist_task(task_id, entry)


# -- endpoints ----------------------------------------------------------------

@app.get("/health")
def health():
    return {"status": "ok"}


@app.post("/interpret/")
async def interpret_endpoint(
    file: UploadFile = File(...),
    hw_overrides: str | None = Form(None),
):
    """
    Upload a .seq file and run the full interpretation pipeline.

    hw_overrides : optional JSON object string of HardwareConstraints attribute
                   overrides (e.g. delay values and TR/RF/START enable flags),
                   applied on top of the service-side hw config.
    """
    overrides: dict = {}
    if hw_overrides:
        try:
            overrides = json.loads(hw_overrides)
            if not isinstance(overrides, dict):
                raise ValueError("hw_overrides must be a JSON object")
        except (ValueError, json.JSONDecodeError) as exc:
            raise HTTPException(status_code=422,
                                detail=f"Invalid hw_overrides: {exc}")

    task_id = str(uuid.uuid4())
    # Store the file under the task_id so parallel uploads never collide
    task_upload_dir = os.path.join(UPLOAD_DIR, task_id)
    os.makedirs(task_upload_dir, exist_ok=True)
    file_path = os.path.join(task_upload_dir, file.filename)
    with open(file_path, "wb") as buf:
        shutil.copyfileobj(file.file, buf)

    entry: dict = {"status": "processing", "result": None}
    _tasks[task_id] = entry
    _persist_task(task_id, entry)
    asyncio.create_task(_run_pipeline(file_path, task_id, overrides))
    return {"status": "accepted", "task_id": task_id,
            "message": f"Processing {file.filename}"}


@app.get("/status/")
def status_endpoint():
    """Return the status of all submitted tasks."""
    return {"tasks": {tid: v["status"] for tid, v in _tasks.items()}}


@app.get("/artifacts/{task_id}")
def artifacts_endpoint(task_id: str):
    """
    Download all output artifacts for a completed task as a zip
    (sync_v2.xml, rf_*.bin, gx/gy/gz.txt, picoscope_params.xml,
    post_request.json). Used by the GUI's Export action in pure-client mode.
    """
    task_dir = _task_dir(task_id)
    if not os.path.isdir(task_dir):
        raise HTTPException(status_code=404,
                            detail=f"No artifacts for task '{task_id}'")
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
        for name in os.listdir(task_dir):
            if name in (_STATUS_FILE, _RESULT_FILE):
                continue
            full = os.path.join(task_dir, name)
            if os.path.isfile(full):
                zf.write(full, arcname=name)
    buf.seek(0)
    return Response(
        content=buf.getvalue(),
        media_type="application/zip",
        headers={"Content-Disposition": f'attachment; filename="{task_id}.zip"'},
    )


@app.get("/result/{task_id}")
def result_endpoint(task_id: str):
    """Return full interpretation result (xml_text, post_json, metadata, waveforms)."""
    entry = _tasks.get(task_id)
    if entry is None:
        # Container may have restarted — try restoring from disk
        entry = _load_task_from_disk(task_id)
        if entry is None:
            raise HTTPException(status_code=404, detail=f"Task '{task_id}' not found")
        _tasks[task_id] = entry   # restore to memory cache
    if entry["status"] == "processing":
        raise HTTPException(status_code=202, detail="Still processing")
    if entry["result"] is None:
        raise HTTPException(status_code=500, detail=entry["status"])
    return entry["result"]