#!/usr/bin/env python3
"""
Alternative Data Acquisition Script
Fetches air pollution data from multiple sources including:
- World Bank (already working)
- WHO Air Quality Database
- AQICN (World Air Quality Index project)
- Generate synthetic/sample data if APIs unavailable
"""

import requests
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from pathlib import Path
import json

# 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")

CITIES = ["New Delhi", "Mumbai", "Bengaluru", "Hyderabad"]
WB_COUNTRY_CODE = "IND"

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)",
    "EN.ATM.NOXE.KT": "Nitrous oxide emissions (thousand metric tons of CO2 equivalent)",
    "EN.ATM.PM25.MC.T1": "PM2.5 pollution, total",
}


def fetch_world_bank_data(country_code, indicators):
    """Fetch historical pollution data from World Bank API."""
    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}")

        url = f"https://api.worldbank.org/v2/country/{country_code}/indicator/{indicator_code}"
        params = {
            "format": "json",
            "per_page": 1000,
            "date": "1990:2024"
        }

        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]
                valid_results = [r for r in results if r.get("value") is not None]
                print(f"  Fetched {len(valid_results)} yearly records with data")

                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 Exception as e:
            print(f"  Error: {e}")

        import time
        time.sleep(0.3)

    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 try_aqicn_api():
    """
    Try to fetch data from AQICN (Air Quality Index China/World) API.
    This is a backup data source.
    """
    print(f"\n{'='*60}")
    print("Attempting to fetch data from AQICN API...")
    print(f"{'='*60}")

    # AQICN API token (using a demo token - users should get their own)
    # Note: This uses the free demo token which has limitations
    token = "demo"  # This is the public demo token from aqicn.org
    base_url = "https://api.waqi.info"

    all_data = []

    city_queries = {
        "New Delhi": "delhi",
        "Mumbai": "mumbai",
        "Bengaluru": "bangalore",
        "Hyderabad": "hyderabad"
    }

    for city_name, city_query in city_queries.items():
        print(f"\nFetching data for {city_name}...")

        url = f"{base_url}/feed/{city_query}/"
        params = {"token": token}

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

            if data.get("status") == "ok" and "data" in data:
                result = data["data"]

                # Extract current readings
                record = {
                    "city": city_name,
                    "station": result.get("city", {}).get("name"),
                    "aqi": result.get("aqi"),
                    "date": result.get("time", {}).get("s"),
                    "pm25": result.get("iaqi", {}).get("pm25", {}).get("v"),
                    "pm10": result.get("iaqi", {}).get("pm10", {}).get("v"),
                    "no2": result.get("iaqi", {}).get("no2", {}).get("v"),
                    "latitude": result.get("city", {}).get("geo", [])[0] if result.get("city", {}).get("geo") else None,
                    "longitude": result.get("city", {}).get("geo", [])[1] if len(result.get("city", {}).get("geo", [])) > 1 else None,
                }
                all_data.append(record)
                print(f"  ✓ Retrieved current AQI data: {record['aqi']}")
            else:
                print(f"  ⚠ No data available (status: {data.get('status')})")

        except Exception as e:
            print(f"  Error: {e}")

        import time
        time.sleep(1)

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


def generate_synthetic_historical_data():
    """
    Generate synthetic historical air quality data for demonstration purposes.
    This is used as a fallback when API data is unavailable.
    Based on known pollution patterns in Indian cities.
    """
    print(f"\n{'='*60}")
    print("Generating synthetic historical data for demonstration...")
    print(f"{'='*60}")

    np.random.seed(42)

    # Generate date range from 2015 to 2024
    start_date = datetime(2015, 1, 1)
    end_date = datetime(2024, 12, 31)
    dates = pd.date_range(start=start_date, end=end_date, freq='D')

    all_data = []

    # Baseline pollution levels (realistic for these cities)
    city_baselines = {
        "New Delhi": {"pm25": 150, "pm10": 250, "no2": 60},
        "Mumbai": {"pm25": 80, "pm10": 120, "no2": 45},
        "Bengaluru": {"pm25": 60, "pm10": 90, "no2": 35},
        "Hyderabad": {"pm25": 70, "pm10": 100, "no2": 40}
    }

    for city, baselines in city_baselines.items():
        print(f"Generating data for {city}...")

        for date in dates:
            # Add seasonal variation (higher in winter)
            month = date.month
            seasonal_factor = 1.4 if month in [11, 12, 1, 2] else 0.8

            # Add weekly variation (lower on weekends)
            weekly_factor = 0.85 if date.weekday() >= 5 else 1.0

            # Add random variation
            for param, baseline in baselines.items():
                value = baseline * seasonal_factor * weekly_factor * np.random.uniform(0.7, 1.3)

                # Add occasional spikes (pollution events)
                if np.random.random() < 0.05:
                    value *= np.random.uniform(1.5, 2.5)

                record = {
                    "city": city,
                    "parameter": param,
                    "value": round(value, 2),
                    "unit": "µg/m³" if param != "no2" else "ppb",
                    "date": date.strftime("%Y-%m-%d"),
                    "location": f"{city} Monitoring Station",
                    "data_source": "synthetic"
                }
                all_data.append(record)

        print(f"  ✓ Generated {len(dates) * len(baselines)} records")

    df = pd.DataFrame(all_data)
    print(f"\n✓ Total synthetic records generated: {len(df):,}")
    return df


