"""
Plot Angular Momentum Budget.

@todo - Figure out and apply correct units to angular momentum.

Volker Hoffmann <volker@cheleb.net>
04 June 2013
"""

import matplotlib.pyplot as plt
from DiskXYZ import Disk
import numpy as np
import argparse
import sys
from math import pi

# Parse Arguments
parser = argparse.ArgumentParser()
parser.add_argument("imin", type=int, help='First Output')
parser.add_argument("imax", type=int, help='Last Output')
parser.add_argument("--show", action="store_true", \
                    help='Show Figure')
parser.add_argument("--save", action="store_true", \
                    help='Save Figure')
parser.add_argument("--disk", action="store_true", \
                    help="Show Total Disk Angular Momentum")
parser.add_argument("--halo", action="store_true", \
                    help="Show Total Halo Angular Momentum")
parser.add_argument("--box", action="store_true", \
                    help="Show Total Box Angular Momentum")
parser.add_argument("--lost_halo", action="store_true", \
                    help="Show (Accumulated) Angular Momentum Lost to Halo")
parser.add_argument("--lost_star", action="store_true", \
                    help="Show (Accumulated) Angular Momentum Lost to Star")
args = parser.parse_args()

# Sanity Checks
if args.imin > args.imax:
    print("Cannot Work With Imin > Imax. Use -h for help.")
    sys.exit(-1)
if not (args.show or args.save):
    print("Specify --save or --show. Use -h for help.")
    sys.exit(-1)
if not (args.disk or args.halo or args.box or args.lost_halo or args.lost_star):
    print("Let me show you something, please? Use -h for help.")
    sys.exit(-1)

# Build Output Range
iouts = range(args.imin, args.imax+1)

# Allocate Arrays
time = np.zeros_like(iouts, dtype='float64')
angular_momentum_total_lost_to_star = np.zeros_like(iouts, dtype='float64')
angular_momentum_total_lost_to_halo = np.zeros_like(iouts, dtype='float64')
angular_momentum_disk = np.zeros_like(iouts, dtype='float64')
angular_momentum_halo = np.zeros_like(iouts, dtype='float64')
angular_momentum_total = np.zeros_like(iouts, dtype='float64')

# Load Data
ii = 0
for iout in iouts:
    print "Loading Data for Output %i." % iout
    disk = Disk(iout)
    disk.load_npz_stats()
    # print disk.info["time"]
    time[ii] = disk.info["time"]
    angular_momentum_total_lost_to_star[ii] = disk.angular_momentum_lost_to_star
    angular_momentum_total_lost_to_halo[ii] = disk.angular_momentum_lost_to_halo
    angular_momentum_disk[ii] = disk.angular_momentum_disk
    angular_momentum_halo[ii] = disk.angular_momentum_halo
    angular_momentum_total[ii] = disk.angular_momentum_total
    ii += 1

# ----------------------------
# Plot Angular Momentum Budget
# ----------------------------
fig1 = plt.figure()
ax1 = fig1.add_subplot(1,1,1)
ax2 = ax1.twiny()
if args.disk:
    ax1.plot(time/2./pi, angular_momentum_disk, 'bs-', label='Disk Angular Momentum')
if args.halo:
    ax1.plot(time/2./pi, angular_momentum_halo, 'rs-', label='Halo Angular Momentum')
if args.box:
    ax1.plot(time/2./pi, angular_momentum_total, 'gs-', label='Box Angular Momentum')
if args.lost_halo:
    ax1.plot(time/2./pi, angular_momentum_total_lost_to_halo, 'ms-', \
             label='Total Angular Momentum Lost to Halo')
if args.lost_star:
    ax1.plot(time/2./pi, angular_momentum_total_lost_to_star, 'ks-', \
             label='Total Angular Momentum Lost to Star')
ax2.plot(time, np.ones_like(time))
ax2.cla(); ax2.set_xlabel('Time [code]')
ax1.grid(); ax1.set_yscale('log')
ax2.set_ylim(ax1.get_ylim()); ax2.set_yscale(ax1.get_yscale())
ax1.set_xlabel('Time [yr]')
ax1.set_ylabel('Angular Momentum [?]')
ax1.legend(loc='best')
fig1.suptitle('Angular Momentum Budget')

if args.save:
    fig1.savefig('angular_momentum_budget.pdf')
if args.show:
    plt.show()