import pylab as pl
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

import astropy #llibreria d'on importarem el model del nostre univers
from astropy.cosmology import FlatLambdaCDM

# tenim un arxiu en local amb les dades copiades del seu article
# l'importem per reading amb el nom 'dades'
with open('./practiques/perlmutter.txt', "r") as dades:
    dat = pd.read_csv(dades, names=['z', 'Magnitud efectiva'], comment='#', sep='\s+')

toprint = dat[:10] #les 10 primeres files de dades
toprint #les imprimim per veure que les nostres dades les hem guardat correctament

def lumdist(z, olambda):
    '''
    Per aquest càlcul prenem:
        Constant de Hubble: 70 (km/s)/Mpc
        Omega matter (densitat de matèria no relativística): 1
        Temperatura de CMB (Cosmic microwave background): 2.725 K
    Paràmetres d'entrada:
        z: redshift
        olambda: la fracció de tota la matèria que es comporta com Dark Energy
    Sortida:
        'Distància lumínica' en Mpc a un redshift donat
    '''
    cosmo = FlatLambdaCDM(H0=70, Om0=1. - olambda, Tcmb0=2.725)
    return  cosmo.luminosity_distance(z)

def effmag(z, olambda, MB):
    '''
    Necessitem convertir la distància en la forma en la qual els astrònoms mesuren la brillantor
    Paràmetres d'entrada:
        z: redshift
        olambda: la fracció de tota la matèria que es comporta com Dark Energy
        MB: magnitud absoluta d'una estrella; B indica que estem referint-nos a la part de
        radiació estelar que s'emet a la part de “blau” de l'espectre
    Sortida:
        Magnitud efectiva, en el cas de SNeIa correspon equivalentment amb la distància (per standard
        candles, quan més tènue veiem un d'aquests objectes, més allunyat es troba de nosaltres)
    '''
    DL = lumdist(z, olambda) #emprem la funció anterior
    return MB + 5. * np.log10(DL.value) #calculem a partir de la magnitud absoluta i la distància lumínica

# plot on eix x és el redshift i l'eix y la lluminositat, equivalent a la distància
# marquem amb 'x' de color negre i amb certa transparència on són les nostres supernoves
pl.plot(dat['z'], dat['Magnitud efectiva'], marker='x', color='k', lw=0.0, alpha=0.6)

#anomenem els nostres eixos
pl.xlabel('Desplaçament al vermell (redshift)')
pl.ylabel('Magnitud efectiva')

#calculate equation for quartic trendline
z = np.polyfit(dat['z'],dat['Magnitud efectiva'] , 5)
p = np.poly1d(z)

x = np.linspace(0.1,0.83,100) #fem les divisions que considerem necessàries entre els valors pels quals tenim dades

# afegim la trendline a la gràfica
plt.plot(x, p(x), 'darkslategrey', label='Línia de tendència')

zs = np.arange(0.01, 0.85, 0.01)
#energia fosca que afegim al nostre univers creat amb Python
pl.plot(zs, effmag(zs, 0.67, 6.), c='r', label='67% Energia Fosca', linestyle='dashdot') 
pl.legend(loc=4, frameon=False) #mostrem una llegenda descrivint el plot amb els labels

#per aquest cas, hem de definir la zona que volem visualitzar de la nostra funció
plt.ylim(13,25)