def generate_summary(openaq_df, wb_df, aqicn_df, synthetic_df, output_file):
    """Generate a comprehensive summary of acquired data."""
    print(f"\n{'='*60}")
    print("Generating Data Acquisition Summary...")
    print(f"{'='*60}")

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

    # OpenAQ/Real-time data summary
    summary_lines.append("\n1. REAL-TIME DATA SOURCES")
    summary_lines.append("-" * 80)

    if not openaq_df.empty:
        summary_lines.append(f"OpenAQ Data: {len(openaq_df):,} records")
        summary_lines.append(f"Date range: {openaq_df['date'].min()} to {openaq_df['date'].max()}")
    else:
        summary_lines.append("OpenAQ Data: Not available (API deprecated/inaccessible)")

    if not aqicn_df.empty:
        summary_lines.append(f"\nAQICN Current Data: {len(aqicn_df):,} records")
        summary_lines.append("Cities covered:")
        for city in aqicn_df['city'].unique():
            city_data = aqicn_df[aqicn_df['city'] == city]
            aqi = city_data.iloc[0]['aqi'] if 'aqi' in city_data.columns else 'N/A'
            summary_lines.append(f"  - {city}: AQI = {aqi}")
    else:
        summary_lines.append("\nAQICN Data: Not available")

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

    if not wb_df.empty:
        summary_lines.append(f"Total Records: {len(wb_df):,}")
        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("\nIndicators:")
        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"    Records: {len(ind_data_clean):,}")
            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 available")

    # Synthetic data summary
    summary_lines.append("\n\n3. SYNTHETIC/DEMONSTRATION DATA")
    summary_lines.append("-" * 80)

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

        summary_lines.append("\nNote: Synthetic data generated based on typical pollution patterns")
        summary_lines.append("for demonstration purposes. Real historical data from World Bank.")
    else:
        summary_lines.append("No synthetic data generated")

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

    if not wb_df.empty:
        wb_df_clean = wb_df[wb_df['value'].notna()]
        if not wb_df_clean.empty:
            summary_lines.append(f"Historical coverage (World Bank): {wb_df_clean['year'].min()} - {wb_df_clean['year'].max()}")

    if not synthetic_df.empty:
        summary_lines.append(f"Detailed synthetic data: 2015 - 2024")

    summary_lines.append("\nData Sources Status:")
    summary_lines.append(f"  ✓ World Bank API: Available")
    summary_lines.append(f"  ✗ OpenAQ API: Unavailable (v2 deprecated, v3 not accessible)")
    summary_lines.append(f"  {'✓' if not aqicn_df.empty else '~'} AQICN API: {'Available (current data only)' if not aqicn_df.empty else 'Limited availability'}")
    summary_lines.append(f"  ✓ Synthetic data: Generated for demonstration")

    summary_lines.append("\n\nRECOMMENDATIONS:")
    summary_lines.append("-" * 80)
    summary_lines.append("1. For real-time data: Consider using CPCB India official data portal")
    summary_lines.append("2. For historical data: World Bank provides country-level annual averages")
    summary_lines.append("3. For city-specific historical data: Manual download from OpenAQ data explorer")
    summary_lines.append("4. Alternative APIs: IQAir, AQICN (requires API key)")

    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("MULTI-SOURCE 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)

    # 1. Fetch World Bank data (historical, country-level)
    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}")

    # 2. Try AQICN API (current data)
    aqicn_df = try_aqicn_api()
    if not aqicn_df.empty:
        aqicn_file = OUTPUT_DIR / "aqicn_current_data.csv"
        aqicn_df.to_csv(aqicn_file, index=False)
        print(f"\n✓ AQICN data saved to: {aqicn_file}")

    # 3. Generate synthetic historical data for demonstration
    synthetic_df = generate_synthetic_historical_data()
    if not synthetic_df.empty:
        synthetic_file = OUTPUT_DIR / "synthetic_daily_data.csv"
        synthetic_df.to_csv(synthetic_file, index=False)
        print(f"\n✓ Synthetic data saved to: {synthetic_file}")

    # Placeholder for OpenAQ (currently unavailable)
    openaq_df = pd.DataFrame()

    # Generate comprehensive summary
    summary_file = RESULTS_DIR / "data_acquisition_summary.txt"
    generate_summary(openaq_df, wb_df, aqicn_df, synthetic_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)
    print("\nData files created:")
    print(f"  1. {OUTPUT_DIR / 'historical_macro_data.csv'} (World Bank)")
    print(f"  2. {OUTPUT_DIR / 'aqicn_current_data.csv'} (AQICN current)")
    print(f"  3. {OUTPUT_DIR / 'synthetic_daily_data.csv'} (Synthetic historical)")
    print(f"\nSummary: {RESULTS_DIR / 'data_acquisition_summary.txt'}")


if __name__ == "__main__":
    main()