import streamlit as st
import random
import numpy as np

def simple_random_sampling(population_size, sample_size, seed=None):
    if seed is not None:
        random.seed(seed)
    population = list(range(1, population_size + 1))
    if sample_size > population_size:
        st.error("Sample size cannot be greater than population size.")
        return None
    return random.sample(population, sample_size)

def stratified_sampling(population_size, sample_size, strata_sizes, seed=None):
    if seed is not None:
        random.seed(seed)
    strata_boundaries = np.cumsum(strata_sizes)
    population = list(range(1, population_size + 1))
    sample = []
    
    for i, size in enumerate(strata_sizes):
        start = 0 if i == 0 else strata_boundaries[i-1]
        end = strata_boundaries[i]
        stratum_population = population[start:end]
        stratum_sample_size = int(size / population_size * sample_size)
        sample.extend(random.sample(stratum_population, stratum_sample_size))
    return sample

def cluster_sampling(clusters, selected_clusters, seed=None):
    if seed is not None:
        random.seed(seed)
    selected = random.sample(clusters, selected_clusters)
    return [item for sublist in selected for item in sublist]

def systematic_sampling(population_size, sample_size, seed=None):
    if seed is not None:
        random.seed(seed)
    step = population_size // sample_size
    start = random.randint(0, step - 1)
    return list(range(start, population_size, step))[:sample_size]

def snowball_sampling(initial_samples, recruit_per_sample, population):
    sample = set(initial_samples)
    while len(sample) < recruit_per_sample:
        recruits = [random.choice(population) for _ in range(recruit_per_sample)]
        sample.update(recruits)
    return list(sample)[:recruit_per_sample]

def main():
    st.title('Advanced Sampling Methods App')
    st.markdown("""
        ## Description
        Select from various sampling methods depending on your research needs.
        Each method has specific inputs tailored to different research scenarios.
    """)

    sampling_method = st.sidebar.selectbox("Choose the sampling method", 
                                           ("Simple Random Sampling", "Stratified Sampling", 
                                            "Cluster Sampling", "Systematic Sampling", 
                                            "Snowball Sampling"))
    
    population_size = st.sidebar.number_input('Enter Population Size', min_value=1, value=1000, step=1)
    sample_size = st.sidebar.number_input('Enter Sample Size', min_value=1, value=50, step=1)
    seed = st.sidebar.number_input('Enter Random Seed (optional)', min_value=0, value=None, step=1, format='%d', key='seed')

    if sampling_method == "Stratified Sampling":
        strata_sizes = st.sidebar.text_input('Enter Strata Sizes (comma-separated, e.g., 10,20,30)')
        strata_sizes = list(map(int, strata_sizes.split(','))) if strata_sizes else []

    if sampling_method == "Cluster Sampling":
        num_clusters = st.sidebar.number_input('Enter Number of Clusters', min_value=1, value=5, step=1)
        selected_clusters = st.sidebar.number_input('Enter Number of Clusters to Sample', min_value=1, value=2, step=1)
        clusters = [list(range(i * 100, (i + 1) * 100)) for i in range(num_clusters)]

    if st.sidebar.button('Generate Sample'):
        sample = []
        if sampling_method == "Simple Random Sampling":
            sample = simple_random_sampling(population_size, sample_size, seed)
        elif sampling_method == "Stratified Sampling":
            sample = stratified_sampling(population_size, sample_size, strata_sizes, seed)
        elif sampling_method == "Cluster Sampling":
            sample = cluster_sampling(clusters, selected_clusters, seed)
        elif sampling_method == "Systematic Sampling":
            sample = systematic_sampling(population_size, sample_size, seed)
        elif sampling_method == "Snowball Sampling":
            initial_samples = [random.randint(1, population_size) for _ in range(5)]
            sample = snowball_sampling(initial_samples, sample_size, range(1, population_size + 1))
        
        if sample:
            st.write(f'### {sampling_method} Output')
            st.write(f'Sample: {sample}')
            st.write(f'Sample Size: {len(sample)}')
            if seed is not None:
                st.write(f'Random Seed Used: {seed}')

if __name__ == '__main__':
    main()