Code Editor for app.py

"""
Data source: https://pokemondb.net/pokedex/national
"""
import streamlit as st
import duckdb
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go

st.title('Pokémon by type')
st.write("Source: editor.ploomber.io/editor/pokemon-plotter-5a01")

conn = duckdb.connect(":memory:")

generation = st.number_input("Generation", min_value=1, max_value=9)

query = """
SELECT unnested_type, COUNT(*) * 100.0 / SUM(COUNT(*)) OVER () AS percentage
FROM (
    SELECT UNNEST(types) as unnested_type
    FROM read_json_auto('pokemon.jsonl')
    WHERE generation = ?
) AS unnested
GROUP BY unnested_type
ORDER BY percentage DESC;
"""
results = conn.execute(query, [generation]).fetchall()


df = pd.DataFrame(results, columns=['Type', 'Percentage'])
print(df["Type"])

# Create a bar chart with colors for each Pokémon type
color_map = {
    'Water': '#6390F0',    # Clean blue
    'Normal': '#A8A878',   # Beige/tan
    'Grass': '#7AC74C',    # Vibrant green
    'Flying': '#A98FF3',   # Light purple/blue
    'Psychic': '#F95587',  # Pink
    'Bug': '#A6B91A',      # Yellow-green
    'Poison': '#A33EA1',   # Purple
    'Fire': '#EE8130',     # Orange-red
    'Ground': '#E2BF65',   # Sandy brown
    'Rock': '#B6A136',     # Rocky brown
    'Fighting': '#C22E28',  # Dark red
    'Dragon': '#6F35FC',   # Deep purple-blue
    'Electric': '#F7D02C', # Yellow
    'Dark': '#705746',     # Dark gray-brown
    'Steel': '#B7B7CE',    # Metallic gray
    'Ghost': '#735797',    # Purple
    'Fairy': '#D685AD',    # Light pink
    'Ice': '#96D9D6'       # Light blue
}

# Map the colors to the DataFrame
df['Color'] = df['Type'].map(color_map)


# Create a bar chart with specified colors using Plotly
fig = px.bar(df, x="Type", y="Percentage", color="Type", 
             color_discrete_map=dict(zip(df['Type'], df['Color'])),
             title="Pokémon Types by Percentage")
st.plotly_chart(fig)


st.title('Compare types among generations')

generations = st.multiselect("Select Generations", options=list(range(1, 10)), default=[1, 2])
generations_ = ','.join(str(gen) for gen in generations)

query_comparison = f"""
SELECT unnested_type as type, generation, COUNT(*) * 100.0 / SUM(COUNT(*)) OVER (PARTITION BY generation) AS percentage
FROM (
    SELECT UNNEST(types) as unnested_type, generation
    FROM read_json_auto('pokemon.jsonl')
    WHERE generation IN ({generations_})
) AS unnested
GROUP BY unnested_type, generation
ORDER BY generation, percentage DESC;
"""
results_comparison = conn.execute(query_comparison).df().sort_values(by='percentage', ascending=False)

fig_comparison = go.Figure()

# Add traces for each generation with corresponding data
for gen in generations:
    data = results_comparison[results_comparison['generation'] == gen]
    fig_comparison.add_trace(go.Bar(
        x=data['type'],
        y=data['percentage'],
        name=f'Generation {gen}',
    ))
# Here we modify the layout to group the bars by generation
fig_comparison.update_layout(barmode='group', xaxis_tickangle=-45,
                              title="Pokémon Types by Percentage for Selected Generations")

st.plotly_chart(fig_comparison)