import pandas as pd
from statistics import NormalDist

def ts_wald_os(data, codes, p0 = 0.5, cc = None):
    '''
    Single Binary Test - One-sample Wald
     
    This function will perform an one-sample Wald test, two-sided
    
    Parameters
    ----------
    data : list or Pandas data series with the data
    codes : list with the two codes to use
    p0 : The hypothesized proportion for the first category (default is 0.5)
    cc : use continuity correction, either None (default), or "yates"
        
    Returns
    -------
    testResults : Pandas dataframe with the test statistic, two-sided significance (p-value) and test used
   
    Notes
    -----
    It uses NormalDist from Python's statistics library
    
    Author
    ------
    Made by P. Stikker
    
    Please visit: https://PeterStatistics.com
    
    YouTube channel: https://www.youtube.com/stikpet
    
    '''
        
    #one-sample score test
    #an approximation using the normal distribution for a one-sample binomial test
    
    #Determine number of successes
    n1 = sum(data==codes[0])
    n2 = sum(data==codes[1])
    n = n1 + n2
    
    minCount = n1
    ExpProp = p0
    if (n2 < n1):
        minCount = n2
        ExpProp = 1 - ExpProp
        
    #Wald approximation
    if cc is None:
        p = minCount / n
        q = 1 - p
        se = (p * (1 - p) / n)**0.5
        Z = (p - ExpProp) / se
        sig2 = 2 * (1 - NormalDist().cdf(abs(Z)))
        testValue = Z
        testUsed = "Wald approximation"
    elif (cc == "yates"):
        #Wald approximation with continuity correction
        p = (minCount + 0.5) / n
        q = 1 - p
        se = (p * (1 - p) / n)**0.5
        Z = (p - ExpProp) / se
        sig2 = 2 * (1 - NormalDist().cdf(abs(Z)))
        testValue = Z
        testUsed = "Wald approximation with Yates continuity correction"
        
    testResults = pd.DataFrame([[testValue, sig2, testUsed]], columns=["statistic", "p-value (2-sided)", "test"])
    
    return (testResults)

#example
dataList = ['Female', 'Male', 'Male', 'Female', 'Male', 'Male', 'Female', 'Female', 'Male', 'Male', 'Male', 'Male', 'Male', 'Male', 'Female', 'Male', 'Female', 'Male', 'Male', 'Female', 'Female', 'Male', 'Male', 'Male', 'Male', 'Male', 'Male', 'Male', 'Female', 'Male', 'Male', 'Male', 'Male', 'Male', 'Male', 'Male', 'Female', 'Male', 'Male', 'Male', 'Male', 'Male', 'Male', 'Male', 'Female', 'Female']
data = pd.Series(dataList)
codes = ['Female', 'Male']

ts_wald_os(data, codes)