#!/usr/bin/env python3
"""
Data Acquisition Script for Air Pollution Data
Fetches historical pollution data from OpenAQ and World Bank APIs
for New Delhi, Mumbai, Bengaluru, and Hyderabad.
"""

import requests
import pandas as pd
import json
import time
from datetime import datetime, timedelta
from pathlib import Path
import sys

# Configuration
API_KEY = "5e9335f14620aa0dcb851537e1a57a6d201210fa06448a3dd0186e7816d5515a"
OUTPUT_DIR = Path("/app/sandbox/session_20251224_200424_4a6cc33ed0fd/workflow/raw_data")
RESULTS_DIR = Path("/app/sandbox/session_20251224_200424_4a6cc33ed0fd/results")

# Target cities with their coordinates (for OpenAQ API)
CITIES = {
    "New Delhi": {"lat": 28.6139, "lon": 77.2090, "country": "IN"},
    "Mumbai": {"lat": 19.0760, "lon": 72.8777, "country": "IN"},
    "Bengaluru": {"lat": 12.9716, "lon": 77.5946, "country": "IN"},
    "Hyderabad": {"lat": 17.3850, "lon": 78.4867, "country": "IN"}
}

# Target parameters
PARAMETERS = ["pm25", "pm10", "no2"]

# World Bank country code for India
WB_COUNTRY_CODE = "IND"

# World Bank indicators for pollution data
WB_INDICATORS = {
    "EN.ATM.PM25.MC.M3": "PM2.5 air pollution, mean annual exposure (micrograms per cubic meter)",
    "EN.ATM.PM25.MC.ZS": "PM2.5 air pollution, population exposed to levels exceeding WHO guideline value (% of total)"
}


def fetch_openaq_data(city_name, city_info, parameters, api_key):
    """
    Fetch air quality data from OpenAQ API for a specific city.

    Args:
        city_name: Name of the city
        city_info: Dictionary with lat, lon, country
        parameters: List of parameters to fetch (pm25, pm10, no2)
        api_key: OpenAQ API key

    Returns:
        DataFrame with the fetched data
    """
    print(f"\n{'='*60}")
    print(f"Fetching OpenAQ data for {city_name}...")
    print(f"{'='*60}")

    base_url = "https://api.openaq.org/v2/measurements"
    headers = {"X-API-Key": api_key}

    all_data = []

    for param in parameters:
        print(f"\nFetching {param.upper()} data for {city_name}...")
        page = 1
        total_fetched = 0

        while True:
            # OpenAQ API parameters
            params = {
                "country": city_info["country"],
                "parameter": param,
                "limit": 1000,  # Max per request
                "page": page,
                "coordinates": f"{city_info['lat']},{city_info['lon']}",
                "radius": 50000,  # 50km radius
                "date_from": "2015-01-01T00:00:00Z",  # Try to get as much history as possible
                "order_by": "datetime"
            }

            try:
                response = requests.get(base_url, headers=headers, params=params, timeout=30)
                response.raise_for_status()
                data = response.json()

                if "results" not in data or len(data["results"]) == 0:
                    print(f"  No more data for {param} (page {page})")
                    break

                results = data["results"]
                total_fetched += len(results)

                # Process results
                for measurement in results:
                    record = {
                        "city": city_name,
                        "parameter": param,
                        "value": measurement.get("value"),
                        "unit": measurement.get("unit"),
                        "date": measurement.get("date", {}).get("utc"),
                        "location": measurement.get("location"),
                        "coordinates": f"{measurement.get('coordinates', {}).get('latitude')},{measurement.get('coordinates', {}).get('longitude')}"
                    }
                    all_data.append(record)

                # Progress update
                if page % 5 == 0:
                    print(f"  Progress: Fetched {total_fetched} measurements for {param} (page {page})")

                # Check if we have more pages
                meta = data.get("meta", {})
                found = meta.get("found", 0)
                if total_fetched >= found:
                    print(f"  Completed: Fetched all {total_fetched} measurements for {param}")
                    break

                page += 1
                time.sleep(0.5)  # Rate limiting

            except requests.exceptions.RequestException as e:
                print(f"  Error fetching data for {param}, page {page}: {e}")
                break
            except Exception as e:
                print(f"  Unexpected error: {e}")
                break

    if all_data:
        df = pd.DataFrame(all_data)
        print(f"\n✓ Total measurements fetched for {city_name}: {len(df)}")
        return df
    else:
        print(f"\n⚠ No data fetched for {city_name}")
        return pd.DataFrame()


