app/api_client.py

"""Generative API clients for double-exposure separation."""

from __future__ import annotations

import os
from dataclasses import dataclass
from typing import List, Optional, Tuple

import numpy as np
from PIL import Image

from app.preprocessing import to_pil


DEFAULT_SEPARATION_MODEL = os.environ.get(
    "REPLICATE_SEPARATION_MODEL",
    "black-forest-labs/flux-dev",
)

DEMO_PERCENTILES = (35.0, 45.0, 50.0, 55.0, 65.0)
REPLICATE_STRENGTHS = (0.55, 0.65, 0.75)


@dataclass
class SeparationResult:
    """Output from a separation attempt."""

    image_a: np.ndarray  # float RGB [0, 1]
    image_b: np.ndarray
    method: str
    message: str
    candidate_id: str = "default"
    scan_analysis: Optional["ScanAnalysis"] = None  # WP-5.1 Fix 7: carry VLM analysis to avoid re-call
    diagnostics: Optional[dict] = None  # WP-6: structured per-source numbers (e.g. DIP init/best loss)


def _luminance(rgb: np.ndarray) -> np.ndarray:
    return (
        0.2126 * rgb[..., 0] + 0.7152 * rgb[..., 1] + 0.0722 * rgb[..., 2]
    )


def _apply_hard_mask(base: np.ndarray, mask: np.ndarray, fill: float = 0.35) -> np.ndarray:
    gray = np.mean(base, axis=-1, keepdims=True)
    out = base * mask[..., np.newaxis] + gray * (~mask[..., np.newaxis]) * fill
    return np.clip(out, 0.0, 1.0).astype(np.float32)


def _demo_separation_percentile(
    positive_rgb: np.ndarray,
    percentile: float,
    candidate_id: str,
    soft: bool = False,
) -> SeparationResult:
    """Split by luminance percentile; optional soft Gaussian boundary."""
    lum = _luminance(positive_rgb)
    threshold = float(np.percentile(lum, percentile))

    if soft:
        sigma = max(threshold * 0.15, 0.02)
        weight_a = np.exp(-((lum - threshold) ** 2) / (2 * sigma**2))
        weight_a = np.where(lum <= threshold, 1.0, weight_a)
        weight_b = 1.0 - weight_a
        image_a = np.clip(positive_rgb * weight_a[..., np.newaxis], 0.0, 1.0)
        image_b = np.clip(positive_rgb * weight_b[..., np.newaxis], 0.0, 1.0)
        strategy = f"soft_p{int(percentile)}"
    else:
        mask_a = lum <= threshold
        image_a = _apply_hard_mask(positive_rgb, mask_a)
        image_b = _apply_hard_mask(positive_rgb, ~mask_a)
        strategy = f"hard_p{int(percentile)}"

    return SeparationResult(
        image_a=image_a.astype(np.float32),
        image_b=image_b.astype(np.float32),
        method=f"demo_{strategy}",
        message="Demo heuristic separation (no API key).",
        candidate_id=candidate_id,
    )


def _demo_separation_spatial(positive_rgb: np.ndarray, axis: str) -> SeparationResult:
    """Split along horizontal or vertical midline with feathered blend."""
    h, w = positive_rgb.shape[:2]
    if axis == "horizontal":
        coord = np.linspace(0, 1, h)[:, np.newaxis]
        coord = np.broadcast_to(coord, (h, w))
        cid = "spatial_h"
    else:
        coord = np.linspace(0, 1, w)[np.newaxis, :]
        coord = np.broadcast_to(coord, (h, w))
        cid = "spatial_v"

    weight_a = np.clip(1.0 - np.abs(coord - 0.5) * 4.0, 0.0, 1.0)
    weight_b = 1.0 - weight_a
    image_a = np.clip(positive_rgb * weight_a[..., np.newaxis], 0.0, 1.0)
    image_b = np.clip(positive_rgb * weight_b[..., np.newaxis], 0.0, 1.0)

    return SeparationResult(
        image_a=image_a.astype(np.float32),
        image_b=image_b.astype(np.float32),
        method=f"demo_{cid}",
        message="Demo heuristic separation (no API key).",
        candidate_id=cid,
    )


