#!/usr/bin/env python3
"""
Grant Landscape Mapping - CDR Funding Search
Searches for significant US-based grants in Carbon Dioxide Removal (CDR)
"""

import json
import requests
import pandas as pd
from datetime import datetime, timedelta
from typing import List, Dict
import time
import sys

# Configuration
KEYWORDS = [
    "Direct Air Capture",
    "DAC",
    "Enhanced Weathering",
    "Ocean Alkalinity",
    "Carbon Dioxide Removal"
]
MIN_AMOUNT = 500000  # $500,000
MONTHS_BACK = 24
OUTPUT_DIR = "/app/sandbox/session_20251212_164706_1df6aaa34e0b/workflow"

def search_nsf_awards(keywords: List[str], start_date: str) -> List[Dict]:
    """
    Search NSF Awards API for grants matching CDR keywords.
    API Documentation: https://www.research.gov/common/webapi/awardapisearch-v1.htm
    """
    print(f"\n🔍 Searching NSF Awards API...")
    base_url = "https://www.research.gov/awardapi-service/v1/awards.json"

    all_awards = []

    for i, keyword in enumerate(keywords, 1):
        print(f"  [{i}/{len(keywords)}] Searching for: '{keyword}'")

        params = {
            'keyword': keyword,
            'dateStart': start_date.replace('-', '/'),
            'printFields': 'id,agency,title,startDate,expDate,fundsObligatedAmt,piFirstName,piLastName,piEmail,abstractText,awardee,awardeeCity,awardeeStateCode,pdPIName'
        }

        try:
            response = requests.get(base_url, params=params, timeout=30)
            if response.status_code == 200:
                data = response.json()
                awards = data.get('response', {}).get('award', [])

                if isinstance(awards, list):
                    all_awards.extend(awards)
                    print(f"    ✓ Found {len(awards)} awards for '{keyword}'")
                elif isinstance(awards, dict):
                    all_awards.append(awards)
                    print(f"    ✓ Found 1 award for '{keyword}'")
                else:
                    print(f"    ℹ No awards found for '{keyword}'")
            else:
                print(f"    ⚠ API returned status {response.status_code} for '{keyword}'")

        except Exception as e:
            print(f"    ⚠ Error searching '{keyword}': {str(e)}")

        time.sleep(0.5)  # Rate limiting

    print(f"\n  Total NSF awards found: {len(all_awards)}")
    return all_awards


def search_doe_sbir_data(keywords: List[str]) -> List[Dict]:
    """
    Search for DOE SBIR/STTR and other DOE funding using available data sources.
    This is a placeholder that returns mock data structure - in production,
    this would query actual DOE databases or scrape DOE websites.
    """
    print(f"\n🔍 Searching DOE and SBIR/STTR databases...")
    print("  ℹ Note: Direct API access requires specific credentials.")
    print("  Using web search to identify recent CDR grants...")

    # This would be replaced with actual API calls or web scraping
    # For now, return empty list to be populated by web search
    return []


def search_arpa_e_data(keywords: List[str]) -> List[Dict]:
    """
    Search ARPA-E project database.
    ARPA-E publishes project data on their website.
    """
    print(f"\n🔍 Searching ARPA-E database...")
    print("  ℹ ARPA-E project data available at arpa-e.energy.gov")

    # This would query ARPA-E database or scrape their website
    # For now, return empty list to be populated by web search
    return []


def filter_grants(grants: List[Dict], min_amount: float) -> List[Dict]:
    """Filter grants by minimum funding amount."""
    filtered = []

    for grant in grants:
        # Handle different amount field names
        amount = None
        if 'fundsObligatedAmt' in grant:
            amount = float(grant.get('fundsObligatedAmt', 0))
        elif 'amount' in grant:
            amount = float(grant.get('amount', 0))
        elif 'awardAmount' in grant:
            amount = float(grant.get('awardAmount', 0))

        if amount and amount >= min_amount:
            filtered.append(grant)

    return filtered


def normalize_grant_data(grants: List[Dict]) -> pd.DataFrame:
    """
    Normalize grant data from different sources into a unified structure.
    Target columns: Awardee, Project Title, Abstract, Amount, Award Date, PI, Funding Agency
    """
    normalized_records = []

    for grant in grants:
        record = {
            'Awardee': None,
            'Project Title': None,
            'Abstract': None,
            'Amount': None,
            'Award Date': None,
            'Principal Investigator (PI)': None,
            'Funding Agency': None
        }

        # NSF format
        if 'agency' in grant:
            record['Funding Agency'] = grant.get('agency', 'NSF')
            record['Awardee'] = grant.get('awardee', '') or grant.get('institution', '')
            record['Project Title'] = grant.get('title', '')
            record['Abstract'] = grant.get('abstractText', '')
            record['Amount'] = float(grant.get('fundsObligatedAmt', 0))
            record['Award Date'] = grant.get('startDate', '')

            # Construct PI name
            pi_first = grant.get('piFirstName', '')
            pi_last = grant.get('piLastName', '')
            pd_pi_name = grant.get('pdPIName', '')

            if pd_pi_name:
                record['Principal Investigator (PI)'] = pd_pi_name
            elif pi_first and pi_last:
                record['Principal Investigator (PI)'] = f"{pi_first} {pi_last}"
            else:
                record['Principal Investigator (PI)'] = pi_first or pi_last

        # DOE/SBIR format (to be implemented)
        elif 'program' in grant:
            record['Funding Agency'] = grant.get('agency', 'DOE')
            record['Awardee'] = grant.get('company', '') or grant.get('recipient', '')
            record['Project Title'] = grant.get('projectTitle', '') or grant.get('title', '')
            record['Abstract'] = grant.get('abstract', '') or grant.get('description', '')
            record['Amount'] = float(grant.get('awardAmount', 0) or grant.get('amount', 0))
            record['Award Date'] = grant.get('awardDate', '') or grant.get('date', '')
            record['Principal Investigator (PI)'] = grant.get('pi', '') or grant.get('principalInvestigator', '')

        # Generic format
        else:
            record['Funding Agency'] = grant.get('agency', 'Unknown')
            record['Awardee'] = grant.get('awardee', '') or grant.get('recipient', '') or grant.get('organization', '')
            record['Project Title'] = grant.get('title', '')
            record['Abstract'] = grant.get('abstract', '') or grant.get('description', '')
            record['Amount'] = float(grant.get('amount', 0))
            record['Award Date'] = grant.get('date', '') or grant.get('awardDate', '')
            record['Principal Investigator (PI)'] = grant.get('pi', '')

        normalized_records.append(record)

    df = pd.DataFrame(normalized_records)

    # Remove duplicates based on title and awardee
    df = df.drop_duplicates(subset=['Project Title', 'Awardee'], keep='first')

    # Sort by amount descending
    df = df.sort_values('Amount', ascending=False)

    return df


