#!/usr/bin/python3

import sys, os
from functools import reduce

import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt

from catalog import Hipparcos

# for additional data not included in catalog
distance_override = { # in pc
    28716: 1800 # chi2 Orionis
    }
mag_override = { # apparent mag
    114971: 3.699 # gamma Piscium
    }

# Reference magnitude: Brighter stars are displayed bigger
Vref = -0.89 # abs mag of Duhe

# Model length scale
lenscale = 0.2 # cm / pc
# Model angular scale
anglescale = 2

def help_exit(msg):
    print("Usage: %s [-c CATFILE] [-o OUTNAME] STARFILE"%sys.argv[0])
    print("CATFILE: catalog file")
    print("OUTNAME: name for output files")
    print("STARFILE: csv with stars in constellation")
    sys.exit(msg)

def main():
    args = sys.argv[1:]
    catalogfile = "hip_main.dat"
    starfile = None
    outputname = None
    while args:
        arg = args.pop(0)
        if arg == '-c':
            cataloguefile = args.pop(0)
        elif arg == '-o':
            outputname = args.pop(0)
        elif starfile is None:
            starfile = arg
        else:
            help_exit("Unrecognized argument %s"%arg)
    if starfile is None:
        help_exit("Missing argument.")
    if outputname is None:
        outputname = starfile

    # read list of stars
    starsHIP = []
    starsLink = [] # [(hip1, hip2)]
    starsName = []
    starsLore = []
    starsInfo = []
    with open(starfile, 'r') as f:
        for line in f:
            if line.startswith('#') or not line.strip():
                continue
            fields = line.split(',')
            hip = int(fields[0])
            starsHIP.append(hip)
            starsLink.extend([(hip, int(s.strip()))
                                  for s in fields[1].split(' ')
                                  if s.strip()])
            starsName.append(fields[2].strip())
            starsLore.append(fields[3].strip())
            starsInfo.append(fields[4].strip())
    print("%d stars in constellation"%len(starsHIP))

    # find stars in catalog
    cat = Hipparcos(catalogfile)
    columns = ['HIP', 'RAdeg', 'e_RAdeg', 'DEdeg', 'e_DEdeg',
               'Plx', 'e_Plx', 'VTmag', 'e_VTmag', 'B-V', 'e_B-V']
    records = []
    for rec in cat.records(columns):
        if rec[0] in starsHIP:
            records.append(rec)
    if len(records) < len(starsHIP):
        print("Not all stars in catalog.")
    data = np.array(records)
    hip = data[:,0]
    ra, raerr = data[:,1], data[:,2]/3600e3 # deg
    de, deerr = data[:,3], data[:,4]/3600e3 # deg
    plx, plxerr = data[:,5], data[:,6] # mas
    V, Verr = data[:,7], data[:,8] # mag
    BV, BVerr = data[:,9], data[:,10] # mag

    # calculate distance from parallax
    d = 1000 / plx
    derr = 1000 / plx**2 * plxerr
    # clean plx and d values
    plx[plx==0.1] = np.nan
    d[plx==0.1] = np.nan

    # put in additional data
    for h, dist in distance_override.items():
        idx = np.where(hip==h)
        d[idx] = dist
        derr[idx] = np.nan
    for h, m in mag_override.items():
        idx = np.where(hip==h)
        V[idx] = m
        Verr[idx] = np.nan

    # calculate absolute magnitudes
    Vabs = V - 5*np.log10(d) + 5
    Vabserr = np.sqrt(Verr**2 + (derr*5/(d*np.log(10)))**2)

    # wrap over stars around 360 degrees
    ra = (ra+180)%360 - 180

    # print table of stars to CSV
    with open(outputname+'.stars.csv', 'w+') as f:
        f.write("Name,RA/deg,DE/deg,Plx/mas,Plx_err/mas,d/pc,d_err/pc,"
                "d scaled/cm,V/mag,V_err/mag,Vabs/mag,Vabs_err/mag,"
                "B-V,(B-V)_err,Lore,Info\n")
        for i, h in enumerate(starsHIP):
            idx = np.where(hip==h)
            cols = [
                ra[idx], de[idx],
                plx[idx], plxerr[idx], d[idx], derr[idx], lenscale*d[idx],
                V[idx], Verr[idx], Vabs[idx], Vabserr[idx], BV[idx], BVerr[idx],
                ]
            f.write(starsName[i]+','
                    +','.join("%f"%np.asscalar(s) for s in cols)
                    +','+starsLore[i]+','+starsInfo[i]+'\n')

    # dot sizes for plots
    area = np.maximum(1, 2*np.maximum(0, 6-V)**2)
    areaabs = np.maximum(1, 2*np.maximum(0, 6-Vabs)**2)
    # 2D equirectangular plot
    fig = plt.figure()
    plt.scatter(ra, de, s=area, color="k", alpha=1)
    for hip1, hip2 in starsLink:
        idxs = np.where(np.logical_or(hip==hip1, hip==hip2))
        plt.plot(ra[idxs], de[idxs], color='b', linewidth=0.5)
    plt.gca().invert_xaxis()
    plt.axis('equal')
    plt.xlabel("RA/deg")
    plt.ylabel("DE/deg")
    plt.title(outputname)
    plt.savefig(outputname+'.equirect.png', dpi=130)
    plt.show()

    # convert to degrees to radians
    ra = ra / 360 * 2*np.pi
    de = de / 360 * 2*np.pi
    # move to cartesian coordinate system
    xx = d * np.cos(ra) * np.cos(de)
    yy = d * np.sin(ra) * np.cos(de)
    zz = d * np.sin(de)

    # simple 3D plot
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(xx, yy, zz, s=areaabs, color="k", marker='.', alpha=1)
    ax.scatter([0], [0], [0], s=10, color="b")
    for i in range(len(hip)): # vertical lines
        ax.plot([xx[i], xx[i]], [yy[i], yy[i]], [0, zz[i]],
                color='k', linewidth=0.3)
    for hip1, hip2 in starsLink: # links
        idxs = np.where(np.logical_or(hip==hip1, hip==hip2))
        ax.plot(xx[idxs], yy[idxs], zz[idxs], color='b', linewidth=0.5)
    # labels
    ax.set_xlabel(r"x / pc")
    ax.set_ylabel(r"y / pc")
    ax.set_zlabel(r"z / pc")
    ax.set_title(outputname+' (not scaled)')
    #plt.savefig(outputname+'.3d.png', dpi=120)
    plt.show()

    # scale to model
    ra2 = ra * anglescale
    de2 = de * anglescale
    xx2 = lenscale * d * np.cos(ra2) * np.cos(de2)
    yy2 = lenscale * d * np.sin(ra2) * np.cos(de2)
    zz2 = lenscale * d * np.sin(de2)

    # print table of link distances
    def dist(idx1, idx2):
        xdist = xx[idx2] - xx[idx1]
        ydist = yy[idx2] - yy[idx1]
        zdist = zz[idx2] - zz[idx1]
        return np.sqrt(xdist**2+ydist**2+zdist**2)
    def dist_scaled(idx1, idx2):
        xdist = xx2[idx2] - xx2[idx1]
        ydist = yy2[idx2] - yy2[idx1]
        zdist = zz2[idx2] - zz2[idx1]
        return np.sqrt(xdist**2+ydist**2+zdist**2)
    with open(outputname+'.links.csv', 'w+') as f:
        f.write("Name 1, Name 2, a / pc, a scaled / cm\n")
        for hip1, hip2 in starsLink:
            idx1 = np.where(hip==hip1)
            idx2 = np.where(hip==hip2)
            f.write(
                starsName[starsHIP.index(hip1)]
                +','+starsName[starsHIP.index(hip2)]
                +',%f'%np.asscalar(dist(idx1,idx2))
                +',%f'%np.asscalar(dist_scaled(idx1, idx2))
                +'\n')

    # find equal scale for all axes, not including earth
    xdiff = np.max(xx2) - np.min(xx2)
    ydiff = np.max(yy2) - np.min(yy2)
    zdiff = np.max(zz2) - np.min(zz2)
    diffmax = max(xdiff, ydiff, zdiff)
    xcenter = (np.min(xx2) + np.max(xx2)) / 2
    ycenter = (np.min(yy2) + np.max(yy2)) / 2
    zcenter = (np.min(zz2) + np.max(zz2)) / 2
    # scaled 3d plot
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter([0], [0], [0], s=10, color="#0000AA") # earth
    ax.scatter(xx2, yy2, zz2, s=areaabs,
               color='k', marker='.', alpha=1)
    for hip1, hip2 in starsLink:
        idxs = np.where(np.logical_or(hip==hip1, hip==hip2))
        ax.plot(xx2[idxs], yy2[idxs], zz2[idxs],
                color='b', linewidth=0.5)
    # limits
    ax.set_xlim([xcenter-diffmax/2, xcenter+diffmax/2])
    ax.set_ylim([ycenter-diffmax/2, ycenter+diffmax/2])
    ax.set_zlim([zcenter-diffmax/2, zcenter+diffmax/2])
    # labels
    ax.set_xlabel("x / cm")
    ax.set_ylabel("y / cm")
    ax.set_zlabel("z / cm")
    ax.set_title(outputname+' (scaled)')
    #plt.savefig(outputname+'.3dscaled.png', dpi=120)
    plt.show()

if __name__ == '__main__':
    main()