import numpy as np
import pandas as pd
from itertools import combinations

#VERSION 1.1
# added three alarm scenarios

TrainerData = pd.read_csv("trainers.csv", index_col = "Name")
MovesetData = pd.read_csv("Bfpalace.csv")
MyPokemon = pd.read_csv("mypokemon.csv")

# USE ONLY LEVEL 100 OR LEVEL 50 FOR NOW
Level = 50

# ADD OR REMOVE PROBLEMATIC MOVES IF YOU LIKE

Alarms1 = set(['Sheer Cold', 'Horn Drill', 'Fissure', 'Guillotine', 'Reversal'])
Alarms2 = set(['Swords Dance', 'Dragon Dance', 'Double Team'])
Alarms3 = set(['Counter', 'Mirror Coat', 'Psych Up'])

ItemAlarmTotal = set(['BrightPowder', 'Lax Incense', 'Quick Claw'])

# Combine alarms
AlarmTotal = Alarms1 | Alarms2 | Alarms3

#==================================#
#=== DON'T TOUCH ANYTHING BELOW ===#
#==================================#

class PokemonSet:
	def __init__(self, Name):
		self.Name = Name
		pokemon_data = MovesetData[MovesetData["Name"] == Name].iloc[0]
		self.Entry = int(pokemon_data["Entry"])
		self.Species = pokemon_data["Species"]
		self.Instance = pokemon_data["Instance"]
		self.Item = pokemon_data["Item"]
		self.Nature = pokemon_data["Nature"]
		self.MoveList = [pokemon_data["Move 1"], pokemon_data["Move 2"], pokemon_data["Move 3"], pokemon_data["Move 4"]]
		self.Speed = int(pokemon_data["Lv " + str(Level) + " Speed"])

		self.Alarm = any(move in AlarmTotal for move in self.MoveList)
		self.Alarm2 = self.Item in ItemAlarmTotal
		self.LocalOdds = 1
		self.GlobalOdds = 1

		self.ATK = pokemon_data["AI_ATK"]
		self.DEF = pokemon_data["AI_DEF"]
		self.SPT = pokemon_data["AI_SPT"]
		self.R_ATK = pokemon_data["R_ATK"]
		self.R_DEF = pokemon_data["R_DEF"]
		self.NOMOVE = pokemon_data["NOMOVE"]
		self.ATK50 = pokemon_data["AI_ATK_50"]
		self.DEF50 = pokemon_data["AI_DEF_50"]
		self.SPT50 = pokemon_data["AI_SPT_50"]
		self.R_ATK50 = pokemon_data["R_ATK_50"]
		self.R_DEF50 = pokemon_data["R_DEF_50"]
		self.NOMOVE50 = pokemon_data["NOMOVE_50"]

		self.Move1Class = pokemon_data["Move 1 Class"]
		self.Move2Class = pokemon_data["Move 2 Class"]
		self.Move3Class = pokemon_data["Move 3 Class"]
		self.Move4Class = pokemon_data["Move 4 Class"]

	def Setprinter(self):
		msg = "!!!" if self.Alarm or self.Alarm2 else ""
		print(f'{msg:>3}', f'{self.Name:>13}', f'{np.round(self.LocalOdds, 6):>9}', f'{self.Item:>13}', f'{self.Nature:>8}',
			  f'{self.MoveList[0]:>13}', f'{self.MoveList[1]:>13}', f'{self.MoveList[2]:>13}', f'{self.MoveList[3]:>13}', f'{self.Speed:>6}')

	def Palaceprinter(self):
		msg = "!!!" if self.Alarm else ""

		# Print move classifications
		print("This Pokemon's moves are classified as:\n")
		print(f'{msg:>3} {self.Name:>13} {self.Item:>13} {self.Nature:>8} '
			  f'{self.Move1Class:>17} {self.Move2Class:>17} {self.Move3Class:>17} {self.Move4Class:>17}')
		print("-" * 113)

