"""Tests for marginal transform utilities."""

from __future__ import annotations

import warnings

import numpy as np
import pandas as pd
import pytest
from scipy.stats import gamma, lognorm, norm, expon

from insurance_copula._marginals import (
    to_pseudo_obs,
    back_transform,
    validate_marginals,
    _pseudo_obs_fallback,
    _PSEUDO_OBS_CLIP,
)
from insurance_copula import HAS_PYVINECOPULIB


# ------------------------------------------------------------------
# to_pseudo_obs
# ------------------------------------------------------------------

class TestToPseudoObs:
    def test_output_shape(self, synthetic_losses):
        u = to_pseudo_obs(synthetic_losses)
        assert u.shape == synthetic_losses.shape

    def test_output_in_unit_interval(self, synthetic_losses):
        u = to_pseudo_obs(synthetic_losses)
        assert (u >= _PSEUDO_OBS_CLIP[0]).all()
        assert (u <= _PSEUDO_OBS_CLIP[1]).all()

    def test_output_clipped(self):
        # Even with extreme values, no exact 0 or 1
        data = pd.DataFrame({
            "a": [1.0, 2.0, 3.0],
            "b": [10.0, 20.0, 30.0],
        })
        u = to_pseudo_obs(data)
        assert (u > 0).all()
        assert (u < 1).all()

    def test_two_column_minimum(self):
        data = pd.DataFrame({
            "x": [1.0, 2.0, 3.0],
            "y": [3.0, 1.0, 2.0],
        })
        u = to_pseudo_obs(data)
        assert u.shape == (3, 2)

    def test_monotone_transform(self):
        # Larger losses should map to larger pseudo-obs
        data = pd.DataFrame({
            "a": [10.0, 20.0, 30.0, 40.0, 50.0],
            "b": [5.0, 5.0, 5.0, 5.0, 5.0],
        })
        u = to_pseudo_obs(data)
        # Column 'a' should be monotone increasing
        assert np.all(np.diff(u[:, 0]) > 0)

    def test_nan_raises(self):
        data = pd.DataFrame({
            "a": [1.0, np.nan, 3.0],
            "b": [1.0, 2.0, 3.0],
        })
        with pytest.raises(ValueError, match="NaN"):
            to_pseudo_obs(data)

    def test_single_observation_raises(self):
        data = pd.DataFrame({"a": [1.0], "b": [2.0]})
        with pytest.raises(ValueError, match="at least 2"):
            to_pseudo_obs(data)

    def test_returns_numpy_array(self, synthetic_losses):
        u = to_pseudo_obs(synthetic_losses)
        assert isinstance(u, np.ndarray)

    def test_consistent_with_fallback(self):
        # Both methods should give same result (up to implementation details)
        data = pd.DataFrame({
            "x": np.arange(1.0, 21.0),
            "y": np.arange(20.0, 0.0, -1.0),
        })
        u_fallback = _pseudo_obs_fallback(data.values)
        u_main = to_pseudo_obs(data)
        # Should be very close (both use rank/(n+1))
        np.testing.assert_allclose(u_main, np.clip(u_fallback, *_PSEUDO_OBS_CLIP),
                                   atol=1e-9)

    def test_large_dataset(self):
        rng = np.random.default_rng(0)
        data = pd.DataFrame(rng.random((1000, 5)), columns=list("abcde"))
        u = to_pseudo_obs(data)
        assert u.shape == (1000, 5)
        assert (u > 0).all()
        assert (u < 1).all()


# ------------------------------------------------------------------
# _pseudo_obs_fallback
# ------------------------------------------------------------------

class TestPseudoObsFallback:
    def test_output_range(self):
        x = np.array([[1.0, 3.0], [2.0, 1.0], [3.0, 2.0]])
        u = _pseudo_obs_fallback(x)
        assert (u > 0).all()
        assert (u < 1).all()

    def test_column_independence(self):
        # Each column is ranked independently
        x = np.column_stack([
            np.array([10.0, 20.0, 30.0]),
            np.array([30.0, 10.0, 20.0]),
        ])
        u = _pseudo_obs_fallback(x)
        # Column 0: 10 < 20 < 30 -> 1/4, 2/4, 3/4
        np.testing.assert_allclose(u[:, 0], [0.25, 0.5, 0.75])
        # Column 1: 30 > 10 < 20 -> 3/4, 1/4, 2/4
        np.testing.assert_allclose(u[:, 1], [0.75, 0.25, 0.5])

    def test_ties_handled(self):
        # Tied values — rankdata handles with average by default
        x = np.array([[1.0, 1.0], [1.0, 2.0], [2.0, 1.0]])
        u = _pseudo_obs_fallback(x)
        assert u.shape == (3, 2)
        assert (u > 0).all()
        assert (u < 1).all()


