from pathlib import Path

import numpy as np

from orisys.hhd import PoissonSolver, nHHD


REPO_ROOT = Path(__file__).resolve().parents[1]


def _circle_mask(shape, radius_ratio=0.35):
    h, w = shape
    yy, xx = np.mgrid[0:h, 0:w]
    cy = (h - 1) / 2.0
    cx = (w - 1) / 2.0
    radius = min(h, w) * radius_ratio
    return ((yy - cy) ** 2 + (xx - cx) ** 2) <= radius ** 2


def _smooth_rhs(shape):
    h, w = shape
    yy, xx = np.mgrid[0:h, 0:w]
    field = (
        np.sin(np.pi * xx / max(1, w - 1))
        + 0.5 * np.cos(2.0 * np.pi * yy / max(1, h - 1))
    )
    return field.astype(np.float32)


def _radial_flow(shape, scale=0.05):
    h, w = shape
    yy, xx = np.mgrid[0:h, 0:w]
    cy = (h - 1) / 2.0
    cx = (w - 1) / 2.0
    flow = np.zeros((h, w, 2), dtype=np.float32)
    flow[:, :, 0] = (xx - cx) * scale
    flow[:, :, 1] = (yy - cy) * scale
    return flow


def _masked_rhs_vector(rhs, mask, pinned):
    vec = rhs[mask].astype(np.float64)
    vec -= vec.mean()
    vec[pinned] = 0.0
    return vec


def _masked_residual_norm(solver, field, rhs, mask):
    cache = solver._masked_cache if solver._masked_cache_matches(mask) else solver._build_masked_cache(mask)
    rhs_vec = _masked_rhs_vector(rhs, mask, cache["pinned"])
    residual = cache["lap"].dot(field[mask].astype(np.float64)) - rhs_vec
    return float(np.linalg.norm(residual))


def test_capfft_solver_improves_over_raw_fft_on_synthetic_mask():
    shape = (40, 48)
    mask = _circle_mask(shape)
    rhs = _smooth_rhs(shape)
    rhs_masked = np.where(mask, rhs, 0).astype(np.float32)

    sparse_solver = PoissonSolver(
        grid=shape,
        spacings=(1.0, 1.0),
        poisson_backend="cpu",
        masked_poisson_solver="sparse",
    )
    capfft_solver = PoissonSolver(
        grid=shape,
        spacings=(1.0, 1.0),
        poisson_backend="cpu",
        masked_poisson_solver="capfft",
        masked_poisson_max_boundary=48,
        masked_poisson_max_iters=3,
    )

    sparse = sparse_solver.solve(rhs_masked, mask=mask)
    capfft = capfft_solver.solve(rhs_masked, mask=mask)
    raw_fft = sparse_solver.solve(rhs_masked, mask=None)

    assert np.allclose(capfft[~mask], 0.0, atol=1e-6)

    err_cap = float(np.linalg.norm((capfft - sparse)[mask]))
    err_raw = float(np.linalg.norm((raw_fft - sparse)[mask]))
    assert err_cap < err_raw

    res_cap = _masked_residual_norm(sparse_solver, capfft, rhs_masked, mask)
    res_raw = _masked_residual_norm(sparse_solver, raw_fft, rhs_masked, mask)
    assert res_cap < res_raw


def test_capfft_hhd_forces_track_sparse_reference():
    shape = (36, 44)
    mask = _circle_mask(shape)
    flow = _radial_flow(shape, scale=0.03)
    flow[~mask] = 0.0

    sparse_hhd = nHHD(
        cuda=False,
        grid=shape,
        spacings=(1.0, 1.0),
        poisson_backend="cpu",
        masked_poisson_solver="sparse",
    )
    capfft_hhd = nHHD(
        cuda=False,
        grid=shape,
        spacings=(1.0, 1.0),
        poisson_backend="cpu",
        masked_poisson_solver="capfft",
        masked_poisson_max_boundary=40,
        masked_poisson_max_iters=3,
    )

    _, sparse_forces = sparse_hhd.get_force(flow, valid_mask=mask)
    _, capfft_forces = capfft_hhd.get_force(flow, valid_mask=mask)

    denom = max(1.0, float(np.linalg.norm(sparse_forces)))
    rel = float(np.linalg.norm(capfft_forces - sparse_forces) / denom)
    assert rel < 0.6

    vnormal = capfft_hhd.get_normal()
    vshear = capfft_hhd.get_shear()
    assert np.allclose(vnormal[~mask], 0.0, atol=1e-5)
    assert np.allclose(vshear[~mask], 0.0, atol=1e-5)


def test_capfft_uv_mask_fixture_improves_over_raw_fft():
    uv_path = REPO_ROOT / "unroll_map_uv.npz"
    if not uv_path.exists():
        return

    full_mask = np.load(uv_path)["valid_mask"].astype(bool)
    mask = full_mask[::24, ::24]
    if mask.sum() < 16:
        mask = full_mask[::16, ::16]

    shape = mask.shape
    rhs = _smooth_rhs(shape)
    rhs_masked = np.where(mask, rhs, 0).astype(np.float32)

    sparse_solver = PoissonSolver(
        grid=shape,
        spacings=(1.0, 1.0),
        poisson_backend="cpu",
        masked_poisson_solver="sparse",
    )
    capfft_solver = PoissonSolver(
        grid=shape,
        spacings=(1.0, 1.0),
        poisson_backend="cpu",
        masked_poisson_solver="capfft",
        masked_poisson_max_boundary=24,
        masked_poisson_max_iters=3,
    )

    sparse = sparse_solver.solve(rhs_masked, mask=mask)
    capfft = capfft_solver.solve(rhs_masked, mask=mask)
    raw_fft = sparse_solver.solve(rhs_masked, mask=None)

    err_cap = float(np.linalg.norm((capfft - sparse)[mask]))
    err_raw = float(np.linalg.norm((raw_fft - sparse)[mask]))
    assert err_cap < err_raw