import math

from mcp_calculator.ops.context import MathError
from mcp_calculator.ops.stats import hypothesis_test, probability, regression, statistics
from tests.support import MathTestCase


class StatisticsTest(MathTestCase):
    def test_basic(self):
        r = statistics([2, 4, 4, 4, 5, 5, 7, 9])
        self.assertEqual(r["mean"], 5)
        self.assertEqual(r["median"], 4.5)
        self.assertEqual(r["mode"], [4.0])
        self.assertApprox(r["std_dev"], 2.13808993529939)

    def test_population(self):
        r = statistics([2, 4, 4, 4, 5, 5, 7, 9], sample=False)
        self.assertEqual(r["std_dev"], 2)
        self.assertTrue(r["kind"].startswith("population"))

    def test_percentiles(self):
        self.assertApprox(statistics(list(range(1, 101)), percentiles=[90])["percentiles"]["90"], 90.1)

    def test_paired(self):
        self.assertApprox(statistics([1, 2, 3, 4, 5], y=[2, 4, 6, 8, 10])["paired"]["pearson_r"], 1)

    def test_single_value(self):
        self.assertEqual(statistics([3])["mean"], 3)

    def test_rejects_empty(self):
        with self.assertRaises(MathError):
            statistics([])


class RegressionTest(MathTestCase):
    def test_linear(self):
        r = regression([1, 2, 3, 4], [3, 5, 7, 9])
        self.assertApprox(r["coefficients"]["intercept"], 1)
        self.assertApprox(r["coefficients"]["slope"], 2)
        self.assertApprox(r["r_squared"], 1)

    def test_prediction(self):
        self.assertApprox(regression([1, 2, 3], [2, 4, 6], predict=[10])["predictions"][0]["y"], 20)

    def test_polynomial(self):
        self.assertApprox(regression([0, 1, 2, 3], [1, 2, 5, 10], model="polynomial", degree=2)["coefficients"]["x^2"], 1)

    def test_exponential(self):
        xs = [0, 1, 2, 3]
        r = regression(xs, [2 * math.exp(0.5 * x) for x in xs], model="exponential")
        self.assertApprox(r["coefficients"]["a"], 2)
        self.assertApprox(r["coefficients"]["b"], 0.5)

    def test_length_mismatch(self):
        with self.assertRaisesRegex(MathError, "same length"):
            regression([1, 2], [1, 2, 3])


class ProbabilityTest(MathTestCase):
    def test_normal_between(self):
        r = probability("normal", "between", {"mean": 100, "std": 15}, lower=85, upper=115)
        self.assertApprox(r["result"], 0.682689492137)

    def test_binomial_discrete_semantics(self):
        r = probability("binomial", "sf", {"n": 10, "p": 0.5}, x=7)
        self.assertApprox(r["result"], 56 / 1024)
        self.assertIn("P(X >= 8)", r["meaning"])

    def test_pmf(self):
        self.assertApprox(probability("poisson", "pdf", {"lambda": 2}, x=0)["result"], math.exp(-2))

    def test_quantile(self):
        self.assertApprox(probability("normal", "quantile", q=0.975)["result"], 1.959963985)

    def test_summary(self):
        self.assertApprox(probability("exponential", "summary", {"rate": 2})["result"]["mean"], 0.5)

    def test_missing_parameter(self):
        with self.assertRaisesRegex(MathError, "needs parameter"):
            probability("binomial", "cdf", {"n": 10}, x=3)

    def test_unknown_parameter(self):
        with self.assertRaisesRegex(MathError, "unknown parameter"):
            probability("normal", "cdf", {"mu": 0}, x=0)


class HypothesisTestTest(MathTestCase):
    def test_one_sample_t(self):
        r = hypothesis_test("one_sample_t", sample=[5.1, 4.9, 5.3, 5.5, 5.0], mu=5)
        self.assertApprox(r["p_value"], 0.2115789946)
        self.assertIs(r["reject_null"], False)

    def test_two_sample_welch(self):
        r = hypothesis_test("two_sample_t", sample=[1, 2, 3, 4, 5], sample2=[6, 7, 8, 9, 10])
        self.assertIs(r["reject_null"], True)
        self.assertEqual(r["variant"], "Welch")

    def test_chi2_independence(self):
        r = hypothesis_test("chi2_independence", table=[[10, 20], [30, 40]])
        self.assertEqual(r["df"], 1)
        self.assertIn("cramers_v", r)

    def test_goodness_of_fit_accepts_probabilities(self):
        self.assertApprox(hypothesis_test("chi2_goodness_of_fit", sample=[50, 50], expected=[0.5, 0.5])["p_value"], 1)

    def test_binomial(self):
        r = hypothesis_test("binomial", sample=[9, 10], p=0.5, alternative="greater")
        self.assertApprox(r["p_value"], 11 / 1024)

    def test_missing_input(self):
        with self.assertRaisesRegex(MathError, "needs 'sample2'"):
            hypothesis_test("paired_t", sample=[1, 2, 3])