"""
Generates a simulated 'galaxy' of stars and gives
positions of those stars relative to some 'Earth'.

Intended to give students a chance to use a mapping
of star positions to determine how far 'Earth' is
from the center.
"""

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


def gen_stars(dx, dy, dz, num):
    """ Generates stars with a disperal from some local center. """
    stars = np.zeros(shape=(num, 3))

    dims = [dx, dy, dz]
    for dim in range(stars.shape[1]):
        v = np.random.normal(0, dims[dim], size=num)
        stars[:,dim] = v

    return stars

def plot_stars(stars):
    fig, axes = plt.subplots(nrows=2)
    axes[0].plot(stars[:,0], stars[:,1], '.')
    axes[0].set_xlabel('X')
    axes[0].set_ylabel('Y')
    axes[1].plot(stars[:,0], stars[:,2], '.')
    axes[1].set_xlabel('X')
    axes[1].set_ylabel('Z')
    plt.show()

def to_equitorial(stars, planet_loc):
    diff = stars - planet_loc
    R = np.linalg.norm(diff, axis=1)
    RA = np.arctan2(diff[:,1], diff[:,0])
    DEC = np.arctan2(diff[:,2], np.linalg.norm(diff[:,:2], axis=1))
    return np.column_stack((R, RA, DEC))

def to_cartesian(df):
    df['X'] = df['R'] * np.cos(df['DEC']) * np.cos(df['RA'])
    df['Y'] = df['R'] * np.cos(df['DEC']) * np.sin(df['RA'])
    df['Z'] = df['R'] * np.sin(df['DEC'])
    return df


if __name__ == '__main__':
    stars = gen_stars(3,2,0.5, 10000)
    plot_stars(stars)
    equi_pts = to_equitorial(stars, np.array([-2,1,0]))
    df = pd.DataFrame(equi_pts, columns=['R', 'RA', 'DEC'])
    df['p_mas'] = 1/(df['R'])
    df['RA_deg'] = (np.degrees(df['RA']) + 360) % 360
    df['DEC_deg'] = np.degrees(df['DEC'])
    df_exp = df.loc[:, ['p_mas', 'RA_deg', 'DEC_deg']]
    df_exp.to_csv('place_in_galaxy.csv')
    df = to_cartesian(df)
    print(df.describe())

    plt.subplot(projection='aitoff')
    plt.plot(df.RA, df.DEC, '.', alpha=0.25)
    plt.grid()
    plt.show()