#!/usr/bin/env python

__build_date__ = "%%fulldate%%"
__program__ = "Punnent Square"
__version__ = "0.0.2"
__filecall__ = "punnentsquare.py"
__author__ = "Alex Brown (\"uh oh spaghetti-o\") <pianoboy3333@juicemedia.tv>"
__tagline__ = "is a python script that does stuff concerning alleles, genotype, and phenotype"
__tagline2__ = "A python script that does stuff concerning alleles, genotype, and phenotype"
__usage__ = __filecall__ + " [OPTIONS]\n" + __program__ + ' ' + __tagline__
__help__ = __usage__ + """
Possible arguments for options:
	<allele>		- outputs information on the given allele
	<genotype>		- outputs information on the given genotype
	<genotype> <genotype>	- gives information on the possible genotypes
					if two species with the given genotypes
					were crossed
	"<genotype> x <genotype>" - same as above
	"<genotype> * <genotype>"
"" indicates a string
<allele> - an allele
<genotype> - a genotype
x - represents crossing them

Notes:
* When crossing genotypes, note the ' x ' in between the two genotypes,
	program will not work if the spaces are absent
* This program is designed very case-insensitive,
	therefore giving a blood genotype such as iAIa will do just fine
* When stating heterozygous genotypes, you must put the
	dominant or big letter first:
	ex: Aa, Bb, Xx, Gg...
	of course this is with the exception of
	heterozygous blood genotypes
	ex: iAiO, iaio, iBiO, ibio, IbiO...
"""
__doc__ = __program__ + ", by " + __author__ + ", Version " + __version__ + __tagline__
__copyright__ = __program__ + " - " + __tagline2__ + "\nCopyright (C) 2006 Alex Brown"
__license__ = "\t" + __program__ + ' - ' + __tagline2__ + """
	Copyright (C) 2006 Alex Brown

	This program is free software; you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation; either version 2 of the License, or
	(at your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License along
	with this program; if not, write to the Free Software Foundation, Inc.,
	51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
"""

import sys
import string
defglobals = False

def case_combos(s):
	if s == '':
		yield ''
		return
	for combo in case_combos(s[1:]):
		yield s[0].upper() + combo
	for combo in case_combos(s[1:]):
		yield s[0].lower() + combo
def list_combos(l):
	combolist = []
	for i in l:
		for combo in case_combos(i):
			combolist += [combo]
	return combolist
def merge_lists(list1, list2):
	i = 0
	if len(list1) > len(list2):
		print "merge_lists(): The length of list1 is greater than the length of list2"
	elif len(list1) < len(list2):
		print "merge_lists(): The length of list1 is less than the length of list2"

	while i < len(list1):
		list1[i] += list2[i]
		i += 1
	return list1

lower_blood_alleles = ['ia', 'ib', 'io']
upper_blood_alleles = ['IA', 'IB', 'IO']
all_blood_alleles = list_combos(['ia', 'ib',  'io'])
a_blood_alleles = list_combos(['ia'])
b_blood_alleles = list_combos(['ib'])
o_blood_alleles = list_combos(['io'])

rec_alleles = list(string.lowercase)
dom_alleles = list(string.uppercase)
all_alleles = rec_alleles + dom_alleles + all_blood_alleles

#Possible genotypes
blood_genotypes = ['iaia', 'iaio', 'ibib', 'ibio', 'iaib', 'ioio']
all_blood_genotypes = list_combos(blood_genotypes)
homorec_genotypes = merge_lists(list(string.lowercase), list(string.lowercase))	# merge the lists of lowercase letters
homodom_genotypes = merge_lists(list(string.uppercase), list(string.uppercase))	# merge the lists of uppercase letters
hetero_genotypes = merge_lists(list(string.uppercase), list(string.lowercase))	# merge the lists of lowercase and uppercase letters
all_genotypes = homorec_genotypes + homodom_genotypes + hetero_genotypes + all_blood_genotypes

def simple_bloodtype(allele):
	if allele in a_blood_alleles:
		return 'ia'
	elif allele in b_blood_alleles:
		return 'ib'
	elif allele in o_blood_alleles:
		return 'io'
	else:
		return
def determine_a_or_g(arg):
	if arg in all_alleles:
		allele_info(arg)
	elif arg in all_genotypes:
		genotype_info(arg)
	else:
		print arg, "is not an allele or genotype"

