# Cochran's Sample Size Calculation
import scipy.stats as stats
def cochran_sample_size(Z, p, e, N):
    n0 = (Z**2 * p * (1 - p)) / (e**2)
    n = n0 / (1 + ((n0 - 1) / N))
    return int(n)

# Parameters
Z = stats.norm.ppf(0.975)   # Calculate the critical z-value for a 95% confidence level in a standard normal distribution
p = 0.5                     # Estimated proportion (use 0.5 for maximum variability, sort of like worst-case scenario)
e = 0.05                    # Margin of error (precision level, more precision requires larger sample size)
N = 25000                   # Finite population size (in this example, student population size)

# Calculate sample size

# Calculate n0 and adjusted sample size for finite population
n0 = (Z**2 * p * (1 - p)) / (e**2)
sample_size = n0 / (1 + ((n0 - 1) / N))
print(f"Ideal sample size (n0): {round(n0)}")
print(f"Adjusted sample size for finite population: {int(sample_size)}")

# Create a distribution chart for each of the three variables after recoding
import matplotlib.pyplot as plt

# Proportion of Gender in the Sample Selection Data and add the proportion on top of each bar
gender_proportions = sample_selection_data['Gender_m'].value_counts(normalize=True).round(2)
gender_proportions.plot(kind='bar', title='Proportion of Gender')
plt.xticks(rotation=0)
# Round to 2 decimal places and convert to percentage
for i, v in enumerate(gender_proportions):
    plt.text(i, v + 0.01, str(round(v * 100, 2)) + '%', ha='center')
plt.show()

# Proportion of Race in the Sample Selection Data
race_proportions = sample_selection_data['Race_m'].value_counts(normalize=True).round(2)
race_proportions.plot(kind='bar', title='Proportion of Race')
plt.xticks(rotation=0)
# Round to 2 decimal places and convert to percentage
for i, v in enumerate(race_proportions):
    plt.text(i, v + 0.01, str(round(v * 100, 2)) + '%', ha='center')
plt.show()

# Proportion of Age Group in the Sample Selection Data
age_proportions = sample_selection_data['Age_Group_2'].value_counts(normalize=True).round(2)
age_proportions.plot(kind='bar', title='Proportion of Age Group')
plt.xticks(rotation=0)
for i, v in enumerate(age_proportions):
    plt.text(i, v + 0.01, str(round(v * 100, 2)) + '%', ha='center')
plt.show()

# This time we are using a different set of parameters. The formula remains the same from Step 2.
# Parameters
Z = stats.norm.ppf(0.975)   # Calculate the critical z-value for
# a 95% confidence level in a standard normal distribution
p = 0.5                     # Estimated proportion (use 0.5 for maximum variability, sort of like worst-case scenario)
e_3 = 0.03                    # Margin of error (precision level, more precision requires larger sample size)
N = 25000                   # Finite population size (in this example, student population size)

# Calculate sample size
# Calculate n0 and adjusted sample size for finite population
n0 = (Z**2 * p * (1 - p)) / (e_3**2)
sample_size_larger_e = n0 / (1 + ((n0 - 1) / N))
print(f"Ideal sample size (n0) for 95% confidence level with 3% margin of error: {round(n0)}")
print(f"Adjusted sample size of finite population for 95% confidence level with 3% margin of error: {int(sample_size_larger_e)}")

# Determine the combined proportions of the strata in the sample selection data with the modified variables
stratum_set_1 = sample_selection_data.groupby(['Gender_m', 'Race_m', 'Age_Group_2'], observed=True).size().reset_index(name='Count')
stratum_set_1['Proportion'] = stratum_set_1['Count'] / stratum_set_1['Count'].sum()

# use the 3% sample size to draw the stratified random sample
stratum_sizes_three_moe = (stratum_set_1['Proportion'] * int(sample_size_larger_e)).round().astype(int)
# Add the grouping variables back to the stratum sizes DataFrame for clarity
stratum_sizes_three_moe = pd.concat([stratum_set_1[['Gender_m', 'Race_m', 'Age_Group_2']], stratum_sizes_three_moe.rename('Sample_Size')], axis=1)
print("Stratum sizes for each subgroup based on 3% margin of error sample size (first 3 rows):")
stratum_sizes_three_moe.head(3)

# Initialize an empty list to hold sampled dataframes
sampled_groups = []
for i, row in stratum_sizes_three_moe.iterrows():
    stratum_filter = (
        (sample_selection_data['Gender_m'] == row['Gender_m']) &
        (sample_selection_data['Race_m'] == row['Race_m']) &
        (sample_selection_data['Age_Group_2'] == row['Age_Group_2'])
    )
    stratum_population = sample_selection_data[stratum_filter]
    # Use the sample size from the current row
    stratum_sample_size = int(row['Sample_Size'])

    # Ensure we don't sample from empty strata or more than available
    n = min(stratum_sample_size, len(stratum_population))
    if n > 0:
        stratum_sample = stratum_population.sample(n=n, replace=False)
        sampled_groups.append(stratum_sample)

# Combine all sampled dataframes into a single dataframe (guard against empty list)
if sampled_groups:
    sampled_dataset = pd.concat(sampled_groups).reset_index(drop=True)
else:
    sampled_dataset = pd.DataFrame(columns=sample_selection_data.columns)