def fetch_world_bank_data(country_code, indicators):
    """
    Fetch historical pollution data from World Bank API.

    Args:
        country_code: World Bank country code (e.g., 'IND' for India)
        indicators: Dictionary of indicator codes and descriptions

    Returns:
        DataFrame with the fetched data
    """
    print(f"\n{'='*60}")
    print(f"Fetching World Bank data for {country_code}...")
    print(f"{'='*60}")

    all_data = []

    for indicator_code, description in indicators.items():
        print(f"\nFetching: {description}")

        # World Bank API endpoint
        url = f"https://api.worldbank.org/v2/country/{country_code}/indicator/{indicator_code}"
        params = {
            "format": "json",
            "per_page": 1000,
            "date": "1990:2024"  # Request all available years
        }

        try:
            response = requests.get(url, params=params, timeout=30)
            response.raise_for_status()
            data = response.json()

            if len(data) > 1 and data[1]:
                results = data[1]
                print(f"  Fetched {len(results)} yearly records")

                for record in results:
                    entry = {
                        "country": record.get("country", {}).get("value"),
                        "country_code": country_code,
                        "indicator": description,
                        "indicator_code": indicator_code,
                        "year": record.get("date"),
                        "value": record.get("value")
                    }
                    all_data.append(entry)
            else:
                print(f"  No data available for {indicator_code}")

        except requests.exceptions.RequestException as e:
            print(f"  Error fetching World Bank data for {indicator_code}: {e}")
        except Exception as e:
            print(f"  Unexpected error: {e}")

        time.sleep(0.3)  # Rate limiting

    if all_data:
        df = pd.DataFrame(all_data)
        print(f"\n✓ Total World Bank records fetched: {len(df)}")
        return df
    else:
        print(f"\n⚠ No World Bank data fetched")
        return pd.DataFrame()


