#!/usr/bin/env python3
"""Minimum numerical reproduction of the acquisition idea behind qOBM/SCqOBM.

This is an educational simulation, not the authors' clinical reconstruction code.
It creates a synthetic, pathology-like phase object, simulates four oblique
illumination captures, and compares physics-based phase recovery from:

1. four captures (two opposing pairs), and
2. one capture from the +x illumination direction.

Only NumPy and Pillow are required.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
from PIL import Image, ImageDraw, ImageFont


def synthetic_phase(size: int, seed: int) -> np.ndarray:
    """Create a deterministic phase phantom with cells and anisotropic tissue."""
    rng = np.random.default_rng(seed)
    axis = np.linspace(-1.0, 1.0, size, endpoint=False)
    x, y = np.meshgrid(axis, axis)

    # Slowly varying tissue background.
    phase = 0.10 * np.exp(-((x / 0.92) ** 6 + (y / 0.82) ** 6))
    phase += 0.035 * np.sin(3.0 * np.pi * x + 0.6 * np.sin(2.0 * np.pi * y))

    # A structure that varies mostly along y. It is deliberately difficult for
    # an illumination direction that measures only the x phase gradient.
    band_center = 0.30 + 0.035 * np.sin(2.0 * np.pi * x)
    phase += 0.30 * np.exp(-((y - band_center) ** 2) / (2.0 * 0.045**2))

    # Elliptical Gaussian nuclei with random orientation and size.
    for _ in range(95):
        cx, cy = rng.uniform(-0.82, 0.82, size=2)
        sx = rng.uniform(0.012, 0.035)
        sy = rng.uniform(0.018, 0.060)
        theta = rng.uniform(0.0, np.pi)
        amplitude = rng.uniform(0.12, 0.42)
        cos_t, sin_t = np.cos(theta), np.sin(theta)
        xr = cos_t * (x - cx) + sin_t * (y - cy)
        yr = -sin_t * (x - cx) + cos_t * (y - cy)
        phase += amplitude * np.exp(-0.5 * ((xr / sx) ** 2 + (yr / sy) ** 2))

    # A few gland-like rings add structures at a larger scale.
    for cx, cy, radius in [(-0.46, -0.38, 0.17), (0.38, -0.34, 0.14)]:
        radial = np.sqrt((x - cx) ** 2 + (y - cy) ** 2)
        phase += 0.18 * np.exp(-((radial - radius) ** 2) / (2.0 * 0.018**2))

    phase -= phase.min()
    phase /= phase.max()
    return phase.astype(np.float64)


def simulate_captures(
    phase: np.ndarray, contrast: float, noise: float, seed: int
) -> dict[str, np.ndarray]:
    """Approximate opposing oblique captures with a weak-phase gradient model."""
    rng = np.random.default_rng(seed + 1)
    grad_y, grad_x = np.gradient(phase)
    gradient_scale = max(float(np.max(np.abs(grad_x))), float(np.max(np.abs(grad_y))))
    grad_x /= gradient_scale
    grad_y /= gradient_scale

    captures: dict[str, np.ndarray] = {}
    for name, gradient, sign in [
        ("plus_x", grad_x, 1.0),
        ("minus_x", grad_x, -1.0),
        ("plus_y", grad_y, 1.0),
        ("minus_y", grad_y, -1.0),
    ]:
        shot_noise = rng.normal(0.0, noise, phase.shape)
        captures[name] = np.clip(1.0 + sign * contrast * gradient + shot_noise, 0.05, 1.95)

    captures["gradient_scale"] = np.asarray(gradient_scale)
    return captures


def integrate_gradients(
    grad_x: np.ndarray, grad_y: np.ndarray, regularization: float
) -> np.ndarray:
    """Least-squares integration of a 2-D gradient field in Fourier space."""
    height, width = grad_x.shape
    fx = np.fft.fftfreq(width)[None, :]
    fy = np.fft.fftfreq(height)[:, None]
    derivative_x = 2j * np.pi * fx
    derivative_y = 2j * np.pi * fy

    numerator = (
        np.conj(derivative_x) * np.fft.fft2(grad_x)
        + np.conj(derivative_y) * np.fft.fft2(grad_y)
    )
    denominator = (
        np.abs(derivative_x) ** 2
        + np.abs(derivative_y) ** 2
        + regularization
    )
    phase_hat = numerator / denominator
    phase_hat[0, 0] = 0.0
    return np.fft.ifft2(phase_hat).real


def affine_align(estimate: np.ndarray, target: np.ndarray) -> np.ndarray:
    """Remove the arbitrary phase offset and scale before comparing results."""
    design = np.column_stack([estimate.ravel(), np.ones(estimate.size)])
    scale, offset = np.linalg.lstsq(design, target.ravel(), rcond=None)[0]
    return scale * estimate + offset


def metrics(estimate: np.ndarray, target: np.ndarray) -> dict[str, float]:
    difference = estimate - target
    dynamic_range = float(np.ptp(target))
    nrmse = float(np.sqrt(np.mean(difference**2)) / dynamic_range)
    correlation = float(np.corrcoef(estimate.ravel(), target.ravel())[0, 1])

    # The row-mean profile isolates structures that vary along y but are nearly
    # constant along x: the ambiguous band for +x-only illumination.
    target_profile = target.mean(axis=1)
    estimate_profile = estimate.mean(axis=1)
    if np.std(estimate_profile) < 1e-12:
        orthogonal_profile_correlation = 0.0
    else:
        orthogonal_profile_correlation = float(
            np.corrcoef(estimate_profile, target_profile)[0, 1]
        )
    return {
        "nrmse": nrmse,
        "pearson_r": correlation,
        "orthogonal_profile_r": orthogonal_profile_correlation,
    }


def colorize(values: np.ndarray, low: float, high: float) -> Image.Image:
    """Apply a compact purple-to-lime colormap without Matplotlib."""
    normalized = np.clip((values - low) / max(high - low, 1e-12), 0.0, 1.0)
    stops = np.asarray(
        [
            [20, 12, 42],
            [72, 38, 117],
            [188, 73, 150],
            [243, 170, 92],
            [220, 252, 130],
        ],
        dtype=np.float64,
    )
    position = normalized * (len(stops) - 1)
    left = np.floor(position).astype(int)
    right = np.minimum(left + 1, len(stops) - 1)
    fraction = (position - left)[..., None]
    rgb = stops[left] * (1.0 - fraction) + stops[right] * fraction
    return Image.fromarray(np.uint8(np.clip(rgb, 0, 255)), mode="RGB")


def gray_image(values: np.ndarray) -> Image.Image:
    low, high = np.percentile(values, [1.0, 99.0])
    normalized = np.clip((values - low) / max(high - low, 1e-12), 0.0, 1.0)
    return Image.fromarray(np.uint8(normalized * 255), mode="L").convert("RGB")


def make_montage(
    phase: np.ndarray,
    captures: dict[str, np.ndarray],
    four_capture: np.ndarray,
    single_capture: np.ndarray,
    output_path: Path,
) -> None:
    phase_low, phase_high = float(phase.min()), float(phase.max())
    dpc_x = (captures["plus_x"] - captures["minus_x"]) / (
        captures["plus_x"] + captures["minus_x"] + 1e-12
    )
    error = np.abs(single_capture - phase)
    panels = [
        (colorize(phase, phase_low, phase_high), "Ground truth phase"),
        (gray_image(captures["plus_x"]), "One oblique capture (+x)"),
        (gray_image(dpc_x), "Opposing-pair DPC (x)"),
        (colorize(four_capture, phase_low, phase_high), "Four-capture recovery"),
        (colorize(single_capture, phase_low, phase_high), "One-capture physics baseline"),
        (colorize(error, 0.0, max(float(np.percentile(error, 99)), 1e-6)), "Single-capture absolute error"),
    ]

    tile_width, tile_height = phase.shape[1], phase.shape[0]
    caption_height = 34
    gap = 12
    canvas = Image.new(
        "RGB",
        (3 * tile_width + 4 * gap, 2 * (tile_height + caption_height) + 3 * gap),
        (10, 13, 18),
    )
    draw = ImageDraw.Draw(canvas)
    try:
        font = ImageFont.load_default(size=14)
    except TypeError:  # Pillow 10.0 compatibility
        font = ImageFont.load_default()
    for index, (image, caption) in enumerate(panels):
        row, column = divmod(index, 3)
        x = gap + column * (tile_width + gap)
        y = gap + row * (tile_height + caption_height + gap)
        canvas.paste(image, (x, y))
        draw.text((x, y + tile_height + 9), caption, fill=(226, 232, 240), font=font)

    output_path.parent.mkdir(parents=True, exist_ok=True)
    canvas.save(output_path, optimize=True)


def run(args: argparse.Namespace) -> dict[str, object]:
    phase = synthetic_phase(args.size, args.seed)
    captures = simulate_captures(phase, args.contrast, args.noise, args.seed)
    gradient_scale = float(captures["gradient_scale"])

    dpc_x = (captures["plus_x"] - captures["minus_x"]) / (
        captures["plus_x"] + captures["minus_x"] + 1e-12
    )
    dpc_y = (captures["plus_y"] - captures["minus_y"]) / (
        captures["plus_y"] + captures["minus_y"] + 1e-12
    )
    grad_x_four = dpc_x * gradient_scale / args.contrast
    grad_y_four = dpc_y * gradient_scale / args.contrast
    four_capture = integrate_gradients(
        grad_x_four, grad_y_four, args.regularization
    )
    four_capture = affine_align(four_capture, phase)

    # A single raw capture mixes a DC background with one directional phase
    # gradient. Mean subtraction is an optimistic estimate of that background.
    grad_x_single = (
        captures["plus_x"] - np.mean(captures["plus_x"])
    ) * gradient_scale / args.contrast
    single_capture = integrate_gradients(
        grad_x_single, np.zeros_like(grad_x_single), args.regularization
    )
    single_capture = affine_align(single_capture, phase)

    result = {
        "configuration": {
            "size": args.size,
            "seed": args.seed,
            "contrast": args.contrast,
            "noise_sigma": args.noise,
            "regularization": args.regularization,
        },
        "four_capture": metrics(four_capture, phase),
        "single_capture_physics_baseline": metrics(single_capture, phase),
        "interpretation": (
            "The one-capture baseline cannot recover phase components that are "
            "constant along x and vary along y. SCqOBM uses a learned tissue prior "
            "to estimate much of this missing information, but the paper reports a "
            "remaining narrow missing frequency band."
        ),
    }

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    make_montage(
        phase,
        captures,
        four_capture,
        single_capture,
        output_dir / "comparison.png",
    )
    (output_dir / "metrics.json").write_text(
        json.dumps(result, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
    )
    return result


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--size", type=int, default=256)
    parser.add_argument("--seed", type=int, default=20260721)
    parser.add_argument("--contrast", type=float, default=0.32)
    parser.add_argument("--noise", type=float, default=0.008)
    parser.add_argument("--regularization", type=float, default=2e-5)
    parser.add_argument("--output-dir", default="scqobm-output")
    return parser.parse_args()


if __name__ == "__main__":
    results = run(parse_args())
    print(json.dumps(results, indent=2, ensure_ascii=False))