#!/usr/bin/env python3
"""
Alternative solution using pygrib for polar stereographic GRIB files
"""

import numpy as np
import pandas as pd

def extract_with_pygrib(grib_file):
    """
    Extract lat/lon and wave data using pygrib library
    """
    try:
        import pygrib
        
        print(f"Opening GRIB file with pygrib: {grib_file}")
        grbs = pygrib.open(grib_file)
        
        # List all messages
        print("Available messages:")
        for i, grb in enumerate(grbs):
            print(f"Message {i+1}: {grb}")
        
        grbs.seek(0)  # Reset to beginning
        
        # Get first message for coordinates
        grb = grbs[1]  # First message
        
        print(f"Grid type: {grb.gridType}")
        print(f"Projection: {getattr(grb, 'projString', 'N/A')}")
        
        # Get lat/lon coordinates
        try:
            lats, lons = grb.latlons()
            print(f"Successfully extracted coordinates")
            print(f"Lat shape: {lats.shape}, range: {lats.min():.3f} to {lats.max():.3f}")
            print(f"Lon shape: {lons.shape}, range: {lons.min():.3f} to {lons.max():.3f}")
            
            # Get data values
            data = grb.values
            print(f"Data shape: {data.shape}, range: {data.min():.3f} to {data.max():.3f}")
            
            # Create flat arrays for lat, lon, data
            lats_flat = lats.flatten()
            lons_flat = lons.flatten()
            data_flat = data.flatten()
            
            # Remove any invalid data points
            valid_mask = ~np.isnan(data_flat) & ~np.isnan(lats_flat) & ~np.isnan(lons_flat)
            
            result_df = pd.DataFrame({
                'latitude': lats_flat[valid_mask],
                'longitude': lons_flat[valid_mask],
                'wave_data': data_flat[valid_mask]
            })
            
            print(f"Created DataFrame with {len(result_df)} valid points")
            
            grbs.close()
            return result_df
            
        except Exception as e:
            print(f"Error extracting lat/lon with pygrib: {e}")
            
            # Try alternative coordinate extraction
            try:
                # Some GRIB files store coordinates differently
                if hasattr(grb, 'latitudeOfFirstGridPointInDegrees'):
                    lat_first = grb.latitudeOfFirstGridPointInDegrees
                    lon_first = grb.longitudeOfFirstGridPointInDegrees
                    print(f"First grid point: lat={lat_first}, lon={lon_first}")
                
                if hasattr(grb, 'Ni') and hasattr(grb, 'Nj'):
                    ni, nj = grb.Ni, grb.Nj
                    print(f"Grid dimensions: {ni} x {nj}")
                    
                # Get projection parameters
                if hasattr(grb, 'projParams'):
                    print(f"Projection parameters: {grb.projParams}")
                    
            except Exception as e2:
                print(f"Error getting grid parameters: {e2}")
            
            grbs.close()
            return None
            
    except ImportError:
        print("pygrib not available. Install with: pip install pygrib")
        return None
    except Exception as e:
        print(f"Error with pygrib: {e}")
        return None

def manual_polar_stereographic_projection(grib_file):
    """
    Manually implement polar stereographic coordinate calculation
    """
    try:
        import pygrib
        
        grbs = pygrib.open(grib_file)
        grb = grbs[1]
        
        # Get polar stereographic parameters
        params = {}
        attr_names = [
            'latitudeOfFirstGridPointInDegrees',
            'longitudeOfFirstGridPointInDegrees', 
            'DxInMetres',
            'DyInMetres',
            'orientationOfTheGridInDegrees',
            'latitudeWhereDxAndDyAreSpecifiedInDegrees',
            'Ni', 'Nj'
        ]
        
        for attr in attr_names:
            if hasattr(grb, attr):
                params[attr] = getattr(grb, attr)
                print(f"{attr}: {params[attr]}")
        
        # Manual coordinate calculation for polar stereographic
        if 'Ni' in params and 'Nj' in params:
            ni, nj = int(params['Ni']), int(params['Nj'])
            
            # Get grid spacing
            dx = params.get('DxInMetres', 25000)  # Default 25km
            dy = params.get('DyInMetres', 25000)
            
            # Create coordinate arrays
            x = np.arange(ni) * dx
            y = np.arange(nj) * dy
            
            X, Y = np.meshgrid(x, y)
            
            # Convert from polar stereographic to lat/lon
            # This is a simplified version - you may need a proper projection library
            lat_origin = params.get('latitudeWhereDxAndDyAreSpecifiedInDegrees', 90.0)
            lon_origin = params.get('orientationOfTheGridInDegrees', 0.0)
            
            print(f"Grid origin: lat={lat_origin}, lon={lon_origin}")
            print(f"Grid spacing: dx={dx}m, dy={dy}m")
            print(f"Grid size: {ni} x {nj}")
            
            # For proper conversion, you'd use pyproj or similar
            try:
                from pyproj import Proj, transform
                
                # Define polar stereographic projection
                proj_polar = Proj(
                    proj='stere',
                    lat_0=lat_origin,
                    lon_0=lon_origin,
                    lat_ts=lat_origin,
                    ellps='sphere'
                )
                
                proj_latlon = Proj(proj='latlong', ellps='sphere')
                
                # Transform coordinates
                lons, lats = transform(proj_polar, proj_latlon, X.flatten(), Y.flatten())
                
                lats = np.array(lats).reshape(nj, ni)
                lons = np.array(lons).reshape(nj, ni)
                
                # Get data
                data = grb.values
                
                # Create DataFrame
                result_df = pd.DataFrame({
                    'latitude': lats.flatten(),
                    'longitude': lons.flatten(),
                    'wave_data': data.flatten()
                })
                
                # Remove invalid points
                valid_mask = ~np.isnan(result_df['wave_data'])
                result_df = result_df[valid_mask]
                
                print(f"Successfully calculated coordinates for {len(result_df)} points")
                
                grbs.close()
                return result_df
                
            except ImportError:
                print("pyproj not available for coordinate transformation")
                grbs.close()
                return None
        
        grbs.close()
        return None
        
    except Exception as e:
        print(f"Error in manual projection: {e}")
        return None

# Example usage
if __name__ == "__main__":
    grib_file = "/tmp/tmpr004q4kw.grib2"  # Your file path
    
    print("=== Trying pygrib extraction ===")
    result1 = extract_with_pygrib(grib_file)
    
    if result1 is not None:
        print(f"Success with pygrib! {len(result1)} points extracted")
        print(result1.head())
    else:
        print("=== Trying manual polar stereographic ===")
        result2 = manual_polar_stereographic_projection(grib_file)
        
        if result2 is not None:
            print(f"Success with manual projection! {len(result2)} points extracted")
            print(result2.head())
        else:
            print("All methods failed. Consider using wgrib2 conversion first.")