%matplotlib inline

from __future__ import division

import quantities as pq
import numpy as np
import matplotlib.pyplot as plt

# Generar un cuadro con versiones de las librerías utilizadas en este notebook
#https://github.com/jrjohansson/version_information
%load_ext version_information
%version_information numpy, matplotlib, quantities

k = 100**(1./5)
print k

m = 0 # Vega es la referencia para la magnitud aparente 0

M = -5 * np.log10(8.1 / 10)
print "M(Vega) = %.2f" %M
print u"módulo de distancia de Vega: %.2f" %(0 - M)

def B(wl,T):
    '''wl es un array de longitudes de onda con unidades de longitud
    T es una temperatura expresada en Kelvin
    el resultado es un array de valores de la radiancia espectral
    con unidades W/(m**2 * nm * sr)
    '''
    I = 2 * pq.constants.h * (pq.c)**2 / wl**5 *  \
        1 / (np.exp((pq.constants.h*pq.c \
        / (wl*pq.constants.k*T)).simplified)-1)
    return I.rescale(pq.watt/(pq.m**2 * pq.nm *pq.sr))

# Definición de las constantes del sistema fotométrico UBV

lambda_u = 365 * pq.nm
delta_u = 68 * pq.nm
lambda_b = 440 * pq.nm
delta_b = 98 * pq.nm
lambda_v = 550 * pq.nm
delta_v = 89 * pq.nm

# Cálculo de Cu-b

T = 42000*pq.kelvin
F = B(lambda_u, T) * delta_u/(B(lambda_b, T)* delta_b)
Cub = -1.19 + 2.5 * np.log10(F)
print Cub


# Cálculo de Cb-v

F = B(lambda_b, T) * delta_b/(B(lambda_v, T)* delta_v)
Cbv = -0.33 + 2.5 * np.log10(F)
print Cbv

def get_UB(T):
    F = B(lambda_u, T) * delta_u/(B(lambda_b, T)* delta_b)
    return -2.5 * np.log10(F) + Cub

def get_BV(T):
    F = B(lambda_b, T) * delta_b/(B(lambda_v, T)* delta_v)
    return -2.5 * np.log10(F) + Cbv


temp = np.arange(2500, 30000, 100)*pq.Kelvin
BmenosV = get_BV(temp)

fig, ax = plt.subplots(figsize=(10, 8))
ax.plot(BmenosV, temp, lw =2)
ax.grid()
ax.set_title(u"Relación Índice de color - Temperatura \n \
     para un cuerpo negro")
ax.title.set_fontsize(18)
ax.set_xlabel(u"Índice de color B-V")
ax.xaxis.label.set_fontsize(15)
ax.set_ylabel("Temperatura superficial en K")
ax.yaxis.label.set_fontsize(15)

# Dibuja tres estrellas representativas

T = np.array([3300, 5800, 22000])*pq.kelvin
BV = [1.85, 0.656, -0.21]
             
ax.scatter(BV[0],T[0], s=600, c='r', marker='*')
ax.scatter(BV[1],T[1], s=600, c='y', marker='*')
ax.scatter(BV[2],T[2], s=600, c='b', marker='*')

# Anotamos los nombres de las tres estrellas

ax.annotate("Betelgeuse", xy=(BV[0],T[0]+1000*pq.kelvin), size=14)
ax.annotate("Sol", xy=(BV[1],T[1]+1000*pq.kelvin), size=14)
ax.annotate("Bellatrix", xy=(BV[2],T[2]+1000*pq.kelvin), size=14);

fig, ax = plt.subplots(figsize=(10, 8))

temp = np.arange(2500,30000,100)*pq.kelvin
BmenosV = get_BV(temp)
UmenosB = get_UB(temp)

ax.scatter(BmenosV, UmenosB)
ax.grid()
ax.set_title(u'Gráfico color-color de los cuerpos negros')
ax.title.set_fontsize(20)
ax.set_xlabel('B-V')
ax.xaxis.label.set_fontsize(15)
ax.set_ylabel('U-B')
ax.yaxis.label.set_fontsize(15)

temp = np.array([3000, 4000, 5000, 7500, 10000, 20000, 30000])*pq.kelvin
BmenosV = get_BV(temp)
UmenosB = get_UB(temp)

ax.scatter(BmenosV, UmenosB, s=150., c='r')

for t in temp:
    ax.annotate("%d K" %t,xy=(get_BV(t)+0.1,get_UB(t)-0.05))