#!/usr/bin/env python3
"""
Demonstrate noise amplification in direct deconvolution.

A sparse pulse signal x is convolved by a Gaussian convolution matrix A:

    y = A x

Then x is reconstructed by the direct linear solve:

    x = solve(A, y)

The clean convolved data are reconstructed almost exactly, but adding noise to
just one data point in y can produce large oscillations in the reconstructed x.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np


def parse_pulses(text: str, n: int) -> list[tuple[int, float]]:
    """Parse pulse specification such as '4:1.0,10:0.7'."""
    pulses: list[tuple[int, float]] = []
    if not text.strip():
        return pulses

    for item in text.split(","):
        item = item.strip()
        if not item:
            continue
        try:
            idx_text, amp_text = item.split(":", maxsplit=1)
            idx = int(idx_text)
            amp = float(amp_text)
        except ValueError as exc:
            raise argparse.ArgumentTypeError(
                f"Invalid pulse specification: {item!r}. Use e.g. 4:1.0,10:0.7"
            ) from exc
        if not (0 <= idx < n):
            raise argparse.ArgumentTypeError(
                f"Pulse index {idx} is out of range. Valid range is 0 to {n - 1}."
            )
        pulses.append((idx, amp))
    return pulses


def make_gaussian_blur_matrix(n: int, sigma: float, dx: float, normalize: str) -> np.ndarray:
    """Return an n x n matrix representing Gaussian convolution on a finite grid."""
    if sigma <= 0:
        raise ValueError("sigma must be positive")
    if dx <= 0:
        raise ValueError("dx must be positive")

    i = np.arange(n)[:, None]
    j = np.arange(n)[None, :]
    distance = (i - j) * dx
    a = np.exp(-0.5 * (distance / sigma) ** 2)

    if normalize == "row":
        # Each observed point is a weighted average of source points.
        # This avoids edge rows having artificially smaller intensity.
        a = a / a.sum(axis=1, keepdims=True)
    elif normalize == "global":
        # Use the same infinite-grid normalization factor for every row.
        # Edge rows then lose intensity because the grid is finite.
        kernel_norm = np.sum(np.exp(-0.5 * ((np.arange(n) - n // 2) * dx / sigma) ** 2))
        a = a / kernel_norm
    elif normalize == "none":
        pass
    else:
        raise ValueError(f"Unknown normalize option: {normalize}")

    return a


def make_pulse_signal(n: int, pulses: list[tuple[int, float]]) -> np.ndarray:
    """Return an n-point signal containing delta-like pulses."""
    x = np.zeros(n, dtype=float)
    for idx, amp in pulses:
        x[idx] += amp
    return x


def save_csv(
    path: Path,
    x_true: np.ndarray,
    y_clean: np.ndarray,
    y_noisy: np.ndarray,
    x_deconv_clean: np.ndarray,
    x_deconv_noisy: np.ndarray,
) -> None:
    """Save the numerical results as CSV."""
    n = len(x_true)
    data = np.column_stack(
        [
            np.arange(n),
            x_true,
            y_clean,
            y_noisy,
            x_deconv_clean,
            x_deconv_noisy,
            x_deconv_noisy - x_deconv_clean,
        ]
    )
    header = (
        "index,x_true,convolved_clean,convolved_with_one_point_noise,"
        "deconv_from_clean,deconv_from_noisy,deconv_noise_effect"
    )
    np.savetxt(path, data, delimiter=",", header=header, comments="")


def plot_results(
    path: Path,
    x_true: np.ndarray,
    y_clean: np.ndarray,
    y_noisy: np.ndarray,
    x_deconv_clean: np.ndarray,
    x_deconv_noisy: np.ndarray,
    cond_a: float,
    sigma: float,
    noise_index: int,
    noise_amplitude: float,
    show: bool,
) -> None:
    """Create a figure comparing clean and noisy direct deconvolution."""
    n = len(x_true)
    grid = np.arange(n)
    noise_effect = x_deconv_noisy - x_deconv_clean

    fig, axes = plt.subplots(3, 1, figsize=(8, 9), sharex=True)

    axes[0].plot(grid, x_true, "o-", label="true pulse signal x")
    axes[0].plot(grid, y_clean, "s-", label="convolved data y = A x")
    axes[0].plot(grid, y_noisy, "x--", label="convolved data + one-point noise")
    axes[0].axvline(noise_index, linestyle=":", label="noise point")
    axes[0].set_ylabel("value")
    axes[0].set_title(f"Gaussian blur: n={n}, sigma={sigma:g}, cond(A)={cond_a:.3e}")
    axes[0].legend(loc="best")
    axes[0].grid(True, alpha=0.3)

    axes[1].plot(grid, x_true, "o-", label="true x")
    axes[1].plot(grid, x_deconv_clean, "s--", label="solve(A, clean y)")
    axes[1].plot(grid, x_deconv_noisy, "x-", label="solve(A, noisy y)")
    axes[1].set_ylabel("deconvolved x")
    axes[1].legend(loc="best")
    axes[1].grid(True, alpha=0.3)

    axes[2].axhline(0.0, linewidth=1.0)
    axes[2].plot(grid, noise_effect, "o-", label="deconvolution change due to one noisy point")
    axes[2].axvline(noise_index, linestyle=":", label=f"noise = {noise_amplitude:g}")
    axes[2].set_xlabel("data index")
    axes[2].set_ylabel("difference")
    axes[2].legend(loc="best")
    axes[2].grid(True, alpha=0.3)

    fig.tight_layout()
    fig.savefig(path, dpi=200)
    if show:
        plt.show()
    plt.close(fig)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Direct deconvolution demo showing noise amplification by numpy.linalg.solve."
    )
    parser.add_argument("--n", type=int, default=15, help="Number of data points. Default: 15")
    parser.add_argument(
        "--sigma",
        type=float,
        default=1.4,
        help="Gaussian blur width in grid units. Larger values make deconvolution more ill-conditioned. Default: 1.4",
    )
    parser.add_argument("--dx", type=float, default=1.0, help="Grid spacing. Default: 1.0")
    parser.add_argument(
        "--pulses",
        type=str,
        default="1:0.0",
#        default="4:1.0,10:0.7",
        help="Pulse list as index:amplitude,index:amplitude. Default: 4:1.0,10:0.7",
    )
    parser.add_argument(
        "--noise-index",
        type=int,
        default=None,
        help="Index where one-point noise is added. Default: center point n//2",
    )
    parser.add_argument(
        "--noise-amplitude",
        type=float,
        default=1.0e-3,
        help="Amplitude of the one-point noise added to convolved data. Default: 1e-3",
    )
    parser.add_argument(
        "--normalize",
        choices=["row", "global", "none"],
        default="row",
        help="Normalization of Gaussian blur matrix. Default: row",
    )
    parser.add_argument(
        "--output-prefix",
        type=str,
        default="deconv_noise_demo",
        help="Output file prefix. PNG and CSV are generated. Default: deconv_noise_demo",
    )
    parser.add_argument("--show", action="store_true", help="Show plot window in addition to saving PNG.")
    parser.add_argument("--print-matrix", action="store_true", help="Print the convolution matrix A.")
    args = parser.parse_args()

    if args.n < 2:
        raise ValueError("--n must be 2 or larger")

    noise_index = args.noise_index
    if noise_index is None:
        noise_index = args.n // 2
    if not (0 <= noise_index < args.n):
        raise ValueError(f"--noise-index must be in the range 0 to {args.n - 1}")

    pulses = parse_pulses(args.pulses, args.n)
    if not pulses:
        raise ValueError("At least one pulse must be specified by --pulses")

    a = make_gaussian_blur_matrix(args.n, args.sigma, args.dx, args.normalize)
    x_true = make_pulse_signal(args.n, pulses)

    y_clean = a @ x_true
    y_noisy = y_clean.copy()
    y_noisy[noise_index] += args.noise_amplitude

    x_deconv_clean = np.linalg.solve(a, y_clean)
    x_deconv_noisy = np.linalg.solve(a, y_noisy)

    cond_a = np.linalg.cond(a)
    clean_error = np.max(np.abs(x_deconv_clean - x_true))
    noise_effect = x_deconv_noisy - x_deconv_clean
    max_noise_effect = np.max(np.abs(noise_effect))
    l2_amplification = np.linalg.norm(noise_effect) / max(abs(args.noise_amplitude), np.finfo(float).eps)
    max_amplification = max_noise_effect / max(abs(args.noise_amplitude), np.finfo(float).eps)

    output_prefix = Path(args.output_prefix)
    png_path = output_prefix.with_suffix(".png")
    csv_path = output_prefix.with_suffix(".csv")

    save_csv(csv_path, x_true, y_clean, y_noisy, x_deconv_clean, x_deconv_noisy)
    plot_results(
        png_path,
        x_true,
        y_clean,
        y_noisy,
        x_deconv_clean,
        x_deconv_noisy,
        cond_a,
        args.sigma,
        noise_index,
        args.noise_amplitude,
        args.show,
    )

    print("=== Direct deconvolution noise amplification demo ===")
    print(f"n                         : {args.n}")
    print(f"Gaussian sigma            : {args.sigma:g}")
    print(f"normalization             : {args.normalize}")
    print(f"pulses                    : {pulses}")
    print(f"condition number cond(A)  : {cond_a:.6e}")
    print(f"max clean reconstruction error |solve(A, A x) - x| : {clean_error:.6e}")
    print(f"noise added to y[{noise_index}]                    : {args.noise_amplitude:.6e}")
    print(f"max |deconvolution change|                         : {max_noise_effect:.6e}")
    print(f"max-amplitude amplification factor                 : {max_amplification:.6e}")
    print(f"L2 amplification factor                            : {l2_amplification:.6e}")
    print(f"saved PNG                  : {png_path}")
    print(f"saved CSV                  : {csv_path}")

    if args.print_matrix:
        np.set_printoptions(precision=4, suppress=True)
        print("\nA =")
        print(a)


if __name__ == "__main__":
    main()