def generate_summary(openaq_df, wb_df, output_file):
    """
    Generate a summary of the acquired data.

    Args:
        openaq_df: DataFrame with OpenAQ data
        wb_df: DataFrame with World Bank data
        output_file: Path to save the summary
    """
    print(f"\n{'='*60}")
    print("Generating Data Acquisition Summary...")
    print(f"{'='*60}")

    summary_lines = []
    summary_lines.append("=" * 80)
    summary_lines.append("DATA ACQUISITION SUMMARY")
    summary_lines.append(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    summary_lines.append("=" * 80)

    # OpenAQ Summary
    summary_lines.append("\n1. OPENAQ DATA")
    summary_lines.append("-" * 80)

    if not openaq_df.empty:
        summary_lines.append(f"Total Records: {len(openaq_df):,}")
        summary_lines.append(f"\nDate Range: {openaq_df['date'].min()} to {openaq_df['date'].max()}")

        summary_lines.append("\nRecords by City:")
        for city in CITIES.keys():
            city_data = openaq_df[openaq_df['city'] == city]
            summary_lines.append(f"  - {city}: {len(city_data):,} measurements")
            if not city_data.empty:
                summary_lines.append(f"    Date range: {city_data['date'].min()} to {city_data['date'].max()}")

        summary_lines.append("\nRecords by Parameter:")
        for param in openaq_df['parameter'].unique():
            param_data = openaq_df[openaq_df['parameter'] == param]
            summary_lines.append(f"  - {param.upper()}: {len(param_data):,} measurements")
    else:
        summary_lines.append("No OpenAQ data acquired.")

    # World Bank Summary
    summary_lines.append("\n\n2. WORLD BANK DATA")
    summary_lines.append("-" * 80)

    if not wb_df.empty:
        summary_lines.append(f"Total Records: {len(wb_df):,}")

        # Filter out None values for year range
        wb_df_clean = wb_df[wb_df['value'].notna()]
        if not wb_df_clean.empty:
            summary_lines.append(f"Year Range: {wb_df_clean['year'].min()} to {wb_df_clean['year'].max()}")

        summary_lines.append("\nRecords by Indicator:")
        for indicator in wb_df['indicator'].unique():
            ind_data = wb_df[wb_df['indicator'] == indicator]
            ind_data_clean = ind_data[ind_data['value'].notna()]
            summary_lines.append(f"  - {indicator}")
            summary_lines.append(f"    Total: {len(ind_data_clean):,} yearly records")
            if not ind_data_clean.empty:
                summary_lines.append(f"    Years: {ind_data_clean['year'].min()} - {ind_data_clean['year'].max()}")
    else:
        summary_lines.append("No World Bank data acquired.")

    # Overall assessment
    summary_lines.append("\n\n3. DATA COVERAGE ASSESSMENT")
    summary_lines.append("-" * 80)

    if not openaq_df.empty:
        openaq_start = pd.to_datetime(openaq_df['date']).min().year
        summary_lines.append(f"OpenAQ earliest data: {openaq_start}")

    if not wb_df.empty:
        wb_df_clean = wb_df[wb_df['value'].notna()]
        if not wb_df_clean.empty:
            wb_start = int(wb_df_clean['year'].min())
            summary_lines.append(f"World Bank earliest data: {wb_start}")

    summary_lines.append("\n" + "=" * 80)

    # Save summary
    summary_text = "\n".join(summary_lines)
    with open(output_file, 'w') as f:
        f.write(summary_text)

    print(summary_text)
    print(f"\n✓ Summary saved to: {output_file}")


def main():
    """Main execution function."""
    print("\n" + "=" * 80)
    print("AIR POLLUTION DATA ACQUISITION")
    print("=" * 80)
    print(f"Start time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

    # Ensure output directories exist
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    RESULTS_DIR.mkdir(parents=True, exist_ok=True)

    # Fetch OpenAQ data for all cities
    openaq_dfs = []
    for city_name, city_info in CITIES.items():
        df = fetch_openaq_data(city_name, city_info, PARAMETERS, API_KEY)
        if not df.empty:
            openaq_dfs.append(df)
        time.sleep(1)  # Pause between cities

    # Combine all OpenAQ data
    if openaq_dfs:
        openaq_combined = pd.concat(openaq_dfs, ignore_index=True)
        openaq_file = OUTPUT_DIR / "openaq_data.csv"
        openaq_combined.to_csv(openaq_file, index=False)
        print(f"\n✓ OpenAQ data saved to: {openaq_file}")
        print(f"  Total records: {len(openaq_combined):,}")
    else:
        openaq_combined = pd.DataFrame()
        print("\n⚠ No OpenAQ data to save")

    # Fetch World Bank data
    wb_df = fetch_world_bank_data(WB_COUNTRY_CODE, WB_INDICATORS)
    if not wb_df.empty:
        wb_file = OUTPUT_DIR / "historical_macro_data.csv"
        wb_df.to_csv(wb_file, index=False)
        print(f"\n✓ World Bank data saved to: {wb_file}")
        print(f"  Total records: {len(wb_df):,}")
    else:
        print("\n⚠ No World Bank data to save")

    # Generate summary
    summary_file = RESULTS_DIR / "data_acquisition_summary.txt"
    generate_summary(openaq_combined, wb_df, summary_file)

    print("\n" + "=" * 80)
    print("DATA ACQUISITION COMPLETE")
    print(f"End time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("=" * 80)


if __name__ == "__main__":
    main()