from rnd import *
import math
import os, sys

def color(str, col, bright=0, reset=1):
	if os.name != "posix":
		"color not supported, am too lazy"
		return str

	ansi = {'R': '31', 'X': '30', 'G': '32', 'Y': '33', 'B': '34', 'M': '35', 'C': '36', 'W': '37'}
	b = '1' if bright else '0'
	a = '\x1b[%s;%sm' % (b, ansi[col.upper()])
	r = ''
	if reset:
		r = '\x1b[0;37m'
		
	return '%s%s%s' % (a, str, r)

class Star(object):
	def __init__(self, lumin, mass, spect, radius, x, y, z):
		self.lumin, self.mass, self.spect, self.radius, self.x, self.y, self.z = \
			lumin, mass, spect, radius, x, y, z
	def dist(self):
		return dist(self.x, self.y, self.z) / 149589000.0 # return value in Astronomical Units 
	def dists(self):
		return [x / 149589000.0  for x in (self.x, self.y, self.z)] # return value in Astronomical Units 
		
	def info(self):
		return "%s-class star at %.3f (%.3f, %.3f, %.3f), mass %.3f, luminosity %.3f, radius %.3f" % \
			tuple([self.spect, self.dist()] + self.dists() + [self.mass, self.lumin, self.radius])
		
class Planet(object):
	def __init__(self, x, y, z, mass, radius, surfarea,
		areainsol, gas_giant, density, atmo,
		albedo, min_temp, max_temp, avg_temp):
		
		self.x, self.y, self.z, self.mass, self.radius, self.surfarea = x, y, z, mass, radius, surfarea
		self.areainsol, self.gas_giant, self.density, self.atmo = areainsol, gas_giant, density, atmo
		self.albedo, self.min_temp, self.max_temp, self.avg_temp = albedo, min_temp, max_temp, avg_temp

	def dist(self):
		return dist(self.x, self.y, self.z) / 149589000.0 # return value in Astronomical Units 
	def dists(self):
		return [x / 149589000.0  for x in (self.x, self.y, self.z)] # return value in Astronomical Units 
	def info(self):
		return "%s%s planet at %.3f (%.3f, %.3f, %.3f) mass %.3e kg, radius %.1f km, density %.3f, atmo %.3e, avg temp %.1f range (%.1f : %.1f), albedo %s" % \
			tuple([self.habitability(), self.type(), self.dist()] + self.dists() + [self.mass, self.radius, self.density, self.atmo,
			self.avg_temp, self.min_temp, self.max_temp, self.albedo])
	def type(self):
		if self.gas_giant:
			return "gas-giant"
		elif self.avg_temp < 60.0:
			return "ice"
		elif self.max_temp >= 700.0:
			return "lava"
		else:
			return "rocky"
	def habitability(self):
		if not self.habitable():
			return color("Uninhabitable ", 'x', 1)
		if self.max_temp < 273.0:
			return color("FREEZING ", 'c', 1)
		elif self.max_temp < 290.0 and self.avg_temp > 273.0:
			return color("COOL ", 'g', 0)
		elif self.min_temp > 290.0 and self.avg_temp > 325.0:
			return color("SCORCHING ", 'y', 1)
		elif self.min_temp > 290.0:
			return color("WARM ", 'g', 0)
		elif self.avg_temp < 273.0:
			return color("COLD ", 'b', 1)
		elif self.avg_temp > 300.0:
			return color("HOT ", 'r', 1)
		elif (self.max_temp - self.min_temp) > 150.0:
			return color("EXTREME ", 'm', 1)
		elif (self.max_temp - self.min_temp) > 60.0:
			return color("EARTHLIKE ", 'g', 0)
		elif (self.max_temp - self.min_temp) > 30.0:
			return color("COMFORTABLE ", 'g', 1)
		else:
			return color("GAIA ", 'g', 1)
	def habitable(self):
		if not self.gas_giant:
			if self.atmo > 1e-6 and (245.0 < self.min_temp < 325.0 or 245.0 < self.max_temp < 325.0):
				return True
		return False
				
	
class Asteroid_Belt(Planet):
	def __init__(self, radius, mass):
		self.radius, self.mass = radius, mass
		self.num_smashed = 2
	
	def dist(self):
		return self.radius
	
	def info(self):
		return "-- Asteroid belt at %.3f, mass %.3e --" % (self.radius, self.mass)
	
	def habitable(self):
		return False		

systemsize = 1e8
stars_randcount = (1.0, 0.5)
planets_randcount = (8.0, 4.0)

