from core import Trainer

import random
import copy
import game_data as gd
import pickle
import matplotlib.pyplot as plt
import numpy as np
import os

#if the results folder does not exist, create it
if not os.path.exists('results'):
    os.makedirs('results')
if not os.path.exists('plots'):
    os.makedirs('plots')

def RandomBattle(trainer,cpu_trainer):
    #print(f"You are in battle against {cpu_trainer.name}!")
    battle_active = True
    
    active_pokemon = trainer.get_active_pokemon()
    cpu_active_pokemon = cpu_trainer.get_active_pokemon()
    #print(f"Your active Pokemon is {active_pokemon.name} and the CPU's active Pokemon is {cpu_active_pokemon.name}.")
    did_player_won = False
    battle_turns_num = 0

    player_hp_evolution = []
    player_selected_moves = []
    player_damage_done = []
    cpu_hp_evolution = []
    cpu_selected_moves = []
    cpu_damage_done = []
    while battle_active:
        player_move_name = random.choice(active_pokemon.get_moves())
        player_selected_moves.append(player_move_name)
        #print(f"{active_pokemon.name} selected {player_move_name} to use against {cpu_active_pokemon.name}:\n")
        try:
            player_damage_done_value = active_pokemon.useMove(player_move_name,cpu_active_pokemon)

            player_damage_done.append(player_damage_done_value)
            player_hp_evolution.append(active_pokemon.get_current_hp())
            cpu_hp_evolution.append(cpu_active_pokemon.get_current_hp())

            battle_turns_num += 1
            #check if the opponent has been defeated
            if not cpu_trainer.active_pokemon_can_battle():
                #print(f"{cpu_trainer.get_active_pokemon().name} has been defeated! You win!")
                did_player_won = True
                battle_active = False
                break

            #if the opponent has not been defeated, then its the CPU turn to play
            cpu_active_pokemon_selected_move_name = random.choice(cpu_active_pokemon.get_moves())
            #print(f"CPU's {cpu_active_pokemon.name} is using {cpu_active_pokemon_selected_move_name} against your {active_pokemon.name}!")
            cpu_selected_moves.append(cpu_active_pokemon_selected_move_name)

            cpu_damage_done_value = cpu_active_pokemon.useMove(cpu_active_pokemon_selected_move_name, active_pokemon)

            cpu_damage_done.append(cpu_damage_done_value)
            player_hp_evolution.append(active_pokemon.get_current_hp())
            cpu_hp_evolution.append(cpu_active_pokemon.get_current_hp())
            battle_turns_num += 1

            #update the player pokemon list
            trainer.update_pokemon_list(active_pokemon)
            #check if the player has been defeated
            if not trainer.active_pokemon_can_battle():
                #print(f"Your {active_pokemon.name} has been defeated!")
                if not trainer.can_battle():
                    #print("You have no more Pokemons that can battle! You lose!")
                    did_player_won=False
                    battle_active = False
                    break
                else:
                    #print("You have other Pokemons that can battle!")
                    continue

        except ValueError as e:
            print(e)

    #manage those cases for which the player wins in the first turn and the cpu does not play any turn
    if len(player_damage_done) == 0:
        player_damage_done = [0]
    if len(cpu_damage_done) == 0:
        cpu_damage_done = [0]

    player_avg_damage_done = np.mean(player_damage_done) #average damage done per turn
    cpu_avg_damage_done = np.mean(cpu_damage_done) #average damage done per turn

    player_residual_hp_percent = trainer.get_active_pokemon().get_current_hp_percent()

    player_history = {'player_hp_evolution':player_hp_evolution,
                      'player_selected_moves':player_selected_moves,
                      'player_damage_done':player_damage_done,
                      'player_avg_damage_done':player_avg_damage_done
    }
    cpu_history = {'cpu_hp_evolution':cpu_hp_evolution,
                   'cpu_selected_moves':cpu_selected_moves,
                   'cpu_damage_done':cpu_damage_done,
                   'cpu_avg_damage_done':cpu_avg_damage_done
    }
    turn_history = {'player_history':player_history,
                    'cpu_history':cpu_history
    }
    battle_results = {
        'player_pokemon':trainer.get_active_pokemon(),
        'encountered_pokemon':cpu_trainer.get_active_pokemon(),
        'result':did_player_won, 
        'battle_turns_num':battle_turns_num,
        'player_residual_hp_percent':player_residual_hp_percent,
        'turn_history':turn_history
    }
    return battle_results

