from abc import abstractmethod
from dataclasses import dataclass
from typing import List

import numpy as np
from scipy.stats import t

from gbstats.shared.models import FrequentistTestResult, Statistic, Uplift
from gbstats.shared.tests import BaseABTest


@dataclass
class FrequentistConfig:
    alpha: float = 0.05
    sequential: bool = False
    sequential_tuning_parameter: float = 5000


class TTest(BaseABTest):
    def __init__(
        self,
        stat_a: Statistic,
        stat_b: Statistic,
        config: FrequentistConfig = FrequentistConfig(),
        test_value: float = 0,
    ):
        """Base class for one- and two-sided T-Tests with unequal variance.
        All values are with respect to relative effects, not absolute effects.
        A result prepared for integration with the stats runner can be
        generated by calling `.compute_result()`

        Args:
            stat_a (Statistic): the "control" or "baseline" statistic
            stat_b (Statistic): the "treatment" or "variation" statistic
            test_value (float, optional): the null hypothesis for the difference in means. Defaults to 0.
            alpha (float, optional): the significance level for the CI construction. Defaults to 0.05.
        """
        super().__init__(stat_a, stat_b)
        self.alpha = config.alpha
        self.test_value = test_value

    @property
    def variance(self) -> float:
        return self.stat_b.variance / (
            pow(self.stat_a.unadjusted_mean, 2) * self.stat_b.n
        ) + self.stat_a.variance * pow(self.stat_b.unadjusted_mean, 2) / (
            pow(self.stat_a.unadjusted_mean, 4) * self.stat_a.n
        )

    @property
    def point_estimate(self) -> float:
        return (self.stat_b.mean - self.stat_a.mean) / self.stat_a.unadjusted_mean

    @property
    def critical_value(self) -> float:
        return (self.point_estimate - self.test_value) / np.sqrt(self.variance)

    @property
    def dof(self) -> float:
        # welch-satterthwaite approx
        return pow(
            self.stat_b.variance / self.stat_b.n + self.stat_a.variance / self.stat_a.n,
            2,
        ) / (
            pow(self.stat_b.variance, 2) / (pow(self.stat_b.n, 2) * (self.stat_b.n - 1))
            + pow(self.stat_a.variance, 2)
            / (pow(self.stat_a.n, 2) * (self.stat_a.n - 1))
        )

    @property
    @abstractmethod
    def p_value(self) -> float:
        pass

    @property
    @abstractmethod
    def confidence_interval(self) -> List[float]:
        pass

    def _default_output(self) -> FrequentistTestResult:
        """Return uninformative output when AB test analysis can't be performed
        adequately
        """
        return FrequentistTestResult(
            expected=0,
            ci=[0, 0],
            p_value=1,
            uplift=Uplift(
                dist="normal",
                mean=0,
                stddev=0,
            ),
        )

    def compute_result(self) -> FrequentistTestResult:
        """Compute the test statistics and return them
        for the main gbstats runner

        Returns:
            FrequentistTestResult -
                note the values are with respect to percent uplift,
                not absolute differences
        """
        if self.stat_a.mean == 0:
            return self._default_output()
        if self._has_zero_variance():
            return self._default_output()
        return FrequentistTestResult(
            expected=self.point_estimate,
            ci=self.confidence_interval,
            p_value=self.p_value,
            uplift=Uplift(
                dist="normal",
                mean=self.point_estimate,
                stddev=np.sqrt(self.variance),
            ),
        )


class TwoSidedTTest(TTest):
    @property
    def p_value(self) -> float:
        return 2 * (1 - t.cdf(abs(self.critical_value), self.dof))

    @property
    def confidence_interval(self) -> List[float]:
        width: float = t.ppf(1 - self.alpha / 2, self.dof) * np.sqrt(self.variance)
        return [self.point_estimate - width, self.point_estimate + width]


class OneSidedTreatmentGreaterTTest(TTest):
    @property
    def p_value(self) -> float:
        return 1 - t.cdf(self.critical_value, self.dof)

    @property
    def confidence_interval(self) -> List[float]:
        width: float = t.ppf(1 - self.alpha, self.dof) * np.sqrt(self.variance)
        return [self.point_estimate - width, np.inf]


class OneSidedTreatmentLesserTTest(TTest):
    @property
    def p_value(self) -> float:
        return t.cdf(self.critical_value, self.dof)

    @property
    def confidence_interval(self) -> List[float]:
        width: float = t.ppf(1 - self.alpha, self.dof) * np.sqrt(self.variance)
        return [-np.inf, self.point_estimate - width]


class SequentialTwoSidedTTest(TTest):
    def __init__(
        self,
        stat_a: Statistic,
        stat_b: Statistic,
        config: FrequentistConfig = FrequentistConfig(),
        test_value: float = 0,
    ):
        super().__init__(stat_a, stat_b, config, test_value)
        self.sequential_tuning_parameter = config.sequential_tuning_parameter

    @property
    def confidence_interval(self) -> List[float]:
        # eq 9 in Waudby-Smith et al. 2023 https://arxiv.org/pdf/2103.06476v7.pdf
        N = self.stat_a.n + self.stat_b.n
        rho = self.rho
        s2 = self.variance * N

        width: float = np.sqrt(s2) * np.sqrt(
            (
                (2 * (N * np.power(rho, 2) + 1))
                * np.log(np.sqrt(N * np.power(rho, 2) + 1) / self.alpha)
                / (np.power(N * rho, 2))
            )
        )
        return [self.point_estimate - width, self.point_estimate + width]

    @property
    def rho(self) -> float:
        # eq 161 in https://arxiv.org/pdf/2103.06476v7.pdf
        return np.sqrt(
            (-2 * np.log(self.alpha) + np.log(-2 * np.log(self.alpha) + 1))
            / self.sequential_tuning_parameter
        )

    @property
    def p_value(self) -> float:
        # eq 155 in https://arxiv.org/pdf/2103.06476v7.pdf
        N = self.stat_a.n + self.stat_b.n
        # slight reparameterization for this quantity below
        st2 = np.power(self.point_estimate - self.test_value, 2) * N / (self.variance)
        tr2p1 = N * np.power(self.rho, 2) + 1
        evalue = np.exp(np.power(self.rho, 2) * st2 / (2 * tr2p1)) / np.sqrt(tr2p1)
        return min(1 / evalue, 1)
