#!/usr/bin/env python3
"""
Aggregate outpatient and inpatient case data to 100m grid cells.

This script:
1. Loads geocoded case data (outpatient + inpatient)
2. Performs spatial join to map each case to its containing grid cell
3. Computes daily aggregates per grid (outpatient_count, inpatient_count)
4. Merges with population data from grid index
5. Computes incidence_rate = total_cases / population
6. Outputs parquet with all grids (including zero-case grids)
"""

import pandas as pd
import geopandas as gpd
from shapely import wkt
import pyarrow as pa
import pyarrow.parquet as pq
from pathlib import Path
import sys

# Paths
PROJECT_ROOT = Path(__file__).parent.parent
CASES_FILE = PROJECT_ROOT / "outputs" / "geocoded_all_cases.csv"
GRID_FILE = PROJECT_ROOT / "processed" / "grid_100m_index.parquet"
OUTPUT_FILE = PROJECT_ROOT / "processed" / "grid_cases_daily.parquet"


def load_cases():
    """Load geocoded case data."""
    print(f"Loading cases from {CASES_FILE}...")
    cases = pd.read_csv(CASES_FILE)
    
    # Filter to valid coordinates
    valid_coords = cases[['latitude', 'longitude']].notnull().all(axis=1)
    cases_valid = cases[valid_coords].copy()
    
    print(f"  Total cases: {len(cases)}")
    print(f"  Cases with valid coordinates: {len(cases_valid)}")
    print(f"  Cases dropped (no coords): {len(cases) - len(cases_valid)}")
    
    # Convert date to datetime
    cases_valid['date'] = pd.to_datetime(cases_valid['date'])
    
    return cases_valid


def load_grid():
    """Load grid index with polygons."""
    print(f"Loading grid from {GRID_FILE}...")
    grid = pd.read_parquet(GRID_FILE)
    
    # Convert WKT strings to shapely geometries
    grid['geometry'] = grid['polygon'].apply(wkt.loads)
    grid_gdf = gpd.GeoDataFrame(grid, geometry='geometry', crs='EPSG:4326')
    
    print(f"  Grid cells: {len(grid_gdf)}")
    return grid_gdf


def spatial_join(cases_gdf, grid_gdf):
    """Perform spatial join to find containing grid for each case."""
    print("Performing spatial join (cases to grids)...")
    
    # Spatial join: find which grid contains each case point
    joined = gpd.sjoin(cases_gdf, grid_gdf[['grid_id', 'geometry', 'center_lon', 'center_lat', 'row', 'col']], 
                       how='left', predicate='within')
    
    print(f"  Cases matched to grids: {joined['grid_id'].notnull().sum()}")
    print(f"  Cases outside grid: {joined['grid_id'].isnull().sum()}")
    
    return joined


def aggregate_cases(joined):
    """Aggregate cases by grid_id and date."""
    print("Aggregating cases by grid and date...")
    
    # Separate by case type
    outpatient = joined[joined['case_type'] == 'outpatient'].copy()
    inpatient = joined[joined['case_type'] == 'inpatient'].copy()
    
    # Aggregate outpatient
    outpatient_agg = outpatient.groupby(['grid_id', 'date']).size().reset_index(name='outpatient_count')
    
    # Aggregate inpatient
    inpatient_agg = inpatient.groupby(['grid_id', 'date']).size().reset_index(name='inpatient_count')
    
    # Full outer join to get all grid-date combinations
    aggregated = outpatient_agg.merge(inpatient_agg, on=['grid_id', 'date'], how='outer')
    
    # Fill NaN with 0
    aggregated['outpatient_count'] = aggregated['outpatient_count'].fillna(0).astype(int)
    aggregated['inpatient_count'] = aggregated['inpatient_count'].fillna(0).astype(int)
    aggregated['total_cases'] = aggregated['outpatient_count'] + aggregated['inpatient_count']
    
    print(f"  Unique grid-date combinations with cases: {len(aggregated)}")
    
    return aggregated