def determine_blood_type(genotype):
	if not genotype in all_genotypes:
		print genotype + " is not a genotype"
		return
	elif not genotype in all_blood_genotypes:
		print genotype + " is a proper genotype, but not a blood genotype"
		return

	if genotype.lower() == 'iaia' or genotype.lower() == 'iaio':
		return genotype[:-2]
	elif genotype.lower() == 'ibib' or genotype.lower() == 'ibio':
		return genotype[:-2]
	elif genotype.lower() == 'iaib':
		return genotype
	elif genotype.lower() == 'ioio':
		return genotype[:-2]
def determine_blood_phenotype(genotype):
	if not genotype in all_genotypes:
		print genotype + " is not a genotype"
		return
	elif not genotype in all_blood_genotypes:
		print genotype + " is a proper genotype, but not a blood genotype"
		return

	if genotype.lower() == 'iaia' or genotype.lower() == 'iaio':
		return 'A'
	elif genotype.lower() == 'ibib' or genotype.lower() == 'ibio':
		return 'B'
	elif genotype.lower() == 'iaib':
		return 'AB'
	elif genotype.lower() == 'ioio':
		return 'O'
def make_genotype(allele1, allele2):
	if not allele1 in all_alleles:
		print allele, "is not a allele"
		return
	if allele1 in all_blood_alleles and allele2 in all_blood_alleles:
		if allele1.lower() == 'ia':
			return allele1 + allele2
		elif allele2.lower() == 'ia':
			return allele2 + allele1
		elif allele1.lower() == 'ib':
			return allele1 + allele2
		elif allele2.lower() == 'ib':
			return allele2 + allele1
		else:
			return allele1 + allele2
	elif (allele1 in all_blood_alleles and not allele2 in all_blood_alleles) or (allele1 not in all_blood_alleles and allele2 in all_blood_alleles):
		print allele1 + " and " + allele2 + " are not both blood alleles"
		return
	# givens could be A, a, or A, A
	# sees that the first is uppercase (dominant)
	# and returns that letter first (A, <other>), because the
	# case of the second given doesn't matter, it will still be correct
	if allele1.isupper():
		return allele1 + allele2
	# opposite of above
	elif allele2.isupper():
		return allele2 + allele1
	# else it must be homo. rec. (aa)
	else:
		return allele1 + allele2
def genotype_info(genotype):
	if not genotype in all_genotypes:
		print genotype + " is not a genotype"
		return False
	else:
		print "Genotype:\t" + genotype
		print "---------------------"

	if genotype in homorec_genotypes:
		print "Type:\t\thomozygous recessive"
		print "Expr. Allele:\t" + genotype[:-1]
	elif genotype in homodom_genotypes:
		print "Type:\t\thomozygous dominant"
		print "Expr. Allele:\t" + genotype[:-1]
	elif genotype in hetero_genotypes:
		print "Type:\t\theterozygous"
		print "Expr. Allele:\t" + genotype[:-1]
	elif genotype in all_blood_genotypes:
		if genotype.lower() == 'ioio':
			print "Type:\t\thomozygous recessive,",
		elif genotype[-2:].lower() == 'io':
			print "Type:\t\theterozygous,",
		elif genotype.lower() == 'iaib':
			print "Type:\t\t",
		else:
			print "Type:\t\thomozygous dominant,",
		print "blood genotype"
		print "Expr. Allele:\t" + determine_blood_type(genotype)
		print "Phenotype:\t" + determine_blood_phenotype(genotype) + " blood type"
		return True

	return True
def allele_info(allele):
	if not allele in all_alleles:
		print allele, "is not an allele"
		return False
	else:
		print "Allele:\t" + allele
		print "---------------------"

	if allele in rec_alleles:
		print "Type:\trecessive"
	elif allele in dom_alleles:
		print "Type:\tdominant"
	elif allele in all_blood_alleles:
		if allele.lower() == 'io':
			print "Type:\t\trecessive",
		else:
			print "Type:\t\tdominant",
		print "blood allele"
	return True
