"""
Generate 100x100m grid summary from geocoded case data
Uses EPSG:4326 coordinates (degrees) directly
"""
import pandas as pd
import numpy as np
from pathlib import Path
from shapely.geometry import Point, box
from shapely.ops import unary_union
import warnings

warnings.filterwarnings('ignore')

# Paths
OUTPUT_DIR = Path("/home/akiba/CA/outputs")
INPUT_FILE = OUTPUT_DIR / "geocoded_all_cases.csv"
OUTPUT_FILE = OUTPUT_DIR / "grid_risk_summary.csv"

# Wuhan bounding box (EPSG:4326 degrees)
WUHAN_BOUNDS = {
    'min_lat': 29.9,
    'max_lat': 31.4,
    'min_lon': 113.6,
    'max_lon': 115.1
}

# Grid resolution: 100m in degrees at Wuhan latitude (~30.5°)
# 1 degree latitude ≈ 111 km
# 1 degree longitude ≈ 111 km * cos(latitude)
GRID_SIZE_LAT = 0.0009  # ~100m latitude
GRID_SIZE_LON = 0.0010  # ~100m longitude at 30.5° latitude


def filter_valid_coordinates(df: pd.DataFrame) -> pd.DataFrame:
    """Filter out invalid coordinates"""
    # Remove null coordinates
    df = df.dropna(subset=['latitude', 'longitude'])
    
    # Filter valid Wuhan bounds
    df = df[
        (df['latitude'] >= WUHAN_BOUNDS['min_lat']) &
        (df['latitude'] <= WUHAN_BOUNDS['max_lat']) &
        (df['longitude'] >= WUHAN_BOUNDS['min_lon']) &
        (df['longitude'] <= WUHAN_BOUNDS['max_lon'])
    ]
    
    # Filter swapped coordinates (lat > 50 or lon > 120 indicates swap)
    df = df[
        (df['latitude'] < 50) &
        (df['longitude'] < 120)
    ]
    
    return df


def create_grid() -> pd.DataFrame:
    """Create 100x100m grid over Wuhan area"""
    grids = []
    grid_id = 0
    
    lat_min = WUHAN_BOUNDS['min_lat']
    lat_max = WUHAN_BOUNDS['max_lat']
    lon_min = WUHAN_BOUNDS['min_lon']
    lon_max = WUHAN_BOUNDS['max_lon']
    
    lat = lat_min
    while lat < lat_max:
        lon = lon_min
        while lon < lon_max:
            center_y = lat + GRID_SIZE_LAT / 2
            center_x = lon + GRID_SIZE_LON / 2
            
            grids.append({
                'grid_id': grid_id,
                'center_y': round(center_y, 6),
                'center_x': round(center_x, 6),
                'lat_min': lat,
                'lat_max': lat + GRID_SIZE_LAT,
                'lon_min': lon,
                'lon_max': lon + GRID_SIZE_LON
            })
            grid_id += 1
            lon += GRID_SIZE_LON
        lat += GRID_SIZE_LAT
    
    return pd.DataFrame(grids)


def aggregate_cases_to_grid(cases_df: pd.DataFrame, grid_df: pd.DataFrame) -> pd.DataFrame:
    """Aggregate cases to grid cells"""
    # Assign each case to a grid cell
    cases_df['grid_lat_idx'] = ((cases_df['latitude'] - WUHAN_BOUNDS['min_lat']) / GRID_SIZE_LAT).astype(int)
    cases_df['grid_lon_idx'] = ((cases_df['longitude'] - WUHAN_BOUNDS['min_lon']) / GRID_SIZE_LON).astype(int)
    cases_df['grid_id'] = cases_df['grid_lat_idx'] * int((WUHAN_BOUNDS['max_lon'] - WUHAN_BOUNDS['min_lon']) / GRID_SIZE_LON) + cases_df['grid_lon_idx']
    
    # Aggregate by grid
    grid_stats = cases_df.groupby('grid_id').agg(
        total_cases=('case_id', 'count'),
        outpatient_cases=('case_type', lambda x: (x == 'outpatient').sum()),
        inpatient_cases=('case_type', lambda x: (x == 'inpatient').sum())
    ).reset_index()
    
    # Merge with grid geometry
    result = grid_df.merge(grid_stats, on='grid_id', how='left')
    
    # Fill NaN with 0 for grids with no cases
    result['total_cases'] = result['total_cases'].fillna(0).astype(int)
    result['outpatient_cases'] = result['outpatient_cases'].fillna(0).astype(int)
    result['inpatient_cases'] = result['inpatient_cases'].fillna(0).astype(int)
    
    # Calculate case density (cases per km²)
    # Grid area = 0.1 km × 0.1 km = 0.01 km²
    result['cases_per_km2'] = result['total_cases'] / 0.01
    
    # Calculate risk index (normalized by max cases)
    max_cases = result['total_cases'].max()
    if max_cases > 0:
        result['risk_index'] = result['total_cases'] / max_cases
    else:
        result['risk_index'] = 0.0
    
    # Assign risk level
    def get_risk_level(risk_index):
        if risk_index >= 0.8:
            return 'high'
        elif risk_index >= 0.6:
            return 'medium_high'
        elif risk_index >= 0.4:
            return 'medium'
        elif risk_index >= 0.2:
            return 'medium_low'
        else:
            return 'low'
    
    result['risk_level'] = result['risk_index'].apply(get_risk_level)
    
    # Select final columns
    result = result[[
        'grid_id', 'center_y', 'center_x',
        'total_cases', 'outpatient_cases', 'inpatient_cases',
        'cases_per_km2', 'risk_index', 'risk_level'
    ]]
    
    return result


def main():
    print(f"Reading geocoded cases from {INPUT_FILE}...")
    df = pd.read_csv(INPUT_FILE)
    print(f"  Total records: {len(df):,}")
    
    print("Filtering valid coordinates...")
    df = filter_valid_coordinates(df)
    print(f"  Valid records: {len(df):,}")
    
    print("Creating 100x100m grid...")
    grid_df = create_grid()
    print(f"  Total grid cells: {len(grid_df):,}")
    
    print("Aggregating cases to grid...")
    result = aggregate_cases_to_grid(df, grid_df)
    
    print(f"Saving to {OUTPUT_FILE}...")
    result.to_csv(OUTPUT_FILE, index=False)
    
    # Summary statistics
    print("\n=== Summary ===")
    print(f"Grid cells with cases: {(result['total_cases'] > 0).sum():,}")
    print(f"Total cases: {result['total_cases'].sum():,}")
    print(f"Max cases in single grid: {result['total_cases'].max():,}")
    print(f"Risk index range: {result['risk_index'].min():.3f} - {result['risk_index'].max():.3f}")
    print(f"Coordinate ranges:")
    print(f"  Latitude: {result['center_y'].min():.4f} to {result['center_y'].max():.4f}")
    print(f"  Longitude: {result['center_x'].min():.4f} to {result['center_x'].max():.4f}")
    print("\nSample row:")
    print(result.iloc[0].to_dict())


if __name__ == "__main__":
    main()