"""Tests for pymeteo.dynamics module."""

import math
import numpy as np
import pytest
import pymeteo.dynamics as dyn
import pymeteo.constants as c


class TestAvgWind:
    """Test avg_wind function."""

    def test_uniform_wind(self):
        u = np.array([10.0, 10.0, 10.0, 10.0])
        v = np.array([5.0, 5.0, 5.0, 5.0])
        z = np.array([0.0, 500.0, 1000.0, 1500.0])
        result = dyn.avg_wind(u, v, z, 0.0, 2000.0)
        assert abs(result[0] - 10.0) < 1e-6
        assert abs(result[1] - 5.0) < 1e-6

    def test_no_data_in_layer(self):
        u = np.array([10.0, 10.0])
        v = np.array([5.0, 5.0])
        z = np.array([0.0, 100.0])
        result = dyn.avg_wind(u, v, z, 500.0, 1000.0)
        assert result == (0.0, 0.0)


class TestWindConversions:
    """Test uv_to_deg and wind_deg_to_uv."""

    def test_uv_to_deg_pure_u(self):
        direction, speed = dyn.uv_to_deg(10.0, 0.0)
        assert abs(speed - 10.0) < 1e-6
        assert abs(direction - 90.0) < 1e-6

    def test_uv_to_deg_pure_v(self):
        direction, speed = dyn.uv_to_deg(0.0, 10.0)
        assert abs(speed - 10.0) < 1e-6
        assert abs(direction) < 1e-6

    def test_wind_deg_to_uv_south_wind(self):
        # 180 degrees, 10 m/s => wind FROM south => v should be negative
        u, v = dyn.wind_deg_to_uv(180.0, 10.0)
        # sin(360) = 0, cos(360) = 1
        assert abs(u) < 1e-6
        assert abs(v - 10.0) < 1e-6


class TestShear:
    """Test wind shear calculation."""

    def test_shear_uniform_wind(self):
        u = np.array([10.0, 10.0, 10.0, 10.0])
        v = np.array([5.0, 5.0, 5.0, 5.0])
        z = np.array([0.0, 1000.0, 2000.0, 3000.0])
        du, dv = dyn.shear(u, v, z, 0.0, 3000.0)
        assert abs(du) < 1e-6
        assert abs(dv) < 1e-6

    def test_shear_linear_profile(self):
        u = np.array([0.0, 5.0, 10.0, 15.0, 20.0])
        v = np.array([0.0, 0.0, 0.0, 0.0, 0.0])
        z = np.array([0.0, 1000.0, 2000.0, 3000.0, 4000.0])
        du, dv = dyn.shear(u, v, z, 0.0, 4000.0)
        assert abs(du - 20.0) < 1e-3
        assert abs(dv) < 1e-6


class TestSRH:
    """Test storm relative helicity."""

    def test_srh_no_shear(self):
        # Uniform wind should give zero SRH
        u = np.array([10.0, 10.0, 10.0, 10.0, 10.0])
        v = np.array([5.0, 5.0, 5.0, 5.0, 5.0])
        z = np.array([0.0, 500.0, 1000.0, 2000.0, 3000.0])
        result = dyn.srh(u, v, z, 0.0, 3000.0, 10.0, 5.0)
        assert abs(result) < 1.0  # essentially zero

    def test_srh_with_veering(self):
        # Veering wind should give positive SRH
        u = np.array([0.0, 5.0, 10.0, 15.0, 20.0])
        v = np.array([10.0, 10.0, 10.0, 10.0, 10.0])
        z = np.array([0.0, 500.0, 1000.0, 2000.0, 3000.0])
        result = dyn.srh(u, v, z, 0.0, 3000.0, 10.0, 5.0)
        # SRH should be nonzero
        assert result != 0


class TestMeanWind:
    """Test mean_wind function."""

    def test_uniform_wind(self):
        u = np.array([10.0, 10.0, 10.0, 10.0])
        v = np.array([5.0, 5.0, 5.0, 5.0])
        z = np.array([0.0, 1000.0, 2000.0, 3000.0])
        uavg, vavg = dyn.mean_wind(u, v, z, 0.0, 3000.0)
        assert abs(uavg - 10.0) < 0.5
        assert abs(vavg - 5.0) < 0.5


