"""
  File:     HRDiagram.py
  Author:   Thomas Mueller, HdA
  
  Description:
  ------------
    Load Hipparcos catalog (binary fits file) and draw a Hertzsprung-Russell
    diagram of absolute versus B-V magnitude.
     
    The catalog data can be downloaded from VizieR 
         <http://vizier.u-strasbg.fr/viz-bin/VizieR>
         
    * Search for "hip" yields "I/234", "I/311", ...
    * Select "I/311/hip2"
    * Set Preferences to 
        - max = unlimited
        - FITS (binary) Table
        - Position in Decimal
    * Check that the following columns are selected:
        "HIP", "Plx", "Hpmag", "B-V"        
    * Press "Submit" to download the data.
    * Rename the file from "asu.fit" to e.g. "I_311_Hipparcos.fit".
"""
     
from astropy.io import fits
import numpy as np
import matplotlib.pyplot as plt

# Gaia DR2 file
filename = 'Gaia DR2-Sterne bis 6m.fit'

# Read fits file
hdulist = fits.open(filename)
hipdata = hdulist[1].data
hdulist.close()

# extract data from Gaia DR2 catalogue
ra = hipdata.field('RA_ICRS')
de = hipdata.field('DE_ICRS')
plx = hipdata.field('Plx')
BmV = hipdata.field('BP-RP')
mag = hipdata.field('Gmag')

# select only those entries, where the parallax is positive
idx = np.where(plx > 0)

# calculate absolute magnitude
AMag = mag[idx] + 5 * (1 + np.log10(plx[idx] * 0.001))

# draw 
fig = plt.figure(figsize = (9, 12))
ax = fig.add_subplot(111)
ax.set_title("H-R diagram of Gaia catalog DR2")
ax.title.set_fontsize(18)

# stars are plotted using scatter-plot
ax.scatter(BmV[idx], AMag, s = 2.2, marker = 'o')

# set x-axis
ax.set_xlim([-1, 4.0])
ax.set_xlabel('G$_{BP}$-G$_{RP}$ [mag]')
ax.xaxis.label.set_fontsize(16)

# set y-axis
ax.set_ylim([8, -8])
ax.set_ylabel('Absolute magnitude M$_G$ [mag]')
ax.yaxis.label.set_fontsize(16)

# gaphic output
ax.grid()
plt.savefig("HRdiagram.png", dpi=100, bbox_inches='tight', pad_inches = 0)
plt.show()