def Main(starter_choice,Nbattles):
    try:
        starter_name = starter_choice
        trainer = Trainer(name='edo', starter_name=starter_name)
        #print(f"{trainer.name} has chosen {trainer.starter} as their starter Pokemon!")
        #print("Adding potions and pokeballs to your inventory...")
        trainer.items = {"potions": 10, "pokeball": 10}
        #print("Your character is ready!")
        #print(f"Welcome {trainer.name}! Your adventure begins now!")
    except ValueError as e:
        print(e)
        print("Something went wrong!")

    game_active_flag = True
    while game_active_flag:
        all_battle_results = []
        for i in range(0,Nbattles):
            #print("You entered the Random Battle mode")
            #print("Battle #",i)
            cpu_trainer = Trainer(name='CPU',starter_name=gd.pokemons_df.index[random.randint(0,len(gd.pokemons_df.index)-1)],is_CPU=True)
            battle_results = RandomBattle(trainer,cpu_trainer)
            all_battle_results.append(battle_results)
            heal_flag = 'yes'
            if heal_flag == 'yes':
                for pokemon in trainer.pokemon_list:
                    pokemon.curr_hp = pokemon.baseStats['hp']
                #print("All your Pokemon have been healed!")

        game_active_flag = False
        #save all battle results
        return all_battle_results

def make_plots(Ngames,Nbattles):
    files_to_analyze = os.listdir('results/')

    avg_damage_per_level = dict()
    for file_name in files_to_analyze:
        pickle_in = open(f"results/{file_name}",'rb')
        all_battle_results = pickle.load(pickle_in)
        pickle_in.close()
        for battle_result in all_battle_results:
            player_pokemon = battle_result['player_pokemon']
            player_pokemon_name = player_pokemon.name
            player_pokemon_level = player_pokemon.level
            player_avg_damage_done = battle_result['turn_history']['player_history']['player_avg_damage_done']

            if player_pokemon_name not in avg_damage_per_level:
                avg_damage_per_level[player_pokemon_name] = {}
            if player_pokemon_level not in avg_damage_per_level[player_pokemon_name]:
                avg_damage_per_level[player_pokemon_name][player_pokemon_level] = []
            avg_damage_per_level[player_pokemon_name][player_pokemon_level].append(player_avg_damage_done)
    for pokemon_name in avg_damage_per_level:
        for level in avg_damage_per_level[pokemon_name]:
            avg_damage_per_level[pokemon_name][level] = np.mean(avg_damage_per_level[pokemon_name][level])

    for pokemon_name in avg_damage_per_level:
        levels = list(avg_damage_per_level[pokemon_name].keys())
        avg_damages = list(avg_damage_per_level[pokemon_name].values())
        plt.figure(figsize=(12, 6))
        plt.bar(levels, avg_damages)
        plt.xlabel('Pokemon Level')
        plt.ylabel('Average Damage Done')
        plt.title(f'Average Damage Done per battle turn for {pokemon_name}')
        plt.xticks(rotation=45)
        plt.tight_layout()
        plt.savefig(f'plots/average_damage_per_turn_{pokemon_name}.pdf')
       
    #pie chart
    available_pokemon_type_distribution = {}
    for pokemon_name in gd.pokemons_dict:
        pokemon_types = gd.pokemons_dict[pokemon_name]['types']
        for pokemon_type in pokemon_types:
            if pokemon_type not in available_pokemon_type_distribution:
                available_pokemon_type_distribution[pokemon_type] = 0
            available_pokemon_type_distribution[pokemon_type] += 1

    encountered_pokemon_types = {}
    for file_name in files_to_analyze:
        pickle_in = open(f"results/{file_name}",'rb')
        all_battle_results = pickle.load(pickle_in)
        pickle_in.close()
        for battle_result in all_battle_results:
            encountered_pokemon = battle_result['encountered_pokemon']
            encountered_pokemon_types_list = encountered_pokemon.types
            for pokemon_type in encountered_pokemon_types_list:
                if pokemon_type not in encountered_pokemon_types:
                    encountered_pokemon_types[pokemon_type] = 0
                encountered_pokemon_types[pokemon_type] += 1

    quantities = list(encountered_pokemon_types.values())
    labels = list(encountered_pokemon_types.keys())
    plt.figure(figsize=(12, 6))
    plt.subplot(1, 2, 1)
    plt.pie(quantities, labels=labels, autopct='%1.1f%%', startangle=140)
    plt.title('Distribution of Encountered Pokemon Types')
    plt.axis('equal')  # Equal aspect ratio ensures that pie is drawn as a circle

    plt.subplot(1, 2, 2)
    quantities_base = list(available_pokemon_type_distribution.values())
    labels_base = list(available_pokemon_type_distribution.keys())
    plt.pie(quantities_base, labels=labels_base, autopct='%1.1f%%', startangle=140)
    plt.title('Distribution of Available Pokemon Types')
    plt.axis('equal')

    plt.savefig('plots/pokemon_types_distribution.pdf')
    

if __name__ == "__main__":
    Ngames = 1000
    Nbattles = 150

    print("Generating data")
    for i in range(Ngames):
        #automatic player creation (choose a random starter pokemon)
        starter_choice = random.choice(gd.available_starters)
        print(f"Generating data for {starter_choice}, game {i}...")
        results = Main(starter_choice = starter_choice, Nbattles=Nbattles)
        pickle_out = open(f"results/{starter_choice}_game_{i}_{Nbattles}_battles.pickle",'wb')
        pickle.dump(results,pickle_out)
        pickle_out.close()
    print("Generating plots")
    make_plots(Ngames, Nbattles)
    print("Finished!")