"""Synthesize a clean calcium-like movie + a noisy version (Poisson + Gaussian).

A few Gaussian blobs ("neurons") brighten and decay over time on a smooth
background. The CLEAN movie is the temporally-correlated signal; the NOISY movie
adds per-frame, temporally-INDEPENDENT shot (Poisson) + read (Gaussian) noise.
This is exactly the regime DeepInterpolation targets: independent noise can be
predicted away from temporal neighbors, correlated signal cannot.
"""
from __future__ import annotations

import numpy as np


def clean_movie(T: int = 120, H: int = 96, W: int = 96, n_neurons: int = 10,
                fps: float = 10.0, seed: int = 0) -> np.ndarray:
    """A clean, noise-free calcium-like movie, scaled to a fluorescence range."""
    rng = np.random.default_rng(seed)
    ys = rng.uniform(0.12, 0.88, n_neurons) * H
    xs = rng.uniform(0.12, 0.88, n_neurons) * W
    radii = rng.uniform(3.0, 6.0, n_neurons)

    yy, xx = np.mgrid[0:H, 0:W]
    footprints = np.stack([
        np.exp(-(((yy - y) ** 2 + (xx - x) ** 2) / (2 * r ** 2)))
        for y, x, r in zip(ys, xs, radii)
    ])  # (N, H, W)

    decay = np.exp(-1.0 / (fps * 0.7))      # tau ~ 0.7 s
    traces = np.zeros((n_neurons, T), dtype=np.float32)
    for n in range(n_neurons):
        rate = rng.uniform(0.04, 0.10)       # spikes / frame
        spikes = rng.random(T) < rate
        amp = rng.uniform(0.9, 1.7)
        c = 0.0
        for t in range(T):
            c = c * decay + (amp if spikes[t] else 0.0)
            traces[n, t] = c

    movie = np.tensordot(traces.T, footprints, axes=(1, 0))  # (T, H, W)
    background = 0.25 + 0.10 * np.sin(yy / H * np.pi)[None]
    movie = movie + background
    movie = np.clip(movie, 0, None)

    movie = (movie / movie.max() * 3000.0 + 300.0).astype(np.float32)
    return movie


def add_noise(clean: np.ndarray, gain: float = 0.6, read_sigma: float = 60.0,
              seed: int = 0) -> np.ndarray:
    """Add per-frame Poisson (shot) + Gaussian (read) noise.

    `gain` scales the Poisson variance (photons per intensity unit); lower gain
    => noisier. `read_sigma` is additive Gaussian read noise (intensity units).
    The noise is independent per pixel and per frame (temporally uncorrelated).
    """
    rng = np.random.default_rng(seed)
    photons = np.clip(clean * gain, 0, None)
    shot = rng.poisson(photons).astype(np.float32) / max(gain, 1e-8)
    read = rng.normal(0.0, read_sigma, clean.shape).astype(np.float32)
    noisy = shot + read
    return noisy.astype(np.float32)