# Script to produce data for figure 2
# Created by Matthew Asker 

import numpy as np
from next_reaction_method_paper import next_reaction
import pandas as pd

a = 0.25
s = 0.1
total_sims = 1000
K_range = [25, 50, 100, 200]

x_th_result = list()
K_result = list()
fix_time = list()
fix_prob = list()

for K in K_range:
    if K <= 100:
        x_th_range = np.linspace(0, 1, K+1)[1:-1]
    else:
        x_th_range = np.linspace(0, 1, 21)[1:-1]
    for x_th in x_th_range:      
        G = next_reaction(s, a, 0.0, 0.0, K, K, x_th, total_sims)
        G.run()
        
        x_th_result.append(x_th)
        K_result.append(K)
        fix_time.append(G.average_time)
        fix_prob.append(G.r_fixation_count / total_sims)
    
raw_data = {'K' : K_result,
            'x_th' : x_th_result,
            'fix_prob' : fix_prob,
            'fix_time' : fix_time}
        
filename = f'fixation probability and time N={K_range}'       
data = pd.DataFrame(raw_data, columns=['x_th', 'K', 'fix_prob', 'fix_time'])
data.to_csv(filename+'.csv', index=False)