import pandas as pd

def process_pokemon_data(species, move1, move2, move3, move4, nature, instance=0, entry=None, item=None, speed_50=None):
    # Load CSV files
    nature_df = pd.read_csv("Nature_Table.csv")
    move_class_df = pd.read_csv("Move_Class.csv")

    # Determine Name
    name = f"{species} {instance}" if instance is not None else f"{species} 0"

    # Fetch nature attributes
    nature_row = nature_df[nature_df["Nature"] == nature].iloc[0]
    attack, defense, support = nature_row[["Attack", "Defense", "Support"]]
    attack_50, defense_50, support_50 = nature_row[["Attack_50", "Defense_50", "Support_50"]]

    # Get move classes
    moves = [move1, move2, move3, move4]
    move_classes = []

    for move in moves:
        move_class_row = move_class_df[move_class_df["Move"] == move]
        move_class = move_class_row["MoveClass"].values[0] if not move_class_row.empty else "None"
        move_classes.append(f"{move} {move_class}")

    # Count occurrences of move classes
    atk_count = sum("ATK" in mc for mc in move_classes)
    def_count = sum("DEF" in mc for mc in move_classes)
    spt_count = sum("SPT" in mc for mc in move_classes)

    contains_atk, contains_def, contains_spt = int(atk_count > 0), int(def_count > 0), int(spt_count > 0)
    num_classes = contains_atk + contains_def + contains_spt

    # Initialize random move probabilities
    r_atk = r_def = nomove = 0
    r_atk_50 = r_def_50 = nomove_50 = 0

    if num_classes == 2:
        if contains_spt == 0:  # ATK & DEF
            if atk_count > def_count:
                r_atk, r_def, nomove = support * 0.5, 0, support * 0.5
            elif atk_count < def_count:
                r_atk, r_def, nomove = 0, support * 0.5, support * 0.5
            else:
                r_atk, r_def, nomove = support * 0.25, support * 0.25, support * 0.5
        elif contains_def == 0:  # ATK & SPT
            r_atk, r_def, nomove = defense * 0.5, 0, defense * 0.5
        elif contains_atk == 0:  # DEF & SPT
            r_atk, r_def, nomove = 0, attack * 0.5, attack * 0.5

        # Repeat for level 50 stats
        if contains_spt == 0:
            if atk_count > def_count:
                r_atk_50, r_def_50, nomove_50 = support_50 * 0.5, 0, support_50 * 0.5
            elif atk_count < def_count:
                r_atk_50, r_def_50, nomove_50 = 0, support_50 * 0.5, support_50 * 0.5
            else:
                r_atk_50, r_def_50, nomove_50 = support_50 * 0.25, support_50 * 0.25, support_50 * 0.5
        elif contains_def == 0:
            r_atk_50, r_def_50, nomove_50 = defense_50 * 0.5, 0, defense_50 * 0.5
        elif contains_atk == 0:
            r_atk_50, r_def_50, nomove_50 = 0, attack_50 * 0.5, attack_50 * 0.5

    elif num_classes == 1:
        if contains_atk:
            r_atk, r_def, nomove = defense * 0.5 + support * 0.5, 0, defense * 0.5 + support * 0.5
        elif contains_def:
            r_atk, r_def, nomove = 0, attack * 0.5 + support * 0.5, attack * 0.5 + support * 0.5
        else:
            r_atk, r_def, nomove = 0, 0, attack + defense

        # Repeat for level 50 stats
        if contains_atk:
            r_atk_50, r_def_50, nomove_50 = defense_50 * 0.5 + support_50 * 0.5, 0, defense_50 * 0.5 + support_50 * 0.5
        elif contains_def:
            r_atk_50, r_def_50, nomove_50 = 0, attack_50 * 0.5 + support_50 * 0.5, attack_50 * 0.5 + support_50 * 0.5
        else:
            r_atk_50, r_def_50, nomove_50 = 0, 0, attack_50 + defense_50

    # AI strategy adjustments
    ai_atk = attack if atk_count > 0 else 0
    ai_def = defense if def_count > 0 else 0
    ai_spt = support if spt_count > 0 else 0
    ai_atk_50 = attack_50 if atk_count > 0 else 0
    ai_def_50 = defense_50 if def_count > 0 else 0
    ai_spt_50 = support_50 if spt_count > 0 else 0

    # Ensure no NaN values by replacing None with 0
    outputs = {
        "Entry": entry or 0,
        "Name": name,
        "Species": species,
        "Instance": instance or 0,
        "Nature": nature,
        "Item": item or 0,
        "Move 1": move1, "Move 2": move2, "Move 3": move3, "Move 4": move4,
        "Lv 50 Speed": speed_50 or 0,
        "Move 1 Class": move_classes[0],
        "Move 2 Class": move_classes[1],
        "Move 3 Class": move_classes[2],
        "Move 4 Class": move_classes[3],
        "ATK_COUNT": atk_count,
        "DEF_COUNT": def_count,
        "SPT_COUNT": spt_count,
        "NumClasses": num_classes,
        "R_ATK": r_atk, "R_DEF": r_def, "NOMOVE": nomove,
        "R_ATK_50": r_atk_50, "R_DEF_50": r_def_50, "NOMOVE_50": nomove_50,
        "AI_ATK": ai_atk, "AI_DEF": ai_def, "AI_SPT": ai_spt,
        "AI_ATK_50": ai_atk_50, "AI_DEF_50": ai_def_50, "AI_SPT_50": ai_spt_50
    }

    print(f"Added {species} with {nature} Nature")

    return outputs



# Define the correct column order
columns_order = [
    "Entry", "Name", "Species", "Instance", "Nature", "Item",
    "Move 1", "Move 2", "Move 3", "Move 4", "Move 1 Class", "Move 2 Class", "Move 3 Class", "Move 4 Class",
    "Possible Ability", "Lv 100 speed", "Lv 50 Speed",
    "AI_ATK", "AI_DEF", "AI_SPT", "R_ATK", "R_DEF", "NOMOVE",
    "AI_ATK_50", "AI_DEF_50", "AI_SPT_50", "R_ATK_50", "R_DEF_50", "NOMOVE_50"
]

# Example list of processed Pokémon data
pokemon_data = [
    process_pokemon_data('Metagross', 'Meteor Mash', 'Earthquake', 'Aerial Ace', 'Shadow Ball', nature = 'Sassy', item = 'Choice Band', speed_50 =106),
    process_pokemon_data('Latios', 'Thunderbolt', 'Psychic', 'Ice Beam', 'Calm Mind', nature = 'Hasty', item = 'Lum Berry', speed_50 =173),
    process_pokemon_data('Swampert', 'Earthquake', 'Rock Slide', 'Curse', 'Rest', nature = 'Brave', item = 'Leftovers', speed_50 =91)
]

# Convert to DataFrame
df = pd.DataFrame(pokemon_data)

# Ensure the DataFrame has the correct column order
df = df.reindex(columns=columns_order)

# Save to CSV
output_csv_path = "mypokemon.csv"
df.to_csv(output_csv_path, index=False)

print(f"CSV file saved to {output_csv_path}")
