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

from BattlePredictor import AI_Logic
ai_predictor = AI_Logic()
# VERSION 1.1
# added three alarm scenarios

TrainerData = pd.read_csv("trainers.csv", index_col = "Name")
MovesetData = pd.read_csv("BFpokemon.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

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.Type1 = pokemon_data["Type1"] # <--- Add this
        self.Type2 = pokemon_data["Type2"] # <--- Add this

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

    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 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



# 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 = [PokemonSet(identifier) for identifier in trainerSets if  (Level != 50 or PokemonSet(identifier).Entry <= 850)]

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)

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))

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))

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


# === START NEW BLOCK ===
print("\n--- AI SWITCH PREDICTION ---")
p_type_input = input('What type is your pokemon? (format: "dragon/flying"): ').strip()
p_levitate_input = input('Does your pokemon have the ability levitate? (y/n): ').strip().lower()

if p_type_input:
    # 1. Parse Inputs
    player_types = [t.strip() for t in p_type_input.split('/')]
    has_levitate = (p_levitate_input == 'y')

    print("Calculating AI Logic...")
    original_count = len(TeamCombinationsAfterSecondPoke)

    # 2. Filter the list
    # We pass the list of valid teams, the switch-in (pokemon2), and your context
    TeamCombinationsAfterSecondPoke = ai_predictor.filter_impossible_teams(
        TeamCombinationsAfterSecondPoke,
        pokemon2,
        player_types,
        has_levitate
    )

    print(f"AI Logic Filtered Teams: {original_count} -> {len(TeamCombinationsAfterSecondPoke)}")
# === END NEW BLOCK ===

# 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()