#!/usr/bin/env python3
"""
Weather interpolation to 100m grid using scipy griddata.
Optimized: vectorized operations, chunked processing, gzip compression.
"""

import pandas as pd
import numpy as np
from scipy.interpolate import griddata
from pathlib import Path
import time
import warnings
warnings.filterwarnings('ignore')

def process_year_fast(year, station_data_dir, grid_parquet_path, output_dir):
    print(f"=== Processing year {year} (optimized) ===")
    t0 = time.time()

    grid_df = pd.read_parquet(grid_parquet_path)
    grid_ids = grid_df['grid_id'].values
    grid_points = grid_df[['center_lon', 'center_lat']].values
    n_grid = len(grid_df)
    print(f"Grid: {n_grid:,} cells")

    station_df = pd.read_parquet(f"{station_data_dir}/daily_wuhan_{year}.parquet")
    station_df['date'] = pd.to_datetime(station_df['date']).dt.date
    dates = sorted(station_df['date'].unique())
    print(f"Days: {len(dates)}")

    pollutants = ['AQI', 'PM25', 'PM10', 'SO2', 'NO2', 'O3', 'CO']

    station_locs = station_df.groupby('station_id').first()[['lat', 'lon', 'district']].reset_index()
    print(f"Stations: {len(station_locs)}")

    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    for poll_idx, poll in enumerate(pollutants):
        print(f"\n[{poll_idx+1}/{len(pollutants)}] {poll}...")
        t1 = time.time()

        all_records = []
        chunk_size = 50

        for chunk_start in range(0, len(dates), chunk_size):
            chunk_dates = dates[chunk_start:chunk_start + chunk_size]
            chunk_records = []

            for date in chunk_dates:
                day_data = station_df[station_df['date'] == date]
                values = day_data.set_index('station_id')[poll]
                merged = station_locs.merge(values.reset_index(), on='station_id', how='inner')

                if len(merged) < 3:
                    continue

                sc = merged[['lon', 'lat']].values
                sv = merged[poll].values
                valid_mask = ~pd.isna(sv)

                if valid_mask.sum() < 3:
                    continue

                result = griddata(sc[valid_mask], sv[valid_mask], grid_points, method='nearest')

                if result is not None and not np.all(np.isnan(result)):
                    valid_result = ~np.isnan(result)
                    if valid_result.any():
                        day_records = pd.DataFrame({
                            'grid_id': grid_ids[valid_result],
                            'date': date,
                            'pollutant': poll,
                            'value': result[valid_result].astype(np.float32)
                        })
                        chunk_records.append(day_records)

            if chunk_records:
                all_records.append(pd.concat(chunk_records, ignore_index=True))

            print(f"  {min(chunk_start + chunk_size, len(dates))}/{len(dates)} days")

        if all_records:
            final_df = pd.concat(all_records, ignore_index=True)
            out_file = output_path / f'grid_weather_{poll}_{year}.parquet'
            final_df.to_parquet(out_file, index=False, compression='gzip')
            size_mb = out_file.stat().st_size / 1024 / 1024
            print(f"  Saved: {len(final_df):,} records, {size_mb:.1f} MB")
        else:
            print(f"  No valid data")

        print(f"  Time: {time.time()-t1:.0f}s")

    print(f"\n=== Total: {time.time()-t0:.0f}s ===")

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--year', type=int, required=True)
    parser.add_argument('--station-data-dir', default='processed/weather')
    parser.add_argument('--grid-parquet', default='processed/grid_100m_index.parquet')
    parser.add_argument('--output-dir', default='processed/weather')
    args = parser.parse_args()
    process_year_fast(args.year, args.station_data_dir, args.grid_parquet, args.output_dir)