class PokemonSetM(PokemonSet):
    def __init__(self, Name):
        self.Name = Name
        # Load from MyPokemon instead of MovesetData
        pokemon_data = MyPokemon[MyPokemon["Name"] == Name]

        if pokemon_data.empty:
            raise ValueError(f"No data found in MyPokemon for Name: {Name}")

        pokemon_data = pokemon_data.iloc[0]  # Extract the first row

        # Initialize attributes using MyPokemon data
        self.Entry = int(pokemon_data["Entry"])
        self.Species = pokemon_data["Species"]
        self.Instance = pokemon_data["Instance"]
        self.Item = pokemon_data["Item"]
        self.Nature = pokemon_data["Nature"]
        self.MoveList = [pokemon_data["Move 1"], pokemon_data["Move 2"],
                         pokemon_data["Move 3"], pokemon_data["Move 4"]]
        self.Speed = int(pokemon_data["Lv " + str(Level) + " Speed"])

        self.Alarm = any(move in AlarmTotal for move in self.MoveList)
        self.Alarm2 = self.Item in ItemAlarmTotal
        self.LocalOdds = 1
        self.GlobalOdds = 1

        self.ATK = pokemon_data["AI_ATK"]
        self.DEF = pokemon_data["AI_DEF"]
        self.SPT = pokemon_data["AI_SPT"]
        self.R_ATK = pokemon_data["R_ATK"]
        self.R_DEF = pokemon_data["R_DEF"]
        self.NOMOVE = pokemon_data["NOMOVE"]
        self.ATK50 = pokemon_data["AI_ATK_50"]
        self.DEF50 = pokemon_data["AI_DEF_50"]
        self.SPT50 = pokemon_data["AI_SPT_50"]
        self.R_ATK50 = pokemon_data["R_ATK_50"]
        self.R_DEF50 = pokemon_data["R_DEF_50"]
        self.NOMOVE50 = pokemon_data["NOMOVE_50"]

        self.Move1Class = pokemon_data["Move 1 Class"]
        self.Move2Class = pokemon_data["Move 2 Class"]
        self.Move3Class = pokemon_data["Move 3 Class"]
        self.Move4Class = pokemon_data["Move 4 Class"]


def SpeciesListing(trainerSpecies, PokemonList):
	# Using dictionary for faster species lookup and summation
	species_odds = {species[0]: 0 for species in trainerSpecies}
	for pokemon in PokemonList:
		if pokemon.Species in species_odds:
			species_odds[pokemon.Species] += pokemon.LocalOdds

	sorted_species = sorted(species_odds.items(), key=lambda pair: pair[1], reverse=True)
	for species, odds in sorted_species:
		print(f'{species.upper():>15}', f'{np.round(odds, 6):>11}')


def printline():
	print("-----------------------------------------------------------------------------------------------------------------")

# Optimizing Team Combination Generation
def generate_unique_teams(PokemonList):
    TeamCombinations = []

    # Step 1: Organize Pokémon into a list of (species, item) pairs
    unique_pokemon = [(pokemon.Species, pokemon.Item, pokemon) for pokemon in PokemonList]
    # print(unique_pokemon)

    # Step 2: Use combinations to select 3 Pokémon at a time
    for comb in combinations(unique_pokemon, 3):
        # print(comb)
        species_set = set()
        item_set = set()
        valid_combination = True

        # Check if species and items are unique
        for species, item, pokemon in comb:
            if species in species_set or item in item_set:
                valid_combination = False
                break
            species_set.add(species)
            item_set.add(item)

        if valid_combination:
            # Extract the original Pokémon objects from the combination
            TeamCombinations.append(tuple(pokemon for _, _, pokemon in comb))

    return TeamCombinations


# Define GT 50 and LT 50 attributes
GT50_attributes = ["ATK", "DEF", "SPT", "R_ATK", "R_DEF", "NOMOVE"]
LT50_attributes = ["ATK50", "DEF50", "SPT50", "R_ATK50", "R_DEF50", "NOMOVE50"]

# Define all four cases
cases = {
    "Case 1 (Both greater than 50% HP)": (GT50_attributes, GT50_attributes),
    "Case 2 (User Pokemon greater than 50% HP, Opponent less than 50% HP)": (GT50_attributes, LT50_attributes),
    "Case 3 (User Pokemon less than 50% HP, Opponent greater than 50% HP)": (LT50_attributes, GT50_attributes),
    "Case 4 (Both less than 50% HP)": (LT50_attributes, LT50_attributes),
}

