import numpy as np
import astropy.constants as const


def calculate_mu(X,Y,Z=None):
	"""
		Calculates mean molecular weight
		Assumes:
			Fully ionized gas
			Composition of Hydrogen, Helium, Metals

		Inputs:
			X 	- Hydrogen mass fraction
			Y 	- Helium mass fraction
			[Z]	- Metals mass fraction
			  	- if Z omitted, assumes X+Y+Z=1
		  	Inputs must sum to 1

		Outputs:
			mu - mean molecular weight

		Warnings:
	"""

	if Z is None:
		Z = 1 - X - Y
	elif (X + Y + Z) != 1:
		print "Incorrect values for X,Y,Z. Must sum to 1"
		raise Exception('calculate_mu incorrect inputs')

	mu = np.power((X / (1./2)) + (Y / (4./3)) + (Z / 2.), -1.)

	return mu

def calculate_mu_e(X,Y,Z=None):
	"""
		Calculates mean molecular weight of electrons
		Assumes:
			Fully ionized gas
			Composition of Hydrogen, Helium, Metals

		Inputs:
			X 	- Hydrogen mass fraction
			Y 	- Helium mass fraction
			[Z]	- Metals mass fraction
			  	- if Z omitted, assumes X+Y+Z=1
		  	Inputs must sum to 1

		Outputs:
			mu_e - mean molecular weight of electrons

		Warnings:
	"""

	if Z is None:
		Z = 1 - X - Y
	elif (X + Y + Z) != 1:
		print "Incorrect values for X,Y,Z. Must sum to 1"
		raise Exception('calculate_mu incorrect inputs')

	mu_e = np.divide(2., (1+X)) #Eq. 4.30, Kippenhan Weigert Weiss 2nd ed.

	return mu_e

def calculate_rad_pressure(T):
	"""
		Calculates radiative pressure of photon gas at temperature T
		Assumes:
			

		Inputs:
			T - temperature [K]

		Outputs:
			P_rad [dyne/cm^2]

		Warnings:
	"""

	if T <= 0:
		print "Temperature must be greater than 0"
		raise Exception('calculate_rad_pressure incorrect inputs')