backend/routers/population.py

from fastapi import APIRouter, Query, HTTPException
from fastapi.responses import JSONResponse
import pandas as pd
from ..utils.pandas_utils import load_data
from ..schemas.population import ContinentStatResponse

router = APIRouter()

df = load_data()

@router.get("/continent")
def get_continent_info():
    continents = df['Continent'].unique()
    num_continents = len(continents)
    return {"number_of_continents": num_continents, "continents": list(continents)}

@router.get("/continents", response_model=ContinentStatResponse)
def get_population_stat_by_continent(continent: str = Query(), data_type: str = Query(), stat: str = Query()):
    valid_stats = ['max', 'min', 'mean', 'sum', 'count']
    if stat not in valid_stats:
        raise HTTPException(status_code=400, detail=f"Invalid stat. Please use one of the following: {valid_stats}.")

    try:
        continent_population_stats = df.groupby('Continent')[data_type].agg(
            Maximum='max', Minimum='min', Average='mean', Total='sum', Number_of_Countries='count'
        )
        continent_countries = df[df['Continent'] == continent]
        
        if continent not in continent_population_stats.index:
            raise HTTPException(status_code=404, detail=f"Continent '{continent}' not found in the data.")

        if stat == 'max':
            population_result = continent_population_stats.loc[continent]['Maximum']
            country_id = continent_countries.loc[continent_countries[data_type].idxmax()]
            country_name = country_id['Country']
            population_value = country_id[data_type]
            return {"result": f"{int(population_result)}. Country: {country_name}, {data_type}: {population_value}"}

        if stat == 'min':
            population_result = continent_population_stats.loc[continent]['Minimum']
            country_id = continent_countries.loc[continent_countries[data_type].idxmin()]
            country_name = country_id['Country']
            population_value = country_id[data_type]
            return {"result": f"{int(population_result)}. Country: {country_name}, {data_type}: {population_value}"}

        if stat == 'mean':
            population_result = continent_population_stats.loc[continent]['Average']
            return {"result": int(population_result)}

        if stat == 'sum':
            population_result = continent_population_stats.loc[continent]['Total']
            return {"result": int(population_result)}

        if stat == 'count':
            population_result = continent_population_stats.loc[continent]['Number_of_Countries']
            return {"result": int(population_result)}

    except KeyError:
        raise HTTPException(status_code=400, detail=f"Data type '{data_type}' not found in the data.")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))