%matplotlib inline

from __future__ import division

import numpy as np
import matplotlib.pyplot as plt
import mechanize
import StringIO
import pandas as pd
import quantities as pq

# 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, pandas, quantities

# URL SQL Search DR10
url = "http://skyserver.sdss3.org/dr10/en/tools/search/sql.aspx"
# Abrimos una instancia de Browser para hacer los queries
br = mechanize.Browser()

def SDSS_select(sql):
    br.open(url)
    br.select_form(name="sql")
    br['cmd'] = sql
    br['format']=['csv']
    response = br.submit()
    file_like = StringIO.StringIO(response.get_data())
    return pd.read_csv(file_like,  skiprows=1)

SQL_V = (
'SELECT '
 'objID, subclass, flags, zwarning, '
 'modelMag_u, modelMag_g, modelMag_r, '
 'modelMag_i, modelMag_z '
'FROM SpecPhoto '
'WHERE '
 'class = "STAR" '
'AND '
 'subclass LIKE "%V%" '
'AND '
 'subclass NOT LIKE "%IV%" '
'AND '
 'subclass NOT LIKE "%VI%" '
'AND '
 'type = 6')

df_V = SDSS_select(SQL_V)
df_V.shape

df_V.columns = ['objID','subclass','flags','zwarning','u','g','r','i','z']
df_V.head()

df_V[df_V['zwarning']!=0].shape

df_V = df_V[df_V['zwarning']==0]

df_V = df_V.drop('zwarning', axis=1)
df_V.shape

df_V = df_V.set_index('objID')
df_V.head()

flg = df_V.loc[1237662640123281773,'flags']
hex(flg)
SELECT objID, dbo.fPhotoFlagsN(flags) FROM SpecPhoto
WHERE
objID= "1237662640123281773" Select * from PhotoFlags
mask = 2**0 + 2**5 + 2**(14+32) + 2**7 + 2**19 + 2**(26+32) + 2**24 + 2**22

f = lambda s: (s & mask) == 0
b =df_V['flags'].map(f)
b.value_counts()

df_V_clean =df_V[b]
df_V_clean.shape

f = lambda s: s[0]
clases = df_V_clean['subclass'].map(f)
clases.value_counts()

b = df_V_clean['subclass'].apply(lambda s : s[0] == 'C')
df_V_clean['subclass'][b].value_counts()

df_V_clean = df_V_clean[~b]
df_V_clean.shape

# Preparamos la query
#(este sería otro ejemplo de como preparar el string con el SQL)
SQL_O = '                            \
SELECT                               \
 objID, subclass, flags,             \
 modelMag_u, modelMag_g, modelMag_r, \
 modelMag_i, modelMag_z              \
FROM SpecPhoto                       \
WHERE                                \
 class = "STAR"                      \
AND                                  \
 subclass LIKE "O%"                  \
AND                                  \
 type = 6                            \
AND                                  \
 zwarning = 0'

# Y la ejecutamos
df_O = SDSS_select(SQL_O)
df_O.columns = ['objID','subclass','flags','u','g','r','i','z']
df_O = df_O.set_index('objID')
df_O.shape

# depuramos los objetos con flags "malos":
f = lambda s: (s & mask) == 0
b =df_O['flags'].map(f)
df_O_clean =df_O[b]
df_O_clean.shape

df = pd.concat([df_V_clean, df_O_clean])
df.shape

df['u-g'] = df['u'] - df['g']
df['g-r'] = df['g'] - df['r']
df['r-i'] = df['r'] - df['i']
df['i-z'] = df['i'] - df['z']

colors = {'O':'#0000FF', 'B':'#CCCCFF', 'A':'#FFFFFF',
          'F':'#FFFFB2', 'G':'#FFFF00', 'K':'#FFA500',
          'M':'#FF0000'}

fig = plt.figure(figsize=(10,10))
ax = fig.add_subplot(111, axisbg='0.6')

ax.set_xlim(-1, 6)
ax.set_ylim(-1, 2.5)
ax.grid()
ax.set_title('Secuencia principal: g-r vs. u-g')

ax.title.set_fontsize(20)
ax.set_xlabel('u-g')
ax.xaxis.label.set_fontsize(20)
ax.set_ylabel('g-r')
ax.yaxis.label.set_fontsize(20)