def create_full_grid_date_index(grid_gdf, aggregated):
    """Create complete grid x date index including zero-case grids."""
    print("Creating full grid x date index...")
    
    # Get date range (2022-2024 matching weather data)
    date_min = pd.Timestamp('2022-01-01')
    date_max = pd.Timestamp('2024-12-31')
    all_dates = pd.date_range(start=date_min, end=date_max, freq='D')
    
    print(f"  Date range: {date_min.date()} to {date_max.date()} ({len(all_dates)} days)")
    
    # Create all grid x date combinations
    grid_ids = grid_gdf['grid_id'].tolist()
    
    # Create multiindex
    full_index = pd.MultiIndex.from_product(
        [grid_ids, all_dates],
        names=['grid_id', 'date']
    )
    full_df = pd.DataFrame(index=full_index).reset_index()
    
    print(f"  Total grid-date combinations: {len(full_df):,}")
    
    # Merge with aggregated data
    result = full_df.merge(aggregated, on=['grid_id', 'date'], how='left')
    
    # Fill NaN with 0 (grids with no cases on that date)
    result['outpatient_count'] = result['outpatient_count'].fillna(0).astype(int)
    result['inpatient_count'] = result['inpatient_count'].fillna(0).astype(int)
    result['total_cases'] = result['total_cases'].fillna(0).astype(int)
    
    print(f"  Grids with at least one case (any date): {result[result['total_cases'] > 0]['grid_id'].nunique()}")
    print(f"  Grids with zero cases (all dates): {result[result['total_cases'] == 0]['grid_id'].nunique()}")
    
    return result


def add_population_and_incidence(result, grid_gdf):
    """Add population data and compute incidence rate."""
    print("Adding population data and computing incidence rate...")
    
    # For now, we don't have population in grid index
    # We'll need to add it from landscan data
    # For this script, we'll set population to 0 as placeholder
    # TODO: Integrate landscan population data
    
    # Extract population from grid if available
    if 'population' in grid_gdf.columns:
        pop_map = grid_gdf[['grid_id', 'population']].set_index('grid_id')['population']
        result['population'] = result['grid_id'].map(pop_map).fillna(0)
    else:
        print("  WARNING: No population column in grid index. Setting population=0 (placeholder)")
        result['population'] = 0
    
    # Compute incidence rate (cases per capita)
    # Avoid division by zero
    result['incidence_rate'] = result.apply(
        lambda row: row['total_cases'] / row['population'] if row['population'] > 0 else 0.0,
        axis=1
    )
    
    return result


def save_output(result, output_file):
    """Save to parquet format."""
    print(f"Saving to {output_file}...")
    
    # Ensure output directory exists
    output_file.parent.mkdir(parents=True, exist_ok=True)
    
    # Convert date to string for parquet compatibility
    result['date'] = result['date'].dt.strftime('%Y-%m-%d')
    
    # Select and order columns
    output_cols = ['grid_id', 'date', 'outpatient_count', 'inpatient_count', 
                   'total_cases', 'population', 'incidence_rate']
    
    result[output_cols].to_parquet(output_file, index=False)
    
    file_size_mb = output_file.stat().st_size / (1024 * 1024)
    print(f"  Saved {len(result):,} rows ({file_size_mb:.1f} MB)")


def main():
    """Main pipeline."""
    print("=" * 60)
    print("Grid Case Aggregation Pipeline")
    print("=" * 60)
    
    # Load data
    cases = load_cases()
    grid = load_grid()
    
    # Convert cases to GeoDataFrame
    print("Converting cases to GeoDataFrame...")
    cases_gdf = gpd.GeoDataFrame(
        cases,
        geometry=gpd.points_from_xy(cases['longitude'], cases['latitude']),
        crs='EPSG:4326'
    )
    
    # Spatial join
    joined = spatial_join(cases_gdf, grid)
    
    # Aggregate
    aggregated = aggregate_cases(joined)
    
    # Create full index
    result = create_full_grid_date_index(grid, aggregated)
    
    # Add population and incidence
    result = add_population_and_incidence(result, grid)
    
    # Save
    save_output(result, OUTPUT_FILE)
    
    print("=" * 60)
    print("Pipeline complete!")
    print(f"Output: {OUTPUT_FILE}")
    print("=" * 60)


if __name__ == "__main__":
    main()