# -*- coding: utf-8 -*-

#AGA0503 - Métodos Numéricos em Astronomia
#Professor: Alex Carciofi
#Aluna: Anne Viegas Rathsam
#No USP: 10873535
#Exercício de Programação 4 - EP4

############################# PACOTES UTILIZADOS #############################

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

################### DETERMINAÇÃO DAS CONSTANTES UTILIZADAS ###################

c = 2.99793E+10    # velocidade da luz em cm/s
k = 1.38065e-16    # constante de Boltzmann em erg/K
h = 6.62607E-27    # constante de Planck em erg.s
sigma = 5.6704e-5  # constante de Stefan-Boltzmann em erg/cm^2sK^4

####################### DEFINIÇÃO DA FUNÇÃO PRINCIPAL #######################

def main():
	''' Dadas medidas de fluxo (com incerteza) por comprimento de onda de uma
	estrela circundada por uma camada de poeira, esse programa calcula a
	temperatura da estrela e da camada de poeira minimizando o chi-quadrado.
	Retorna, além dos valores da temperatura e do chi-quadrado mínimo, gráficos
	de chi-quadrado com temperaturas da estrela e da poeira fixos, mapa de 
	chi-quadrado, do espectro com o ajuste realizado e dos resíduos do ajuste.
	Utiliza as funções chi2, corpo_negro e modelo_planck.
	'''

	# Passo 1: Ler a tabela de dados em arrays específicos

	dados = np.loadtxt("dados_Anne.txt", skiprows=1) 
	# Pular a primeira linha do arquivo

	lm = dados[:,0] * 1.0E-8 # Convertendo para cm
	flux = dados[:,1] * 1.0E8 # Convertendo para erg/(cm *s)**2
	sigma_flux = dados[:,2] * 1.0E8

	# Passo 2: Gerar grade de N valores de T1 e T2 e calcular os N x N valores do 
	# chi-quadrado correspondentes. 

	T_min_estrela = float(input("Digite o valor mínimo para a temperatura da estrela, em Kelvin: "))
	T_max_estrela = float(input("Digite o valor máximo para a temperatura da estrela, em Kelvin: "))
	T_min_poeira = float(input("Digite o valor mínimo para a temperatura da poeira, em Kelvin: "))
	T_max_poeira = float(input("Digite o valor máximo para a temperatura da poeira, em Kelvin: "))
	n = int(input("Digite a quantidade de valores de temperaturas (n): "))
	
	print("\n-----------------------------------------------------------")
	
	print("\nCalculando...")
	
	print("\n-----------------------------------------------------------")
	
	# Criar vetores com os valores de temperatura da estrela (T1) e da poeira (T2)
	
	T1 = np.linspace(T_min_estrela, T_max_estrela, n)
	T2 = np.linspace(T_min_poeira, T_max_poeira, n)
	
	# Calcular os n x n valores de chi-quadrado
	
	chi_2 = []
	
	# Vamos criar 3 laços, pois precisamos calcular o chi_quadrado para cada
	# temperatura da poeira, da estrela e comprimento de onda
	
	j=0 # j é uma variável auxiliar para guardar o índice de chi_2 que está
	# sendo calculado
	
	for t1 in T1:
		for t2 in T2:
			chi_i = 0
			for i in range(0, len(lm), 1):
				modelo = modelo_planck(lm[i], t1, t2)
				chi_i += chi2(modelo, flux[i], sigma_flux[i])
			chi_2.append(chi_i)
			
	# Passo 3: Encontrar o menor valor de chi-quadrado e os valores de t1 e t2
	# correspondentes. 
	
	# Podemos fazer isso enquanto os chi-quadrados são calculados para facilitar
			
			if chi_2[j] == min(chi_2):
				temp_estrela = t1
				temp_poeira = t2
			j+=1
				
	# Encontramos as temperaturas t1 e t2 para as quais o chi-quadrado é mínimo
	
	print("\nO menor valor de chi-quadrado encontrado foi", min(chi_2))
	print("\nA temperatura da estrela é, em Kelvin,", temp_estrela)
	print("\nA temperatura da camada de poeira é, em Kelvin,", temp_poeira)
	
	# Passo 4: Fazer gráficos de chi^2 vs. T1 e chi^2 vs. T2
	
	# Com T2 fixa:
	
	chi_temp_poeira = []
	
	for t in T1:
		chi_i = 0
		for i in range(0, len(lm), 1):
			modelo = modelo_planck(lm[i], t, temp_poeira)
			chi_i += chi2(modelo, flux[i], sigma_flux[i])
		chi_temp_poeira.append(chi_i)
		
	# Plotar o gráfico
	
	plt.plot(T1, chi_temp_poeira)
	plt.grid()
	plt.title(r"$\chi^{2}$ por temperatura da estrela")
	plt.xlabel("Temperatura da estrela (K)")
	plt.ylabel(r"$\chi^{2}$")
	plt.show()
	#plt.savefig("chi_quadrado_t2fixa.png")
	
	# Com T1 fixa:
	
	chi_temp_estrela = []
	
	for t in T2:
		chi_i = 0
		for i in range(0, len(lm), 1):
			modelo = modelo_planck(lm[i], t, temp_estrela)
			chi_i += chi2(modelo, flux[i], sigma_flux[i])
		chi_temp_estrela.append(chi_i)
		
	# Plotar o gráfico
	
	plt.plot(T2, chi_temp_estrela, color='g')
	plt.grid()
	plt.title(r"$\chi^{2}$ por temperatura da poeira")
	plt.xlabel("Temperatura da poeira (K)")
	plt.ylabel(r"$\chi^{2}$")
	plt.show()
	#plt.savefig("chi_quadrado_t1fixa.png")
	
	# Passo 5: Fazer um gráfico 2D do chi-quadrado para visualizar se o mínimo 
	# encontrado é único, ou se há outros mínimos possíveis em outros lugares 
	# do espaço de parâmetros.
	
	X,Y = np.meshgrid(chi_temp_poeira, chi_temp_estrela)
	z = X + Y 
	
	plt.imshow(z, interpolation='nearest', cmap='Spectral')
	plt.colorbar()
	plt.title(r"Mapa de $\chi^{2}$")
	plt.xlabel("Temperatura da estrela")
	plt.ylabel("Temperatura da poeira")
	plt.xticks([]) 
	plt.yticks([])
	plt.show()
	#plt.savefig("mapa_chi_quadrado.png")

	# Passo 6: Faça um gráfico dos dados e do melhor modelo, bem como dos 
	#resíduos.
	
	# Primeiro precisamos calcular o vetor dos dados do modelo mod
	
	mod = []
	
	for i in range(0, len(lm), 1):
		mod_i = modelo_planck(lm[i], temp_estrela, temp_poeira)
		mod.append(mod_i)
		
	# Agora precisamos calcular os resíduos
	
	res = []
	
	for i in range(0, len(lm), 1):
		res_i = ((flux[i] - mod[i])/sigma_flux[i])
		res.append(res_i)
		
	# Gráfico dos dados e modelo
	
	fig, ax = plt.subplots()
	ax.errorbar(lm, flux, yerr=sigma_flux, color='k', linestyle='', marker='.')
	plt.plot(lm, mod, color='m', linestyle='-', marker='')
	plt.xscale('log')
	plt.yscale('log')
	plt.title("Espectro")
	plt.xlabel(r"$\lambda$ (cm)")
	plt.ylabel("Fluxo (erg/$cm^{2}$ $s^{2}$)")
	plt.gca().legend(('Modelo','Dados'))
	plt.grid()
	plt.show()
	#plt.savefig("espectro.png")
	
	# Gráfico dos resíduos
	
	plt.plot(lm, res, color='k', marker='.', linestyle='')
	plt.title("Resíduos")
	plt.xlabel("res x")
	plt.ylabel("res y")
	plt.grid()
	plt.show()
	#plt.savefig("residuos.png")

