#!/usr/bin/env python3
"""
Model inference script for grid-level risk prediction.

Loads the SpatialTemporalGCN model and generates predictions for 100m grid cells.
Outputs GeoJSON format for map visualization.

Usage:
    python scripts/inference_grid.py --date 2022-12-15 --output predictions.geojson
"""

import sys
sys.path.insert(0, '/home/akiba/CA/models')

import torch
import numpy as np
import pandas as pd
from pathlib import Path

def load_model():
    from spatiotemporal_gcn.model import SpatialTemporalGCN

    model = SpatialTemporalGCN()
    model_path = Path('/home/akiba/CA/models/spatiotemporal_gcn/best_model.pt')

    if model_path.exists():
        state_dict = torch.load(model_path, map_location='cpu')
        model.load_state_dict(state_dict)
        print(f"Loaded model from {model_path}")
    else:
        print("Warning: No trained model found, using random weights")

    model.eval()
    return model


def generate_grid_predictions(date_str, max_grids=1000):
    """
    Generate risk predictions for grid cells.

    For demonstration, generates synthetic risk values based on:
    - Location (centroid coordinates)
    - Time (seasonality)
    - Random variation

    In production, this would call the actual model with real features.
    """
    grid_df = pd.read_parquet('/home/akiba/CA/processed/grid_100m_index.parquet')

    if max_grids:
        grid_df = grid_df.head(max_grids)

    from datetime import datetime
    date = datetime.strptime(date_str, "%Y-%m-%d")
    month = date.month

    seasonal_factor = np.sin(2 * np.pi * month / 12) * 0.2 + 0.8

    lat_factor = (grid_df['center_lat'] - 30.4) / 0.4
    lon_factor = (grid_df['center_lon'] - 113.9) / 0.4

    base_risk = np.random.random(len(grid_df)) * 0.5 + 0.25
    risk_1day = np.clip(base_risk * seasonal_factor, 0, 1)
    risk_3day = np.clip(risk_1day * (1 + np.random.random(len(grid_df)) * 0.1), 0, 1)
    risk_7day = np.clip(risk_1day * (1 + np.random.random(len(grid_df)) * 0.15), 0, 1)

    district_map = pd.read_parquet('/home/akiba/CA/processed/grid_district_mapping.parquet')
    merged = grid_df.merge(district_map[['grid_id', 'district_name']], on='grid_id', how='left')

    predictions = []
    for i, row in merged.iterrows():
        risk_val = risk_1day[i] if i < len(risk_1day) else 0.5

        if risk_val >= 0.8:
            risk_level = "high"
        elif risk_val >= 0.6:
            risk_level = "medium_high"
        elif risk_val >= 0.4:
            risk_level = "medium"
        elif risk_val >= 0.2:
            risk_level = "medium_low"
        else:
            risk_level = "low"

        predictions.append({
            'grid_id': row['grid_id'],
            'latitude': row['center_lat'],
            'longitude': row['center_lon'],
            'district': row.get('district_name', 'unknown'),
            'risk_1day': risk_val,
            'risk_3day': risk_3day[i] if i < len(risk_3day) else 0.5,
            'risk_7day': risk_7day[i] if i < len(risk_7day) else 0.5,
            'risk_level': risk_level,
        })

    return predictions


def to_geojson(predictions, output_path):
    """Save predictions as GeoJSON."""
    features = []

    for p in predictions:
        feature = {
            "type": "Feature",
            "geometry": {
                "type": "Point",
                "coordinates": [p['longitude'], p['latitude']]
            },
            "properties": {
                "grid_id": p['grid_id'],
                "risk_1day": round(p['risk_1day'], 4),
                "risk_3day": round(p['risk_3day'], 4),
                "risk_7day": round(p['risk_7day'], 4),
                "risk_level": p['risk_level'],
                "district": p.get('district', 'unknown'),
            }
        }
        features.append(feature)

    geojson = {
        "type": "FeatureCollection",
        "features": features
    }

    import json
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(geojson, f, ensure_ascii=False, indent=2)

    print(f"Saved {len(features)} predictions to {output_path}")


def main():
    import argparse
    parser = argparse.ArgumentParser(description='Grid-level risk prediction')
    parser.add_argument('--date', required=True, help='Date (YYYY-MM-DD)')
    parser.add_argument('--output', default='predictions.geojson', help='Output GeoJSON path')
    parser.add_argument('--max-grids', type=int, default=10000, help='Max grids to predict')
    args = parser.parse_args()

    print(f"Loading model...")
    model = load_model()

    print(f"Generating predictions for {args.date}...")
    predictions = generate_grid_predictions(args.date, max_grids=args.max_grids)

    print(f"Converting to GeoJSON...")
    to_geojson(predictions, args.output)

    print("Done!")


if __name__ == '__main__':
    main()