def main():
    """Main execution function."""
    print("=" * 70)
    print("CDR GRANT LANDSCAPE MAPPING")
    print("=" * 70)
    print(f"\nSearch Criteria:")
    print(f"  Keywords: {', '.join(KEYWORDS)}")
    print(f"  Minimum Amount: ${MIN_AMOUNT:,}")
    print(f"  Timeframe: Last {MONTHS_BACK} months")

    # Calculate date range
    end_date = datetime.now()
    start_date = end_date - timedelta(days=MONTHS_BACK * 30)
    start_date_str = start_date.strftime('%m-%d-%Y')

    print(f"  Date Range: {start_date_str} to {end_date.strftime('%m-%d-%Y')}")

    # Search multiple databases
    all_grants = []

    # 1. Search NSF
    nsf_grants = search_nsf_awards(KEYWORDS, start_date_str)
    all_grants.extend(nsf_grants)

    # 2. Search DOE/SBIR (placeholder for now)
    doe_grants = search_doe_sbir_data(KEYWORDS)
    all_grants.extend(doe_grants)

    # 3. Search ARPA-E (placeholder for now)
    arpa_grants = search_arpa_e_data(KEYWORDS)
    all_grants.extend(arpa_grants)

    print(f"\n📊 Total grants retrieved: {len(all_grants)}")

    # Filter by minimum amount
    print(f"\n🔍 Filtering grants with amount >= ${MIN_AMOUNT:,}...")
    filtered_grants = filter_grants(all_grants, MIN_AMOUNT)
    print(f"  ✓ {len(filtered_grants)} grants meet criteria")

    # Save raw results
    raw_output_path = f"{OUTPUT_DIR}/raw_grants.json"
    print(f"\n💾 Saving raw results to: {raw_output_path}")
    with open(raw_output_path, 'w') as f:
        json.dump({
            'search_criteria': {
                'keywords': KEYWORDS,
                'min_amount': MIN_AMOUNT,
                'months_back': MONTHS_BACK,
                'start_date': start_date_str,
                'end_date': end_date.strftime('%m-%d-%Y')
            },
            'total_grants_found': len(all_grants),
            'grants_after_filtering': len(filtered_grants),
            'grants': filtered_grants
        }, f, indent=2)
    print(f"  ✓ Saved {len(filtered_grants)} grants to raw_grants.json")

    # Process into structured CSV
    print(f"\n📋 Processing grants into structured CSV...")
    df = normalize_grant_data(filtered_grants)

    csv_output_path = f"{OUTPUT_DIR}/grants_cleaned.csv"
    df.to_csv(csv_output_path, index=False)
    print(f"  ✓ Saved cleaned data to: grants_cleaned.csv")
    print(f"  ✓ Columns: {', '.join(df.columns.tolist())}")
    print(f"  ✓ Total records: {len(df)}")

    # Display summary statistics
    print(f"\n📈 Summary Statistics:")
    print(f"  Total Funding: ${df['Amount'].sum():,.2f}")
    print(f"  Average Award: ${df['Amount'].mean():,.2f}")
    print(f"  Median Award: ${df['Amount'].median():,.2f}")
    print(f"  Largest Award: ${df['Amount'].max():,.2f}")
    print(f"  Smallest Award: ${df['Amount'].min():,.2f}")

    if not df.empty:
        print(f"\n🏆 Top 5 Awards by Amount:")
        top_5 = df.nlargest(5, 'Amount')[['Awardee', 'Project Title', 'Amount', 'Funding Agency']]
        for idx, row in top_5.iterrows():
            print(f"  {idx+1}. {row['Awardee']}: ${row['Amount']:,.2f} ({row['Funding Agency']})")
            print(f"     {row['Project Title'][:80]}...")

    print(f"\n✅ Grant landscape mapping completed!")
    print(f"   Raw data: {raw_output_path}")
    print(f"   Cleaned data: {csv_output_path}")
    print("=" * 70)


if __name__ == "__main__":
    main()