# Function to get user input for Pokémon level group
def get_pokemon_group(pokemon_name):
    while True:
        try:
            user_input = input(f"Enter HP for {pokemon_name} (1 for greater than half HP, 2 for less than half HP): ")
            if user_input == '':
                exit()
            user_input = int(user_input)
            if user_input in [1, 2]:
                return user_input == 1  # Return True if GT 50, False if LT 50
            else:
                print("Invalid input. Please enter 1 or 2.")
        except ValueError:
            print("Invalid input. Please enter a valid number (1 or 2).")


# Function to calculate probability table for a given case
def calculate_probability_table(case_name, attrs_user, attrs_opp, mypokemon, opponent_pokemon):
    case_data = {}
    for attr1 in attrs_user:
        for attr2 in attrs_opp:
            prob = getattr(mypokemon, attr1) * getattr(opponent_pokemon, attr2)
            if prob > 0:  # Only keep non-zero probabilities
                case_data[f"{attr1} x {attr2}"] = prob

    # Convert to DataFrame, sort in descending order, and return
    return pd.DataFrame(case_data.items(), columns=["Combination", "Probability"])\
             .sort_values(by="Probability", ascending=False)

def display_cases(selected_case, cases, mypokemon, opponent_pokemon):
    # Calculate and display the remaining three cases
    print("\n*** Other Cases ***")
    for case_name, (attrs_user, attrs_opp) in cases.items():
        if case_name != selected_case:
            print(f"\n{case_name}\n")
            other_table = calculate_probability_table(case_name, attrs_user, attrs_opp, mypokemon, opponent_pokemon)
            print(other_table)

    # Calculate and display the selected case first
    print(f"\n*** Selected Case: {selected_case}, Scroll Up For Other Cases ***\n")
    selected_table = calculate_probability_table(selected_case, *cases[selected_case], mypokemon, opponent_pokemon)
    print(selected_table)


def determine_selected_case(mypokemon, pokemon3):
        # Get user input for each Pokémon
        mypokemon_GT50 = get_pokemon_group(mypokemon.Name)
        pokemon3_GT50 = get_pokemon_group(pokemon3.Name)

        # Identify the selected case based on user input
        if mypokemon_GT50 and pokemon3_GT50:
            return "Case 1 (Both greater than 50% HP)"
        elif mypokemon_GT50 and not pokemon3_GT50:
            return "Case 2 (User Pokemon greater than 50% HP, Opponent less than 50% HP)"
        elif not mypokemon_GT50 and pokemon3_GT50:
            return "Case 3 (User Pokemon less than 50% HP, Opponent greater than 50% HP)"
        else:
            return "Case 4 (Both less than 50% HP)"

# List to store PokemonSet objects
my_pokemon_list = []

# Iterate through each row and create PokemonSet objects
print("Reading User Pokemon")
printline()
for _, row in MyPokemon.iterrows():
    pokemon = PokemonSetM(row["Name"])  # Assuming Name is the identifier
    my_pokemon_list.append(pokemon)
    print(f"Added {pokemon.Species}")


printline()

# Main execution flow
printline()
print("What trainer are you battling?")
Trainer = input()
if Trainer == '':
    exit()

Trainer = Trainer.upper()
printline()

trainerSets = list(TrainerData.loc[[Trainer]].values[0])[1:]  # Exclude Name from the list
trainerSets = [x for x in trainerSets if pd.notna(x)]

PokemonList = []
for identifier in trainerSets:
    p = PokemonSet(identifier)
    if p.Entry > 1000:
        PokemonList.append(p)
    elif Level != 50 or p.Entry <= 850:
        PokemonList.append(p)
		
totalalarms = 0
totalitemalarms=0
for Pokemon in PokemonList:
    if Pokemon.Alarm == True:
        totalalarms += 1
    if Pokemon.Alarm2 == True:
        totalitemalarms += 1