def generate_demo_candidates(
    positive_rgb: np.ndarray,
    num_candidates: int = 3,
) -> List[SeparationResult]:
    """Generate diverse demo candidates without API access."""
    pool: List[SeparationResult] = []

    for p in DEMO_PERCENTILES:
        pool.append(
            _demo_separation_percentile(
                positive_rgb, p, candidate_id=f"hard_p{int(p)}", soft=False
            )
        )
    pool.append(
        _demo_separation_percentile(
            positive_rgb, 50.0, candidate_id="soft_p50", soft=True
        )
    )
    pool.extend([
        _demo_separation_spatial(positive_rgb, "horizontal"),
        _demo_separation_spatial(positive_rgb, "vertical"),
    ])

    return pool[: max(1, min(num_candidates, len(pool)))]


def _replicate_separation(
    positive_rgb: np.ndarray,
    model: str = DEFAULT_SEPARATION_MODEL,
    prompt_strength: float = 0.65,
    candidate_id: str = "replicate_0",
) -> SeparationResult:
    """Call Replicate for generative separation via image-to-image prompting."""
    import replicate

    pil = to_pil(positive_rgb)
    prompt = (
        "Separate this double-exposed photograph into the two distinct original "
        "scenes. Recover clear, photorealistic details from each exposure."
    )

    output = replicate.run(
        model,
        input={
            "prompt": prompt,
            "image": pil,
            "prompt_strength": prompt_strength,
            "num_inference_steps": 28,
            "guidance": 3.5,
        },
    )

    if isinstance(output, list):
        url = str(output[0])
    else:
        url = str(output)

    from io import BytesIO
    import urllib.request

    with urllib.request.urlopen(url) as resp:
        gen_img = Image.open(BytesIO(resp.read())).convert("RGB")

    gen_rgb = np.asarray(gen_img, dtype=np.float32) / 255.0
    if gen_rgb.shape[:2] != positive_rgb.shape[:2]:
        gen_pil = gen_img.resize(
            (positive_rgb.shape[1], positive_rgb.shape[0]),
            Image.Resampling.LANCZOS,
        )
        gen_rgb = np.asarray(gen_pil, dtype=np.float32) / 255.0

    lum_gen = _luminance(gen_rgb)
    lum_orig = _luminance(positive_rgb)
    lum_b = np.clip(lum_orig - lum_gen * 0.5, 0.0, 1.0)
    scale = lum_b[..., np.newaxis] / np.clip(lum_orig[..., np.newaxis], 1e-4, 1.0)
    image_b = np.clip(positive_rgb * scale, 0.0, 1.0).astype(np.float32)

    return SeparationResult(
        image_a=gen_rgb.astype(np.float32),
        image_b=image_b,
        method=f"replicate:{model}@s{prompt_strength:.2f}",
        message=f"Generative separation via {model} (strength={prompt_strength:.2f})",
        candidate_id=candidate_id,
    )


def generate_replicate_candidates(
    positive_rgb: np.ndarray,
    num_candidates: int = 3,
    model: str = DEFAULT_SEPARATION_MODEL,
) -> List[SeparationResult]:
    """Generate multiple Replicate separations with varied prompt strength."""
    strengths = REPLICATE_STRENGTHS[: max(1, num_candidates)]
    results: List[SeparationResult] = []
    for i, strength in enumerate(strengths):
        try:
            results.append(
                _replicate_separation(
                    positive_rgb,
                    model=model,
                    prompt_strength=strength,
                    candidate_id=f"replicate_s{int(strength * 100)}",
                )
            )
        except Exception as exc:
            results.append(
                SeparationResult(
                    image_a=positive_rgb.copy(),
                    image_b=positive_rgb.copy(),
                    method="replicate_error",
                    message=f"Candidate failed: {exc}",
                    candidate_id=f"replicate_fail_{i}",
                )
            )
    return results


def _append_one_demix(candidates, positive_rgb, h_total, confidence_mask, img2img, vlm, method):
    """WP-5.1 Fix 8 helper to dedup append logic."""
    if h_total is None or confidence_mask is None:
        return
    try:
        from app.demix import analyze_scan, residual_demix, DemixConfig
        analysis = analyze_scan(positive_rgb, vlm=vlm)
        cfg = DemixConfig(iterations=2, strength=0.55, use_instruct_edit=("instruct" in method))
        d = residual_demix(positive_rgb, h_total, confidence_mask, img2img, analysis, cfg, method=method)
        if "instruct" in method:
            d.candidate_id = f"demix_instruct_k{analysis.k_judgment:.1f}_i{cfg.iterations}"
        candidates.append(d)
    except Exception:
        pass


