#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Oct 29 04:44:24 2018

@author: danial
"""

# import packages
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patheffects as PathEffects
from matplotlib import rc
rc('font',**{'family':'sans-serif','sans-serif':['Helvetica']})
## for Palatino and other serif fonts use:
#rc('font',**{'family':'serif','serif':['Palatino']})
rc('text', usetex=True)

file = open("list-stars.txt")
listoflines = file.readlines()

lines = []
i = 1

for i in range(6,len(listoflines)):
    line = listoflines[i]
    lines.append(line)
    i+=1

# define lists of attributes

names = []
spectralclasses = []
BVcolours = []
UBcolours = []
effTemps = []
absolmags = []
luminosities = []

i = 1

for i in range(len(lines)):
    line = lines[i]
    linesplit = line.split()

    name = linesplit[0]
    names.append(name)

    spectralclass = linesplit[1]
    spectralclasses.append(spectralclass)

    BVcolour = linesplit[2]
    BVcolours.append(BVcolour)

    UBcolour = linesplit[3]
    UBcolours.append(UBcolour)

    effTemp = linesplit[4]
    effTemps.append(effTemp)

    absolmag = linesplit[5]
    absolmags.append(absolmag)


    luminosity = linesplit[6]
    luminosities.append(luminosity)

# if you want to check, if all your lists are right, uncomment this passage
'''
print("names \n", names, "\n")
print("spectralclasses \n", spectralclasses,"\n")
print("BVcolours \n", BVcolours, "\n")
print("UBcolours \n", UBcolours, "\n")
print("effTemps \n", effTemps, "\n")
print("absolmags \n", absolmags,"\n")
print("luminosities \n", luminosities,"\n")
'''
X = []
Y = []
i = 0

for i in range(23):

    x = float(effTemps[i])
    X.append(x)

    y = np.log10(float(luminosities[i]))
    #y = float(luminosities[i])
    Y.append(y)

    i+=1

print(X,"\n")
print(Y)


fig, ax = plt.subplots(figsize=(5,5), dpi = 300, facecolor = 'black')

plt.rc('text', usetex=True)
plt.rc('font', family='serif')


ax.spines['bottom'].set_color('white')
ax.spines['left'].set_color('white')

ax.title.set_color('white')
ax.yaxis.label.set_color('white')
ax.xaxis.label.set_color('white')
ax.tick_params(axis='x', colors='white')
ax.tick_params(axis='y', colors='white')

# names of axes
ax.set_title(r'Hertzsprung-Russel Diagramm')
ax.set_xlabel(r'$ \displaystyle T_{eff} [\mathrm{K}]$', fontsize = 10)
ax.set_ylabel(r'$ \displaystyle \log(L/L_{\odot}[\mathrm{mag}])$', fontsize = 10)
ax.invert_xaxis()

# make some annotations
i = 0

spectralcolor = []

for i in range(len(names)):

    if spectralclasses[i] == 'O':
        c = 'darkblue'
        spectralcolor.append(c)

    elif spectralclasses[i] == 'B':
        c = 'cyan'
        spectralcolor.append(c)

    elif spectralclasses[i] == 'A':
        c = 'ivory'
        spectralcolor.append(c)

    elif spectralclasses[i] == 'F':
        c = 'beige'
        spectralcolor.append(c)

    elif spectralclasses[i] == 'G':
        c = 'wheat'
        spectralcolor.append(c)

    elif spectralclasses[i] == 'K':
        c = 'orange'
        spectralcolor.append(c)

    elif spectralclasses[i] == 'M':
        c = 'red'
        spectralcolor.append(c)

    txt = ax.annotate( r'$ \displaystyle '+names[i]+'$', (X[i],Y[i]), color ='white', fontsize=7)
    txt.set_path_effects([PathEffects.withStroke(linewidth=1, foreground='black')])
    i+=1

ax.set_facecolor('black')
ax.scatter(X,Y,color = spectralcolor)

# close and save file
plt.savefig("HRD-plot.png", bbox_inches=0,  pad_inches = 0,  dpi = 300)