"""tksynthetic.py

最小二乗プログラムの動作確認用の内部生成データユーティリティ。
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Mapping, Optional, Sequence, Union
import csv
import numpy as np


ArrayLike = Sequence[float] | np.ndarray


@dataclass
class SyntheticData:
    """合成データ一式。"""

    x: np.ndarray
    y: np.ndarray
    y_clean: np.ndarray
    y_noise: np.ndarray
    true_params: dict
    noise_std: float
    seed: Optional[int] = None


def make_x_grid(
    xmin: float = 0.0,
    xmax: float = 1.0,
    n: int = 50,
    *,
    kind: str = "linear",
) -> np.ndarray:
    """線形または対数 x グリッドを作る。"""
    if kind == "linear":
        return np.linspace(float(xmin), float(xmax), int(n))
    if kind == "log":
        if xmin <= 0 or xmax <= 0:
            raise ValueError("log grid requires xmin > 0 and xmax > 0")
        return np.geomspace(float(xmin), float(xmax), int(n))
    raise ValueError("kind must be 'linear' or 'log'")


def generate_noisy_data(
    model_func: Callable[[np.ndarray, Mapping[str, float]], ArrayLike],
    x: ArrayLike,
    true_params: Mapping[str, float],
    *,
    noise_std: float = 0.05,
    seed: Optional[int] = 0,
    noise: str = "normal",
    relative_noise: bool = False,
) -> SyntheticData:
    """モデル関数からノイズ付きデータを生成する。

    model_func:
        y_clean = model_func(x_array, true_params)
    """
    rng = np.random.default_rng(seed)
    x_arr = np.asarray(x, dtype=float).reshape(-1)
    p = {k: float(v) for k, v in true_params.items()}

    y_clean = np.asarray(model_func(x_arr, p), dtype=float).reshape(-1)
    if y_clean.size != x_arr.size:
        raise ValueError("model_func output size and x size mismatch")

    if noise == "normal":
        eps = rng.normal(0.0, float(noise_std), size=x_arr.size)
    elif noise == "uniform":
        half_width = np.sqrt(3.0) * float(noise_std)
        eps = rng.uniform(-half_width, half_width, size=x_arr.size)
    else:
        raise ValueError("noise must be 'normal' or 'uniform'")

    if relative_noise:
        y_noise = eps * np.maximum(np.abs(y_clean), np.finfo(float).eps)
    else:
        y_noise = eps

    y = y_clean + y_noise

    return SyntheticData(
        x=x_arr,
        y=y,
        y_clean=y_clean,
        y_noise=y_noise,
        true_params=p,
        noise_std=float(noise_std),
        seed=seed,
    )


def generate_replicates(
    model_func: Callable[[np.ndarray, Mapping[str, float]], ArrayLike],
    x: ArrayLike,
    true_params: Mapping[str, float],
    *,
    noise_std: float = 0.05,
    seed: Optional[int] = 0,
    n_replicates: int = 3,
) -> SyntheticData:
    """同じ x に対する繰り返し測定風データを生成する。"""
    rng = np.random.default_rng(seed)
    x_arr = np.asarray(x, dtype=float).reshape(-1)
    p = {k: float(v) for k, v in true_params.items()}
    y_clean = np.asarray(model_func(x_arr, p), dtype=float).reshape(-1)

    xs = []
    ys = []
    ycs = []
    yns = []

    for _ in range(int(n_replicates)):
        eps = rng.normal(0.0, float(noise_std), size=x_arr.size)
        xs.append(x_arr)
        ys.append(y_clean + eps)
        ycs.append(y_clean)
        yns.append(eps)

    return SyntheticData(
        x=np.concatenate(xs),
        y=np.concatenate(ys),
        y_clean=np.concatenate(ycs),
        y_noise=np.concatenate(yns),
        true_params=p,
        noise_std=float(noise_std),
        seed=seed,
    )


def save_xy_csv(
    path: Union[str, Path],
    data: SyntheticData,
    *,
    include_clean: bool = True,
) -> None:
    """合成データを CSV に保存する。"""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    fieldnames = ["x", "y"]
    if include_clean:
        fieldnames += ["y_clean", "y_noise"]

    with path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()

        for i in range(data.x.size):
            row = {
                "x": float(data.x[i]),
                "y": float(data.y[i]),
            }
            if include_clean:
                row["y_clean"] = float(data.y_clean[i])
                row["y_noise"] = float(data.y_noise[i])
            writer.writerow(row)


def synthetic_summary(data: SyntheticData) -> str:
    """合成データの概要を文字列化する。"""
    lines = [
        "SyntheticData",
        f"  N         = {data.x.size}",
        f"  x range   = [{np.nanmin(data.x):.6g}, {np.nanmax(data.x):.6g}]",
        f"  y range   = [{np.nanmin(data.y):.6g}, {np.nanmax(data.y):.6g}]",
        f"  noise_std = {data.noise_std:.6g}",
        f"  seed      = {data.seed}",
        "  true_params:",
    ]

    for k, v in data.true_params.items():
        lines.append(f"    {k:16s} = {v:.10g}")

    return "\n".join(lines)