#include "../squiggle.h"
#include "../squiggle_more.h"
#include <stdio.h>
#include <stdlib.h>

// Estimate functions
double sample_ais_1(uint64_t* seed)
{
    double num_arxiv_ml_authors_2024 = 7379; // Number of authors who published in the stats.ML category on arxiv in 2023
    double fraction_of_ml = sample_beta(7.41986324742243, 114.487997692331, seed); // fraction they are of the field. 0.03 to 0.1. https://nunosempere.com/blog/2023/03/15/fit-beta/
    double fraction_of_their_research_thats_relevant = sample_beta(0.8277362357555023, 25.259989675532076, seed); // fraction of their research that is safety relevant, 0.001 to 0.1
    double academia_adjustment = sample_beta(1.9872200324266, 6.36630125578423, seed); // 0.05 0.5 adjustment because they are from academia
    
    return num_arxiv_ml_authors_2024 * fraction_of_their_research_thats_relevant * academia_adjustment / fraction_of_ml;
}

double sample_ais_2(uint64_t* seed)
{
    double num_arxiv_ml_authors_2024 = 7379; // Number of authors who published in the stats.ML category on arxiv in 2023
    double fraction_of_ml = sample_beta(7.41986324742243, 114.487997692331, seed); // fraction they are of the field. 0.03 to 0.1. https://nunosempere.com/blog/2023/03/15/fit-beta/
    double fraction_of_their_research_thats_relevant = sample_beta(3.28962721497463, 17.7686162987246, seed); // fraction of their research that is safety relevant, 0.001 to 0.1
    double academia_adjustment = sample_beta(2.23634269185645, 3.73532102339597, seed); // 0.05 0.5 adjustment because they are from academia
    
    return num_arxiv_ml_authors_2024 * fraction_of_their_research_thats_relevant * academia_adjustment / fraction_of_ml; 
}

int main()
{
    // set randomness seed
    uint64_t* seed = malloc(sizeof(uint64_t));
    *seed = 1000; // xorshift can't start with 0

    int n_samples = 10 * MILLION;

    printf("# AIS 1\n");
    double* xs = malloc(sizeof(double) * (size_t)n_samples);
    sampler_parallel(sample_ais_1, xs, 16, n_samples);
    printf("# Stats\n");
    array_print_stats(xs, n_samples);
    printf("\n# Histogram\n");
    array_print_histogram(xs, n_samples, 23);

    printf("# AIS 2\n");
    sampler_parallel(sample_ais_2, xs, 16, n_samples);
    printf("# Stats\n");
    array_print_stats(xs, n_samples);
    printf("\n# Histogram\n");
    array_print_histogram(xs, n_samples, 23);

    free(seed);
}