class TestBRN:
    """Test Bulk Richardson Number."""

    def test_brn_positive_cape(self):
        # Create a wind profile with shear
        u = np.array([0.0, 5.0, 10.0, 15.0, 20.0, 25.0, 30.0])
        v = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
        z = np.array([0.0, 1000.0, 2000.0, 3000.0, 4000.0, 5000.0, 6000.0])
        cape = 2000.0
        result = dyn.brn(u, v, z, cape)
        assert result > 0


class TestStormMotion:
    """Test storm motion calculations."""

    def test_bunkers_returns_4_values(self):
        u = np.array([0.0, 5.0, 10.0, 15.0, 20.0, 25.0, 30.0])
        v = np.array([0.0, 2.0, 5.0, 7.0, 8.0, 9.0, 10.0])
        z = np.array([0.0, 1000.0, 2000.0, 3000.0, 4000.0, 5000.0, 6000.0])
        result = dyn.storm_motion_bunkers(u, v, z)
        assert len(result) == 4

    def test_rasmussen_returns_4_values(self):
        u = np.array([0.0, 5.0, 10.0, 15.0, 20.0, 25.0, 30.0])
        v = np.array([0.0, 2.0, 5.0, 7.0, 8.0, 9.0, 10.0])
        z = np.array([0.0, 1000.0, 2000.0, 3000.0, 4000.0, 5000.0, 6000.0])
        result = dyn.storm_motion_rasmussen(u, v, z)
        assert len(result) == 4

    def test_bunkers_zero_shear_no_crash(self):
        """Bunkers should not crash when du=0 (previously divided by zero)."""
        u = np.array([10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0])
        v = np.array([5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0])
        z = np.array([0.0, 1000.0, 2000.0, 3000.0, 4000.0, 5000.0, 6000.0])
        result = dyn.storm_motion_bunkers(u, v, z)
        assert len(result) == 4
        for val in result:
            assert np.isfinite(val)

    def test_rasmussen_zero_shear_no_crash(self):
        """Rasmussen should not crash when du=0 (previously divided by zero)."""
        u = np.array([10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0])
        v = np.array([5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0])
        z = np.array([0.0, 1000.0, 2000.0, 3000.0, 4000.0, 5000.0, 6000.0])
        result = dyn.storm_motion_rasmussen(u, v, z)
        assert len(result) == 4
        for val in result:
            assert np.isfinite(val)


class TestCirculation:
    """Test circulation calculation."""

    def test_circulation_with_missing_values(self):
        n = 4
        x = np.array([c.missingval, 1.0, 1.0, 0.0])
        y = np.array([0.0, 0.0, 1.0, 1.0])
        z = np.array([0.0, 0.0, 0.0, 0.0])
        u = np.array([1.0, 1.0, 1.0, 1.0])
        v = np.array([0.0, 0.0, 0.0, 0.0])
        w = np.array([0.0, 0.0, 0.0, 0.0])
        C, _, _ = dyn.circulation(u, v, w, x, y, z)
        assert C == c.missingval


class TestIntegralBdz:
    """Test integral of buoyancy over height."""

    def test_zero_buoyancy(self):
        n = 5
        th = np.array([300.0, 301.0, 302.0, 303.0, 304.0])
        thp = np.array([300.0, 301.0, 302.0, 303.0, 304.0])
        z = np.array([0.0, 1.0, 2.0, 3.0, 4.0])  # km
        # th == thp -> division by zero in formula, but let's check missing
        result = dyn.integral_Bdz(th, thp, z)
        # when th == thp, th_avg - thp_avg = 0, which causes inf/nan
        # This tests that the function handles it without crashing
        assert isinstance(result, (float, np.floating))

    def test_missing_values(self):
        th = np.array([c.missingval, 301.0])
        thp = np.array([300.0, 301.0])
        z = np.array([0.0, 1.0])
        result = dyn.integral_Bdz(th, thp, z)
        assert result == c.missingval


class TestIntegralDt:
    """Test integral over time."""

    def test_constant_integrand(self):
        # Integral of 1.0 from 0 to 10 should be 10
        i = np.array([1.0, 1.0, 1.0, 1.0, 1.0])
        t = np.array([0.0, 2.5, 5.0, 7.5, 10.0])
        result = dyn.integral_dt(i, t)
        assert abs(result - 10.0) < 1e-3

    def test_linear_integrand(self):
        # Integral of x from 0 to 4 should be 8
        i = np.array([0.0, 1.0, 2.0, 3.0, 4.0])
        t = np.array([0.0, 1.0, 2.0, 3.0, 4.0])
        result = dyn.integral_dt(i, t)
        assert abs(result - 8.0) < 0.1