def _append_deep_prior(
    candidates,
    positive_rgb,
    density,
    log_exposure,
    confidence_mask,
    film_curve,
    dip_policy=None,
):
    """WP-6 helper: append the Double-DIP candidate (soft-fail, lazy import).

    Skipped silently when density/log_exposure are unavailable (same rule as demix:
    never optimize against the legacy circular objective).
    WP-13.1: dip_policy defaults to DEFAULT_POLICY (bench-safe); app passes APP_POLICY.
    """
    if density is None or log_exposure is None or confidence_mask is None:
        return
    try:
        from baselines.double_dip import double_dip_separate, DoubleDIPConfig
        from scoring_policy import DEFAULT_POLICY
        if film_curve is None:
            from film_physics import get_film_curve
            film_curve = get_film_curve("Generic")
        cfg = DoubleDIPConfig(policy=dip_policy or DEFAULT_POLICY)
        res = double_dip_separate(
            positive_rgb, log_exposure, density, confidence_mask, film_curve, config=cfg
        )
        if res is not None:
            candidates.append(res)
    except Exception:
        pass


def generate_candidates(
    positive_rgb: np.ndarray,
    num_candidates: int = 3,
    api_token: Optional[str] = None,
    model: str = DEFAULT_SEPARATION_MODEL,
    h_total: Optional[np.ndarray] = None,
    confidence_mask: Optional[np.ndarray] = None,
    density: Optional[np.ndarray] = None,
    log_exposure=None,
    include_deep_prior: bool = False,
    film_curve=None,
    dip_policy=None,
) -> Tuple[List[SeparationResult], str]:
    """
    Generate multiple separation candidates for ranking.

    h_total / confidence_mask (WP-5): when provided, demix source is registered
    (stub in demo; replicate+ instruct variant when token).
    density / log_exposure / include_deep_prior / film_curve (WP-6): when the flag
    is set and the density path is live, the Double-DIP source is registered.
    film_curve should be the user's selected stock so DIP optimizes the right physics.

    Returns:
        Tuple of (candidate list, mode description string).
    """
    num_candidates = max(1, min(int(num_candidates), 5))
    token = api_token or os.environ.get("REPLICATE_API_TOKEN", "").strip()

    if not token:
        candidates = generate_demo_candidates(positive_rgb, num_candidates)
        if h_total is not None and confidence_mask is not None:
            # WP-5.1 Fix 4+8: demo uses helper, vlm=None
            from app.demix import stub_cleanup
            _append_one_demix(candidates, positive_rgb, h_total, confidence_mask, stub_cleanup, None, "demix_stub")
        if include_deep_prior:
            _append_deep_prior(
                candidates, positive_rgb, density, log_exposure, confidence_mask, film_curve,
                dip_policy=dip_policy,
            )
        return candidates, "demo"

    os.environ["REPLICATE_API_TOKEN"] = token
    try:
        candidates = generate_replicate_candidates(positive_rgb, num_candidates, model)
        if all(c.method == "replicate_error" for c in candidates):
            raise RuntimeError("All Replicate candidates failed")
        # WP-5.1 Fix 3+4+8: append via helper; gate instruct
        if h_total is not None and confidence_mask is not None:
            from app.demix import replicate_img2img as _rep_img2img
            vlm = (lambda: __import__("app.demix", fromlist=["anthropic_vlm"]).anthropic_vlm if os.environ.get("ANTHROPIC_API_KEY") else None)()
            _append_one_demix(candidates, positive_rgb, h_total, confidence_mask, _rep_img2img, vlm, "demix_replicate")
            if num_candidates >= 3:
                _append_one_demix(candidates, positive_rgb, h_total, confidence_mask, _rep_img2img, vlm, "demix_instruct")
        if include_deep_prior:
            _append_deep_prior(
                candidates, positive_rgb, density, log_exposure, confidence_mask, film_curve,
                dip_policy=dip_policy,
            )
        return candidates, "replicate"
    except Exception as exc:
        candidates = generate_demo_candidates(positive_rgb, num_candidates)
        for c in candidates:
            c.message = f"Replicate unavailable ({exc}). Using demo candidate."
        if h_total is not None and confidence_mask is not None:
            # WP-5.1 Fix 4+8: fallback demo via helper, vlm=None
            from app.demix import stub_cleanup
            _append_one_demix(candidates, positive_rgb, h_total, confidence_mask, stub_cleanup, None, "demix_stub")
        if include_deep_prior:
            _append_deep_prior(
                candidates, positive_rgb, density, log_exposure, confidence_mask, film_curve,
                dip_policy=dip_policy,
            )
        return candidates, "demo_fallback"


def result_to_pil_pair(result: SeparationResult) -> Tuple[Image.Image, Image.Image]:
    """Convert separation arrays to PIL images for display."""
    return to_pil(result.image_a), to_pil(result.image_b)