import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

import utilities

# Open the file "/hipparcos-voidmain.csv" for reading
f = open("./hipparcos-voidmain.csv", "r")

dataset = pd.read_csv(f, sep=",", header=0)

# Isolate the data we want to plot
inputData = np.array(dataset[["B-V", "e_B-V", "Vmag", "Plx", "e_Plx"]], dtype=float)

# Filter the data
# Get rid of the NaN values
inputData = inputData[~np.isnan(inputData).any(axis=1)]
# Apply the following conditions:
#   - e_B-V should be smaller than 0.1
#   - Plx should be larger than 0.1
#   - e_Plx should be smaller than 5
inputData = inputData[(inputData[:, 1] < 0.1) & (inputData[:, 3] > 0.1) & (inputData[:, 4] < 5)]

# Compute the distance
distance = 1000 / inputData[:, 3] # unit: parsec
# Compute the absolute magnitude
absoluteMag = inputData[:, 2] - 5 * np.log10(distance / 10)
bmv = inputData[:, 0]

# Map B-V to RGB
rgb = np.array([utilities.bv2rgb(item) for item in bmv])

# Plot the Hertzsprung-Russell diagram
plt.figure(figsize=(10, 10))
plt.scatter(bmv, absoluteMag, s=0.05, c=rgb, alpha=0.75)
plt.xlim(-0.5, 2.5)
plt.ylim(-15, 15)
plt.xlabel("B-V")
plt.ylabel("Vmag")
plt.title("Hertzsprung-Russell diagram")
plt.gca().set_facecolor((0, 0, 0))
plt.gca().invert_yaxis()
plt.show()

# Close the file
f.close()