# -*- coding: utf-8 -*-
"""
Code to generate Figure S2 of the Supplementary Information.

Created on Fri Apr 12 16:43:08 2024
@author: Andrei Sontag
"""

import numpy as np
import seaborn as sn
import pandas as pd
import os, fnmatch

npoints = 1001

alf_vals = np.linspace(0,1,npoints)
kap_vals = np.linspace(0,1.15,npoints)

stats = np.zeros((npoints,npoints))

for i,alfa in enumerate(alf_vals):
    for j,kappa in enumerate(kap_vals):
        delta = (1-alfa)**2 - 4*kappa*alfa
        z = np.sqrt(delta)
        if delta >= 0 and z <= 1:
            stats[i,j] = True
        else:
            stats[i,j] = False
            
            
import matplotlib.pylab as pylab
params = {'legend.fontsize': 25,
          'figure.figsize': (10, 10),
         'axes.labelsize': 25,
         'axes.titlesize': 25,
         'xtick.labelsize': 25,
         'ytick.labelsize': 25}
pylab.rcParams.update(params)

cols = ['{0:0.2f}'.format(val) for val in kap_vals]
idxs = ['{0:0.2f}'.format(val) for val in alf_vals]
stats = pd.DataFrame(stats,columns = cols, index = idxs)
ax = sn.heatmap(stats,cmap=['#E8E9E8','cornflowerblue'],cbar=False)
ax.set_xlabel(r'$\kappa$')
ax.set_ylabel(r'$\alpha$')

os.chdir(r'.\With Abstentions\fitted_rates')

# Find data files
fileOfDirectory = os.listdir('.')
pattern = r'rates_*'
files = []
for filename in fileOfDirectory:
        if fnmatch.fnmatch(filename, pattern):
                files.append(filename)

# Show the name of the files
print(files)

for file in files:
    
    params = np.loadtxt(file,delimiter=',')
    
    arl = params[0]
    brl = params[1]
    crl = params[2]
    ars = params[3]
    brs = params[4]
    crs = params[5]
    asr = params[6]
    bsr = params[7]
    csr = params[8]
    
    alf = (2*arl+ars)/(bsr-brs)
    kap = asr/crs
    print(alf)
    print(kap)
    ax.scatter((kap/kap_vals[-1])*npoints-1,alf*npoints-1,marker='X',s=300,linewidth=2,ec='white',fc='black')

for axis in ['top','bottom','left','right']:
    ax.spines[axis].set_linewidth(1.5)
    ax.spines[axis].set_linewidth(1.5)