%matplotlib inline

from __future__ import division

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

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

colspecs = [(0,4), (4,14), (25,31), (31,37), (41,42), (43,44), (51,60), 
            (90,96), (96,102), (102,107), (109,114), (115,120), (127,147)]

labels =['HR', 'Name', 'HD', 'SAO', 'IRflag', 'Multiple', 'VarID', 
         'GLON', 'GLAT', 'Vmag', 'B-V', 'U-B', 'SpType']

bsc = pd.read_fwf('../datos/bsc5.dat', header= None, colspecs=colspecs,
                    names=labels, index_col=0)

bsc.head()

bsc.tail()

bsc['SpType'].nunique()

pd.isnull(bsc['SpType']).sum()

b = pd.notnull(bsc['SpType'])
bsc = bsc[b]
bsc.info()

b1 = bsc['SpType'].str.contains('V')
b2 = bsc['SpType'].str.contains('IV')
b3 = bsc['SpType'].str.contains('VI')

# Luminosity class "V", not ("IV" or "VI" or "VII"9
b = np.logical_and(b1,np.logical_and(np.logical_not(b2),np.logical_not(b3)))
main = bsc[b]
main.shape


main[0:10]

main['SpType'].map(lambda s:s[0]).unique()

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

fig, ax = plt.subplots(figsize=(10, 10))
ax.set_axis_bgcolor('0.6')

ax.grid()
ax.set_title('Color-color plot: main sequence BSC catalog')

ax.title.set_fontsize(20)
ax.set_xlabel('B-V')
ax.xaxis.label.set_fontsize(20)
ax.set_ylabel('U-B')
ax.yaxis.label.set_fontsize(20)

for cls in 'OBAFGKM':
    f = lambda s: s[0] == cls
    b = main['SpType'].astype('string').map(f)
    x = main[b]['B-V']
    y = main[b]['U-B']
    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')

b =  bsc['SpType'].str.contains('VI')
bsc[b].shape

b =  bsc['SpType'].str.contains('I')
giants = bsc[b]
giants.shape

giants['SpType'].map(lambda s:s[0]).unique()

colors = {'O':'#0000FF', 'B':'#CCCCFF', 'A':'#FFFFFF', 'F':'#FFFFB2', 'G':'#FFFF00', 'K':'#FFA500', 'M':'#FF0000',
          'W':'#000099' ,'S':'#B80000', 'C':'#780000'}

fig, ax = plt.subplots()
ax.set_axis_bgcolor('0.6')
ax.grid()
ax.set_title('Color-color plot: exotic classes BSC catalog')

ax.title.set_fontsize(15)
ax.set_xlabel('B-V')
ax.xaxis.label.set_fontsize(12)
ax.set_ylabel('U-B')
ax.yaxis.label.set_fontsize(12)

for cls in 'WSC':
    f = lambda s: s[0] == cls
    b = giants['SpType'].astype('string').map(f)
    x = giants[b]['B-V']
    y = giants[b]['U-B']
    c = colors[cls]
    ax.scatter(x, y, c = c, s=30,
               edgecolors='None', label = cls)
  
legend = ax.legend(scatterpoints=1,markerscale = 2, shadow=True)
frame = legend.get_frame()
frame.set_facecolor('0.90')

