"""Tests for beam geometry utilities."""

import numpy as np
import pytest

from oncosim.physics.beam_geometry import (
    generate_beam_profile,
    compute_beam_intersection,
    create_tumor_mask,
    create_oar_mask,
    compute_dvh,
)


class TestBeamProfile:
    def test_output_shape(self):
        profile = generate_beam_profile(0.0, grid_shape=(64, 64))
        assert profile.shape == (64, 64)

    def test_values_zero_to_one(self):
        profile = generate_beam_profile(45.0)
        assert np.all(profile >= 0)
        assert np.all(profile <= 1.0 + 1e-10)

    def test_different_angles(self):
        p0 = generate_beam_profile(0.0)
        p90 = generate_beam_profile(90.0)
        assert not np.allclose(p0, p90)

    def test_normalization(self):
        profile = generate_beam_profile(0.0)
        assert np.max(profile) > 0.9  # Peak should be close to 1


class TestMasks:
    def test_tumor_mask_shape(self):
        mask = create_tumor_mask((32, 32), 10.0, (64, 64))
        assert mask.shape == (64, 64)

    def test_tumor_mask_binary(self):
        mask = create_tumor_mask((32, 32), 10.0)
        unique = np.unique(mask)
        assert all(v in [0.0, 1.0] for v in unique)

    def test_tumor_mask_area(self):
        mask = create_tumor_mask((32, 32), 10.0, (64, 64))
        area = np.sum(mask)
        expected = np.pi * 10.0**2
        assert abs(area - expected) < expected * 0.15  # Within 15%

    def test_oar_mask_same_as_tumor(self):
        t = create_tumor_mask((20, 20), 5.0)
        o = create_oar_mask((20, 20), 5.0)
        assert np.array_equal(t, o)

    def test_mask_center_is_inside(self):
        mask = create_tumor_mask((32, 32), 10.0, (64, 64))
        assert mask[32, 32] == 1.0

    def test_mask_far_from_center_is_outside(self):
        mask = create_tumor_mask((32, 32), 5.0, (64, 64))
        assert mask[0, 0] == 0.0

    def test_small_radius(self):
        mask = create_tumor_mask((32, 32), 1.0, (64, 64))
        assert np.sum(mask) > 0
        assert np.sum(mask) < 20

    def test_multiple_oar_masks_no_overlap(self):
        oar1 = create_oar_mask((10, 10), 5.0, (64, 64))
        oar2 = create_oar_mask((50, 50), 5.0, (64, 64))
        overlap = np.sum(oar1 * oar2)
        assert overlap == 0.0


class TestBeamIntersection:
    def test_intersection_positive(self):
        profile = generate_beam_profile(0.0, grid_shape=(64, 64))
        mask = create_tumor_mask((32, 32), 10.0, (64, 64))
        intersection = compute_beam_intersection(profile, mask)
        assert intersection > 0

    def test_intersection_zero_no_overlap(self):
        profile = generate_beam_profile(0.0, grid_shape=(64, 64), beam_width_pixels=2.0)
        mask = create_tumor_mask((5, 5), 2.0, (64, 64))
        # Very narrow beam far from small target - should have minimal intersection
        # (May not be exactly zero due to gaussian tails)
        intersection = compute_beam_intersection(profile, mask)
        assert intersection >= 0


class TestDVH:
    def test_dvh_shape(self):
        dose = np.random.rand(64, 64)
        mask = create_tumor_mask((32, 32), 10.0, (64, 64))
        dvh = compute_dvh(dose, mask)
        assert dvh.shape == (100,)

    def test_dvh_starts_at_one(self):
        dose = np.random.rand(64, 64) + 0.1
        mask = create_tumor_mask((32, 32), 10.0, (64, 64))
        dvh = compute_dvh(dose, mask)
        assert dvh[0] == pytest.approx(1.0)

    def test_dvh_monotonic_decrease(self):
        dose = np.random.rand(64, 64)
        mask = create_tumor_mask((32, 32), 10.0, (64, 64))
        dvh = compute_dvh(dose, mask)
        for i in range(len(dvh) - 1):
            assert dvh[i] >= dvh[i + 1] - 1e-10

    def test_dvh_empty_mask(self):
        dose = np.random.rand(64, 64)
        mask = np.zeros((64, 64))
        dvh = compute_dvh(dose, mask)
        assert np.all(dvh == 0)