trainerSpecies = list({pokemon.Species: 0 for pokemon in PokemonList}.items())  # Unique species count

# Generate all unique team combinations
TeamCombinations = generate_unique_teams(PokemonList)
# print(len(TeamCombinations))

alarmcounter = 0
itemalarmcounter = 0

alarm1counter = 0
alarm2counter = 0
alarm3counter = 0

# Loop over all team combinations
for pokemon1, pokemon2, pokemon3 in TeamCombinations:
    # Check if the combination of the 3 Pokémon has a move that triggers an alarm
    teammoves = pokemon1.MoveList + pokemon2.MoveList + pokemon3.MoveList
    movealarm = any(move in AlarmTotal for move in teammoves)
    movealarm1 = any(move in Alarms1 for move in teammoves)
    movealarm2 = any(move in Alarms2 for move in teammoves)
    movealarm3 = any(move in Alarms3 for move in teammoves)
    if movealarm:
        alarmcounter += 1
    if movealarm1:
        alarm1counter += 1
    if movealarm2:
        alarm2counter += 1
    if movealarm3:
        alarm3counter += 1

    # Check if the combination of the 3 Pokémon has an item that triggers an alarm
    teamitems = {pokemon1.Item, pokemon2.Item, pokemon3.Item}  # Use a set to avoid counting duplicate items
    itemalarm = any(item in ItemAlarmTotal for item in teamitems)
    if itemalarm:
        itemalarmcounter += 1

print("This trainer has", len(PokemonList), "different sets, which constitute", int(len(TeamCombinations)), "different teams.")
print("There are", totalitemalarms, "alarming item sets, and there is a", round(itemalarmcounter/len(TeamCombinations),4), "chance of running into at least one of them.")
print("There are", totalalarms, "alarming movesets, and there is a", round(alarmcounter/len(TeamCombinations),4), "chance of running into at least one of them.")
print("There is", round(alarm1counter/len(TeamCombinations),4), "chance of finding Alarm1:", Alarms1)
print("There is", round(alarm2counter/len(TeamCombinations),4), "chance of finding Alarm2:", Alarms2)
print("There is", round(alarm3counter/len(TeamCombinations),4), "chance of finding Alarm3:", Alarms3)
print("There is", round(itemalarmcounter/len(TeamCombinations),4), "chance of finding an Alarming Item:", ItemAlarmTotal)
printline()

# Compute and update probabilities for Pokémon based on their participation in team combinations
for pokemon in PokemonList:
    counter = sum(1 for team in TeamCombinations if pokemon in team)
    pokemon.LocalOdds = counter / (3 * len(TeamCombinations))  # Divide by 3 because each team has 3 Pokémon

# Sort and display Pokémon probabilities
PokemonList = sorted(PokemonList, key=lambda pokemon: pokemon.Instance)
# printline()

# for pokemon in PokemonList:
# 	pokemon.Setprinter()


printline()
SpeciesListing(trainerSpecies, PokemonList)
printline()
print("Please enter a Pokemon you have seen.")
temp = input()
if temp == '':
	exit()
temp = temp.capitalize()
# temp2 = temp.upper()
printline()
print("Possible sets:")
odds = 0
for pokemon in PokemonList:
	if pokemon.Species == temp:
		odds += pokemon.LocalOdds
for pokemon in PokemonList:
	if pokemon.Species == temp:
		pokemon.LocalOdds /= odds
		pokemon.Setprinter()
printline()
print("Which one was it? Enter a number.")
number = input()
if number == '':
	exit()
printline()
pokemon1 = PokemonSet(temp + ' ' + str(number))


printline()


print("List of User Pokemon")
printline()
for pokemon in my_pokemon_list:
	pokemon.Setprinter()
printline()
print("Which Pokemon are you battling with?")
temp_mine = input()
if temp_mine == '':
	exit()
temp_mine = temp_mine.capitalize()

my_pokemon_name = ''
for pokemon in my_pokemon_list:
	if pokemon.Species == temp_mine:
		my_pokemon_name = pokemon.Name


mypokemon = PokemonSetM(my_pokemon_name)
printline()