# Create an alias expected by downstream cells and align column naming
sampled_data = sampled_dataset.copy()
if 'Age_Group' not in sampled_data.columns and 'Age_Group_2' in sampled_data.columns:
    sampled_data['Age_Group'] = sampled_data['Age_Group_2']
# View the first three rows of the sampled dataset
print("Sampled dataset (first 3 rows):")
sampled_data[['StudentID', 'Gender_m', 'Race_m', 'Age_Group']].head(3)

import pandas as pd
data_folder = './Data Files/'
# Import the respondents data file (synthetic respondents data based on a survey)
respondents_data = pd.read_excel(data_folder + 'respondents_file.xlsx')
# Make the output view smaller for presentation
respondents_data[:3]

# Recall the propotions from the sampled data
sampled_data_proportions = {
    "gender_sampled": sampled_dataset['Gender_m'].value_counts(normalize=True).round(3).to_dict(),
    "race_sampled": sampled_dataset['Race_m'].value_counts(normalize=True).round(3).to_dict(),
    "age_group_sampled": sampled_dataset['Age_Group_2'].value_counts(normalize=True).round(3).to_dict()
}

## Calculate the respondents demographic proportions of the three variables: Gender, Race, and Age Group and create a dictionary by the propotion and round to three decimal places
respondent_proportions = {
    "Gender_r": respondents_data['Q1_Gender'].value_counts(normalize=True).round(3).to_dict(),
    "Race_r": respondents_data['Q2_Race'].value_counts(normalize=True).round(3).to_dict(),
    "Age_Group_r": respondents_data['Q3_AgeGroup'].value_counts(normalize=True).round(3).to_dict()
}

# Define the two dictionaries for population and respondents proportions
popl_dict = sampled_data_proportions
resp_dict = respondent_proportions

# Using each of the labels in the variables, then the population or respondent proportions can be retrieved from the dictionaries.
# For example, if Q1_Gender is 'Female', then the population proportion of 0.574 will be retrieved.
# Then, the respondent proportion, 'Gender_r' of 0.631 will be retrieved for the denominator.

# Now we can calculate the weights for each respondent based on the three variables
def calculate_weight(row, popl_dict, resp_dict):
    gender_weight = popl_dict['gender_sampled'].get(row['Q1_Gender'], 0) / resp_dict['Gender_r'].get(row['Q1_Gender'], 1)
    race_weight = popl_dict['race_sampled'].get(row['Q2_Race'], 0) / resp_dict['Race_r'].get(row['Q2_Race'], 1)
    age_group_weight = popl_dict['age_group_sampled'].get(row['Q3_AgeGroup'], 0) / resp_dict['Age_Group_r'].get(row['Q3_AgeGroup'], 1)
    total_weight = gender_weight * race_weight * age_group_weight
    return total_weight

# Apply the weight calculation to each row in the respondents_data DataFrame
respondents_data['O_Weight'] = respondents_data.apply(calculate_weight, axis=1, args=(popl_dict, resp_dict))
# View the first three rows with the calculated weights
print("Respondents data with calculated weights (first 3 rows, rounded to 3 decimal places):")  
respondents_data[['Q1_Gender', 'Q2_Race', 'Q3_AgeGroup', 'O_Weight']].head(3).round(3)

# Apply the weight calculation to each row in the respondents_data DataFrame
respondents_data['Q4_with_weights'] = respondents_data['Q4_Traffic_Accidents'] * respondents_data['O_Weight']
# View the first three rows with the calculated weighted scores
print("Respondents data with weighted Q4_Traffic_Accidents scores (first 3 rows, rounded to 3 decimal places):")  
respondents_data[['Q4_Traffic_Accidents', 'O_Weight', 'Q4_with_weights']].head(3).round(3)

traffic_accident_response = respondents_data[['Q4_Traffic_Accidents', 'Q4_with_weights']].copy()
non_weighted_mean = traffic_accident_response['Q4_Traffic_Accidents'].mean()
weighted_mean = traffic_accident_response['Q4_with_weights'].mean()
print(f"Non-weighted mean of Q4_Traffic_Accidents: {non_weighted_mean:.3f}")
print(f"Weighted mean of Q4_Traffic_Accidents: {weighted_mean:.3f}")

# Visual comparison of non-weighted vs weighted means with histogram
import matplotlib.pyplot as plt

plt.figure(figsize=(8, 5))
plt.clf()  # Clear any previous plots

# Histogram of non-weighted responses
n, bins, patches = plt.hist(
    respondents_data['Q4_Traffic_Accidents'],
    bins=6,
    alpha=0.5,
    label='Non-weighted',
    color='blue',
    edgecolor='black'
)

# Histogram of weighted responses
n_w, bins_w, patches_w = plt.hist(
    respondents_data['Q4_with_weights'],
    bins=6,
    alpha=0.5,
    label='Weighted',
    color='orange',
    edgecolor='black'
)

plt.xlabel('Q4_Traffic_Accidents (0=Never, 5=Very Often)')
plt.ylabel('Frequency')
plt.title('Distribution of Q4_Traffic_Accidents Responses')
plt.legend()
plt.tight_layout()
plt.show()

respondents_data[['Q4_Traffic_Accidents', 'Q4_with_weights']].describe().round(3)