"""
This module contains functions for performing one-sample Z-tests and calculating
various statistics such as Cohen's d, Z-score, p-value, confidence intervals, and
standard errors.

The module includes the following functions:
- calculate_central_ci_from_cohens_d_one_sample: Calculate the confidence intervals
and standard error for Cohen's d effect size in a one-sample Z-test.
- One_Sample_ZTests: A class containing static methods for performing one-sample
Z-tests from Z-score, parameters, and data.
"""

import numpy as np
from scipy.stats import norm

from src.stats.results import CohenD


def calculate_central_ci_from_cohens_d_one_sample(
    cohens_d, sample_size, confidence_level
):
    """
    Calculate the confidence intervals and standard error for Cohen's d effect size in a
    one-sample Z-test.

    This function calculates the confidence intervals of the effect size (Cohen's d) for a
    one-sample Z-test or two dependent samples test using the Hedges and Olkin (1985)
    formula to estimate the standard error.

    Parameters
    ----------
    cohens_d : float
        The calculated Cohen's d effect size
    sample_size : int
        The size of the sample
    confidence_level : float
        The confidence level as a decimal (e.g., 0.95 for 95%)

    Returns
    -------
    tuple
        A tuple containing:
        - ci_lower (float): Lower bound of the confidence interval
        - ci_upper (float): Upper bound of the confidence interval
        - standard_error_es (float): Standard error of the effect size

    Notes
    -----
    Since the effect size in the population and its standard deviation are unknown,
    we estimate it based on the sample using the Hedges and Olkin (1985) formula
    to estimate the standard deviation of the effect size.
    """
    standard_error_es = np.sqrt((1 / sample_size) + ((cohens_d**2 / (2 * sample_size))))
    z_critical_value = norm.ppf(confidence_level + ((1 - confidence_level) / 2))
    ci_lower, ci_upper = (
        cohens_d - standard_error_es * z_critical_value,
        cohens_d + standard_error_es * z_critical_value,
    )
    return ci_lower, ci_upper, standard_error_es


class OneSampleZResults:
    """
    A class to store results from one-sample Z statistical tests.

    This class contains attributes to store various statistical measures including:
    - Effect size (Cohen's d)
    - Z-score and p-value
    - Standard error of the mean
    - Confidence intervals for Cohen's d
    - Standard error of the effect size
    """

    def __init__(self, z_score: float, p_value: float, cohens_d: CohenD) -> None:
        self.cohens_d: CohenD = cohens_d
        self.z_score: float = z_score
        self.p_value: float = p_value

        self.standard_error_of_the_mean: float | None = None
        self.sample_mean: float | None = None
        self.population_mean: float | None = None
        self.population_sd: float | None = None
        self.sample_sd: float | None = None
        self.difference_between_means: float | None = None
        self.sample_size: float | None = None


def one_sample_from_z_score(
    z_score: float, sample_size: float, confidence_level: float = 0.95
) -> OneSampleZResults:
    """
    Calculate the one-sample Z-test results from a given Z-score.

    """

    p_value = min(float(norm.sf((abs(z_score))) * 2), 0.99999)
    cohens_d = z_score / np.sqrt(sample_size)
    ci_lower, ci_upper, standard_error_es = (
        calculate_central_ci_from_cohens_d_one_sample(
            cohens_d, sample_size, confidence_level
        )
    )

    cohens_d = CohenD(
        value=round(cohens_d, 4),
        ci_lower=round(ci_lower, 4),
        ci_upper=round(ci_upper, 4),
        standard_error=round(standard_error_es, 4),
    )

    results = OneSampleZResults(z_score, p_value, cohens_d)

    return results


def one_sample_from_parameters(
    sample_mean: float,
    sample_size: int,
    population_mean: int,
    population_sd,
    confidence_level=0.95,
) -> OneSampleZResults:
    """
    Calculate the one-sample Z-test results from given population and sample parameters.
    """

    mean_standard_error = population_sd / np.sqrt(sample_size)
    z_score = (population_mean - sample_mean) / mean_standard_error
    cohens_d = (population_mean - sample_mean) / population_sd
    p_value = min(float(norm.sf((abs(z_score))) * 2), 0.99999)
    ci_lower, ci_upper, standard_error_es = (
        calculate_central_ci_from_cohens_d_one_sample(
            cohens_d, sample_size, confidence_level
        )
    )

    cohens_d = CohenD(
        value=round(cohens_d, 4),
        ci_lower=round(ci_lower, 4),
        ci_upper=round(ci_upper, 4),
        standard_error=round(standard_error_es, 4),
    )

    # Create results object
    results = OneSampleZResults(z_score, p_value, cohens_d)
    results.standard_error_of_the_mean = round(mean_standard_error, 4)
    results.sample_mean = sample_mean
    results.population_mean = population_mean
    results.population_sd = population_sd
    return results


def one_sample_from_data(
    column: list,
    population_mean: float,
    population_sd: float,
    confidence_level: float = 0.95,
) -> OneSampleZResults:
    """
    Calculate the one-sample Z-test results from given sample data.
    """

    sample_mean = np.mean(column)
    sample_sd = np.std(column, ddof=1)
    diff_mean = population_mean - sample_mean
    sample_size = len(column)
    standard_error = population_sd / (np.sqrt(sample_size))
    z_score = diff_mean / standard_error
    cohens_d = (
        diff_mean
    ) / population_sd  # This is the effect size for one sample z-test Cohen's d
    p_value = min(float(norm.sf((abs(z_score))) * 2), 0.99999)
    ci_lower, ci_upper, standard_error_es = (
        calculate_central_ci_from_cohens_d_one_sample(
            cohens_d, sample_size, confidence_level
        )
    )

    cohens_d = CohenD(
        value=cohens_d,
        ci_lower=ci_lower,
        ci_upper=ci_upper,
        standard_error=standard_error_es,
    )

    results = OneSampleZResults(z_score, p_value, cohens_d)
    results.sample_size = sample_size
    results.population_mean = population_mean
    results.sample_mean = sample_mean
    results.sample_sd = sample_sd
    results.standard_error_of_the_mean = standard_error
    results.difference_between_means = diff_mean

    return results
