"""Tests for analytical dose calculation module."""

import numpy as np
import pytest

from oncosim.physics.dose_calc import calculate_beam_dose, calculate_plan_dose


class TestBeamDose:
    def test_output_shape(self):
        dose = calculate_beam_dose(0.0, grid_shape=(64, 64))
        assert dose.shape == (64, 64)

    def test_output_shape_custom(self):
        dose = calculate_beam_dose(90.0, grid_shape=(32, 32))
        assert dose.shape == (32, 32)

    def test_values_non_negative(self):
        dose = calculate_beam_dose(45.0)
        assert np.all(dose >= 0)

    def test_values_bounded(self):
        dose = calculate_beam_dose(0.0)
        assert np.max(dose) <= 1.0 + 1e-10

    def test_attenuation_with_depth(self):
        dose = calculate_beam_dose(0.0, grid_shape=(64, 64))
        center_row = 32
        # Beam from right (0 deg): dose should decrease left to right at center
        left_dose = dose[center_row, 10]
        right_dose = dose[center_row, 50]
        # At least one side should have dose (beam enters from some direction)
        assert np.max(dose) > 0

    def test_different_angles_produce_different_distributions(self):
        dose_0 = calculate_beam_dose(0.0)
        dose_90 = calculate_beam_dose(90.0)
        assert not np.allclose(dose_0, dose_90)

    def test_dose_symmetry_opposite_angles(self):
        dose_0 = calculate_beam_dose(0.0, grid_shape=(64, 64))
        dose_180 = calculate_beam_dose(180.0, grid_shape=(64, 64))
        # Opposite beams should have different patterns
        assert not np.allclose(dose_0, dose_180)

    def test_beam_width_effect(self):
        narrow = calculate_beam_dose(0.0, beam_width_mm=10.0)
        wide = calculate_beam_dose(0.0, beam_width_mm=40.0)
        # Wider beam should have more spread (higher sum)
        assert np.sum(wide) > np.sum(narrow)

    def test_zero_dose_outside_beam(self):
        dose = calculate_beam_dose(0.0, grid_shape=(64, 64), beam_width_mm=5.0)
        # Far from beam axis should have very low dose
        assert dose[0, 0] < 0.01 or dose[63, 0] < 0.01

    def test_high_attenuation(self):
        dose = calculate_beam_dose(0.0, mu=0.1)
        assert np.max(dose) <= 1.0

    def test_low_attenuation(self):
        low = calculate_beam_dose(0.0, mu=0.001)
        high = calculate_beam_dose(0.0, mu=0.01)
        assert np.sum(low) > np.sum(high)

    def test_multiple_angles_coverage(self):
        for angle in [0, 45, 90, 135, 180, 225, 270, 315]:
            dose = calculate_beam_dose(float(angle))
            assert dose.shape == (64, 64)
            assert np.max(dose) > 0

    def test_pixel_size_parameter(self):
        dose_small = calculate_beam_dose(0.0, pixel_size_mm=1.0)
        dose_large = calculate_beam_dose(0.0, pixel_size_mm=4.0)
        assert dose_small.shape == dose_large.shape

    def test_dtype_float64(self):
        dose = calculate_beam_dose(0.0)
        assert dose.dtype == np.float64


class TestPlanDose:
    def test_plan_dose_single_beam(self):
        single = calculate_beam_dose(0.0)
        plan = calculate_plan_dose([0.0], None)
        assert np.allclose(single, plan, atol=1e-10)

    def test_plan_dose_uniform_weights(self):
        plan = calculate_plan_dose([0.0, 90.0], None)
        assert plan.shape == (64, 64)
        assert np.all(plan >= 0)

    def test_plan_dose_custom_weights(self):
        plan = calculate_plan_dose([0.0, 90.0], [0.7, 0.3])
        assert plan.shape == (64, 64)

    def test_plan_dose_superposition(self):
        plan = calculate_plan_dose([0.0, 90.0, 180.0, 270.0], None)
        single = calculate_beam_dose(0.0)
        # Plan with 4 beams should have higher total dose than single beam
        assert np.sum(plan) > np.sum(single) * 0.5

    def test_plan_dose_grid_size(self):
        plan = calculate_plan_dose([0.0, 90.0], None, grid_shape=(32, 32))
        assert plan.shape == (32, 32)