def cross(genotype1, genotype2):
	if (genotype1 in all_blood_genotypes) == True and (genotype2 in all_blood_genotypes) == True:
		pm1 = genotype1[:2]; pm2 = genotype1[2:]
		pf1 = genotype2[:2]; pf2 = genotype2[2:]
		blood = True
	elif (genotype1 in all_genotypes) == True and (genotype2 in all_genotypes) == True:
		pm1 = genotype1[0]; pm2 = genotype1[1]
		pf1 = genotype2[0]; pf2 = genotype2[1]
		blood = False
	else:
		print genotype1 + " and " + genotype2 + " cannot be crossed, they are not proper genotypes"
		return False

	# Diag of output
	"""
	PM = Parent Male
	PF = Parent Female
	B# = box, number (possible genotype)
	   | M | M |
	---|---|---|---
	 F | 1 | 2 |
	---|---|---|---
	 F | 3 | 4 |
	---|---|---|---
	   |   |   |

	    | io | io |
	----|----|----|----
	 ia | 1  | 2  |
	----|----|----|----
	 ia | 3  | 4  |
	----|----|----|----
	    |    |    |
	"""

	print "\t" + genotype1 + ' x ' + genotype2
	print "---------------------"

	if blood == False:
		if genotype1.lower() != genotype2.lower():
			print genotype1 + " and " + genotype2 + " cannot be crossed, they are not alleles for the same trait"
			return
		b1 = make_genotype(pm1, pf1)
		b2 = make_genotype(pm2, pf1)
		b3 = make_genotype(pm1, pf2)
		b4 = make_genotype(pm2, pf2)
		print "    | " + pm1 + "  | " + pm2 + "  |"
		print "----|----|----|----"
		print "  " + pf1 + " | " + b1 + " | " + b2 + " |"
		print "----|----|----|----"
		print "  " + pf2 + " | " + b3 + " | " + b4 + " |"
		print "----|----|----|----"
		print "    |    |    |"
		print

		homodom, homorec, hetero = 0, 0, 0
		for i in [b1, b2, b3, b4]:
			if i in homodom_genotypes:
				homodom += 1
			elif i in homorec_genotypes:
				homorec += 1
			elif i in hetero_genotypes:
				hetero += 1
		#print homodom, homorec, hetero
		print "Possible Genotypes:"
		print		"\tHomozygous Dominant (" + genotype1.lower() + "):\t" + str(int((homodom / 4.0) * 100)) + "%"
		print		"\tHomozygous Recessive (" + genotype1.upper() + "):\t" + str(int((homorec / 4.0) * 100)) + "%"
		print		"\tHeterozygous (" + pm1.upper() + pm1.lower() + "):\t\t" + str(int((hetero / 4.0) * 100)) + "%"
		print "Possible Phenotypes:"
		print		"\tDominant Trait (" + pm1.upper() + "):\t\t" + str(int((homodom + hetero / 4.0) * 100)) + "%"
		print		"\tRecessive Trait (" + pm1.lower() + "):\t\t" + str(int((homorec / 4.0) * 100)) + "%"
	elif blood == True:
		def sp(s):
			if s.lower() == 'ab': return ' '
			else: return '  '
		b1 = determine_blood_phenotype(make_genotype(pm1, pf1))
		b2 = determine_blood_phenotype(make_genotype(pm2, pf1))
		b3 = determine_blood_phenotype(make_genotype(pm1, pf2))
		b4 = determine_blood_phenotype(make_genotype(pm2, pf2))
		
		print "    | " + pm1 + " | " + pm2 + " |"
		print "----|----|----|----"
		print " " + pf1 + " | " + b1 + sp(b1) + "| " + b2 + sp(b2) + "|"
		print "----|----|----|----"
		print " " + pf2 + " | " + b3 + sp(b3) + "| " + b4 + sp(b4) + "|"
		print "----|----|----|----"
		print "    |    |    |"
		print

	return True
def main(toparse, *args):
	"""
	main(*args) -> True if successful, else False
	args - list of arguments to the program, if __name__ == "__main__",
		then args is set to sys.argv which is
		the args passed to it at execution
	"""
	args = toparse
	# Make sure correct number of arguments
	if "--help" in args or "-h" in args:
		print __help__
		return True
	elif len(args) > 3:
		sys.exit("Too many arguments\nTry `" + __filecall__ + " --help' for more information,\n\tor see " + __filecall__ + ".__help__")
		return False
	elif len(args) < 2:
		sys.exit("Too little arguments\nTry `" + __filecall__ + " --help' for more information,\n\tor see " + __filecall__ + ".__help__")
		return False

	# Parse output
	if len(args) == 2:
		arg = args[1]
		if ' x ' in arg != -1:
			split = arg.split(' x ')
			if len(split) == 2:
				cross(split[0], split[1])
			else:
				sys.exit("Too many arguments\nTry `" + __filecall__ +\
					 " --help' for more information,\n\tor see " + __filecall__ + ".__help__")
			return False
		elif ' * ' in arg:
			split = arg.split(' * ')
			if len(split) == 2:
				cross(split[0], split[1])
			else:
				sys.exit("Too many arguments\nTry `" + __filecall__ +\
					 " --help' for more information,\n\tor see " + __filecall__ + ".__help__")
				return False
		else:
			determine_a_or_g(arg)
	elif len(args) == 3:
		split = args[1:]
		cross(split[0], split[1])
	else:
		print "main(): Error"
		return False

	return True

if __name__ == '__main__':
	main(sys.argv)