#!/usr/bin/python
#
# Poketypes allows you to list the types your team has, and prints
# how many / which types you have effectiveness, weakness, normal damage, or immunity towards

import argparse


type_index = ['normal', 'fire', 'water', 'electric', 'grass', 'ice', 'fighting', 'poison', 'ground', 'flying', 'psychic', 'bug', 'rock', 'ghost', 'dragon', 'dark', 'steel']

type_chart =[
    [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.5, 0.0, 1.0, 1.0, 0.5], # Normal
    [1.0, 0.5, 0.5, 1.0, 2.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0, 2.0, 0.5, 1.0, 0.5, 1.0, 2.0], # Fire
    [1.0, 2.0, 0.5, 1.0, 0.5, 1.0, 1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 2.0, 1.0, 0.5, 1.0, 1.0], # Water
    [1.0, 1.0, 2.0, 0.5, 0.5, 1.0, 1.0, 1.0, 0.0, 2.0, 1.0, 1.0, 1.0, 1.0, 0.5, 1.0, 1.0], # Electric
    [1.0, 0.5, 2.0, 1.0, 0.5, 1.0, 1.0, 0.5, 2.0, 0.5, 1.0, 0.5, 2.0, 1.0, 0.5, 1.0, 0.5], # Grass
    [1.0, 0.5, 0.5, 1.0, 2.0, 0.5, 1.0, 1.0, 2.0, 2.0, 1.0, 1.0, 1.0, 1.0, 2.0, 1.0, 0.5], # Ice
    [2.0, 1.0, 1.0, 1.0, 1.0, 2.0, 1.0, 0.5, 1.0, 0.5, 0.5, 0.5, 2.0, 0.0, 1.0, 2.0, 2.0], # Fighting
    [1.0, 1.0, 1.0, 1.0, 2.0, 1.0, 1.0, 0.5, 0.5, 1.0, 1.0, 1.0, 0.5, 0.5, 1.0, 1.0, 0.0], # Poison
    [1.0, 2.0, 1.0, 2.0, 0.5, 1.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.5, 2.0, 1.0, 1.0, 1.0, 2.0], # Ground
    [1.0, 1.0, 1.0, 0.5, 2.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 2.0, 0.5, 1.0, 1.0, 1.0, 0.5], # Flying
    [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 2.0, 2.0, 1.0, 1.0, 0.5, 1.0, 1.0, 1.0, 1.0, 0.0, 0.5], # Psychic
    [1.0, 0.5, 1.0, 1.0, 2.0, 1.0, 0.5, 0.5, 1.0, 0.5, 2.0, 1.0, 1.0, 0.5, 1.0, 2.0, 0.5], # Bug
    [1.0, 2.0, 1.0, 1.0, 1.0, 2.0, 0.5, 1.0, 0.5, 2.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 0.5], # Rock
    [0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 2.0, 1.0, 1.0, 2.0, 1.0, 0.5, 0.5], # Ghost
    [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 2.0, 1.0, 0.5], # Dragon
    [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.5, 1.0, 1.0, 1.0, 2.0, 1.0, 1.0, 2.0, 1.0, 0.5, 0.5], # Dark
    [1.0, 0.5, 0.5, 0.5, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 0.5], # Steel
    ]


def calculate_attack_values(types):
    atks = [0.0]*len(type_index)

    for t in types:
        t_index = type_index.index(t)

        for i in range(len(type_index)):
            value = type_chart[t_index][i]
            if value > atks[i]:
                atks[i] = value

    return atks


def calculate_defense_values(types):
    defs = [2.0]*len(type_index)

    for t in types:
        t_index = type_index.index(t)

        for i in range(len(type_index)):
            value = type_chart[i][t_index]
            if value < defs[i]:
                defs[i] = value

    return defs


def get_types_gt(strengths, value):
    """ Given a list of strengths, returns a list of all the types that are > value"""
    
    ret = []
    for i in range(len(strengths)):
        if strengths[i] > value:
            ret.append(type_index[i])

    return ret


def get_types_lt(strengths, value):
    """ Given a list of strengths, returns a list of all the types that are < value"""
    
    ret = []
    for i in range(len(strengths)):
        if strengths[i] < value:
            ret.append(type_index[i])

    return ret


def get_types_eq(strengths, value):
    """ Given a list of strengths, returns a list of all the types that are == value"""
    
    ret = []
    for i in range(len(strengths)):
        if strengths[i] == value:
            ret.append(type_index[i])

    return ret


def print_entry(src_list):
    try:
        print '{:10}'.format(src_list.pop(0)),
    except IndexError:
        print ' '*10,


             
def main():
    type_list = parse_args().types

    print '\n{space}Attack\n{space}======\n'.format(space=' '*16)
    print 'Strong     Normal     Weak       Immune'
    print '------     ------     ----       ------'

    strengths = calculate_attack_values(type_list)
    strong_list = get_types_gt(strengths, 1.0)
    normal_list = get_types_eq(strengths, 1.0)
    weak_list   = list(set(get_types_lt(strengths, 1.0)) - set(get_types_eq(strengths, 0.0)))
    immune_list = get_types_eq(strengths, 0.0)

    while strong_list or normal_list or weak_list or immune_list:
        print_entry(strong_list)
        print_entry(normal_list)
        print_entry(weak_list)
        print_entry(immune_list)
        print


    strengths = calculate_defense_values(type_list)

    strong_list = list(set(get_types_lt(strengths, 1.0)) - set(get_types_eq(strengths, 0.0)))
    normal_list = get_types_eq(strengths, 1.0)
    weak_list = get_types_gt(strengths, 1.0)
    immune_list = get_types_eq(strengths, 0.0)

    print '\n\n{space}Defend\n{space}======\n'.format(space=' '*16)
    print 'Strong     Normal     Weak       Immune'
    print '------     ------     ----       ------'

    while strong_list or normal_list or weak_list or immune_list:
        print_entry(strong_list)
        print_entry(normal_list)
        print_entry(weak_list)
        print_entry(immune_list)
        print

    print


def parse_args():
    parser = argparse.ArgumentParser(description='Check your team\'s type strengths and weaknesses', formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument('types', nargs=argparse.REMAINDER, help="A list of the types your team is composed of")
    return parser.parse_args()


if __name__ == '__main__': main()