import pandas as pd

def load_data(pokemon_moves_file: str, move_class_file: str):
    # Load the CSV files
    pokemon_moves = pd.read_csv(pokemon_moves_file)
    move_class = pd.read_csv(move_class_file)
    return pokemon_moves, move_class

def get_pokemon_list(pokemon_moves: pd.DataFrame):
    # Return list of all Pokémon names
    return pokemon_moves['Pokemon'].dropna().tolist()

def get_moves_for_pokemon(pokemon_moves: pd.DataFrame, pokemon_name: str):
    # Get all non-NaN moves for the selected Pokémon
    row = pokemon_moves[pokemon_moves['Pokemon'] == pokemon_name]
    if row.empty:
        return []
    moves = row.iloc[0, 1:].dropna().tolist()
    return moves

def classify_moves(moves: list, move_class: pd.DataFrame):
    # Create a move -> class mapping
    move_class_dict = dict(zip(move_class['Move'], move_class['MoveClass']))

    # Separate into DEF and SPT categories
    def_moves = []
    spt_moves = []
    for move in moves:
        category = move_class_dict.get(move)
        if category == "DEF":
            def_moves.append(move)
        elif category == "SPT":
            spt_moves.append(move)
    return def_moves, spt_moves

def main():
    # File paths
    pokemon_moves_file = "pokemon_moves.csv"
    move_class_file = "Move_Class.csv"


    # Load data
    pokemon_moves, move_class = load_data(pokemon_moves_file, move_class_file)

    # Display available Pokémon
    pokemons = get_pokemon_list(pokemon_moves)
    print("Available Pokémon:")
    for idx, pkmn in enumerate(pokemons, start=1):
        print(f"{idx}. {pkmn}")

    # Ask user to pick a Pokémon
    choice = input("\nEnter the Pokémon name exactly as listed: ").strip()

    if choice not in pokemons:
        print(f"Error: '{choice}' is not in the list.")
        return

    # Get and classify moves
    moves = get_moves_for_pokemon(pokemon_moves, choice)
    def_moves, spt_moves = classify_moves(moves, move_class)

    # Display results
    print(f"\n{choice} - DEF Moves:")
    if def_moves:
        for move in def_moves:
            print(f"- {move}")
    else:
        print("No DEF moves found.")

    print(f"\n{choice} - SPT Moves:")
    if spt_moves:
        for move in spt_moves:
            print(f"- {move}")
    else:
        print("No SPT moves found.")

if __name__ == "__main__":
    main()