import sys

import matplotlib.pyplot as plt
from math import *
import numpy as np


# Reading the fits file with the frequencies:

def read_sample(filename):

	from astropy.io import fits

	# read all hdus:
	hdulist = fits.open(filename)
	
	# PARAM is numpy array of parameters (radius and beta) for each point in the sample:
	# PARAM = [..., [ radius in M, beta ], [..., ...], ...]
	param   = hdulist['PARAM'].data
	
	# FREQ-RE is a numpy array of real parts of the eigenfrequencies
	freq_re = hdulist['FREQ-RE'].data

	# FREQ-IM is a numpy array of imaginary parts of the eigenfrequencies
	freq_im = hdulist['FREQ-IM'].data

	# SPIN is the black hole spin
	a = hdulist[0].header['SPIN']

	# AWNUM is the azimuthal wavenumber
	m = hdulist[0].header['AWNUM']

	# other informations:
	print(hdulist[0].header['PINDEX'])		# polytropic index
	print(hdulist[0].header['R0MIN'])		# minimal central radius
	print(hdulist[0].header['R0MAX'])		# maximal central radius
	print(hdulist[0].header['ROUTMAX'])		# maximal outer radius of the torus
		
	return (a, m, param, freq_re, freq_im) 


(a, m, param, freq_re, freq_im) = read_sample('sample/RAD+_a00_m+0.fits')


# this plots the sample in radius - beta plane:

plt.figure(1)
plt.plot(param[:,0], param[:,1], ls='', marker='x', color='black')


from scipy.interpolate import griddata

# this plots the radial mode frequencies for different betas:

plt.figure(2)
betas = np.linspace(0., 0.5, 10)
radii = np.linspace(6.02, 20., 100)
for b in betas:
	points = np.asarray([[r, b] for r in radii])
	frequencies = griddata(param, freq_re, points, method='linear')
	plt.plot([points[i,0] for i in range(len(points))], frequencies, ls='-', marker='')

plt.xlabel(r'$r[M]$')
plt.ylabel(r'$\omega[M^{-1}]$')

plt.show()