mypokemon.Palaceprinter()
pokemon1.Palaceprinter()

printline()
selected_case = determine_selected_case(mypokemon, pokemon1)
printline()
display_cases(selected_case, cases, mypokemon, pokemon1)

print("")
input("Press Enter to continue...")
# printline()


printline()
TeamCombinationsAfterFirstPoke = [
    team for team in TeamCombinations
    if all(pokemon.Species != pokemon1.Species and pokemon.Item != pokemon1.Item for pokemon in team)
]

# print(len(TeamCombinationsAfterFirstPoke))

alarmcounter = 0
itemalarmcounter = 0

alarm1counter = 0
alarm2counter = 0
alarm3counter = 0

# Loop over all team combinations
for pokemon1, pokemon2, pokemon3 in TeamCombinationsAfterFirstPoke:
    # Check if the combination of the 3 Pokémon has a move that triggers an alarm
    teammoves = pokemon2.MoveList + pokemon3.MoveList
    movealarm = any(move in AlarmTotal for move in teammoves)
    movealarm1 = any(move in Alarms1 for move in teammoves)
    movealarm2 = any(move in Alarms2 for move in teammoves)
    movealarm3 = any(move in Alarms3 for move in teammoves)
    if movealarm:
        alarmcounter += 1
    if movealarm1:
        alarm1counter += 1
    if movealarm2:
        alarm2counter += 1
    if movealarm3:
        alarm3counter += 1

    # Check if the combination of the 3 Pokémon has an item that triggers an alarm
    teamitems = {pokemon2.Item, pokemon3.Item}  # Use a set to avoid counting duplicate items
    itemalarm = any(item in ItemAlarmTotal for item in teamitems)
    if itemalarm:
        itemalarmcounter += 1

for pokemon in PokemonList:
	counter = 0
	for team in TeamCombinationsAfterFirstPoke:
		if pokemon in team:
			counter += 1
	pokemon.LocalOdds = counter/(3*len(TeamCombinationsAfterFirstPoke))

PokemonList = sorted(PokemonList, key=lambda pokemon: pokemon.Instance)

# for pokemon in PokemonList:
# 	pokemon.Setprinter()

print("Updated Pokemon Probabilities")
printline()
SpeciesListing(trainerSpecies, PokemonList)
printline()
print("There is", round(alarm1counter/len(TeamCombinationsAfterFirstPoke),4), "chance of finding Alarm1:", Alarms1)
print("There is", round(alarm2counter/len(TeamCombinationsAfterFirstPoke),4), "chance of finding Alarm2:", Alarms2)
print("There is", round(alarm3counter/len(TeamCombinationsAfterFirstPoke),4), "chance of finding Alarm3:", Alarms3)
print("There is", round(itemalarmcounter/len(TeamCombinationsAfterFirstPoke),4), "chance of finding an Alarming Item:", ItemAlarmTotal)
printline()
print("Please enter another Pokemon you have seen.")
temp = input()
if temp == '':
	exit()
temp = temp.capitalize()
# temp2 = temp.upper()
printline()
print("Possible sets:")
odds = 0
for pokemon in PokemonList:
	if pokemon.Species == temp:
		odds += pokemon.LocalOdds
for pokemon in PokemonList:
	if pokemon.Species == temp:
		pokemon.LocalOdds /= odds
		pokemon.Setprinter()
printline()
print("Which one was it? Enter a number.")
number = input()
if number == '':
	exit()
printline()

pokemon2 = PokemonSet(temp + ' ' + str(number))

printline()
print("List of User Pokemon")
printline()
for pokemon in my_pokemon_list:
	pokemon.Setprinter()
printline()
print("Which Pokemon are you battling with?")
temp_mine = input()
if temp_mine == '':
	exit()
temp_mine = temp_mine.capitalize()

my_pokemon_name = ''
for pokemon in my_pokemon_list:
	if pokemon.Species == temp_mine:
		my_pokemon_name = pokemon.Name


mypokemon = PokemonSetM(my_pokemon_name)
printline()


mypokemon.Palaceprinter()
pokemon2.Palaceprinter()