star_spectra = ['O', 'B', 'A', 'F', 'G', 'K', 'M']
star_lumin = [(500000.0, 30000.0), (30000.0, 25.0), (25.0, 5.0), (5.0, 1.5), (1.5, 0.6), (0.6, 0.008), (0.08, 0.0008)]
star_masses = [(40.0, 16.0), (16.0, 2.1), (2.1, 1.4), (1.4, 1.04), (1.04, 0.8), (0.8, 0.45), (0.45, 0.1)]
star_radii = [(22.0, 8.0), (7.0, 1.8), (1.8, 1.4), (1.4, 1.15), (1.15, 0.96), (0.96, 0.7), (0.7, 0.25)]
star_prob = [0.0, 0.05, 0.35, 0.6, 0.8, 0.97, 1.0]
gen_counter = 0
starsys_counter = 0

def dist(x, y, z):
	return math.sqrt(x**2 + y**2 + z**2)

def gen_system(force_habitable=False):
	global gen_counter
	single = True
	system_habitable = False
	while single or (force_habitable and not system_habitable):
		gen_counter += 1
		single = False
		num_stars = round(gaussian(*stars_randcount))
		if num_stars < 1.0:
			num_stars = 1.0
		num_stars = int(num_stars)
		
		num_planets = round(gaussian(*planets_randcount))
		if num_planets < 0.0:
			num_planets = 0.0
		num_planets = int(num_planets)
		
		stars = []
		x = y = z = 0.0
		
		while len(stars) < num_stars:
			r = random()
			for i, e in enumerate(star_prob):
				if r <= e:
					si = i
					break
			r = random()
			slumin = (star_lumin[si][1] * 0.8) + (star_masses[si][0] * 1.2 * r)
			r2 = r + (random() * 0.3) - 0.15
			smass = (star_masses[si][1] * 0.8) + (star_masses[si][0] * 1.2 * r)
			r2 = r + (random() * 0.3) - 0.15
			sradius = (star_radii[si][1] * 0.8) + (star_radii[si][0] * 1.2 * r)
		
			stars += [Star(slumin, smass, star_spectra[si], sradius, x, y, z)]
		
			# first star is always at 0.0, 0.0, 0.0, next star gets random coords
			x = gaussian(0.0, systemsize)
			y = gaussian(0.0, systemsize)
			if len(stars) > 1:
				z = gaussian(0.0, systemsize / 10e3)
		
		
		
		planets = []
		sys_lumin = sum([x.lumin for x in stars])
		sys_size = math.sqrt(sys_lumin / 2.0)
		
		while len(planets) < num_planets:
		
			if random() >= 0.5:
				loc = 2e7
				loczone = 20.0
			else:
				loc = 2e7
				loczone = 1.0
			
			x = gaussian(loc * randsign(), systemsize * sys_size * loczone)
			y = gaussian(loc * randsign(), systemsize * sys_size * loczone)
			z = gaussian(0, systemsize * sys_size * loczone/ 10e0)
		
			reject = False
			for star in stars:
				d = dist(x-star.x,y-star.y,z-star.z)
		
				deathzone = star.radius + (star.radius * 0.25 * math.sqrt(star.lumin))
				if d < deathzone:
					reject = True
		
			if reject:
				continue
		
			r = random()
			rocky = False
			gas_giant = False
			if r < 0.6:
				# rock planet
				rocky = True
				massmin = 1.0e19
				massmax = 8.8e25
				mass = 0.0
				while mass < massmin or mass > massmax:
					massexp = gaussian(24.5, 1.25)
					mass = 10.0 ** massexp
				density = 0.0
				while density < 0.75:
					density = gaussian(4.5, 1.0)
				mag = -1.0
				while mag < 0.0:
					mag = gaussian(density, 3.0)
		
			else:
				# gas giant
				gas_giant = True
				massmin = 1e25
				massmax = 4e28
				mass = 0.0
				while mass < massmin or mass > massmax:
					massexp = gaussian(26.5, 1.0)
					mass = 10.0 ** massexp
				density = 0.0
				while density < 0.3:
					density = gaussian(1.5, 0.75)
				mag = -1.0
				while mag < 2.0:
					mag = gaussian(7.0, 3.0)
		
		
			volume = mass / density / 1e12
			radius = pow(volume * 0.75 / math.pi, 1.0/3.0)
			circsurfarea = math.pi * radius**2
			surfarea = circsurfarea * 4.0
			insol = 0.0
			for star in stars:
				d = dist(x-star.x,y-star.y,z-star.z)
		
				spower = 3.8e26 * star.lumin
				distsphere = 4.0 * math.pi * d**2
				insol += (circsurfarea / distsphere) * spower
		
			areainsol = insol / surfarea / 1e6
			if areainsol > 75000:
				continue
		
			atmo = None
			atmo_count = 0
			while atmo == None or atmo > density:
				atmo_count += 1
				if gas_giant:
					atmo_exp = randrange(-5.0, 12.0)
				else:
					atmo_exp = randrange(-9.0, 2.0)
				
				atmo = 1.3 ** atmo_exp
				
				if atmo_count > 50 and atmo > density:
					print "ATMO FAIL for gasgiant:%s, wanted less than %f, kept getting values like %f due to exponent %f" % (gas_giant, density, atmo, atmo_exp)
					atmo = 1.3 ** (atmo_exp - 10.0)
					break
				
			albedo = -1.0
			while albedo <= 0.01 or albedo >= 0.99:
				albedo = gaussian(0.35, 0.25)
		
			emissivity = 1.0 - albedo
			#emissivity = randrange(0.00001, 1.0)
		
			emiss_diff = (((10.0 - pow(atmo, 1.0/2.0)) / 10.0) ** 100.0) / 2.0
			
			emissivity_min = emissivity - (emissivity * emiss_diff)
			emissivity_max = emissivity + (emissivity * emiss_diff)
		
			avg_temp = pow(((1.0 - albedo) * areainsol) / (4.0 * emissivity * 5.67e-8), 1.0/4.0)
			min_temp = pow(((1.0 - albedo) * areainsol) / (4.0 * emissivity_max * 5.67e-8), 1.0/4.0)
			max_temp = pow(((1.0 - albedo) * areainsol) / (4.0 * emissivity_min * 5.67e-8), 1.0/4.0)
			
			planets.append(Planet(x, y, z, mass, radius, surfarea,
				areainsol, gas_giant, density, atmo, 
				albedo, min_temp, max_temp, avg_temp))
		
		
		planets.sort(cmp=lambda a, b: cmp(dist(a.x, a.y, a.z), dist(b.x, b.y, b.z)))
		
		dels = []
		inserts = []
		lpl = None
		for i in range(0, len(planets)-1):
			pl = planets[i]
			npl = planets[i+1]
			
			if (npl.dist() - pl.dist()) < math.sqrt((npl.mass + pl.mass) / 1e22) * 0.001:
				"BOOM. asteroids"
				if i in dels:
					# ooh, three or more planets smashing together! exciting!
					# this one has already been deleted by the previous one, so we only need to add the next delete
					# also update the old asteroid belt to include this planet
					dels += [i+1]
					belt = inserts[-1][1]
					belt.mass += npl.mass
					belt.radius = ((belt.radius * float(belt.num_smashed)) + npl.dist()) /  (belt.num_smashed + 1)
					belt.num_smashed += 1
				else:
					inserts.append((i - len(dels), Asteroid_Belt((npl.dist() + pl.dist()) / 2.0, npl.mass + pl.mass)))
					dels += [i, i+1]
		
		dels.reverse()
		for dd in dels:
			del planets[dd]
		
		inserts.reverse()
		for ii in inserts:
			planets.insert(ii[0], ii[1])
	
		system_habitable = any([x.habitable() for x in planets])
		
	return (stars, planets)