######################### DEFINIÇÃO DA FUNÇÃO CHI2 ##########################	

def chi2(mod, obs, inc):
	''' Essa função calcula uma componente da somatória do chi-quadrado de uma
	observação obs com incerteza inc, que esperamos que siga o modelo mod.
	'''
	
	chi = ((mod - obs)/inc)**2
	
	return chi
	
#################### DEFINIÇÃO DA FUNÇÃO DE CORPO NEGRO #####################	

def corpo_negro(lm, T):
	''' Essa função calcula o valor da função de corpo negro, dados valores do
	comprimento de onda (lm) e da temperatura (T) do corpo estudado.
	'''
	
	numerador = 2*math.pi*h*(c**2)
	exponencial = math.exp((h*c)/(lm*k*T))  
	denominador = (lm**5) * (exponencial - 1)
	
	B = numerador/denominador # B é a função de corpo negro
	
	return B

####################### DEFINIÇÃO DA FUNÇÃO DE PLANCK #######################	

def modelo_planck(lm, t1, t2):
	''' Essa função calcula o fluxo do espectro de uma estrela com temperatura
	t1 circundada por uma camada de poeira com temperatura t2, avaliado em um 
	comprimento de onda lm.
	Utiliza a função corpo_negro.
	'''
	
	S1 = corpo_negro(lm, t1)/(sigma * (t1**4))
	S2 = corpo_negro(lm, t2)/(sigma * (t2**4))
	
	S = S1 + S2
	
	return S
	
main()