printline()
selected_case = determine_selected_case(mypokemon, pokemon2)
printline()
display_cases(selected_case, cases, mypokemon, pokemon2)

printline()
print("")
input("Press Enter to continue...")
printline()


TeamCombinationsAfterSecondPoke = [
    team for team in TeamCombinationsAfterFirstPoke
    if all(pokemon.Species != pokemon2.Species and pokemon.Item != pokemon2.Item for pokemon in team)
]

# print(len(TeamCombinationsAfterSecondPoke))


alarmcounter = 0
itemalarmcounter = 0

alarm1counter = 0
alarm2counter = 0
alarm3counter = 0

# Loop over all team combinations
for pokemon1, pokemon2, pokemon3 in TeamCombinationsAfterSecondPoke:
    # Check if the combination of the 3 Pokémon has a move that triggers an alarm
    teammoves = pokemon3.MoveList
    movealarm = any(move in AlarmTotal for move in teammoves)
    movealarm1 = any(move in Alarms1 for move in teammoves)
    movealarm2 = any(move in Alarms2 for move in teammoves)
    movealarm3 = any(move in Alarms3 for move in teammoves)
    if movealarm:
        alarmcounter += 1
    if movealarm1:
        alarm1counter += 1
    if movealarm2:
        alarm2counter += 1
    if movealarm3:
        alarm3counter += 1

    # Check if the combination of the 3 Pokémon has an item that triggers an alarm
    teamitems = {pokemon3.Item}  # Use a set to avoid counting duplicate items
    itemalarm = any(item in ItemAlarmTotal for item in teamitems)
    if itemalarm:
        itemalarmcounter += 1

for pokemon in PokemonList:
	counter = 0
	for team in TeamCombinationsAfterSecondPoke:
		if pokemon in team:
			counter += 1
	pokemon.LocalOdds = counter/(3*len(TeamCombinationsAfterSecondPoke))

PokemonList = sorted(PokemonList, key=lambda pokemon: pokemon.Instance)

# for pokemon in PokemonList:
# 	pokemon.Setprinter()


print("Updated Pokemon Probabilities")
printline()
SpeciesListing(trainerSpecies, PokemonList)
printline()
print("There is", round(alarm1counter/len(TeamCombinationsAfterSecondPoke),4), "chance of finding Alarm1:", Alarms1)
print("There is", round(alarm2counter/len(TeamCombinationsAfterSecondPoke),4), "chance of finding Alarm2:", Alarms2)
print("There is", round(alarm3counter/len(TeamCombinationsAfterSecondPoke),4), "chance of finding Alarm3:", Alarms3)
print("There is", round(itemalarmcounter/len(TeamCombinationsAfterSecondPoke),4), "chance of finding an Alarming Item:", ItemAlarmTotal)
printline()
print("Please enter the last Pokemon.")
temp = input()
if temp == '':
	exit()
temp = temp.capitalize()
# temp2 = temp.upper()
printline()
print("Possible sets:")
odds = 0
for pokemon in PokemonList:
	if pokemon.Species == temp:
		odds += pokemon.LocalOdds
for pokemon in PokemonList:
	if pokemon.Species == temp:
		pokemon.LocalOdds /= odds
		pokemon.Setprinter()

printline()
print("Which one was it? Enter a number.")
number = input()

if number == '':
	exit()

printline()

pokemon3 = PokemonSet(temp + ' ' + str(number))

printline()
print("List of User Pokemon")
printline()
for pokemon in my_pokemon_list:
	pokemon.Setprinter()
printline()
print("Which Pokemon are you battling with?")
temp_mine = input()
if temp_mine == '':
	exit()
temp_mine = temp_mine.capitalize()

my_pokemon_name = ''
for pokemon in my_pokemon_list:
	if pokemon.Species == temp_mine:
		my_pokemon_name = pokemon.Name


mypokemon = PokemonSetM(my_pokemon_name)
printline()

mypokemon.Palaceprinter()
pokemon3.Palaceprinter()

printline()
selected_case = determine_selected_case(mypokemon, pokemon3)
printline()
display_cases(selected_case, cases, mypokemon, pokemon3)

print("")
#printline()