def print_system(syst, show_hab=True):
	stars, planets = syst
	print "\nStars:\n"
	
	for star in stars:
		print star.info()
	
	print "\nPlanets:\n"
		
	hab = False
	for i, pl in enumerate(planets):
		if pl.habitable():
			hab = True
		print "#%i: %s" % (i+1, pl.info())
	
	if show_hab:
		print ""
		if hab:
			print "System is habitable!"
		else:
			print "System UNINHABITABLE!"
	
def main():
	if len(sys.argv) > 1:
		runs = int(sys.argv[1])
	else:
		runs = 1
	
	fl = False
	systems = []
	dead_systems = []
	for i in range(runs):
		syst = gen_system(fl)
		fl = True
		if any([x.habitable() for x in syst[1]]):
			systems.append(syst)
		else:
			dead_systems.append(syst)
	
	print color("\n    ***************** INHABITABLE SYSTEMS (%d) *****************\n", 'm', 1) % (len(systems),)
	for syst in systems:
		print_system(syst, False)
	
	ansi_cancel = ""
	if os.name == "posix":
		ansi_cancel = "\x1b[0;37m"
	print color("\n    ***************** UNINHABITABLE SYSTEMS (%d) *****************%s\n", 'm', 1) % (len(dead_systems), ansi_cancel)
	for syst in dead_systems:
		print_system(syst, False)
		
		
if __name__ == '__main__':
	main()