for cls in 'OBAFGKM':
    b = df['subclass'].map(lambda s: s[0] ==cls)
    x = df[b]['u-g']
    y = df[b]['g-r']
    c = colors[cls]
    ax.scatter(x, y, c = c, s=6, edgecolors='none', label = cls)
legend = ax.legend(scatterpoints=1,markerscale = 6, shadow=True)
frame = legend.get_frame()
frame.set_facecolor('0.90')

def plot_gr_ug(df, cls, axes):
    x = df['u-g']
    y = df['g-r']
    axes.scatter(x, y, c = colors[cls], s=6, edgecolors='none')
    axes.set_xlim(-1, 6)
    axes.set_ylim(-1,2.5)
    axes.set_xlabel('u-g')
    axes.set_ylabel('g-r')
    axes.grid()
    axes.set_title('Secuencia principal: class %s' %cls);

# Definición de las constantes del sistema fotométrico del SDSS
lambda_u = 3540 * pq.angstrom
delta_u = 599 * pq.angstrom
lambda_g = 4770 * pq.angstrom
delta_g = 1379 * pq.angstrom
lambda_r = 6222 * pq.angstrom
delta_r = 1382 * pq.angstrom
lambda_i = 7632 * pq.angstrom
delta_i = 1535 * pq.angstrom
lambda_z = 9049 * pq.angstrom
delta_u = 1370 * pq.angstrom

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))

# El flujo recibido lo aproximamos como el producto de la radiancia espectral
# B(lambda, T) por la anchura delta lambda

def get_UG(T):    
    F = B(lambda_g, T) * delta_g/(B(lambda_u, T)* delta_u)
    return 2.5 * np.log10(F)

def get_GR(T):    # el flujo recibido lo aproximamos como el producto de la radiancia espectral el lambda por la anchura delta lambda
    F = B(lambda_r, T) * delta_r/(B(lambda_g, T)* delta_g)
    return 2.5 * np.log10(F)

def get_RI(T):
    F = B(lambda_i, T) * delta_i/(B(lambda_r, T)* delta_r)
    return 2.5 * np.log10(F)

fig = plt.figure(figsize=(10,20))
c ='OBAFGKM'
for i in range(0,7):   
    ax1 = fig.add_subplot(4,2,i+1, axisbg='0.4')
    cls = c[i]
    b = df['subclass'].map(lambda s: s[0] ==cls)
    plot_gr_ug(df[b],cls=cls,axes=ax1)
    T1 = 2000 * pq.kelvin
    T2 = 15000 * pq.kelvin
    GmenosR_1 = get_GR(T1)
    UmenosG_1 = get_UG(T1)
    GmenosR_2 = get_GR(T2)
    UmenosG_2 = get_UG(T2)
    # Trazamos la línea "blackbody"
    ax1.plot([UmenosG_1, UmenosG_2], [GmenosR_1, GmenosR_2], c='k',ls='--')

def plot_ri_gr(df, cls, axes):
    x = df['g-r']
    y = df['r-i']
    axes.scatter(x, y, c = colors[cls], s=6, edgecolors='none')
    axes.set_xlim(-1, 2.5)
    axes.set_ylim(-1, 2)
    axes.set_xlabel('g-r')
    axes.set_ylabel('r-i')
    axes.grid()
    axes.set_title('Secuencia principal: class %s' %cls);

fig = plt.figure(figsize=(10,20))
for i in range(0,7):
    c ='OBAFGKM'
    ax2 = fig.add_subplot(4,2,i+1, axisbg='0.4')
    cls = c[i]
    b = df['subclass'].map(lambda s: s[0] ==cls)
    plot_ri_gr(df[b],cls=cls,axes=ax2)
    T1 = 2000 * pq.kelvin
    T2 = 15000 * pq.kelvin
    GmenosR_1 = get_GR(T1)
    RmenosI_1 = get_RI(T1)
    GmenosR_2 = get_GR(T2)
    RmenosI_2 = get_RI(T2)

    ax2.plot([GmenosR_1, GmenosR_2], [RmenosI_1, RmenosI_2], c='k', ls='--')