def B(wl,T):
    '''wl is an array of wavelengths with units of length
    T is a temperature in Kelvin
    the result is an array of spectral radiances
    with units 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))

#  Defining the UBV photometric system constants
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

# Function to calculate 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)

# Function to calculate Cb-v
F = B(lambda_b, T) * delta_b/(B(lambda_v, T)* delta_v)
Cbv = -0.33 + 2.5 * np.log10(F)

# Functions to calculate blackbody color indices
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

fig, ax = plt.subplots(figsize=(10, 10))
ax.set_axis_bgcolor('0.6')
ax.grid()
ax.set_title('Color-color main sequence and giants BSC catalog')

ax.title.set_fontsize(20)
ax.set_xlabel('B-V')
ax.xaxis.label.set_fontsize(20)
ax.set_ylabel('U-B')
ax.yaxis.label.set_fontsize(20)


for cls in 'OBAFGKM':
    f = lambda s: s[0] == cls
    b = giants['SpType'].astype('string').map(f)
    x = giants[b]['B-V']
    y = giants[b]['U-B']
    c = colors[cls]
    ax.scatter(x, y, c = c, s=50, alpha = 0.5,
               edgecolors='None', label = cls)

x = main['B-V']
y = main['U-B']
ax.scatter(x, y, c = 'black', s=6, edgecolors='none',
           label='Main')    
    
legend = ax.legend(scatterpoints=1, markerscale = 2,shadow=True)
frame = legend.get_frame()
frame.set_facecolor('0.90')

T1 = 2500 * pq.kelvin
T2 = 25000 * pq.kelvin

BminusV_1 = get_BV(T1)
UminusB_1 = get_UB(T1)
BminusV_2 = get_BV(T2)
UminusB_2 = get_UB(T2)

# Just two points are sufficient to determine the line
ax.plot([BminusV_1, BminusV_2], [UminusB_1, UminusB_2], lw = 3, c='k') 
ax.text(0.9,0.4,'blackbody', fontsize=20, rotation = 35);

# dataframe with only class IV stars
bIV = giants['SpType'].str.contains('IV')
giantsIV = giants[bIV]
giantsIV.shape

# dataframe with only class III stars
bIII = giants['SpType'].str.contains('III')
giantsIII = giants[bIII]
giantsIII.shape

# dataframe with only class II stars
bII = giants['SpType'].str.contains('II')
b = np.logical_and(bII,np.logical_not(bIII))
giantsII = giants[b]
giantsII.shape

# dataframe with only class I stars
bI = giants['SpType'].str.contains('I')
b = np.logical_and(bI,np.logical_not(bII))
giantsI = giants[b]
giantsI.shape

fig = plt.figure(figsize=(10,15))
ax1 = fig.add_subplot(321)
x = main['B-V']
y = main['U-B']
ax1.scatter(x, y, c = 'black', s=0.1)
ax1.set_xlim(-0.5, 2.5)
ax1.set_ylim(-1.5, 3)
ax1.set_xlabel('B-V')
ax1.set_ylabel('U-B')
ax1.grid()
ax1.set_title('main sequence (MK V)');

ax2 = fig.add_subplot(322)
x = giantsIV['B-V']
y = giantsIV['U-B']
ax2.scatter(x, y, c = 'black', s=0.1)
ax2.set_xlim(-0.5, 2.5)
ax2.set_ylim(-1.5, 3)
ax2.set_xlabel('B-V')
ax2.set_ylabel('U-B')
ax2.grid()
ax2.set_title('Subgiants (MK IV)');

ax3 = fig.add_subplot(323)
x = giantsIII['B-V']
y = giantsIII['U-B']
ax3.scatter(x, y, c = 'black', s=0.1)
ax3.set_xlim(-0.5, 2.5)
ax3.set_ylim(-1.5, 3)
ax3.set_xlabel('B-V')
ax3.set_ylabel('U-B')
ax3.grid()
ax3.set_title('Giants (MK III)');

ax4 = fig.add_subplot(324)
x = giantsII['B-V']
y = giantsII['U-B']
ax4.scatter(x, y, c = 'black', s=0.1)
ax4.set_xlim(-0.5, 2.5)
ax4.set_ylim(-1.5, 3)
ax4.set_xlabel('B-V')
ax4.set_ylabel('U-B')
ax4.grid()
ax4.set_title('Bright giants (MK II)');

ax5 = fig.add_subplot(325)
x = giantsI['B-V']
y = giantsI['U-B']
ax5.scatter(x, y, c = 'black', s=0.1)
ax5.set_xlim(-0.5, 2.5)
ax5.set_ylim(-1.5, 3)
ax5.set_xlabel('B-V')
ax5.set_ylabel('U-B')
ax5.grid()
ax5.set_title('Supergiants (MK I)');