# ------------------------------------------------------------------
# back_transform
# ------------------------------------------------------------------

class TestBackTransform:
    def test_output_shape(self):
        u = np.array([[0.1, 0.5], [0.5, 0.9], [0.9, 0.2]])
        marginals = {"flood": gamma(a=0.8, scale=5000), "storm": gamma(a=1.2, scale=800)}
        result = back_transform(u, ["flood", "storm"], marginals)
        assert result.shape == (3, 2)
        assert list(result.columns) == ["flood", "storm"]

    def test_output_positive(self):
        u = np.random.default_rng(0).random((100, 2))
        marginals = {"flood": gamma(a=0.8, scale=5000), "storm": gamma(a=1.2, scale=800)}
        result = back_transform(u, ["flood", "storm"], marginals)
        assert (result.values >= 0).all()

    def test_column_order_matches_perils(self):
        u = np.array([[0.5, 0.5]])
        marginals = {
            "storm": gamma(a=1.0, scale=100),
            "flood": gamma(a=2.0, scale=1000),
        }
        result = back_transform(u, ["flood", "storm"], marginals)
        assert list(result.columns) == ["flood", "storm"]

    def test_missing_peril_raises(self):
        u = np.array([[0.5, 0.5]])
        marginals = {"flood": gamma(a=0.8, scale=5000)}
        with pytest.raises(ValueError, match="storm"):
            back_transform(u, ["flood", "storm"], marginals)

    def test_column_count_mismatch_raises(self):
        u = np.array([[0.5, 0.5, 0.5]])
        marginals = {"flood": gamma(a=0.8, scale=5000), "storm": gamma(a=1.2, scale=800)}
        with pytest.raises(ValueError, match="perils"):
            back_transform(u, ["flood", "storm"], marginals)

    def test_returns_dataframe(self):
        u = np.array([[0.5, 0.5]])
        marginals = {"a": gamma(a=1.0, scale=100), "b": gamma(a=1.0, scale=200)}
        result = back_transform(u, ["a", "b"], marginals)
        assert isinstance(result, pd.DataFrame)

    def test_with_lognorm(self):
        u = np.linspace(0.01, 0.99, 50).reshape(25, 2)
        marginals = {
            "flood": lognorm(s=1.5, scale=1000),
            "storm": lognorm(s=1.2, scale=500),
        }
        result = back_transform(u, ["flood", "storm"], marginals)
        assert result.shape == (25, 2)
        assert (result.values > 0).all()

    def test_with_normal(self):
        u = np.array([[0.5, 0.5]])
        marginals = {"a": norm(loc=0, scale=1), "b": norm(loc=100, scale=10)}
        result = back_transform(u, ["a", "b"], marginals)
        assert result.shape == (1, 2)

    def test_single_peril_pair(self):
        u = np.array([[0.25], [0.75]])
        marginals = {"flood": gamma(a=1.0, scale=5000)}
        result = back_transform(u, ["flood"], marginals)
        assert result.shape == (2, 1)


# ------------------------------------------------------------------
# validate_marginals
# ------------------------------------------------------------------

class TestValidateMarginals:
    def test_valid_marginals_passes(self):
        marginals = {
            "flood": gamma(a=0.8, scale=5000),
            "storm": gamma(a=1.2, scale=800),
        }
        # Should not raise
        validate_marginals(["flood", "storm"], marginals)

    def test_missing_peril_raises(self):
        marginals = {"flood": gamma(a=0.8, scale=5000)}
        with pytest.raises(ValueError, match="storm"):
            validate_marginals(["flood", "storm"], marginals)

    def test_no_ppf_raises(self):
        class BadMarginal:
            pass

        with pytest.raises(TypeError, match=".ppf"):
            validate_marginals(["flood"], {"flood": BadMarginal()})  # type: ignore[dict-item]

    def test_broken_ppf_raises(self):
        class BrokenMarginal:
            def ppf(self, q):
                raise RuntimeError("broken")

        with pytest.raises(ValueError, match="PPF of marginal"):
            validate_marginals(["flood"], {"flood": BrokenMarginal()})  # type: ignore[dict-item]

    def test_six_perils_valid(self, marginals_6perils):
        perils = ["flood", "subsidence", "storm", "eow", "fire", "theft"]
        validate_marginals(perils, marginals_6perils)  # Should not raise
